fix kwargs

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

@@ -81,7 +81,7 @@ class DynamicTenantScheduler(PersistentScheduler):
                 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}")

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

@@ -22,119 +22,217 @@ BEAT_EXPIRES_DEFAULT = 15 * 60  # 15 minutes (in seconds)
 cloud_tasks_to_schedule = [
     {
         "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-for-indexing",
-        "task": OnyxCeleryTask.CLOUD_CHECK_FOR_INDEXING,
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
         "schedule": timedelta(seconds=15),
         "options": {
             "priority": OnyxCeleryPriority.HIGHEST,
             "expires": BEAT_EXPIRES_DEFAULT,
         },
-    },
-]
-
-# tasks that run in either self-hosted on cloud
-tasks_to_schedule = [
-    {
-        "name": "check-for-vespa-sync",
-        "task": OnyxCeleryTask.CHECK_FOR_VESPA_SYNC_TASK,
-        "schedule": timedelta(seconds=20),
-        "options": {
-            "priority": OnyxCeleryPriority.MEDIUM,
-            "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,
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-for-vespa-sync",
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
+        "schedule": timedelta(seconds=20),
+        "options": {
+            "priority": OnyxCeleryPriority.HIGHEST,
+            "expires": BEAT_EXPIRES_DEFAULT,
+        },
+        "kwargs": {
+            "task_name": OnyxCeleryTask.CHECK_FOR_VESPA_SYNC_TASK,
+        },
+    },
+    {
+        "name": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check-for-prune",
+        "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_PRUNING,
+        },
     },
     {
-        "name": "kombu-message-cleanup",
-        "task": OnyxCeleryTask.KOMBU_MESSAGE_CLEANUP_TASK,
-        "schedule": timedelta(seconds=3600),
-        "options": {
-            "priority": OnyxCeleryPriority.LOWEST,
-            "expires": BEAT_EXPIRES_DEFAULT,
-        },
-    },
-    {
-        "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),
-        "options": {
-            "priority": OnyxCeleryPriority.LOW,
-            "expires": BEAT_EXPIRES_DEFAULT,
-            "queue": OnyxCeleryQueues.MONITORING,
-        },
-    },
-    {
-        "name": "check-for-doc-permissions-sync",
-        "task": OnyxCeleryTask.CHECK_FOR_DOC_PERMISSIONS_SYNC,
+        "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.MEDIUM,
+            "priority": OnyxCeleryPriority.HIGHEST,
             "expires": BEAT_EXPIRES_DEFAULT,
         },
+        "kwargs": {
+            "task_name": OnyxCeleryTask.CHECK_FOR_DOC_PERMISSIONS_SYNC,
+        },
     },
     {
-        "name": "check-for-external-group-sync",
-        "task": OnyxCeleryTask.CHECK_FOR_EXTERNAL_GROUP_SYNC,
+        "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": f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_monitor-background-processes",
+        "task": OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
+        "schedule": timedelta(minutes=5),
+        "options": {
+            "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

backend/onyx/background/celery/tasks/indexing/tasks.py

@@ -15,7 +15,6 @@ from redis import Redis
 from redis.lock import Lock as RedisLock
 
 from onyx.background.celery.apps.app_base import task_logger
-from onyx.background.celery.tasks.beat_schedule import BEAT_EXPIRES_DEFAULT
 from onyx.background.celery.tasks.indexing.utils import _should_index
 from onyx.background.celery.tasks.indexing.utils import get_unfenced_index_attempt_ids
 from onyx.background.celery.tasks.indexing.utils import IndexingCallback
@@ -26,15 +25,12 @@ from onyx.background.indexing.run_indexing import run_indexing_entrypoint
 from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT
 from onyx.configs.constants import CELERY_INDEXING_LOCK_TIMEOUT
 from onyx.configs.constants import CELERY_TASK_WAIT_FOR_FENCE_TIMEOUT
-from onyx.configs.constants import ONYX_CLOUD_TENANT_ID
-from onyx.configs.constants import OnyxCeleryPriority
 from onyx.configs.constants import OnyxCeleryTask
 from onyx.configs.constants import OnyxRedisLocks
 from onyx.configs.constants import OnyxRedisSignals
 from onyx.db.connector import mark_ccpair_with_indexing_trigger
 from onyx.db.connector_credential_pair import fetch_connector_credential_pairs
 from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
-from onyx.db.engine import get_all_tenant_ids
 from onyx.db.engine import get_session_with_tenant
 from onyx.db.enums import IndexingMode
 from onyx.db.index_attempt import get_index_attempt
@@ -786,61 +782,3 @@ def connector_indexing_proxy_task(
         f"search_settings={search_settings_id}"
     )
     return
-
-
-@shared_task(
-    name=OnyxCeleryTask.CLOUD_CHECK_FOR_INDEXING,
-    trail=False,
-    bind=True,
-)
-def cloud_check_for_indexing(self: Task) -> bool | None:
-    """a lightweight task used to kick off individual check tasks for each tenant."""
-    time_start = time.monotonic()
-
-    redis_client = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)
-
-    lock_beat: RedisLock = redis_client.lock(
-        OnyxRedisLocks.CLOUD_CHECK_INDEXING_BEAT_LOCK,
-        timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,
-    )
-
-    # these tasks should never overlap
-    if not lock_beat.acquire(blocking=False):
-        return None
-
-    last_lock_time = time.monotonic()
-
-    try:
-        tenant_ids = get_all_tenant_ids()
-        for tenant_id in tenant_ids:
-            current_time = time.monotonic()
-            if current_time - last_lock_time >= (CELERY_GENERIC_BEAT_LOCK_TIMEOUT / 4):
-                lock_beat.reacquire()
-                last_lock_time = current_time
-
-            self.app.send_task(
-                OnyxCeleryTask.CHECK_FOR_INDEXING,
-                kwargs=dict(
-                    tenant_id=tenant_id,
-                ),
-                priority=OnyxCeleryPriority.HIGH,
-                expires=BEAT_EXPIRES_DEFAULT,
-            )
-    except SoftTimeLimitExceeded:
-        task_logger.info(
-            "Soft time limit exceeded, task is being terminated gracefully."
-        )
-    except Exception:
-        task_logger.exception("Unexpected exception during cloud indexing check")
-    finally:
-        if lock_beat.owned():
-            lock_beat.release()
-        else:
-            task_logger.error("cloud_check_for_indexing - Lock not owned on completion")
-            redis_lock_dump(lock_beat, redis_client)
-
-    time_elapsed = time.monotonic() - time_start
-    task_logger.info(
-        f"cloud_check_for_indexing finished: num_tenants={len(tenant_ids)} elapsed={time_elapsed:.2f}"
-    )
-    return True

backend/onyx/background/celery/tasks/shared/tasks.py

@@ -1,15 +1,22 @@
+import time
 from http import HTTPStatus
 
 import httpx
 from celery import shared_task
 from celery import Task
 from celery.exceptions import SoftTimeLimitExceeded
+from redis.lock import Lock as RedisLock
 from tenacity import RetryError
 
 from onyx.access.access import get_access_for_document
 from onyx.background.celery.apps.app_base import task_logger
+from onyx.background.celery.tasks.beat_schedule import BEAT_EXPIRES_DEFAULT
 from onyx.background.celery.tasks.shared.RetryDocumentIndex import RetryDocumentIndex
+from onyx.configs.constants import CELERY_GENERIC_BEAT_LOCK_TIMEOUT
+from onyx.configs.constants import ONYX_CLOUD_TENANT_ID
+from onyx.configs.constants import OnyxCeleryPriority
 from onyx.configs.constants import OnyxCeleryTask
+from onyx.configs.constants import OnyxRedisLocks
 from onyx.db.document import delete_document_by_connector_credential_pair__no_commit
 from onyx.db.document import delete_documents_complete__no_commit
 from onyx.db.document import fetch_chunk_count_for_document
@@ -18,10 +25,13 @@ from onyx.db.document import get_document_connector_count
 from onyx.db.document import mark_document_as_modified
 from onyx.db.document import mark_document_as_synced
 from onyx.db.document_set import fetch_document_sets_for_document
+from onyx.db.engine import get_all_tenant_ids
 from onyx.db.engine import get_session_with_tenant
 from onyx.document_index.document_index_utils import get_both_index_names
 from onyx.document_index.factory import get_default_document_index
 from onyx.document_index.interfaces import VespaDocumentFields
+from onyx.redis.redis_pool import get_redis_client
+from onyx.redis.redis_pool import redis_lock_dump
 from onyx.server.documents.models import ConnectorCredentialPairIdentifier
 
 DOCUMENT_BY_CC_PAIR_CLEANUP_MAX_RETRIES = 3
@@ -199,3 +209,72 @@ def document_by_cc_pair_cleanup_task(
         return False
 
     return True
+
+
+@shared_task(
+    name=OnyxCeleryTask.CLOUD_BEAT_TASK_GENERATOR,
+    trail=False,
+    bind=True,
+)
+def cloud_beat_task_generator(
+    self: Task,
+    task_name: str,
+    queue: str = OnyxCeleryTask.DEFAULT,
+    priority: int = OnyxCeleryPriority.MEDIUM,
+) -> bool | None:
+    """a lightweight task used to kick off individual beat tasks per tenant."""
+    time_start = time.monotonic()
+
+    redis_client = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)
+
+    lock_beat: RedisLock = redis_client.lock(
+        OnyxRedisLocks.CLOUD_CHECK_INDEXING_BEAT_LOCK,
+        timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,
+    )
+
+    # these tasks should never overlap
+    if not lock_beat.acquire(blocking=False):
+        return None
+
+    last_lock_time = time.monotonic()
+
+    try:
+        tenant_ids = get_all_tenant_ids()
+        for tenant_id in tenant_ids:
+            current_time = time.monotonic()
+            if current_time - last_lock_time >= (CELERY_GENERIC_BEAT_LOCK_TIMEOUT / 4):
+                lock_beat.reacquire()
+                last_lock_time = current_time
+
+            self.app.send_task(
+                task_name,
+                kwargs=dict(
+                    tenant_id=tenant_id,
+                ),
+                queue=queue,
+                priority=priority,
+                expires=BEAT_EXPIRES_DEFAULT,
+            )
+    except SoftTimeLimitExceeded:
+        task_logger.info(
+            "Soft time limit exceeded, task is being terminated gracefully."
+        )
+    except Exception:
+        task_logger.exception("Unexpected exception during cloud_beat_task_generator")
+    finally:
+        if not lock_beat.owned():
+            task_logger.error(
+                "cloud_beat_task_generator - Lock not owned on completion"
+            )
+            redis_lock_dump(lock_beat, redis_client)
+        else:
+            lock_beat.release()
+
+    time_elapsed = time.monotonic() - time_start
+    task_logger.info(
+        f"cloud_beat_task_generator finished: "
+        f"task={task_name} "
+        f"num_tenants={len(tenant_ids)} "
+        f"elapsed={time_elapsed:.2f}"
+    )
+    return True

backend/onyx/configs/constants.py

@@ -295,6 +295,7 @@ class OnyxRedisLocks:
     ANONYMOUS_USER_ENABLED = "anonymous_user_enabled"
 
     CLOUD_CHECK_INDEXING_BEAT_LOCK = "da_lock:cloud_check_indexing_beat"
+    CLOUD_BEAT_TASK_GENERATOR_LOCK = "da_lock:cloud_beat_task_generator"
 
 
 class OnyxRedisSignals:
@@ -317,6 +318,10 @@ ONYX_CLOUD_TENANT_ID = "cloud"
 
 
 class OnyxCeleryTask:
+    DEFAULT = "celery"
+
+    CLOUD_BEAT_TASK_GENERATOR = f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_generate_beat_tasks"
+
     CHECK_FOR_CONNECTOR_DELETION = "check_for_connector_deletion_task"
     CHECK_FOR_VESPA_SYNC_TASK = "check_for_vespa_sync_task"
     CHECK_FOR_INDEXING = "check_for_indexing"
@@ -324,8 +329,10 @@ class OnyxCeleryTask:
     CHECK_FOR_DOC_PERMISSIONS_SYNC = "check_for_doc_permissions_sync"
     CHECK_FOR_EXTERNAL_GROUP_SYNC = "check_for_external_group_sync"
     CHECK_FOR_LLM_MODEL_UPDATE = "check_for_llm_model_update"
+
     MONITOR_VESPA_SYNC = "monitor_vespa_sync"
     MONITOR_BACKGROUND_PROCESSES = "monitor_background_processes"
+
     KOMBU_MESSAGE_CLEANUP_TASK = "kombu_message_cleanup_task"
     CONNECTOR_PERMISSION_SYNC_GENERATOR_TASK = (
         "connector_permission_sync_generator_task"
@@ -343,8 +350,6 @@ class OnyxCeleryTask:
     CHECK_TTL_MANAGEMENT_TASK = "check_ttl_management_task"
     AUTOGENERATE_USAGE_REPORT_TASK = "autogenerate_usage_report_task"
 
-    CLOUD_CHECK_FOR_INDEXING = f"{ONYX_CLOUD_CELERY_TASK_PREFIX}_check_for_indexing"
-
 
 REDIS_SOCKET_KEEPALIVE_OPTIONS = {}
 REDIS_SOCKET_KEEPALIVE_OPTIONS[socket.TCP_KEEPINTVL] = 15