improved timeouts + worker configs

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

@@ -21,11 +21,11 @@ from danswer.background.celery.celery_redis import RedisDocumentSet
 from danswer.background.celery.celery_redis import RedisUserGroup
 from danswer.background.celery.celery_utils import celery_is_worker_primary
 from danswer.configs.constants import DanswerRedisLocks
-from danswer.db.engine import get_all_tenant_ids
 from danswer.redis.redis_pool import get_redis_client
 from danswer.utils.logger import ColoredFormatter
 from danswer.utils.logger import PlainFormatter
 from danswer.utils.logger import setup_logger
+from shared_configs.configs import MULTI_TENANT
 from shared_configs.configs import SENTRY_DSN
 
 
@@ -173,44 +173,30 @@ def wait_for_redis(sender: Any, **kwargs: Any) -> None:
 
 
 def on_secondary_worker_init(sender: Any, **kwargs: Any) -> None:
+    logger.info("Running as a secondary celery worker.")
+
+    # Exit early if multi-tenant since primary worker check not needed
+    if MULTI_TENANT:
+        return
+
+    # Set up variables for waiting on primary worker
     WAIT_INTERVAL = 5
     WAIT_LIMIT = 60
-
-    logger.info("Running as a secondary celery worker.")
-    logger.info("Waiting for all tenant primary workers to be ready...")
+    r = get_redis_client(tenant_id=None)
     time_start = time.monotonic()
 
+    logger.info("Waiting for primary worker to be ready...")
     while True:
-        tenant_ids = get_all_tenant_ids()
-        # Check if we have a primary worker lock for each tenant
-        all_tenants_ready = all(
-            get_redis_client(tenant_id=tenant_id).exists(
-                DanswerRedisLocks.PRIMARY_WORKER
-            )
-            for tenant_id in tenant_ids
-        )
-
-        if all_tenants_ready:
+        if r.exists(DanswerRedisLocks.PRIMARY_WORKER):
             break
 
         time_elapsed = time.monotonic() - time_start
-        ready_tenants = sum(
-            1
-            for tenant_id in tenant_ids
-            if get_redis_client(tenant_id=tenant_id).exists(
-                DanswerRedisLocks.PRIMARY_WORKER
-            )
-        )
-
         logger.info(
-            f"Not all tenant primary workers are ready yet. "
-            f"Ready tenants: {ready_tenants}/{len(tenant_ids)} "
-            f"elapsed={time_elapsed:.1f} timeout={WAIT_LIMIT:.1f}"
+            f"Primary worker is not ready yet. elapsed={time_elapsed:.1f} timeout={WAIT_LIMIT:.1f}"
         )
-
         if time_elapsed > WAIT_LIMIT:
             msg = (
-                f"Not all tenant primary workers were ready within the timeout "
+                f"Primary worker was not ready within the timeout. "
                 f"({WAIT_LIMIT} seconds). Exiting..."
             )
             logger.error(msg)
@@ -218,7 +204,7 @@ def on_secondary_worker_init(sender: Any, **kwargs: Any) -> None:
 
         time.sleep(WAIT_INTERVAL)
 
-    logger.info("All tenant primary workers are ready. Continuing...")
+    logger.info("Wait for primary worker completed successfully. Continuing...")
     return
 
 
@@ -230,26 +216,20 @@ def on_worker_shutdown(sender: Any, **kwargs: Any) -> None:
     if not celery_is_worker_primary(sender):
         return
 
-    if not hasattr(sender, "primary_worker_locks"):
+    if not sender.primary_worker_lock:
         return
 
