From 008bf0bb5492bbf07c1e10550ee18caea1645379 Mon Sep 17 00:00:00 2001
From: Cristhian Garcia <crisgarta8@gmail.com>
Date: Fri, 4 Oct 2024 16:36:18 -0500
Subject: [PATCH] feat: allow to setup multiple celery queues

---
 tutor/hooks/catalog.py                        |  3 ++
 tutor/plugins/openedx.py                      | 43 +++++++++++++++++++
 tutor/templates/local/docker-compose.prod.yml | 38 ++++++----------
 3 files changed, 58 insertions(+), 26 deletions(-)

diff --git a/tutor/hooks/catalog.py b/tutor/hooks/catalog.py
index b1bfcdbf7c..be4fac7fea 100644
--- a/tutor/hooks/catalog.py
+++ b/tutor/hooks/catalog.py
@@ -514,6 +514,9 @@ def your_filter_callback(some_data):
     #: :param str file_path: The path to the file being checked.
     IS_FILE_RENDERED: Filter[bool, [str]] = Filter()
 
+    #: List of workers
+    CELERY_WORKERS_CONFIG: Filter[dict[str, dict[str, Any]], []] = Filter()
+
 
 class Contexts:
     """
diff --git a/tutor/plugins/openedx.py b/tutor/plugins/openedx.py
index 2984d60792..010b9a7142 100644
--- a/tutor/plugins/openedx.py
+++ b/tutor/plugins/openedx.py
@@ -142,3 +142,46 @@ def is_directory_mounted(image_name: str, dirname: str) -> bool:
 hooks.Filters.ENV_TEMPLATE_VARIABLES.add_item(
     ("iter_mounted_directories", iter_mounted_directories)
 )
+
+
+@hooks.Filters.CELERY_WORKERS_CONFIG.add(priority=hooks.priorities.HIGH)
+def _add_default_celery_queues(worker_configs):
+    worker_configs = {
+        "lms": {
+            "default": {
+                "concurrency": 1,
+                "pool": "prefork",
+            },
+        },
+        "cms": {
+            "default": {
+                "concurrency": 1,
+                "pool": "prefork",
+            },
+        },
+    }
+    return worker_configs
+
+
+@hooks.lru_cache
+def get_celery_workers_config() -> dict[str, dict[str, t.Any]]:
+    """
+    This function is cached for performance.
+    """
+    return hooks.Filters.CELERY_WORKERS_CONFIG.apply({})
+
+
+def iter_celery_workers_config() -> dict[str, dict[str, t.Any]]:
+    """
+    Yield:
+
+        (name, dict)
+    """
+    return {name: config for name, config in get_celery_workers_config().items()}
+
+
+hooks.Filters.ENV_TEMPLATE_VARIABLES.add_items(
+    [
+        ("iter_celery_workers_config", iter_celery_workers_config),
+    ]
+)
diff --git a/tutor/templates/local/docker-compose.prod.yml b/tutor/templates/local/docker-compose.prod.yml
index aa08f24121..e13e53dc2f 100644
--- a/tutor/templates/local/docker-compose.prod.yml
+++ b/tutor/templates/local/docker-compose.prod.yml
@@ -26,42 +26,28 @@ services:
     {% endif %}
 
   ############# LMS and CMS workers
-  lms-worker:
+  {% for service, variants in iter_celery_workers_config().items() %}
+  {% for variant, config in variants.items() %}
+  {% set deployment = service + "-" + "worker" + "-" + variant.replace("_", "-") %}
+  {{deployment}}:
     image: {{ DOCKER_IMAGE_OPENEDX }}
     environment:
-      SERVICE_VARIANT: lms
-      DJANGO_SETTINGS_MODULE: lms.envs.tutor.production
-    command: celery --app=lms.celery worker --loglevel=info --hostname=edx.lms.core.default.%%h --max-tasks-per-child=100 --exclude-queues=edx.cms.core.default
+      SERVICE_VARIANT: {{service}}
+      DJANGO_SETTINGS_MODULE: {{service}}.envs.tutor.production
+    command: celery --app={{service}}.celery worker --loglevel=info --hostname=edx.{{service}}.core.default.%%h --max-tasks-per-child=100 --queues=edx.{{service}}.core.{{variant}}
     restart: unless-stopped
     volumes:
       - ../apps/openedx/settings/lms:/openedx/edx-platform/lms/envs/tutor:ro
       - ../apps/openedx/settings/cms:/openedx/edx-platform/cms/envs/tutor:ro
       - ../apps/openedx/config:/openedx/config:ro
-      - ../../data/lms:/openedx/data
+      - ../../data/{{service}}:/openedx/data
       - ../../data/openedx-media:/openedx/media
-      {%- for mount in iter_mounts(MOUNTS, "openedx", "lms-worker") %}
+      {%- for mount in iter_mounts(MOUNTS, "openedx", service + "-worker") %}
       - {{ mount }}
       {%- endfor %}
     depends_on:
-      - lms
-
-  cms-worker:
-    image: {{ DOCKER_IMAGE_OPENEDX }}
-    environment:
-      SERVICE_VARIANT: cms
-      DJANGO_SETTINGS_MODULE: cms.envs.tutor.production
-    command: celery --app=cms.celery worker --loglevel=info --hostname=edx.cms.core.default.%%h --max-tasks-per-child 100 --exclude-queues=edx.lms.core.default
-    restart: unless-stopped
-    volumes:
-      - ../apps/openedx/settings/lms:/openedx/edx-platform/lms/envs/tutor:ro
-      - ../apps/openedx/settings/cms:/openedx/edx-platform/cms/envs/tutor:ro
-      - ../apps/openedx/config:/openedx/config:ro
-      - ../../data/cms:/openedx/data
-      - ../../data/openedx-media:/openedx/media
-      {%- for mount in iter_mounts(MOUNTS, "openedx", "cms-worker") %}
-      - {{ mount }}
-      {%- endfor %}
-    depends_on:
-      - cms
+      - {{service}}
+    {% endfor %}
+    {% endfor %}
 
   {{ patch("local-docker-compose-prod-services")|indent(2) }}