Merge pull request #457 from Koed00/resource_limits

Resource limits: max rss for workers

django_q/__init__.py

@@ -3,10 +3,4 @@ VERSION = (1, 3, 1)
 default_app_config = "django_q.apps.DjangoQConfig"
 
 
-__all__ = ["conf", "cluster", "models", "tasks", "croniter"]
-
-# Optional Imports
-try:
-    from croniter import croniter
-except ImportError:
-    croniter = None
+__all__ = ["conf", "cluster", "models", "tasks"]

django_q/admin.py

@@ -2,10 +2,9 @@
 from django.contrib import admin
 from django.utils.translation import gettext_lazy as _
 
-from django_q.conf import Conf
+from django_q.conf import Conf, croniter
 from django_q.models import Success, Failure, Schedule, OrmQ
 from django_q.tasks import async_task
-from django_q import croniter
 
 
 class TaskAdmin(admin.ModelAdmin):

django_q/cluster.py

@@ -10,6 +10,7 @@ from time import sleep
 
 # External
 import arrow
+
 # Django
 from django import db
 from django.conf import settings
@@ -19,14 +20,21 @@ from django.utils.translation import gettext_lazy as _
 # Local
 import django_q.tasks
 from django_q.brokers import get_broker, Broker
-from django_q.conf import Conf, logger, psutil, get_ppid, error_reporter
+from django_q.conf import (
+    Conf,
+    logger,
+    psutil,
+    get_ppid,
+    error_reporter,
+    croniter,
+    resource,
+)
 from django_q.humanhash import humanize
 from django_q.models import Task, Success, Schedule
 from django_q.queues import Queue
 from django_q.signals import pre_execute
 from django_q.signing import SignedPackage, BadSignature
 from django_q.status import Stat, Status
-from django_q import croniter
 
 
 class Cluster:
@@ -102,10 +110,10 @@ class Cluster:
     @property
     def is_stopping(self) -> bool:
         return (
-                self.stop_event
-                and self.start_event
-                and self.start_event.is_set()
-                and self.stop_event.is_set()
+            self.stop_event
+            and self.start_event
+            and self.start_event.is_set()
+            and self.stop_event.is_set()
         )
 
     @property
@@ -115,13 +123,13 @@ class Cluster:
 
 class Sentinel:
     def __init__(
-            self,
-            stop_event,
-            start_event,
-            cluster_id,
-            broker=None,
-            timeout=Conf.TIMEOUT,
-            start=True,
+        self,
+        stop_event,
+        start_event,
+        cluster_id,
+        broker=None,
+        timeout=Conf.TIMEOUT,
+        start=True,
     ):
         # Make sure we catch signals for the pool
         signal.signal(signal.SIGINT, signal.SIG_IGN)
@@ -376,7 +384,7 @@ def monitor(result_queue: Queue, broker: Broker = None):
 
 
 def worker(
-        task_queue: Queue, result_queue: Queue, timer: Value, timeout: int = Conf.TIMEOUT
+    task_queue: Queue, result_queue: Queue, timer: Value, timeout: int = Conf.TIMEOUT
 ):
     """
     Takes a task from the task queue, tries to execute it and puts the result back in the result queue
@@ -433,7 +441,7 @@ def worker(
             result_queue.put(task)
             timer.value = -1  # Idle
             # Recycle
-            if task_count == Conf.RECYCLE:
+            if task_count == Conf.RECYCLE or rss_check():
                 timer.value = -2  # Recycled
                 break
     logger.info(_(f"{name} stopped doing work"))
@@ -551,9 +559,9 @@ def scheduler(broker: Broker = None):
     try:
         with db.transaction.atomic(using=Schedule.objects.db):
             for s in (
-                    Schedule.objects.select_for_update()
-                            .exclude(repeats=0)
-                            .filter(next_run__lt=timezone.now())
+                Schedule.objects.select_for_update()
+                .exclude(repeats=0)
+                .filter(next_run__lt=timezone.now())
             ):
                 args = ()
                 kwargs = {}
@@ -692,3 +700,11 @@ def set_cpu_affinity(n: int, process_ids: list, actual: bool = not Conf.TESTING)
             if actual:
                 p.cpu_affinity(affinity)
             logger.info(_(f"{pid} will use cpu {affinity}"))
+
+
+def rss_check():
+    if Conf.MAX_RSS and resource:
+        return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss >= Conf.MAX_RSS
+    elif Conf.MAX_RSS and psutil:
+        return psutil.Process().memory_info().rss >= Conf.MAX_RSS * 1024
+    return False

django_q/conf.py

@@ -16,6 +16,16 @@ try:
 except ImportError:
     psutil = None
 
+try:
+    from croniter import croniter
+except ImportError:
+    croniter = None
+
+try:
+    import resource
+except ModuleNotFoundError:
+    resource = None
+
 
 class Conf:
     """
