Check queue size to launch the heavy processing worker (#4288)

Author: marcoacierno

backend/pycon/celery.py

@@ -23,7 +23,10 @@ def setup_periodic_tasks(sender, **kwargs):
     )
     from schedule.tasks import process_schedule_items_videos_to_upload
     from files_upload.tasks import delete_unused_files
-    from pycon.tasks import check_for_idle_heavy_processing_workers
+    from pycon.tasks import (
+        check_for_idle_heavy_processing_workers,
+        check_pending_heavy_processing_work,
+    )
 
     add = sender.add_periodic_task
 
@@ -47,3 +50,8 @@ def setup_periodic_tasks(sender, **kwargs):
         check_for_idle_heavy_processing_workers,
         name="Check for idle heavy processing workers",
     )
+    add(
+        timedelta(minutes=1),
+        check_pending_heavy_processing_work,
+        name="Check pending heavy processing work",
+    )

backend/pycon/tasks.py

@@ -3,7 +3,7 @@
 import boto3
 from django.conf import settings
 from pycon.celery_utils import OnlyOneAtTimeTask
-
+from redis import Redis
 from collections import Counter
 from pycon.celery import app
 from django.core.cache import cache
@@ -153,3 +153,14 @@ def check_for_idle_heavy_processing_workers():
 
 def build_idle_worker_cache_key(worker_name: str) -> str:
     return f"celery-workers-idle:{worker_name}"
+
+
+@app.task(base=OnlyOneAtTimeTask)
+def check_pending_heavy_processing_work():
+    redis = Redis.from_url(settings.CELERY_BROKER_URL)
+    tasks_in_queue = redis.llen("heavy_processing")
+
+    if tasks_in_queue == 0:
+        return
+
+    launch_heavy_processing_worker.delay()

backend/pycon/tests/test_tasks.py

@@ -3,6 +3,7 @@
     build_idle_worker_cache_key,
     launch_heavy_processing_worker,
     check_for_idle_heavy_processing_workers,
+    check_pending_heavy_processing_work,
 )
 
 
@@ -383,3 +384,23 @@ def test_dont_shutdown_heavy_processing_worker_if_last_check_is_not_5_mins_old(m
 
     mock_cache.set.assert_not_called()
     mock_broadcast.assert_not_called()
+
+
+def test_check_pending_heavy_processing_work(mocker):
+    mock_launch = mocker.patch("pycon.tasks.launch_heavy_processing_worker")
+    redis_mock = mocker.patch("pycon.tasks.Redis", return_value=mocker.MagicMock())
+    redis_mock.from_url.return_value.llen.return_value = 10
+
+    check_pending_heavy_processing_work()
+
+    mock_launch.delay.assert_called()
+
+
+def test_check_pending_heavy_processing_work_with_no_work(mocker):
+    mock_launch = mocker.patch("pycon.tasks.launch_heavy_processing_worker")
+    redis_mock = mocker.patch("pycon.tasks.Redis", return_value=mocker.MagicMock())
+    redis_mock.from_url.return_value.llen.return_value = 0
+
+    check_pending_heavy_processing_work()
+
+    mock_launch.delay.assert_not_called()

backend/video_uploads/admin.py

@@ -1,8 +1,8 @@
-from pycon.tasks import launch_heavy_processing_worker
 from django.utils.safestring import mark_safe
 from django.contrib import admin
 from django.db import transaction
 
+from pycon.tasks import check_pending_heavy_processing_work
 from video_uploads.tasks import process_wetransfer_to_s3_transfer_request
 from video_uploads.models import WetransferToS3TransferRequest
 
@@ -16,7 +16,7 @@ def _on_commit():
         process_wetransfer_to_s3_transfer_request.apply_async(
             args=[request_obj.id], queue="heavy_processing"
         )
-        launch_heavy_processing_worker.delay()
+        check_pending_heavy_processing_work.delay()
 
     transaction.on_commit(_on_commit)
 

backend/video_uploads/tests/test_admin.py

@@ -12,8 +12,8 @@ def test_queue_wetransfer_to_s3_transfer_request(
     mock_process_wetransfer_to_s3_transfer_request = mocker.patch(
         "video_uploads.admin.process_wetransfer_to_s3_transfer_request"
     )
-    mock_launch_heavy_processing_worker = mocker.patch(
-        "video_uploads.admin.launch_heavy_processing_worker"
+    mock_check_pending_heavy_processing_work = mocker.patch(
+        "video_uploads.admin.check_pending_heavy_processing_work"
     )
 
     request = WetransferToS3TransferRequestFactory()
@@ -29,7 +29,7 @@ def test_queue_wetransfer_to_s3_transfer_request(
     mock_process_wetransfer_to_s3_transfer_request.apply_async.assert_called_once_with(
         args=[request.id], queue="heavy_processing"
     )
-    mock_launch_heavy_processing_worker.delay.assert_called_once()
+    mock_check_pending_heavy_processing_work.delay.assert_called_once()
 
 
 def test_retry_transfer():