-    for tenant_id, lock in sender.primary_worker_locks.items():
-        try:
-            if lock and lock.owned():
-                logger.debug(f"Attempting to release lock for tenant {tenant_id}")
-                try:
-                    lock.release()
-                    logger.debug(f"Successfully released lock for tenant {tenant_id}")
-                except Exception as e:
-                    logger.error(
-                        f"Failed to release lock for tenant {tenant_id}. Error: {str(e)}"
-                    )
-                finally:
-                    sender.primary_worker_locks[tenant_id] = None
-        except Exception as e:
-            logger.error(
-                f"Error checking lock status for tenant {tenant_id}. Error: {str(e)}"
-            )
+    logger.info("Releasing primary worker lock.")
+    lock = sender.primary_worker_lock
+    try:
+        if lock.owned():
+            try:
+                lock.release()
+                sender.primary_worker_lock = None
+            except Exception as e:
+                logger.error(f"Failed to release primary worker lock: {e}")
+    except Exception as e:
+        logger.error(f"Failed to check if primary worker lock is owned: {e}")
 
 
 def on_setup_logging(

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

@@ -3,26 +3,144 @@ from typing import Any
 
 from celery import Celery
 from celery import signals
+from celery.beat import PersistentScheduler  # type: ignore
 from celery.signals import beat_init
 
 import danswer.background.celery.apps.app_base as app_base
-from danswer.configs.constants import DanswerCeleryPriority
 from danswer.configs.constants import POSTGRES_CELERY_BEAT_APP_NAME
 from danswer.db.engine import get_all_tenant_ids
 from danswer.db.engine import SqlEngine
 from danswer.utils.logger import setup_logger
+from danswer.utils.variable_functionality import fetch_versioned_implementation
 
-logger = setup_logger()
+logger = setup_logger(__name__)
 
 celery_app = Celery(__name__)
 celery_app.config_from_object("danswer.background.celery.configs.beat")
 
 
+class DynamicTenantScheduler(PersistentScheduler):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        logger.info("Initializing DynamicTenantScheduler")
+        super().__init__(*args, **kwargs)
+        self._reload_interval = timedelta(minutes=2)
+        self._last_reload = self.app.now() - self._reload_interval
+        # Let the parent class handle store initialization
+        self.setup_schedule()
+        self._update_tenant_tasks()
+        logger.info(f"Set reload interval to {self._reload_interval}")
+
+    def setup_schedule(self) -> None:
+        logger.info("Setting up initial schedule")
+        super().setup_schedule()
+        logger.info("Initial schedule setup complete")
+
+    def tick(self) -> float:
+        retval = super().tick()
+        now = self.app.now()
+        if (
+            self._last_reload is None
+            or (now - self._last_reload) > self._reload_interval
+        ):
+            logger.info("Reload interval reached, initiating tenant task update")
+            self._update_tenant_tasks()
+            self._last_reload = now
+            logger.info("Tenant task update completed, reset reload timer")
+        return retval
+
+    def _update_tenant_tasks(self) -> None:
+        logger.info("Starting tenant task update process")
+        try:
+            logger.info("Fetching all tenant IDs")
+            tenant_ids = get_all_tenant_ids()
+            logger.info(f"Found {len(tenant_ids)} tenants")
+
+            logger.info("Fetching tasks to schedule")
+            tasks_to_schedule = fetch_versioned_implementation(
+                "danswer.background.celery.tasks.beat_schedule", "get_tasks_to_schedule"
+            )
+
+            new_beat_schedule: dict[str, dict[str, Any]] = {}
+
+            current_schedule = self.schedule.items()
+
+            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 tenants in schedule")
+
+            for tenant_id in tenant_ids:
+                if tenant_id not in existing_tenants:
+                    logger.info(f"Processing new tenant: {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):
+                logger.info(
+                    "Schedule update required",
+                    extra={
+                        "new_tasks": len(new_beat_schedule),
+                        "current_tasks": len(current_schedule),
+                    },
+                )
+
+                # 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", {}),
+                    )
+
+                # Update the schedule using the scheduler's methods
+                self.schedule.clear()
+                self.schedule.update(entries)
+
+                # Ensure changes are persisted
+                self.sync()
+
+                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 tenant tasks: {str(e)}")
+
+    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
+
+
 @beat_init.connect
 def on_beat_init(sender: Any, **kwargs: Any) -> None:
     logger.info("beat_init signal received.")
 
-    # celery beat shouldn't touch the db at all. But just setting a low minimum here.
+    # Celery beat shouldn't touch the db at all. But just setting a low minimum here.
     SqlEngine.set_app_name(POSTGRES_CELERY_BEAT_APP_NAME)
     SqlEngine.init_engine(pool_size=2, max_overflow=0)
     app_base.wait_for_redis(sender, **kwargs)
@@ -35,68 +153,4 @@ def on_setup_logging(
     app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)
 
 
-#####
-# Celery Beat (Periodic Tasks) Settings
-#####
-
-tenant_ids = get_all_tenant_ids()
-
-tasks_to_schedule = [
-    {
-        "name": "check-for-vespa-sync",
-        "task": "check_for_vespa_sync_task",
-        "schedule": timedelta(seconds=5),
-        "options": {"priority": DanswerCeleryPriority.HIGH},
-    },
-    {
-        "name": "check-for-connector-deletion",
-        "task": "check_for_connector_deletion_task",
-        "schedule": timedelta(seconds=60),
-        "options": {"priority": DanswerCeleryPriority.HIGH},
-    },
-    {
-        "name": "check-for-indexing",
-        "task": "check_for_indexing",
-        "schedule": timedelta(seconds=10),
-        "options": {"priority": DanswerCeleryPriority.HIGH},
-    },
-    {
-        "name": "check-for-prune",
-        "task": "check_for_pruning",
-        "schedule": timedelta(seconds=10),
-        "options": {"priority": DanswerCeleryPriority.HIGH},
-    },
-    {
-        "name": "kombu-message-cleanup",
-        "task": "kombu_message_cleanup_task",
-        "schedule": timedelta(seconds=3600),
-        "options": {"priority": DanswerCeleryPriority.LOWEST},
-    },
-    {
-        "name": "monitor-vespa-sync",
-        "task": "monitor_vespa_sync",
-        "schedule": timedelta(seconds=5),
-        "options": {"priority": DanswerCeleryPriority.HIGH},
-    },
-]
-
-
-# Build the celery beat schedule dynamically
-beat_schedule = {}
-
-for tenant_id in tenant_ids:
-    for task in tasks_to_schedule:
-        task_name = f"{task['name']}-{tenant_id}"  # Unique name for each scheduled task
-        beat_schedule[task_name] = {
-            "task": task["task"],
-            "schedule": task["schedule"],
-            "options": task["options"],
-            "kwargs": {"tenant_id": tenant_id},  # Must pass tenant_id as an argument
-        }
-
-# Include any existing beat schedules
-existing_beat_schedule = celery_app.conf.beat_schedule or {}
-beat_schedule.update(existing_beat_schedule)
-
-# Update the Celery app configuration once
-celery_app.conf.beat_schedule = beat_schedule
+celery_app.conf.beat_scheduler = DynamicTenantScheduler

backend/danswer/background/celery/apps/light.py

@@ -85,5 +85,6 @@ celery_app.autodiscover_tasks(
     [
         "danswer.background.celery.tasks.shared",
         "danswer.background.celery.tasks.vespa",
+        "danswer.background.celery.tasks.connector_deletion",
     ]
 )

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

@@ -24,10 +24,10 @@ from danswer.background.celery.celery_utils import celery_is_worker_primary
 from danswer.configs.constants import CELERY_PRIMARY_WORKER_LOCK_TIMEOUT
 from danswer.configs.constants import DanswerRedisLocks
 from danswer.configs.constants import POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME
-from danswer.db.engine import get_all_tenant_ids
 from danswer.db.engine import SqlEngine
 from danswer.redis.redis_pool import get_redis_client
 from danswer.utils.logger import setup_logger
+from shared_configs.configs import MULTI_TENANT
 
 
 logger = setup_logger()
@@ -79,91 +79,90 @@ def on_worker_init(sender: Any, **kwargs: Any) -> None:
 
     logger.info("Running as the primary celery worker.")
 
-    sender.primary_worker_locks = {}
+    if MULTI_TENANT:
+        return
 
     # This is singleton work that should be done on startup exactly once
-    # by the primary worker
-    tenant_ids = get_all_tenant_ids()
-    for tenant_id in tenant_ids:
-        r = get_redis_client(tenant_id=tenant_id)
+    # by the primary worker. This is unnecessary in the multi tenant scenario
+    r = get_redis_client(tenant_id=None)
 
-        # For the moment, we're assuming that we are the only primary worker
-        # that should be running.
-        # TODO: maybe check for or clean up another zombie primary worker if we detect it
-        r.delete(DanswerRedisLocks.PRIMARY_WORKER)
+    # For the moment, we're assuming that we are the only primary worker
+    # that should be running.
+    # TODO: maybe check for or clean up another zombie primary worker if we detect it
+    r.delete(DanswerRedisLocks.PRIMARY_WORKER)
 
-        # this process wide lock is taken to help other workers start up in order.
-        # it is planned to use this lock to enforce singleton behavior on the primary
-        # worker, since the primary worker does redis cleanup on startup, but this isn't
-        # implemented yet.
-        lock = r.lock(
-            DanswerRedisLocks.PRIMARY_WORKER,
-            timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT,
-        )
+    # this process wide lock is taken to help other workers start up in order.
+    # it is planned to use this lock to enforce singleton behavior on the primary
+    # worker, since the primary worker does redis cleanup on startup, but this isn't
+    # implemented yet.
+    lock = r.lock(
+        DanswerRedisLocks.PRIMARY_WORKER,
+        timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT,
+    )
 
-        logger.info("Primary worker lock: Acquire starting.")
-        acquired = lock.acquire(blocking_timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 2)
-        if acquired:
-            logger.info("Primary worker lock: Acquire succeeded.")
-        else:
-            logger.error("Primary worker lock: Acquire failed!")
-            raise WorkerShutdown("Primary worker lock could not be acquired!")
+    logger.info("Primary worker lock: Acquire starting.")
+    acquired = lock.acquire(blocking_timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 2)
+    if acquired:
+        logger.info("Primary worker lock: Acquire succeeded.")
+    else:
+        logger.error("Primary worker lock: Acquire failed!")
+        raise WorkerShutdown("Primary worker lock could not be acquired!")
 
-        # tacking on our own user data to the sender
-        sender.primary_worker_locks[tenant_id] = lock
+    # tacking on our own user data to the sender
+    sender.primary_worker_lock = lock
 
-        # As currently designed, when this worker starts as "primary", we reinitialize redis
-        # to a clean state (for our purposes, anyway)
-        r.delete(DanswerRedisLocks.CHECK_VESPA_SYNC_BEAT_LOCK)
-        r.delete(DanswerRedisLocks.MONITOR_VESPA_SYNC_BEAT_LOCK)
+    # As currently designed, when this worker starts as "primary", we reinitialize redis
+    # to a clean state (for our purposes, anyway)
+    r.delete(DanswerRedisLocks.CHECK_VESPA_SYNC_BEAT_LOCK)
+    r.delete(DanswerRedisLocks.MONITOR_VESPA_SYNC_BEAT_LOCK)
 
-        r.delete(RedisConnectorCredentialPair.get_taskset_key())
-        r.delete(RedisConnectorCredentialPair.get_fence_key())
+    r.delete(RedisConnectorCredentialPair.get_taskset_key())
+    r.delete(RedisConnectorCredentialPair.get_fence_key())
 
-        for key in r.scan_iter(RedisDocumentSet.TASKSET_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisDocumentSet.TASKSET_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisDocumentSet.FENCE_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisDocumentSet.FENCE_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisUserGroup.TASKSET_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisUserGroup.TASKSET_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisUserGroup.FENCE_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisUserGroup.FENCE_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisConnectorDeletion.TASKSET_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisConnectorDeletion.TASKSET_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisConnectorDeletion.FENCE_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisConnectorDeletion.FENCE_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisConnectorPruning.TASKSET_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisConnectorPruning.TASKSET_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisConnectorPruning.GENERATOR_COMPLETE_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisConnectorPruning.GENERATOR_COMPLETE_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisConnectorPruning.GENERATOR_PROGRESS_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisConnectorPruning.GENERATOR_PROGRESS_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisConnectorPruning.FENCE_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisConnectorPruning.FENCE_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisConnectorIndexing.TASKSET_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisConnectorIndexing.TASKSET_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisConnectorIndexing.GENERATOR_COMPLETE_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisConnectorIndexing.GENERATOR_COMPLETE_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisConnectorIndexing.GENERATOR_PROGRESS_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisConnectorIndexing.GENERATOR_PROGRESS_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisConnectorIndexing.FENCE_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisConnectorIndexing.FENCE_PREFIX + "*"):
+        r.delete(key)
 
-        for key in r.scan_iter(RedisConnectorStop.FENCE_PREFIX + "*"):
-            r.delete(key)
+    for key in r.scan_iter(RedisConnectorStop.FENCE_PREFIX + "*"):
+        r.delete(key)
 
 
 @worker_ready.connect
@@ -216,52 +215,36 @@ class HubPeriodicTask(bootsteps.StartStopStep):
             if not celery_is_worker_primary(worker):
                 return
 
-            if not hasattr(worker, "primary_worker_locks"):
+            if not hasattr(worker, "primary_worker_lock"):
                 return
 
-            # Retrieve all tenant IDs
-            tenant_ids = get_all_tenant_ids()
+            lock = worker.primary_worker_lock
 
-            for tenant_id in tenant_ids:
-                lock = worker.primary_worker_locks.get(tenant_id)
-                if not lock:
-                    continue  # Skip if no lock for this tenant
+            r = get_redis_client(tenant_id=None)
 
-                r = get_redis_client(tenant_id=tenant_id)
+            if lock.owned():
+                task_logger.debug("Reacquiring primary worker lock.")
+                lock.reacquire()
+            else:
+                task_logger.warning(
+                    "Full acquisition of primary worker lock. "
+                    "Reasons could be worker restart or lock expiration."
+                )
+                lock = r.lock(
+                    DanswerRedisLocks.PRIMARY_WORKER,
+                    timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT,
+                )
 
-                if lock.owned():
-                    task_logger.debug(
-                        f"Reacquiring primary worker lock for tenant {tenant_id}."
-                    )
-                    lock.reacquire()
+                task_logger.info("Primary worker lock: Acquire starting.")
+                acquired = lock.acquire(
+                    blocking_timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 2
+                )
+                if acquired:
+                    task_logger.info("Primary worker lock: Acquire succeeded.")
+                    worker.primary_worker_lock = lock
                 else:
-                    task_logger.warning(
-                        f"Full acquisition of primary worker lock for tenant {tenant_id}. "
-                        "Reasons could be worker restart or lock expiration."
-                    )
-                    lock = r.lock(
-                        DanswerRedisLocks.PRIMARY_WORKER,
-                        timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT,
-                    )
-
-                    task_logger.info(
-                        f"Primary worker lock for tenant {tenant_id}: Acquire starting."
-                    )
-                    acquired = lock.acquire(
-                        blocking_timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 2
-                    )
-                    if acquired:
-                        task_logger.info(
-                            f"Primary worker lock for tenant {tenant_id}: Acquire succeeded."
-                        )
-                        worker.primary_worker_locks[tenant_id] = lock
-                    else:
-                        task_logger.error(
-                            f"Primary worker lock for tenant {tenant_id}: Acquire failed!"
-                        )
-                        raise TimeoutError(
-                            f"Primary worker lock for tenant {tenant_id} could not be acquired!"
-                        )
+                    task_logger.error("Primary worker lock: Acquire failed!")
+                    raise TimeoutError("Primary worker lock could not be acquired!")
 
         except Exception:
             task_logger.exception("Periodic task failed.")

backend/danswer/background/celery/apps/scheduler.py

@@ -0,0 +1,96 @@
+from datetime import timedelta
+from typing import Any
+
+from celery.beat import PersistentScheduler  # type: ignore
+from celery.utils.log import get_task_logger
+
+from danswer.db.engine import get_all_tenant_ids
+from danswer.utils.variable_functionality import fetch_versioned_implementation
+
+logger = get_task_logger(__name__)
+
+
+class DynamicTenantScheduler(PersistentScheduler):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._reload_interval = timedelta(minutes=1)
+        self._last_reload = self.app.now() - self._reload_interval
+
+    def setup_schedule(self) -> None:
+        super().setup_schedule()
+
+    def tick(self) -> float:
+        retval = super().tick()
+        now = self.app.now()
+        if (
+            self._last_reload is None
+            or (now - self._last_reload) > self._reload_interval
+        ):
+            logger.info("Reloading schedule to check for new tenants...")
+            self._update_tenant_tasks()
+            self._last_reload = now
+        return retval
+
+    def _update_tenant_tasks(self) -> None:
+        logger.info("Checking for tenant task updates...")
+        try:
+            tenant_ids = get_all_tenant_ids()
+            tasks_to_schedule = fetch_versioned_implementation(
+                "danswer.background.celery.tasks.beat_schedule", "get_tasks_to_schedule"
+            )
+
+            new_beat_schedule: dict[str, dict[str, Any]] = {}
+
+            current_schedule = getattr(self, "_store", {"entries": {}}).get(
+                "entries", {}
+            )
+
+            existing_tenants = set()
+            for task_name in current_schedule.keys():
+                if "-" in task_name:
+                    existing_tenants.add(task_name.split("-")[-1])
+
+            for tenant_id in tenant_ids:
+                if tenant_id not in existing_tenants:
+                    logger.info(f"Found new tenant: {tenant_id}")
+
+                for task in tasks_to_schedule():
+                    task_name = f"{task['name']}-{tenant_id}"
+                    new_task = {
+                        "task": task["task"],
+                        "schedule": task["schedule"],
+                        "kwargs": {"tenant_id": tenant_id},
+                    }
+                    if options := task.get("options"):
+                        new_task["options"] = options
+                    new_beat_schedule[task_name] = new_task
+
+            if self._should_update_schedule(current_schedule, new_beat_schedule):
+                logger.info(
+                    "Updating schedule",
+                    extra={
+                        "new_tasks": len(new_beat_schedule),
+                        "current_tasks": len(current_schedule),
+                    },
+                )
+                if not hasattr(self, "_store"):
+                    self._store: dict[str, dict] = {"entries": {}}
+                self.update_from_dict(new_beat_schedule)
+                logger.info(f"New schedule: {new_beat_schedule}")
+
+                logger.info("Tenant tasks updated successfully")
+            else:
+                logger.debug("No schedule updates needed")
+
+        except (AttributeError, KeyError):
+            logger.exception("Failed to process task configuration")
+        except Exception:
+            logger.exception("Unexpected error updating tenant tasks")
+
+    def _should_update_schedule(
+        self, current_schedule: dict, new_schedule: dict
+    ) -> bool:
+        """Compare schedules to determine if an update is needed."""
+        current_tasks = set(current_schedule.keys())
+        new_tasks = set(new_schedule.keys())
+        return current_tasks != new_tasks

backend/danswer/background/celery/tasks/beat_schedule.py

@@ -0,0 +1,48 @@
+from datetime import timedelta
+from typing import Any
+
+from danswer.configs.constants import DanswerCeleryPriority
+
+
+tasks_to_schedule = [
+    {
+        "name": "check-for-vespa-sync",
+        "task": "check_for_vespa_sync_task",
+        "schedule": timedelta(seconds=5),
+        "options": {"priority": DanswerCeleryPriority.HIGH},
+    },
+    {
+        "name": "check-for-connector-deletion",
+        "task": "check_for_connector_deletion_task",
+        "schedule": timedelta(seconds=20),
+        "options": {"priority": DanswerCeleryPriority.HIGH},
+    },
+    {
+        "name": "check-for-indexing",
+        "task": "check_for_indexing",
+        "schedule": timedelta(seconds=10),
+        "options": {"priority": DanswerCeleryPriority.HIGH},
+    },
+    {
+        "name": "check-for-prune",
+        "task": "check_for_pruning",
+        "schedule": timedelta(seconds=10),
+        "options": {"priority": DanswerCeleryPriority.HIGH},
+    },
+    {
+        "name": "kombu-message-cleanup",
+        "task": "kombu_message_cleanup_task",
+        "schedule": timedelta(seconds=3600),
+        "options": {"priority": DanswerCeleryPriority.LOWEST},
+    },
+    {
+        "name": "monitor-vespa-sync",
+        "task": "monitor_vespa_sync",
+        "schedule": timedelta(seconds=5),
+        "options": {"priority": DanswerCeleryPriority.HIGH},
+    },
+]
+
+
+def get_tasks_to_schedule() -> list[dict[str, Any]]:
+    return tasks_to_schedule

backend/danswer/background/celery/versioned_apps/beat.py

@@ -1,8 +1,6 @@
 """Factory stub for running celery worker / celery beat."""
-from danswer.utils.variable_functionality import fetch_versioned_implementation
+from danswer.background.celery.apps.beat import celery_app
 from danswer.utils.variable_functionality import set_is_ee_based_on_env_variable
 
 set_is_ee_based_on_env_variable()
-app = fetch_versioned_implementation(
-    "danswer.background.celery.apps.beat", "celery_app"
-)
+app = celery_app

backend/ee/danswer/background/celery/apps/beat.py

@@ -1,52 +0,0 @@
-#####
-# Celery Beat (Periodic Tasks) Settings
-#####
-from datetime import timedelta
-
-from danswer.background.celery.apps.beat import celery_app
-from danswer.db.engine import get_all_tenant_ids
-
-
-tenant_ids = get_all_tenant_ids()
-
-tasks_to_schedule = [
-    {
-        "name": "sync-external-doc-permissions",
-        "task": "check_sync_external_doc_permissions_task",
-        "schedule": timedelta(seconds=30),  # TODO: optimize this
-    },
-    {
-        "name": "sync-external-group-permissions",
-        "task": "check_sync_external_group_permissions_task",
-        "schedule": timedelta(seconds=60),  # TODO: optimize this
-    },
-    {
-        "name": "autogenerate_usage_report",
-        "task": "autogenerate_usage_report_task",
-        "schedule": timedelta(days=30),  # TODO: change this to config flag
-    },
-    {
-        "name": "check-ttl-management",
-        "task": "check_ttl_management_task",
-        "schedule": timedelta(hours=1),
-    },
-]
-
-# Build the celery beat schedule dynamically
-beat_schedule = {}
-
-for tenant_id in tenant_ids:
-    for task in tasks_to_schedule:
-        task_name = f"{task['name']}-{tenant_id}"  # Unique name for each scheduled task
-        beat_schedule[task_name] = {
-            "task": task["task"],
-            "schedule": task["schedule"],
-            "kwargs": {"tenant_id": tenant_id},  # Must pass tenant_id as an argument
-        }
-
-# Include any existing beat schedules
-existing_beat_schedule = celery_app.conf.beat_schedule or {}
-beat_schedule.update(existing_beat_schedule)
-
-# Update the Celery app configuration
-celery_app.conf.beat_schedule = beat_schedule

backend/ee/danswer/background/celery/tasks/beat_schedule.py

@@ -0,0 +1,33 @@
+from datetime import timedelta
+from typing import Any
+
+from danswer.background.celery.tasks.beat_schedule import (
+    tasks_to_schedule as base_tasks_to_schedule,
+)
+
+ee_tasks_to_schedule = [
+    {
+        "name": "sync-external-doc-permissions",
+        "task": "check_sync_external_doc_permissions_task",
+        "schedule": timedelta(seconds=30),  # TODO: optimize this
+    },
+    {
+        "name": "sync-external-group-permissions",
+        "task": "check_sync_external_group_permissions_task",
+        "schedule": timedelta(seconds=60),  # TODO: optimize this
+    },
+    {
+        "name": "autogenerate_usage_report",
+        "task": "autogenerate_usage_report_task",
+        "schedule": timedelta(days=30),  # TODO: change this to config flag
+    },
+    {
+        "name": "check-ttl-management",
+        "task": "check_ttl_management_task",
+        "schedule": timedelta(hours=1),
+    },
+]
+
+
+def get_tasks_to_schedule() -> list[dict[str, Any]]:
+    return ee_tasks_to_schedule + base_tasks_to_schedule

deployment/cloud_kubernetes/hpa/workers_hpa.yaml

@@ -0,0 +1,56 @@
+apiVersion: autoscaling/v2
+kind: HorizontalPodAutoscaler
+metadata:
+  name: celery-worker-heavy-hpa
+spec:
+  scaleTargetRef:
+    apiVersion: apps/v1
+    kind: Deployment
+    name: celery-worker-heavy
+  minReplicas: 1
+  maxReplicas: 5
+  metrics:
+    - type: Resource
+      resource:
+        name: cpu
+        target:
+          type: Utilization
+          averageUtilization: 60
+---
+apiVersion: autoscaling/v2
+kind: HorizontalPodAutoscaler
+metadata:
+  name: celery-worker-light-hpa
+spec:
+  scaleTargetRef:
+    apiVersion: apps/v1
+    kind: Deployment
+    name: celery-worker-light
+  minReplicas: 1
+  maxReplicas: 10
+  metrics:
+    - type: Resource
+      resource:
+        name: cpu
+        target:
+          type: Utilization
+          averageUtilization: 70
+---
+apiVersion: autoscaling/v2
+kind: HorizontalPodAutoscaler
+metadata:
+  name: celery-worker-indexing-hpa
+spec:
+  scaleTargetRef:
+    apiVersion: apps/v1
+    kind: Deployment
+    name: celery-worker-indexing
+  minReplicas: 1
+  maxReplicas: 10
+  metrics:
+    - type: Resource
+      resource:
+        name: cpu
+        target:
+          type: Utilization
+          averageUtilization: 70

deployment/cloud_kubernetes/workers/beat.yaml

@@ -0,0 +1,44 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: celery-beat
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: celery-beat
+  template:
+    metadata:
+      labels:
+        app: celery-beat
+    spec:
+      containers:
+        - name: celery-beat
+          image: danswer/danswer-backend:v0.11.0-cloud.beta.3
+          imagePullPolicy: Always
+          command:
+            [
+              "celery",
+              "-A",
+              "danswer.background.celery.versioned_apps.beat",
+              "beat",
+              "--loglevel=INFO",
+            ]
+          env:
+            - name: REDIS_PASSWORD
+              valueFrom:
+                secretKeyRef:
+                  name: danswer-secrets
+                  key: redis_password
+            - name: DANSWER_VERSION
+              value: "v0.11.0-cloud.beta.3"
+          envFrom:
+            - configMapRef:
+                name: env-configmap
+          resources:
+            requests:
+              cpu: "250m"
+              memory: "512Mi"
+            limits:
+              cpu: "500m"
+              memory: "1Gi"

deployment/cloud_kubernetes/workers/heavy_worker.yaml

@@ -0,0 +1,60 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: celery-worker-heavy
+spec:
+  replicas: 2
+  selector:
+    matchLabels:
+      app: celery-worker-heavy
+  template:
+    metadata:
+      labels:
+        app: celery-worker-heavy
+    spec:
+      containers:
+        - name: celery-worker-heavy
+          image: danswer/danswer-backend:v0.11.0-cloud.beta.3
+          imagePullPolicy: Always
+          command:
+            [
+              "celery",
+              "-A",
+              "danswer.background.celery.versioned_apps.heavy",
+              "worker",
+              "--loglevel=INFO",
+              "--hostname=heavy@%n",
+              "-Q",
+              "connector_pruning",
+            ]
+          env:
+            - name: REDIS_PASSWORD
+              valueFrom:
+                secretKeyRef:
+                  name: danswer-secrets
+                  key: redis_password
+            - name: DANSWER_VERSION
+              value: "v0.11.0-cloud.beta.3"
+          envFrom:
+            - configMapRef:
+                name: env-configmap
+          volumeMounts:
+            - name: vespa-certificates
+              mountPath: "/app/certs"
+              readOnly: true
+          resources:
+            requests:
+              cpu: "1000m"
+              memory: "2Gi"
+            limits:
+              cpu: "2000m"
+              memory: "4Gi"
+      volumes:
+        - name: vespa-certificates
+          secret:
+            secretName: vespa-certificates
+            items:
+              - key: cert.pem
+                path: cert.pem
+              - key: key.pem
+                path: key.pem

deployment/cloud_kubernetes/workers/indexing_worker.yaml

@@ -0,0 +1,60 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: celery-worker-indexing
+spec:
+  replicas: 3
+  selector:
+    matchLabels:
+      app: celery-worker-indexing
+  template:
+    metadata:
+      labels:
+        app: celery-worker-indexing
+    spec:
+      containers:
+        - name: celery-worker-indexing
+          image: danswer/danswer-backend:v0.11.0-cloud.beta.3
+          imagePullPolicy: Always
+          command:
+            [
+              "celery",
+              "-A",
+              "danswer.background.celery.versioned_apps.indexing",
+              "worker",
+              "--loglevel=INFO",
+              "--hostname=indexing@%n",
+              "-Q",
+              "connector_indexing",
+            ]
+          env:
+            - name: REDIS_PASSWORD
+              valueFrom:
+                secretKeyRef:
+                  name: danswer-secrets
+                  key: redis_password
+            - name: DANSWER_VERSION
+              value: "v0.11.0-cloud.beta.3"
+          envFrom:
+            - configMapRef:
+                name: env-configmap
+          volumeMounts:
+            - name: vespa-certificates
+              mountPath: "/app/certs"
+              readOnly: true
+          resources:
+            requests:
+              cpu: "500m"
+              memory: "1Gi"
+            limits:
+              cpu: "1000m"
+              memory: "2Gi"
+      volumes:
+        - name: vespa-certificates
+          secret:
+            secretName: vespa-certificates
+            items:
+              - key: cert.pem
+                path: cert.pem
+              - key: key.pem
+                path: key.pem

deployment/cloud_kubernetes/workers/light_worker.yaml

@@ -0,0 +1,60 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: celery-worker-light
+spec:
+  replicas: 2
+  selector:
+    matchLabels:
+      app: celery-worker-light
+  template:
+    metadata:
+      labels:
+        app: celery-worker-light
+    spec:
+      containers:
+        - name: celery-worker-light
+          image: danswer/danswer-backend:v0.11.0-cloud.beta.3
+          imagePullPolicy: Always
+          command:
+            [
+              "celery",
+              "-A",
+              "danswer.background.celery.versioned_apps.light",
+              "worker",
+              "--loglevel=INFO",
+              "--hostname=light@%n",
+              "-Q",
+              "vespa_metadata_sync,connector_deletion",
+            ]
+          env:
+            - name: REDIS_PASSWORD
+              valueFrom:
+                secretKeyRef:
+                  name: danswer-secrets
+                  key: redis_password
+            - name: DANSWER_VERSION
+              value: "v0.11.0-cloud.beta.3"
+          envFrom:
+            - configMapRef:
+                name: env-configmap
+          volumeMounts:
+            - name: vespa-certificates
+              mountPath: "/app/certs"
+              readOnly: true
+          resources:
+            requests:
+              cpu: "500m"
+              memory: "1Gi"
+            limits:
+              cpu: "1000m"
+              memory: "2Gi"
+      volumes:
+        - name: vespa-certificates
+          secret:
+            secretName: vespa-certificates
+            items:
+              - key: cert.pem
+                path: cert.pem
+              - key: key.pem
+                path: key.pem

deployment/cloud_kubernetes/workers/primary.yaml

@@ -0,0 +1,60 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: celery-worker-primary
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: celery-worker-primary
+  template:
+    metadata:
+      labels:
+        app: celery-worker-primary
+    spec:
+      containers:
+        - name: celery-worker-primary
+          image: danswer/danswer-backend:v0.11.0-cloud.beta.3
+          imagePullPolicy: Always
+          command:
+            [
+              "celery",
+              "-A",
+              "danswer.background.celery.versioned_apps.primary",
+              "worker",
+              "--loglevel=INFO",
+              "--hostname=primary@%n",
+              "-Q",
+              "celery,periodic_tasks,vespa_metadata_sync",
+            ]
+          env:
+            - name: REDIS_PASSWORD
+              valueFrom:
+                secretKeyRef:
+                  name: danswer-secrets
+                  key: redis_password
+            - name: DANSWER_VERSION
+              value: "v0.11.0-cloud.beta.3"
+          envFrom:
+            - configMapRef:
+                name: env-configmap
+          volumeMounts:
+            - name: vespa-certificates
+              mountPath: "/app/certs"
+              readOnly: true
+          resources:
+            requests:
+              cpu: "500m"
+              memory: "1Gi"
+            limits:
+              cpu: "1000m"
+              memory: "2Gi"
+      volumes:
+        - name: vespa-certificates
+          secret:
+            secretName: vespa-certificates
+            items:
+              - key: cert.pem
+                path: cert.pem
+              - key: key.pem
+                path: key.pem