@@ -104,6 +114,10 @@ class Conf:
     # Number of tasks each worker can handle before it gets recycled. Useful for releasing memory
     RECYCLE = conf.get("recycle", 500)
 
+    # The maximum resident set size in kilobytes before a worker will recycle. Useful for limiting memory usage
+    # Not available on all platforms
+    MAX_RSS = conf.get("max_rss", None)
+
     # Number of seconds to wait for a worker to finish.
     TIMEOUT = conf.get("timeout", None)
 
@@ -211,7 +225,7 @@ if Conf.ERROR_REPORTER:
         # and instantiate an ErrorReporter using the provided config
         for name, conf in error_conf.items():
             for entry in pkg_resources.iter_entry_points(
-                    "djangoq.errorreporters", name
+                "djangoq.errorreporters", name
             ):
                 Reporter = entry.load()
                 reporters.append(Reporter(**conf))

django_q/models.py

@@ -13,8 +13,8 @@ from picklefield import PickledObjectField
 from picklefield.fields import dbsafe_decode
 
 # Local
+from django_q.conf import croniter
 from django_q.signing import SignedPackage
-from django_q import croniter
 
 
 class Task(models.Model):

django_q/tests/test_cluster.py

@@ -338,6 +338,40 @@ def test_recycle(broker, monkeypatch):
     assert Success.objects.count() == Conf.SAVE_LIMIT
     broker.delete_queue()
 
+@pytest.mark.django_db
+def test_max_rss(broker, monkeypatch):
+    # set up the Sentinel
+    broker.list_key = 'test_max_rss_test:q'
+    async_task('django_q.tests.tasks.multiply', 2, 2, broker=broker)
+    start_event = Event()
+    stop_event = Event()
+    cluster_id = uuidlib.uuid4()
+    # override settings
+    monkeypatch.setattr(Conf, 'MAX_RSS', 40000)
+    monkeypatch.setattr(Conf, 'WORKERS', 1)
+    # set a timer to stop the Sentinel
+    threading.Timer(3, stop_event.set).start()
+    s = Sentinel(stop_event, start_event, cluster_id=cluster_id, broker=broker)
+    assert start_event.is_set()
+    assert s.status() == Conf.STOPPED
+    assert s.reincarnations == 1
+    async_task('django_q.tests.tasks.multiply', 2, 2, broker=broker)
+    task_queue = Queue()
+    result_queue = Queue()
+    # push the task
+    pusher(task_queue, stop_event, broker=broker)
+    # worker should exit on recycle
+    worker(task_queue, result_queue, Value('f', -1))
+    # check if the work has been done
+    assert result_queue.qsize() == 1
+    # save_limit test
+    monkeypatch.setattr(Conf, 'SAVE_LIMIT', 1)
+    result_queue.put('STOP')
+    # run monitor
+    monitor(result_queue)
+    assert Success.objects.count() == Conf.SAVE_LIMIT
+    broker.delete_queue()
+
 
 @pytest.mark.django_db
 def test_bad_secret(broker, monkeypatch):

