Bugfix/beat redux (#3639)

* WIP

* WIP

* try spinning out check for indexing into a system task

* check for the correct delimiter

* use constants

---------

Co-authored-by: Richard Kuo (Danswer) <rkuo@onyx.app>
Co-authored-by: Richard Kuo <rkuo@rkuo.com>

backend/onyx/background/celery/apps/app_base.py

@@ -20,6 +20,7 @@ from sqlalchemy.orm import Session
 from onyx.background.celery.apps.task_formatters import CeleryTaskColoredFormatter
 from onyx.background.celery.apps.task_formatters import CeleryTaskPlainFormatter
 from onyx.background.celery.celery_utils import celery_is_worker_primary
+from onyx.configs.constants import ONYX_CLOUD_CELERY_TASK_PREFIX
 from onyx.configs.constants import OnyxRedisLocks
 from onyx.db.engine import get_sqlalchemy_engine
 from onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client
@@ -100,6 +101,10 @@ def on_task_postrun(
     if not task_id:
         return
 
+    if task.name.startswith(ONYX_CLOUD_CELERY_TASK_PREFIX):
+        # this is a cloud / all tenant task ... no postrun is needed
+        return
+
     # Get tenant_id directly from kwargs- each celery task has a tenant_id kwarg
     if not kwargs:
         logger.error(f"Task {task.name} (ID: {task_id}) is missing kwargs")

backend/onyx/background/celery/apps/beat.py

@@ -1,5 +1,6 @@
 from datetime import timedelta
 from typing import Any
+from typing import cast
 
 from celery import Celery
 from celery import signals
@@ -7,12 +8,14 @@ from celery.beat import PersistentScheduler  # type: ignore
 from celery.signals import beat_init
 
 import onyx.background.celery.apps.app_base as app_base
+from onyx.configs.constants import ONYX_CLOUD_CELERY_TASK_PREFIX
 from onyx.configs.constants import POSTGRES_CELERY_BEAT_APP_NAME
 from onyx.db.engine import get_all_tenant_ids
 from onyx.db.engine import SqlEngine
 from onyx.utils.logger import setup_logger
 from onyx.utils.variable_functionality import fetch_versioned_implementation
 from shared_configs.configs import IGNORED_SYNCING_TENANT_LIST
+from shared_configs.configs import MULTI_TENANT
 
 logger = setup_logger(__name__)
 
@@ -28,7 +31,7 @@ class DynamicTenantScheduler(PersistentScheduler):
         self._last_reload = self.app.now() - self._reload_interval
         # Let the parent class handle store initialization
         self.setup_schedule()
-        self._update_tenant_tasks()
+        self._try_updating_schedule()
         logger.info(f"Set reload interval to {self._reload_interval}")
 
     def setup_schedule(self) -> None:
@@ -44,105 +47,154 @@ class DynamicTenantScheduler(PersistentScheduler):
             or (now - self._last_reload) > self._reload_interval
         ):
             logger.info("Reload interval reached, initiating task update")
-            self._update_tenant_tasks()
+            try:
+                self._try_updating_schedule()
+            except (AttributeError, KeyError) as e:
+                logger.exception(f"Failed to process task configuration: {str(e)}")
+            except Exception as e:
+                logger.exception(f"Unexpected error updating tasks: {str(e)}")
+
             self._last_reload = now
             logger.info("Task update completed, reset reload timer")
         return retval
 
-    def _update_tenant_tasks(self) -> None:
-        logger.info("Starting task update process")
-        try:
-            logger.info("Fetching all IDs")
-            tenant_ids = get_all_tenant_ids()
-            logger.info(f"Found {len(tenant_ids)} IDs")
+    def _generate_schedule(
+        self, tenant_ids: list[str] | list[None]
+    ) -> dict[str, dict[str, Any]]:
+        """Given a list of tenant id's, generates a new beat schedule for celery."""
+        logger.info("Fetching tasks to schedule")
 
-            logger.info("Fetching tasks to schedule")
-            tasks_to_schedule = fetch_versioned_implementation(
-                "onyx.background.celery.tasks.beat_schedule", "get_tasks_to_schedule"
+        new_schedule: dict[str, dict[str, Any]] = {}
+
+        if MULTI_TENANT:
+            # cloud tasks only need the single task beat across all tenants
+            get_cloud_tasks_to_schedule = fetch_versioned_implementation(
+                "onyx.background.celery.tasks.beat_schedule",
+                "get_cloud_tasks_to_schedule",
             )
 
-            new_beat_schedule: dict[str, dict[str, Any]] = {}
+            cloud_tasks_to_schedule: list[
+                dict[str, Any]
+            ] = get_cloud_tasks_to_schedule()
+            for task in cloud_tasks_to_schedule:
+                task_name = task["name"]
+                cloud_task = {
+                    "task": task["task"],
+                    "schedule": task["schedule"],
+                    "kwargs": {},
+                }
+                if options := task.get("options"):
+                    logger.debug(f"Adding options to task {task_name}: {options}")
+                    cloud_task["options"] = options
+                new_schedule[task_name] = cloud_task
 
-            current_schedule = self.schedule.items()
+        # regular task beats are multiplied across all tenants
+        get_tasks_to_schedule = fetch_versioned_implementation(
+            "onyx.background.celery.tasks.beat_schedule", "get_tasks_to_schedule"
+        )
 
-            existing_tenants = set()
-            for task_name, _ in current_schedule:
-                if "-" in task_name:
-                    existing_tenants.add(task_name.split("-")[-1])
-            logger.info(f"Found {len(existing_tenants)} existing items in schedule")
+        tasks_to_schedule: list[dict[str, Any]] = get_tasks_to_schedule()
 
-            for tenant_id in tenant_ids:
-                if (
-                    IGNORED_SYNCING_TENANT_LIST
-                    and tenant_id in IGNORED_SYNCING_TENANT_LIST
-                ):
-                    logger.info(
-                        f"Skipping tenant {tenant_id} as it is in the ignored syncing list"
-                    )
-                    continue
-
-                if tenant_id not in existing_tenants:
-                    logger.info(f"Processing new item: {tenant_id}")
-
-                for task in tasks_to_schedule():
-                    task_name = f"{task['name']}-{tenant_id}"
-                    logger.debug(f"Creating task configuration for {task_name}")
-                    new_task = {
-                        "task": task["task"],
-                        "schedule": task["schedule"],
-                        "kwargs": {"tenant_id": tenant_id},
-                    }
-                    if options := task.get("options"):
-                        logger.debug(f"Adding options to task {task_name}: {options}")
-                        new_task["options"] = options
-                    new_beat_schedule[task_name] = new_task
-
-            if self._should_update_schedule(current_schedule, new_beat_schedule):
+        for tenant_id in tenant_ids:
+            if IGNORED_SYNCING_TENANT_LIST and tenant_id in IGNORED_SYNCING_TENANT_LIST:
                 logger.info(
-                    "Schedule update required",
-                    extra={
-                        "new_tasks": len(new_beat_schedule),
-                        "current_tasks": len(current_schedule),
-                    },
+                    f"Skipping tenant {tenant_id} as it is in the ignored syncing list"
                 )
+                continue
 
-                # Create schedule entries
-                entries = {}
-                for name, entry in new_beat_schedule.items():
-                    entries[name] = self.Entry(
-                        name=name,
-                        app=self.app,
-                        task=entry["task"],
-                        schedule=entry["schedule"],
-                        options=entry.get("options", {}),
-                        kwargs=entry.get("kwargs", {}),
+            for task in tasks_to_schedule:
+                task_name = task["name"]
+                tenant_task_name = f"{task['name']}-{tenant_id}"
+
+                logger.debug(f"Creating task configuration for {tenant_task_name}")
+                tenant_task = {
+                    "task": task["task"],
+                    "schedule": task["schedule"],
+                    "kwargs": {"tenant_id": tenant_id},
+                }
+                if options := task.get("options"):
+                    logger.debug(
+                        f"Adding options to task {tenant_task_name}: {options}"
                     )
+                    tenant_task["options"] = options
+                new_schedule[tenant_task_name] = tenant_task
 
-                # Update the schedule using the scheduler's methods
-                self.schedule.clear()
-                self.schedule.update(entries)
+        return new_schedule
 
-                # Ensure changes are persisted
-                self.sync()
+    def _try_updating_schedule(self) -> None:
+        """Only updates the actual beat schedule on the celery app when it changes"""
 
-                logger.info("Schedule update completed successfully")
-            else:
-                logger.info("Schedule is up to date, no changes needed")
-        except (AttributeError, KeyError) as e:
-            logger.exception(f"Failed to process task configuration: {str(e)}")
-        except Exception as e:
-            logger.exception(f"Unexpected error updating tasks: {str(e)}")
+        logger.info("_try_updating_schedule starting")
 
-    def _should_update_schedule(
-        self, current_schedule: dict, new_schedule: dict
-    ) -> bool:
-        """Compare schedules to determine if an update is needed."""
-        logger.debug("Comparing current and new schedules")
-        current_tasks = set(name for name, _ in current_schedule)
-        new_tasks = set(new_schedule.keys())
-        needs_update = current_tasks != new_tasks
-        logger.debug(f"Schedule update needed: {needs_update}")
-        return needs_update
+        tenant_ids = get_all_tenant_ids()
+        logger.info(f"Found {len(tenant_ids)} IDs")
+
+        # get current schedule and extract current tenants
+        current_schedule = self.schedule.items()
+
+        current_tenants = set()
+        for task_name, _ in current_schedule:
+            task_name = cast(str, task_name)
+            if task_name.startswith(ONYX_CLOUD_CELERY_TASK_PREFIX):
+                continue
+
+            if "_" in task_name:
+                # example: "check-for-condition-tenant_12345678-abcd-efgh-ijkl-12345678"
+                # -> "12345678-abcd-efgh-ijkl-12345678"
+                current_tenants.add(task_name.split("_")[-1])
+        logger.info(f"Found {len(current_tenants)} existing items in schedule")
+
+        for tenant_id in tenant_ids:
+            if tenant_id not in current_tenants:
+                logger.info(f"Processing new tenant: {tenant_id}")
+
+        new_schedule = self._generate_schedule(tenant_ids)
+
+        if DynamicTenantScheduler._compare_schedules(current_schedule, new_schedule):
+            logger.info(
+                "_try_updating_schedule: Current schedule is up to date, no changes needed"
+            )
+            return
+
+        logger.info(
+            "Schedule update required",
+            extra={
+                "new_tasks": len(new_schedule),
+                "current_tasks": len(current_schedule),
+            },
+        )
+
+        # Create schedule entries
+        entries = {}
+        for name, entry in new_schedule.items():
+            entries[name] = self.Entry(
+                name=name,
+                app=self.app,
+                task=entry["task"],
+                schedule=entry["schedule"],
+                options=entry.get("options", {}),
+                kwargs=entry.get("kwargs", {}),
+            )
+
+        # Update the schedule using the scheduler's methods
+        self.schedule.clear()
+        self.schedule.update(entries)
+
+        # Ensure changes are persisted
+        self.sync()
+
+        logger.info("_try_updating_schedule: Schedule updated successfully")
+
+    @staticmethod
+    def _compare_schedules(schedule1: dict, schedule2: dict) -> bool:
+        """Compare schedules to determine if an update is needed.
+        True if equivalent, False if not."""
+        current_tasks = set(name for name, _ in schedule1)
+        new_tasks = set(schedule2.keys())
+        if current_tasks != new_tasks:
+            return False
+
+        return True
 
 
 @beat_init.connect

backend/onyx/background/celery/apps/primary.py

@@ -17,7 +17,7 @@ from redis.lock import Lock as RedisLock
 import onyx.background.celery.apps.app_base as app_base
 from onyx.background.celery.apps.app_base import task_logger
 from onyx.background.celery.celery_utils import celery_is_worker_primary
-from onyx.background.celery.tasks.indexing.tasks import (
+from onyx.background.celery.tasks.indexing.utils import (
     get_unfenced_index_attempt_ids,
 )
 from onyx.configs.constants import CELERY_PRIMARY_WORKER_LOCK_TIMEOUT