Feature/more celery fanout (#3740)

* WIP * migrate most beat tasks to fan out strategy * fix kwargs * migrate EE tasks * lock on the task_name level * typo fix --------- Co-authored-by: Richard Kuo (Danswer) <[email protected]>
onyx-dot-app · Jan 23, 2025 · a2d8e81 · a2d8e81
1 parent b1e05bb
commit a2d8e81
Show file tree

Hide file tree

Showing 6 changed files with 299 additions and 137 deletions.
diff --git a/backend/ee/onyx/background/celery/tasks/beat_schedule.py b/backend/ee/onyx/background/celery/tasks/beat_schedule.py
@@ -1,30 +1,70 @@
 from datetime import timedelta
 from typing import Any
 
+from onyx.background.celery.tasks.beat_schedule import BEAT_EXPIRES_DEFAULT
 from onyx.background.celery.tasks.beat_schedule import (
     cloud_tasks_to_schedule as base_cloud_tasks_to_schedule,
 )
 from onyx.background.celery.tasks.beat_schedule import (
     tasks_to_schedule as base_tasks_to_schedule,
 )
+from onyx.configs.constants import ONYX_CLOUD_CELERY_TASK_PREFIX
+from onyx.configs.constants import OnyxCeleryPriority
 from onyx.configs.constants import OnyxCeleryTask
+from shared_configs.configs import MULTI_TENANT
 
-ee_tasks_to_schedule = [
+ee_cloud_tasks_to_schedule = [
     {
-        "name": "autogenerate-usage-report",
-        "task": OnyxCeleryTask.AUTOGENERATE_USAGE_REPORT_TASK,
-        "schedule": timedelta(days=30),  # TODO: change this to config flag
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_autogenerate-usage-report",
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
+        "schedule": timedelta(days=30),
+        "options": {
+            "priority": OnyxCeleryPriority.HIGHEST,
+            "expires": BEAT_EXPIRES_DEFAULT,
+        },
+        "kwargs": {
+            "task_name": OnyxCeleryTask.AUTOGENERATE_USAGE_REPORT_TASK,
+        },
     },
     {
-        "name": "check-ttl-management",
-        "task": OnyxCeleryTask.CHECK_TTL_MANAGEMENT_TASK,
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-ttl-management",
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
         "schedule": timedelta(hours=1),
+        "options": {
+            "priority": OnyxCeleryPriority.HIGHEST,
+            "expires": BEAT_EXPIRES_DEFAULT,
+        },
+        "kwargs": {
+            "task_name": OnyxCeleryTask.CHECK_TTL_MANAGEMENT_TASK,
+        },
     },
 ]
 
+if not MULTI_TENANT:
+    ee_tasks_to_schedule = [
+        {
+            "name": "autogenerate-usage-report",
+            "task": OnyxCeleryTask.AUTOGENERATE_USAGE_REPORT_TASK,
+            "schedule": timedelta(days=30),  # TODO: change this to config flag
+            "options": {
+                "priority": OnyxCeleryPriority.MEDIUM,
+                "expires": BEAT_EXPIRES_DEFAULT,
+            },
+        },
+        {
+            "name": "check-ttl-management",
+            "task": OnyxCeleryTask.CHECK_TTL_MANAGEMENT_TASK,
+            "schedule": timedelta(hours=1),
+            "options": {
+                "priority": OnyxCeleryPriority.MEDIUM,
+                "expires": BEAT_EXPIRES_DEFAULT,
+            },
+        },
+    ]
+
 
 def get_cloud_tasks_to_schedule() -> list[dict[str, Any]]:
-    return base_cloud_tasks_to_schedule
+    return ee_cloud_tasks_to_schedule + base_cloud_tasks_to_schedule
 
 
 def get_tasks_to_schedule() -> list[dict[str, Any]]:

diff --git a/backend/onyx/background/celery/apps/beat.py b/backend/onyx/background/celery/apps/beat.py
@@ -81,7 +81,7 @@ def _generate_schedule(
                 cloud_task = {
                     "task": task["task"],
                     "schedule": task["schedule"],
-                    "kwargs": {},
+                    "kwargs": task.get("kwargs", {}),
                 }
                 if options := task.get("options"):
                     logger.debug(f"Adding options to task {task_name}: {options}")

diff --git a/backend/onyx/background/celery/tasks/beat_schedule.py b/backend/onyx/background/celery/tasks/beat_schedule.py
@@ -17,134 +17,234 @@
 BEAT_EXPIRES_DEFAULT = 15 * 60  # 15 minutes (in seconds)
 
 # tasks that only run in the cloud
-# the name attribute must start with ONYX_CELERY_CLOUD_PREFIX = "cloud" to be filtered
+# the name attribute must start with ONYX_CLOUD_CELERY_TASK_PREFIX = "cloud" to be filtered
 # by the DynamicTenantScheduler
 cloud_tasks_to_schedule = [
-    {
-        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-for-indexing",
-        "task": OnyxCeleryTask.CLOUD_CHECK_FOR_INDEXING,
-        "schedule": timedelta(seconds=15),
-        "options": {
-            "priority": OnyxCeleryPriority.HIGHEST,
-            "expires": BEAT_EXPIRES_DEFAULT,
-        },
-    },
+    # cloud specific tasks
     {
         "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-alembic",
         "task": OnyxCeleryTask.CLOUD_CHECK_ALEMBIC,
         "schedule": timedelta(hours=1),
         "options": {
+            "queue": OnyxCeleryQueues.MONITORING,
             "priority": OnyxCeleryPriority.HIGH,
             "expires": BEAT_EXPIRES_DEFAULT,
-            "queue": OnyxCeleryQueues.MONITORING,
         },
     },
-]
-
-# tasks that run in either self-hosted on cloud
-tasks_to_schedule = [
+    # remaining tasks are cloud generators for per tenant tasks
     {
-        "name": "check-for-vespa-sync",
-        "task": OnyxCeleryTask.CHECK_FOR_VESPA_SYNC_TASK,
-        "schedule": timedelta(seconds=20),
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-for-indexing",
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
+        "schedule": timedelta(seconds=15),
         "options": {
-            "priority": OnyxCeleryPriority.MEDIUM,
+            "priority": OnyxCeleryPriority.HIGHEST,
             "expires": BEAT_EXPIRES_DEFAULT,
         },
+        "kwargs": {
+            "task_name": OnyxCeleryTask.CHECK_FOR_INDEXING,
+        },
     },
     {
-        "name": "check-for-connector-deletion",
-        "task": OnyxCeleryTask.CHECK_FOR_CONNECTOR_DELETION,
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-for-connector-deletion",
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
         "schedule": timedelta(seconds=20),
         "options": {
-            "priority": OnyxCeleryPriority.MEDIUM,
+            "priority": OnyxCeleryPriority.HIGHEST,
             "expires": BEAT_EXPIRES_DEFAULT,
         },
+        "kwargs": {
+            "task_name": OnyxCeleryTask.CHECK_FOR_CONNECTOR_DELETION,
+        },
     },
     {
-        "name": "check-for-prune",
-        "task": OnyxCeleryTask.CHECK_FOR_PRUNING,
-        "schedule": timedelta(seconds=15),
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-for-vespa-sync",
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
+        "schedule": timedelta(seconds=20),
         "options": {
-            "priority": OnyxCeleryPriority.MEDIUM,
+            "priority": OnyxCeleryPriority.HIGHEST,
             "expires": BEAT_EXPIRES_DEFAULT,
         },
+        "kwargs": {
+            "task_name": OnyxCeleryTask.CHECK_FOR_VESPA_SYNC_TASK,
+        },
     },
     {
-        "name": "kombu-message-cleanup",
-        "task": OnyxCeleryTask.KOMBU_MESSAGE_CLEANUP_TASK,
-        "schedule": timedelta(seconds=3600),
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-for-prune",
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
+        "schedule": timedelta(seconds=15),
         "options": {
-            "priority": OnyxCeleryPriority.LOWEST,
+            "priority": OnyxCeleryPriority.HIGHEST,
             "expires": BEAT_EXPIRES_DEFAULT,
         },
+        "kwargs": {
+            "task_name": OnyxCeleryTask.CHECK_FOR_PRUNING,
+        },
     },
     {
-        "name": "monitor-vespa-sync",
-        "task": OnyxCeleryTask.MONITOR_VESPA_SYNC,
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor-vespa-sync",
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
         "schedule": timedelta(seconds=5),
         "options": {
-            "priority": OnyxCeleryPriority.MEDIUM,
+            "priority": OnyxCeleryPriority.HIGHEST,
             "expires": BEAT_EXPIRES_DEFAULT,
         },
+        "kwargs": {
+            "task_name": OnyxCeleryTask.MONITOR_VESPA_SYNC,
+        },
     },
     {
-        "name": "monitor-background-processes",
-        "task": OnyxCeleryTask.MONITOR_BACKGROUND_PROCESSES,
-        "schedule": timedelta(minutes=5),
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-for-doc-permissions-sync",
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
+        "schedule": timedelta(seconds=30),
         "options": {
-            "priority": OnyxCeleryPriority.LOW,
+            "priority": OnyxCeleryPriority.HIGHEST,
             "expires": BEAT_EXPIRES_DEFAULT,
-            "queue": OnyxCeleryQueues.MONITORING,
+        },
+        "kwargs": {
+            "task_name": OnyxCeleryTask.CHECK_FOR_DOC_PERMISSIONS_SYNC,
         },
     },
     {
-        "name": "check-for-doc-permissions-sync",
-        "task": OnyxCeleryTask.CHECK_FOR_DOC_PERMISSIONS_SYNC,
-        "schedule": timedelta(seconds=30),
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-for-external-group-sync",
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
+        "schedule": timedelta(seconds=20),
         "options": {
-            "priority": OnyxCeleryPriority.MEDIUM,
+            "priority": OnyxCeleryPriority.HIGHEST,
             "expires": BEAT_EXPIRES_DEFAULT,
         },
+        "kwargs": {
+            "task_name": OnyxCeleryTask.CHECK_FOR_EXTERNAL_GROUP_SYNC,
+        },
     },
     {
-        "name": "check-for-external-group-sync",
-        "task": OnyxCeleryTask.CHECK_FOR_EXTERNAL_GROUP_SYNC,
-        "schedule": timedelta(seconds=20),
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor-background-processes",
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
+        "schedule": timedelta(minutes=5),
         "options": {
-            "priority": OnyxCeleryPriority.MEDIUM,
+            "priority": OnyxCeleryPriority.HIGHEST,
             "expires": BEAT_EXPIRES_DEFAULT,
         },
+        "kwargs": {
+            "task_name": OnyxCeleryTask.MONITOR_BACKGROUND_PROCESSES,
+            "queue": OnyxCeleryQueues.MONITORING,
+            "priority": OnyxCeleryPriority.LOW,
+        },
     },
 ]
 
-if not MULTI_TENANT:
-    tasks_to_schedule.append(
+if LLM_MODEL_UPDATE_API_URL:
+    cloud_tasks_to_schedule.append(
         {
-            "name": "check-for-indexing",
-            "task": OnyxCeleryTask.CHECK_FOR_INDEXING,
-            "schedule": timedelta(seconds=15),
+            "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-for-llm-model-update",
+            "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
+            "schedule": timedelta(hours=1),  # Check every hour
             "options": {
-                "priority": OnyxCeleryPriority.MEDIUM,
+                "priority": OnyxCeleryPriority.HIGHEST,
                 "expires": BEAT_EXPIRES_DEFAULT,
             },
+            "kwargs": {
+                "task_name": OnyxCeleryTask.CHECK_FOR_LLM_MODEL_UPDATE,
+                "priority": OnyxCeleryPriority.LOW,
+            },
         }
     )
 
-# Only add the LLM model update task if the API URL is configured
-if LLM_MODEL_UPDATE_API_URL:
-    tasks_to_schedule.append(
-        {
-            "name": "check-for-llm-model-update",
-            "task": OnyxCeleryTask.CHECK_FOR_LLM_MODEL_UPDATE,
-            "schedule": timedelta(hours=1),  # Check every hour
-            "options": {
-                "priority": OnyxCeleryPriority.LOW,
-                "expires": BEAT_EXPIRES_DEFAULT,
+# tasks that run in either self-hosted on cloud
+tasks_to_schedule: list[dict] = []
+
+if not MULTI_TENANT:
+    tasks_to_schedule.extend(
+        [
+            {
+                "name": "check-for-indexing",
+                "task": OnyxCeleryTask.CHECK_FOR_INDEXING,
+                "schedule": timedelta(seconds=15),
+                "options": {
+                    "priority": OnyxCeleryPriority.MEDIUM,
+                    "expires": BEAT_EXPIRES_DEFAULT,
+                },
             },
-        }
+            {
+                "name": "check-for-connector-deletion",
+                "task": OnyxCeleryTask.CHECK_FOR_CONNECTOR_DELETION,
+                "schedule": timedelta(seconds=20),
+                "options": {
+                    "priority": OnyxCeleryPriority.MEDIUM,
+                    "expires": BEAT_EXPIRES_DEFAULT,
+                },
+            },
+            {
+                "name": "check-for-vespa-sync",
+                "task": OnyxCeleryTask.CHECK_FOR_VESPA_SYNC_TASK,
+                "schedule": timedelta(seconds=20),
+                "options": {
+                    "priority": OnyxCeleryPriority.MEDIUM,
+                    "expires": BEAT_EXPIRES_DEFAULT,
+                },
+            },
+            {
+                "name": "check-for-pruning",
+                "task": OnyxCeleryTask.CHECK_FOR_PRUNING,
+                "schedule": timedelta(hours=1),
+                "options": {
+                    "priority": OnyxCeleryPriority.MEDIUM,
+                    "expires": BEAT_EXPIRES_DEFAULT,
+                },
+            },
+            {
+                "name": "monitor-vespa-sync",
+                "task": OnyxCeleryTask.MONITOR_VESPA_SYNC,
+                "schedule": timedelta(seconds=5),
+                "options": {
+                    "priority": OnyxCeleryPriority.MEDIUM,
+                    "expires": BEAT_EXPIRES_DEFAULT,
+                },
+            },
+            {
+                "name": "check-for-doc-permissions-sync",
+                "task": OnyxCeleryTask.CHECK_FOR_DOC_PERMISSIONS_SYNC,
+                "schedule": timedelta(seconds=30),
+                "options": {
+                    "priority": OnyxCeleryPriority.MEDIUM,
+                    "expires": BEAT_EXPIRES_DEFAULT,
+                },
+            },
+            {
+                "name": "check-for-external-group-sync",
+                "task": OnyxCeleryTask.CHECK_FOR_EXTERNAL_GROUP_SYNC,
+                "schedule": timedelta(seconds=20),
+                "options": {
+                    "priority": OnyxCeleryPriority.MEDIUM,
+                    "expires": BEAT_EXPIRES_DEFAULT,
+                },
+            },
+            {
+                "name": "monitor-background-processes",
+                "task": OnyxCeleryTask.MONITOR_BACKGROUND_PROCESSES,
+                "schedule": timedelta(minutes=5),
+                "options": {
+                    "priority": OnyxCeleryPriority.LOW,
+                    "expires": BEAT_EXPIRES_DEFAULT,
+                    "queue": OnyxCeleryQueues.MONITORING,
+                },
+            },
+        ]
     )
 
+    # Only add the LLM model update task if the API URL is configured
+    if LLM_MODEL_UPDATE_API_URL:
+        tasks_to_schedule.append(
+            {
+                "name": "check-for-llm-model-update",
+                "task": OnyxCeleryTask.CHECK_FOR_LLM_MODEL_UPDATE,
+                "schedule": timedelta(hours=1),  # Check every hour
+                "options": {
+                    "priority": OnyxCeleryPriority.LOW,
+                    "expires": BEAT_EXPIRES_DEFAULT,
+                },
+            }
+        )
+
 
 def get_cloud_tasks_to_schedule() -> list[dict[str, Any]]:
     return cloud_tasks_to_schedule