docs/configure.rst

@@ -51,6 +51,13 @@ recycle
 
 The number of tasks a worker will process before recycling . Useful to release memory resources on a regular basis. Defaults to ``500``.
 
+max_rss
+~~~~~~~
+
+The maximum resident set size in kilobytes before a worker will recycle and release resources. Useful for limiting memory usage.
+Only supported on platforms that implement the python resource module or install the :ref:`psutil<psutil_package>` module.
+Defaults to ``None``.
+
 .. _timeout:
 
 timeout

poetry.lock

@@ -71,10 +71,10 @@ description = "The AWS SDK for Python"
 name = "boto3"
 optional = true
 python-versions = "*"
-version = "1.14.16"
+version = "1.14.19"
 
 [package.dependencies]
-botocore = ">=1.17.16,<1.18.0"
+botocore = ">=1.17.19,<1.18.0"
 jmespath = ">=0.7.1,<1.0.0"
 s3transfer = ">=0.3.0,<0.4.0"
 
@@ -84,7 +84,7 @@ description = "Low-level, data-driven core of boto 3."
 name = "botocore"
 optional = true
 python-versions = "*"
-version = "1.17.16"
+version = "1.17.19"
 
 [package.dependencies]
 docutils = ">=0.10,<0.16"
@@ -550,12 +550,12 @@ blessed = [
     {file = "blessed-1.17.8.tar.gz", hash = "sha256:7671d057b2df6ddbefd809009fb08feb2f8d2d163d240b5e765088a90519b2f1"},
 ]
 boto3 = [
-    {file = "boto3-1.14.16-py2.py3-none-any.whl", hash = "sha256:c2a223f4b48782e8b160b2130265e2a66081df111f630a5a384d6909e29a5aa9"},
-    {file = "boto3-1.14.16.tar.gz", hash = "sha256:ce5a4ab6af9e993d1864209cbbb6f4812f65fbc57ad6b95e5967d8bf38b1dcfb"},
+    {file = "boto3-1.14.19-py2.py3-none-any.whl", hash = "sha256:1776200c04152dd308e297cd18abb183d93383587dca358a768155507f15965b"},
+    {file = "boto3-1.14.19.tar.gz", hash = "sha256:944c02dbf96dbaf52498acdbea60d4eb2aa717f40a797f86f8d47d8905e02441"},
 ]
 botocore = [
-    {file = "botocore-1.17.16-py2.py3-none-any.whl", hash = "sha256:99d995ef99cf77458a661f3fc64e0c3a4ce77ca30facfdf0472f44b2953dd856"},
-    {file = "botocore-1.17.16.tar.gz", hash = "sha256:fe0c4f7cd6b67eff3b7cb8dff6709a65d6fca10b7b7449a493b2036915e98b4c"},
+    {file = "botocore-1.17.19-py2.py3-none-any.whl", hash = "sha256:7890e83cd28967f854fd54d6c8bdb009868aec1c31a71dcc82e4577562f1affb"},
+    {file = "botocore-1.17.19.tar.gz", hash = "sha256:4eef7d38de1bee3bb60a66d53a73d47ec5ea30d3d43befa90840ba1d52791971"},
 ]
 certifi = [
     {file = "certifi-2020.6.20-py2.py3-none-any.whl", hash = "sha256:8fc0819f1f30ba15bdb34cceffb9ef04d99f420f68eb75d901e9560b8749fc41"},

requirements.txt

@@ -7,8 +7,8 @@
 arrow==0.15.7             # via -r requirements.in
 asgiref==3.2.10           # via django
 blessed==1.17.8           # via -r requirements.in
-boto3==1.14.16            # via -r requirements.in
-botocore==1.17.16         # via boto3, s3transfer
+boto3==1.14.19            # via -r requirements.in
+botocore==1.17.19         # via boto3, s3transfer
 certifi==2020.6.20        # via requests
 chardet==3.0.4            # via requests
 croniter==0.3.34          # via -r requirements.in