From 13be391a8d9498149341579d4d832ea946ae8210 Mon Sep 17 00:00:00 2001 From: ali Date: Mon, 1 Jun 2026 12:23:11 +0530 Subject: [PATCH 1/3] UN-3497 [FEAT] Introduce queue_backend seam and migrate Celery call sites MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wrap @shared_task and current_app.send_task behind a single workers/queue_backend/ module so a later phase can route specific tasks through PG Queue without touching call sites. Pure passthrough today — behavioural no-op over Celery, locked in by 28 equivalence and characterisation tests. Co-Authored-By: Claude Opus 4.7 (1M context) --- workers/api-deployment/tasks.py | 9 +- workers/api-deployment/worker.py | 3 +- workers/callback/tasks.py | 9 +- workers/callback/worker.py | 3 +- workers/executor/tasks.py | 4 +- workers/executor/worker.py | 3 +- .../file_processing/structure_tool_task.py | 3 +- workers/file_processing/tasks.py | 13 +- workers/file_processing/worker.py | 3 +- workers/general/tasks.py | 8 +- workers/general/worker.py | 3 +- workers/ide_callback/tasks.py | 13 +- workers/ide_callback/worker.py | 58 ++++ workers/log_consumer/tasks.py | 6 +- workers/log_consumer/worker.py | 3 +- workers/notification/tasks.py | 12 +- workers/notification/worker.py | 3 +- workers/pyproject.toml | 7 +- workers/queue_backend/__init__.py | 22 ++ workers/queue_backend/decorator.py | 38 +++ workers/queue_backend/dispatch.py | 46 +++ workers/run-worker.sh | 271 ++++++++++++++---- workers/scheduler/tasks.py | 15 +- workers/scheduler/worker.py | 3 +- .../shared/patterns/notification/helper.py | 4 +- .../test_dispatch_sites_characterisation.py | 150 +++++----- workers/tests/test_queue_backend_seam.py | 229 +++++++++++++++ 27 files changed, 758 insertions(+), 183 deletions(-) create mode 100644 workers/ide_callback/worker.py create mode 100644 workers/queue_backend/__init__.py create mode 100644 workers/queue_backend/decorator.py create mode 100644 workers/queue_backend/dispatch.py create mode 100644 workers/tests/test_queue_backend_seam.py diff --git a/workers/api-deployment/tasks.py b/workers/api-deployment/tasks.py index 6406a864e6..4a0e2683d6 100644 --- a/workers/api-deployment/tasks.py +++ b/workers/api-deployment/tasks.py @@ -7,6 +7,7 @@ import time from typing import Any +from queue_backend import worker_task from shared.api import InternalAPIClient from shared.enums.status_enums import PipelineStatus from shared.enums.task_enums import TaskName @@ -289,7 +290,7 @@ def _unified_api_execution( logger.warning(f"Failed to cleanup StateStore context: {cleanup_error}") -@app.task( +@worker_task( bind=True, name=TaskName.ASYNC_EXECUTE_BIN_API, autoretry_for=(Exception,), @@ -348,7 +349,7 @@ def async_execute_bin_api( ) -@app.task( +@worker_task( bind=True, name=TaskName.ASYNC_EXECUTE_BIN, autoretry_for=(Exception,), @@ -993,7 +994,7 @@ def _calculate_manual_review_decisions_for_batch_api( return [False] * len(batch) -@app.task(bind=True) +@worker_task(bind=True) @monitor_performance @retry(max_attempts=3, base_delay=2.0) @with_execution_context @@ -1064,7 +1065,7 @@ def api_deployment_status_check( raise -@app.task(bind=True) +@worker_task(bind=True) @monitor_performance @with_execution_context def api_deployment_cleanup( diff --git a/workers/api-deployment/worker.py b/workers/api-deployment/worker.py index 8d29899d56..74dee8df90 100644 --- a/workers/api-deployment/worker.py +++ b/workers/api-deployment/worker.py @@ -4,6 +4,7 @@ Handles API deployment, cleanup, and status checking tasks. """ +from queue_backend import worker_task from shared.enums.worker_enums import WorkerType from shared.infrastructure.config.builder import WorkerBuilder from shared.infrastructure.config.registry import WorkerRegistry @@ -59,7 +60,7 @@ def check_api_deployment_health(): ) -@app.task(bind=True) +@worker_task(bind=True) def healthcheck(self): """Health check task for monitoring systems.""" return { diff --git a/workers/callback/tasks.py b/workers/callback/tasks.py index 42599f0659..80ba1f0540 100644 --- a/workers/callback/tasks.py +++ b/workers/callback/tasks.py @@ -7,8 +7,7 @@ import time from typing import Any -# Use Celery current_app to avoid circular imports -from celery import current_app as app +from queue_backend import worker_task # Import shared worker infrastructure from shared.api import InternalAPIClient @@ -1512,7 +1511,7 @@ def _process_batch_callback_core( logger.warning("api_client.close() failed during callback cleanup: %s", e) -@app.task( +@worker_task( bind=True, name=TaskName.PROCESS_BATCH_CALLBACK, max_retries=0, # Match Django backend pattern @@ -1536,7 +1535,7 @@ def process_batch_callback(self, results, *args, **kwargs) -> dict[str, Any]: return _process_batch_callback_core(self, results, *args, **kwargs) -@app.task( +@worker_task( bind=True, name="process_batch_callback_api", autoretry_for=(Exception,), @@ -1932,7 +1931,7 @@ def _publish_final_workflow_ui_logs_api( ) -@app.task( +@worker_task( bind=True, name="workflow_manager.workflow_v2.file_execution_tasks.process_batch_callback", max_retries=0, diff --git a/workers/callback/worker.py b/workers/callback/worker.py index 10f18e034a..a46f990fb1 100644 --- a/workers/callback/worker.py +++ b/workers/callback/worker.py @@ -3,6 +3,7 @@ Celery worker for processing file processing callbacks and status updates. """ +from queue_backend import worker_task from shared.enums.worker_enums import WorkerType from shared.infrastructure.config.builder import WorkerBuilder from shared.infrastructure.config.registry import WorkerRegistry @@ -61,7 +62,7 @@ def check_callback_health(): ) -@app.task(bind=True) +@worker_task(bind=True) def healthcheck(self): """Health check task for monitoring systems.""" return { diff --git a/workers/executor/tasks.py b/workers/executor/tasks.py index a30243bca3..97faf2bf7f 100644 --- a/workers/executor/tasks.py +++ b/workers/executor/tasks.py @@ -5,7 +5,7 @@ ExecutionOrchestrator, and returns an ExecutionResult dict. """ -from celery import shared_task +from queue_backend import worker_task from shared.clients import UsageAPIClient from shared.enums.task_enums import TaskName from shared.infrastructure.config import WorkerConfig @@ -27,7 +27,7 @@ ) -@shared_task( +@worker_task( bind=True, name=TaskName.EXECUTE_EXTRACTION, autoretry_for=(ConnectionError, TimeoutError, OSError), diff --git a/workers/executor/worker.py b/workers/executor/worker.py index 9b92466aa5..611f074e5a 100644 --- a/workers/executor/worker.py +++ b/workers/executor/worker.py @@ -7,6 +7,7 @@ import logging import os +from queue_backend import worker_task from shared.enums.worker_enums import WorkerType from shared.infrastructure.config.builder import WorkerBuilder from shared.infrastructure.config.registry import WorkerRegistry @@ -67,7 +68,7 @@ def check_executor_health(): ) -@app.task(bind=True) +@worker_task(bind=True) def healthcheck(self): """Health check task for monitoring systems.""" return { diff --git a/workers/file_processing/structure_tool_task.py b/workers/file_processing/structure_tool_task.py index 3fcf1999c4..d89d1c9d6f 100644 --- a/workers/file_processing/structure_tool_task.py +++ b/workers/file_processing/structure_tool_task.py @@ -23,6 +23,7 @@ from typing import Any from file_processing.worker import app +from queue_backend import worker_task from shared.enums.task_enums import TaskName from shared.infrastructure.context import StateStore @@ -190,7 +191,7 @@ def _should_skip_extraction_for_smart_table( # ----------------------------------------------------------------------- -@app.task(bind=True, name=str(TaskName.EXECUTE_STRUCTURE_TOOL)) +@worker_task(bind=True, name=str(TaskName.EXECUTE_STRUCTURE_TOOL)) def execute_structure_tool(self, params: dict) -> dict: """Execute structure tool as a Celery task. diff --git a/workers/file_processing/tasks.py b/workers/file_processing/tasks.py index 78dbcc3ede..614266376a 100644 --- a/workers/file_processing/tasks.py +++ b/workers/file_processing/tasks.py @@ -10,6 +10,8 @@ import time from typing import Any +from queue_backend import worker_task + # Import shared worker infrastructure from shared.api import InternalAPIClient @@ -40,9 +42,6 @@ ) from shared.processing.files.processor import FileProcessor -# Import manual review service with WorkflowUtil access -from worker import app - from unstract.core.data_models import ( ExecutionStatus, FileBatchData, @@ -255,7 +254,7 @@ def _process_file_batch_core( return _compile_batch_result(context) -@app.task( +@worker_task( bind=True, name=TaskName.PROCESS_FILE_BATCH, max_retries=0, # Match Django backend pattern @@ -1507,7 +1506,7 @@ def _process_file( ) -@app.task( +@worker_task( bind=True, name=TaskName.PROCESS_FILE_BATCH_API, max_retries=0, # Match Django backend @@ -1993,7 +1992,7 @@ def resilient_executor(func): # Resilient file processor -@app.task(bind=True) +@worker_task(bind=True) @resilient_executor @with_execution_context def process_file_batch_resilient( @@ -2030,7 +2029,7 @@ def process_file_batch_resilient( # Register the same task function with the old Django task names for compatibility -@app.task( +@worker_task( bind=True, name="workflow_manager.workflow_v2.file_execution_tasks.process_file_batch", max_retries=0, diff --git a/workers/file_processing/worker.py b/workers/file_processing/worker.py index 80794a35b3..ff51161a99 100644 --- a/workers/file_processing/worker.py +++ b/workers/file_processing/worker.py @@ -4,6 +4,7 @@ Handles file uploads, text extraction, and processing workflows. """ +from queue_backend import worker_task from shared.enums.worker_enums import WorkerType from shared.infrastructure.config.builder import WorkerBuilder from shared.infrastructure.config.registry import WorkerRegistry @@ -59,7 +60,7 @@ def check_file_processing_health(): ) -@app.task(bind=True) +@worker_task(bind=True) def healthcheck(self): """Health check task for monitoring systems.""" return { diff --git a/workers/general/tasks.py b/workers/general/tasks.py index addf18b9a6..aa784c81fa 100644 --- a/workers/general/tasks.py +++ b/workers/general/tasks.py @@ -7,7 +7,7 @@ import time from typing import Any -from celery import shared_task +from queue_backend import worker_task from scheduler.tasks import execute_pipeline_task_v2 # Import shared worker infrastructure using new structure @@ -133,7 +133,7 @@ def _log_batch_creation_statistics( ) -@app.task( +@worker_task( bind=True, name=TaskName.ASYNC_EXECUTE_BIN_GENERAL, autoretry_for=(Exception,), @@ -1440,7 +1440,7 @@ def _validate_batch_data_integrity_dataclass( ) -@app.task( +@worker_task( bind=True, name="async_execute_bin", autoretry_for=(Exception,), @@ -1536,7 +1536,7 @@ def async_execute_bin( logger.info("✅ Registered scheduler tasks in general worker for backward compatibility") -@shared_task(name="scheduler.tasks.execute_pipeline_task", bind=True) +@worker_task(name="scheduler.tasks.execute_pipeline_task", bind=True) def execute_pipeline_task( self, workflow_id: Any, diff --git a/workers/general/worker.py b/workers/general/worker.py index 1d9920c8d2..036d9b1d49 100644 --- a/workers/general/worker.py +++ b/workers/general/worker.py @@ -4,6 +4,7 @@ processing, and workflow orchestration. """ +from queue_backend import worker_task from shared.enums.worker_enums import WorkerType from shared.infrastructure.config.builder import WorkerBuilder from shared.infrastructure.config.registry import WorkerRegistry @@ -59,7 +60,7 @@ def check_general_worker_health(): ) -@app.task(bind=True) +@worker_task(bind=True) def healthcheck(self): """Health check task for monitoring systems.""" return { diff --git a/workers/ide_callback/tasks.py b/workers/ide_callback/tasks.py index c6494e036f..0902f4609b 100644 --- a/workers/ide_callback/tasks.py +++ b/workers/ide_callback/tasks.py @@ -15,6 +15,7 @@ from typing import Any from celery import current_app as app +from queue_backend import worker_task from shared.clients.prompt_studio_client import PromptStudioAPIClient logger = logging.getLogger(__name__) @@ -140,7 +141,7 @@ def _track_subscription_usage(org_id: str, run_id: str) -> None: # ------------------------------------------------------------------ -@app.task(name="ide_index_complete") +@worker_task(name="ide_index_complete") def ide_index_complete( result_dict: dict[str, Any], callback_kwargs: dict[str, Any] | None = None, @@ -299,7 +300,7 @@ def ide_index_complete( raise -@app.task(name="ide_index_error") +@worker_task(name="ide_index_error") def ide_index_error( failed_task_id: str, callback_kwargs: dict[str, Any] | None = None, @@ -345,7 +346,7 @@ def ide_index_error( logger.exception("ide_index_error callback failed") -@app.task(name="ide_prompt_complete") +@worker_task(name="ide_prompt_complete") def ide_prompt_complete( result_dict: dict[str, Any], callback_kwargs: dict[str, Any] | None = None, @@ -480,7 +481,7 @@ def ide_prompt_complete( raise -@app.task(name="ide_prompt_error") +@worker_task(name="ide_prompt_error") def ide_prompt_error( failed_task_id: str, callback_kwargs: dict[str, Any] | None = None, @@ -533,7 +534,7 @@ def _get_extraction_client(): return ExtractionAPIClient() -@app.task(name="extraction_complete") +@worker_task(name="extraction_complete") def extraction_complete( result_dict: dict[str, Any], callback_kwargs: dict[str, Any] | None = None, @@ -625,7 +626,7 @@ def extraction_complete( raise -@app.task(name="extraction_error") +@worker_task(name="extraction_error") def extraction_error( failed_task_id: str, callback_kwargs: dict[str, Any] | None = None, diff --git a/workers/ide_callback/worker.py b/workers/ide_callback/worker.py new file mode 100644 index 0000000000..cb1bb4200c --- /dev/null +++ b/workers/ide_callback/worker.py @@ -0,0 +1,58 @@ +"""IDE Callback Worker + +Celery worker for Prompt Studio IDE post-execution callbacks. +Serves the ``ide_callback`` queue defined in shared/infrastructure/config/registry.py. +""" + +from shared.enums.worker_enums import WorkerType +from shared.infrastructure.config.builder import WorkerBuilder +from shared.infrastructure.config.registry import WorkerRegistry +from shared.infrastructure.logging import WorkerLogger + +# Setup worker - this executes when module is imported by Celery +logger = WorkerLogger.setup(WorkerType.IDE_CALLBACK) +app, config = WorkerBuilder.build_celery_app(WorkerType.IDE_CALLBACK) + + +def check_ide_callback_health(): + """Custom health check for IDE callback worker.""" + from shared.infrastructure.monitoring.health import HealthCheckResult, HealthStatus + + try: + from shared.utils.api_client_singleton import get_singleton_api_client + + client = get_singleton_api_client(config) + api_healthy = client is not None + + if api_healthy: + return HealthCheckResult( + name="ide_callback_health", + status=HealthStatus.HEALTHY, + message="IDE callback worker is healthy", + details={ + "worker_type": "ide_callback", + "api_client": "healthy", + "queue": "ide_callback", + }, + ) + else: + return HealthCheckResult( + name="ide_callback_health", + status=HealthStatus.DEGRADED, + message="IDE callback worker partially functional", + details={"api_client": "unhealthy"}, + ) + + except Exception as e: + return HealthCheckResult( + name="ide_callback_health", + status=HealthStatus.DEGRADED, + message=f"Health check failed: {e}", + details={"error": str(e)}, + ) + + +# Register health check +WorkerRegistry.register_health_check( + WorkerType.IDE_CALLBACK, "ide_callback_health", check_ide_callback_health +) diff --git a/workers/log_consumer/tasks.py b/workers/log_consumer/tasks.py index 3eaa280858..aa0d7aec3d 100644 --- a/workers/log_consumer/tasks.py +++ b/workers/log_consumer/tasks.py @@ -8,7 +8,7 @@ from urllib.parse import quote import socketio -from celery import shared_task +from queue_backend import worker_task from shared.infrastructure.config import WorkerConfig from shared.infrastructure.logging import WorkerLogger from shared.utils.api_client_singleton import get_singleton_api_client @@ -68,7 +68,7 @@ ) -@shared_task(name=LogProcessingTask.TASK_NAME) +@worker_task(name=LogProcessingTask.TASK_NAME) def logs_consumer(**kwargs: Any) -> None: """Task to process logs from log publisher. @@ -115,7 +115,7 @@ def logs_consumer(**kwargs: Any) -> None: # Health check task for monitoring -@shared_task(name="log_consumer_health_check") +@worker_task(name="log_consumer_health_check") def health_check() -> dict[str, Any]: """Health check task for log consumer worker. diff --git a/workers/log_consumer/worker.py b/workers/log_consumer/worker.py index df90ca276b..5b56391637 100644 --- a/workers/log_consumer/worker.py +++ b/workers/log_consumer/worker.py @@ -3,6 +3,7 @@ Celery worker for processing execution logs and WebSocket emissions. """ +from queue_backend import worker_task from shared.enums.worker_enums import WorkerType from shared.infrastructure.config.builder import WorkerBuilder from shared.infrastructure.config.registry import WorkerRegistry @@ -58,7 +59,7 @@ def check_log_consumer_health(): ) -@app.task(bind=True) +@worker_task(bind=True) def healthcheck(self): """Health check task for monitoring systems.""" return { diff --git a/workers/notification/tasks.py b/workers/notification/tasks.py index 32a684ad9c..c092fbc05d 100644 --- a/workers/notification/tasks.py +++ b/workers/notification/tasks.py @@ -8,7 +8,6 @@ import os from typing import Any -from celery import shared_task from notification.enums import PlatformType from notification.providers.base_provider import ( DeliveryError, @@ -22,6 +21,7 @@ log_notification_failure, log_notification_success, ) +from queue_backend import worker_task from shared.infrastructure.config import WorkerConfig from shared.infrastructure.logging import WorkerLogger @@ -64,7 +64,7 @@ def _get_webhook_provider_for_url(url: str): return WebhookProvider() -@shared_task(name="process_notification") +@worker_task(name="process_notification") def process_notification( notification_type: str, priority: bool = False, **kwargs: Any ) -> dict[str, Any]: @@ -164,7 +164,7 @@ def process_notification( } -@shared_task(bind=True, name="send_webhook_notification") +@worker_task(bind=True, name="send_webhook_notification") def send_webhook_notification( self, url: str, @@ -281,7 +281,7 @@ def send_webhook_notification( return None -@shared_task(name="send_batch_notifications") +@worker_task(name="send_batch_notifications") def send_batch_notifications( notifications: list[dict[str, Any]], batch_id: str | None = None, @@ -367,7 +367,7 @@ def send_batch_notifications( return results -@shared_task(name="priority_notification") +@worker_task(name="priority_notification") def priority_notification(notification_type: str, **kwargs: Any) -> dict[str, Any]: """High-priority notification processor. @@ -387,7 +387,7 @@ def priority_notification(notification_type: str, **kwargs: Any) -> dict[str, An return process_notification(notification_type, priority=True, **kwargs) -@shared_task(name="notification_health_check") +@worker_task(name="notification_health_check") def notification_health_check() -> dict[str, Any]: """Health check task for notification worker.""" try: diff --git a/workers/notification/worker.py b/workers/notification/worker.py index 112a4b649b..4fca079950 100644 --- a/workers/notification/worker.py +++ b/workers/notification/worker.py @@ -3,6 +3,7 @@ Celery worker for processing notifications including webhooks, emails, SMS. """ +from queue_backend import worker_task from shared.enums.worker_enums import WorkerType from shared.infrastructure.config.builder import WorkerBuilder from shared.infrastructure.config.registry import WorkerRegistry @@ -62,7 +63,7 @@ def check_notification_health(): ) -@app.task(bind=True) +@worker_task(bind=True) def healthcheck(self): """Health check task for monitoring systems.""" return { diff --git a/workers/pyproject.toml b/workers/pyproject.toml index 3803c9416f..4d5bef1230 100644 --- a/workers/pyproject.toml +++ b/workers/pyproject.toml @@ -116,7 +116,8 @@ known_first_party = [ "notification", "scheduler", "log_consumer", - "ide_callback" + "ide_callback", + "queue_backend" ] [tool.mypy] @@ -163,6 +164,7 @@ addopts = [ "--cov=scheduler", "--cov=log_consumer", "--cov=ide_callback", + "--cov=queue_backend", "--cov-report=term-missing", "--cov-report=html", "--cov-report=xml" @@ -185,7 +187,8 @@ source = [ "notification", "scheduler", "log_consumer", - "ide_callback" + "ide_callback", + "queue_backend" ] omit = [ "*/tests/*", diff --git a/workers/queue_backend/__init__.py b/workers/queue_backend/__init__.py new file mode 100644 index 0000000000..16ba181f6d --- /dev/null +++ b/workers/queue_backend/__init__.py @@ -0,0 +1,22 @@ +"""Queue-backend seam for workers. + +This module is the single place where the choice of queue substrate +(Celery+RabbitMQ today; PG Queue in the future) lives. + +Today both entry points are no-op aliases over Celery primitives: + +* ``dispatch(task_name, args, kwargs, queue)`` -> ``current_app.send_task(...)`` +* ``@worker_task`` -> ``@shared_task`` + +In a later PR these will gain per-task routing: when a task name appears +in ``WORKER_PG_QUEUE_ENABLED_TASKS``, dispatch routes through the PG Queue +backend instead. Default (env var empty) keeps 100% of traffic on Celery. + +Call sites should migrate to this module so the eventual substrate switch +is a single-flag operation rather than a codebase-wide rewrite. +""" + +from .decorator import worker_task +from .dispatch import dispatch + +__all__ = ["dispatch", "worker_task"] diff --git a/workers/queue_backend/decorator.py b/workers/queue_backend/decorator.py new file mode 100644 index 0000000000..f62232b714 --- /dev/null +++ b/workers/queue_backend/decorator.py @@ -0,0 +1,38 @@ +"""Task registration decorator. + +Today: a transparent wrapper over ``celery.shared_task``. +Future: registers the task body with whichever substrates are enabled +(Celery + optionally PG Queue), so a single ``@worker_task`` definition +can be served by either consumer. + +Accepts both Celery decorator forms — ``shared_task`` handles them +internally, so a pass-through ``*args, **kwargs`` is enough: + + @worker_task + def healthcheck(self): ... + + @worker_task(bind=True, name="my.task") + def my_task(self, payload): ... +""" + +from __future__ import annotations + +from typing import Any + +from celery import shared_task + + +def worker_task(*args: Any, **kwargs: Any) -> Any: + """Register a function as a worker task via the queue_backend seam. + + Today this is a one-line passthrough to ``celery.shared_task``. The + indirection is the seam: when PR #15 adds PG Queue routing, the + consumer-registration logic lands here without touching call sites. + + The return type is ``Any`` because ``shared_task`` returns different + objects depending on call form — a ``PromiseProxy`` for the bare + ``@worker_task`` form and a decorator factory for the parameterised + ``@worker_task(name=...)`` form. Pinning a tighter type would lock + out future routing variants without buying real safety today. + """ + return shared_task(*args, **kwargs) diff --git a/workers/queue_backend/dispatch.py b/workers/queue_backend/dispatch.py new file mode 100644 index 0000000000..6edcb3b95e --- /dev/null +++ b/workers/queue_backend/dispatch.py @@ -0,0 +1,46 @@ +"""Transport-agnostic task dispatch. + +Today: thin pass-through to ``celery.current_app.send_task``. +Future: per-task routing between Celery and PG Queue based on +``WORKER_PG_QUEUE_ENABLED_TASKS``. + +The signature intentionally exposes only what the current two raw +``current_app.send_task`` sites actually use (args, kwargs, queue). +More Celery options can be added when a real call site needs them — +not before. +""" + +from __future__ import annotations + +from typing import Any + +from celery import current_app + + +def dispatch( + task_name: str, + *, + args: list[Any] | None = None, + kwargs: dict[str, Any] | None = None, + queue: str | None = None, +) -> Any: + """Enqueue a task by name. + + Args: + task_name: Registered task name (e.g. "send_webhook_notification"). + args: Positional task args. Defaults to ``[]``. + kwargs: Keyword task args. Defaults to ``{}``. + queue: Target queue name. Defaults to the task's bound queue. + + Returns: + Whatever the underlying transport returns (today: a Celery + ``AsyncResult``). Callers should not assume a particular type + beyond truthiness — this is the seam point that lets PG Queue + return a different handle later. + """ + return current_app.send_task( + task_name, + args=args if args is not None else [], + kwargs=kwargs if kwargs is not None else {}, + queue=queue, + ) diff --git a/workers/run-worker.sh b/workers/run-worker.sh index 27d9fc8893..4a16d11247 100755 --- a/workers/run-worker.sh +++ b/workers/run-worker.sh @@ -41,6 +41,8 @@ declare -A WORKERS=( ["scheduler"]="scheduler" ["schedule"]="scheduler" ["${EXECUTOR_WORKER_TYPE}"]="${EXECUTOR_WORKER_TYPE}" + ["ide-callback"]="ide_callback" + ["ide_callback"]="ide_callback" ["all"]="all" ) @@ -57,6 +59,7 @@ declare -A WORKER_QUEUES=( ["notification"]="notifications,notifications_webhook,notifications_email,notifications_sms,notifications_priority" ["scheduler"]="scheduler" ["${EXECUTOR_WORKER_TYPE}"]="celery_executor_legacy" + ["ide_callback"]="ide_callback" ) # Worker health ports @@ -69,6 +72,7 @@ declare -A WORKER_HEALTH_PORTS=( ["notification"]="8085" ["scheduler"]="8087" ["${EXECUTOR_WORKER_TYPE}"]="8088" + ["ide_callback"]="8089" ) # Function to display usage @@ -87,6 +91,7 @@ WORKER_TYPE: notification, notify Run notification worker scheduler, schedule Run scheduler worker (scheduled pipeline tasks) executor Run executor worker (extraction execution tasks) + ide-callback Run IDE callback worker (Prompt Studio post-execution callbacks) all Run all workers (in separate processes, includes auto-discovered pluggable workers) Note: Pluggable workers in pluggable_worker/ directory are automatically discovered and can be run by name. @@ -101,7 +106,11 @@ OPTIONS: -P, --pool TYPE Set Celery pool type (threads, prefork, gevent, solo, eventlet) -n, --hostname NAME Set custom worker hostname/name -k, --kill Kill running workers and exit + -r, --restart Kill matching worker(s) then relaunch (with WORKER_TYPE, + restarts only that worker; without it, restarts all) -s, --status Show status of running workers + -L, --logs [WORKER] Live-tail worker log files (all if WORKER omitted) + -C, --clear-logs Delete worker .log files created by -d / 'all' runs -h, --help Show this help message EXAMPLES: @@ -134,6 +143,21 @@ EXAMPLES: # Kill all running workers $0 -k + # Live-tail all worker logs + $0 -L + + # Tail just one worker's log + $0 -L general + + # Wipe out old log files + $0 -C + + # Restart all workers + $0 -r -P prefork + + # Restart just one worker + $0 -r general -l DEBUG + ENVIRONMENT: The script will load environment variables from .env file if present. Required variables: @@ -146,16 +170,10 @@ ENVIRONMENT: See sample.env for full configuration options. HEALTH CHECKS: - Each worker exposes a health check endpoint: - - API Deployment: http://localhost:8080/health - - General: http://localhost:8081/health - - File Processing: http://localhost:8082/health - - Callback: http://localhost:8083/health - - Log Consumer: http://localhost:8084/health - - Notification: http://localhost:8085/health - - Scheduler: http://localhost:8087/health - - Executor: http://localhost:8088/health - - Pluggable workers: http://localhost:8090+/health (auto-assigned ports) + Workers can optionally bind an HTTP health server on the port assigned + via {WORKER}_HEALTH_PORT env vars (8080-8089 by default). It's used by + K8s liveness probes in production but is NOT load-bearing locally — + 'run-worker.sh -s' only reports whether the Celery process is alive. EOF } @@ -275,32 +293,142 @@ validate_env() { } # Function to get worker PIDs +# Anchors the match on the hostname value's start (preceded by a non-word +# character) so e.g. "callback" doesn't also match "ide_callback". get_worker_pids() { local worker_type=$1 - pgrep -f "uv run celery.*worker.*$worker_type" || true + pgrep -f -- "[^[:alnum:]_]${worker_type}-worker-" 2>/dev/null || true } -# Function to kill workers -kill_workers() { - print_status $YELLOW "Killing all running workers..." - - for worker in "${!WORKERS[@]}"; do - if [[ "$worker" == "all" ]]; then +# Function to resolve canonical worker dir names from the WORKERS map +# (skips aliases and "all"). Includes discovered pluggable workers. +list_core_worker_dirs() { + local seen="" + for key in "${!WORKERS[@]}"; do + local value="${WORKERS[$key]}" + if [[ "$value" == "all" ]]; then continue fi + if [[ "$seen" == *" $value "* ]]; then + continue + fi + seen="$seen $value " + echo "$value" + done +} - local worker_dir="${WORKERS[${worker}]}" - local pids=$(pgrep -f "uv run celery.*worker" || true) +list_pluggable_worker_dirs() { + for key in "${!PLUGGABLE_WORKERS[@]}"; do + local value="${PLUGGABLE_WORKERS[$key]}" + if [[ "$key" == "$value" ]]; then + echo "$value" + fi + done +} - if [[ -n "$pids" ]]; then - print_status $YELLOW "Killing worker processes: $pids" - echo "$pids" | xargs kill -TERM 2>/dev/null || true - sleep 2 - # Force kill if still running - echo "$pids" | xargs kill -KILL 2>/dev/null || true +# Function to resolve the log file path for a given canonical worker dir. +# Core workers live at WORKERS_DIR/$dir/$dir.log; pluggable workers +# live at WORKERS_DIR/pluggable_worker/$dir/$dir.log. +resolve_log_file() { + local worker_dir=$1 + local core_path="$WORKERS_DIR/$worker_dir/$worker_dir.log" + local pluggable_path="$WORKERS_DIR/pluggable_worker/$worker_dir/$worker_dir.log" + if [[ -f "$core_path" ]]; then + echo "$core_path" + elif [[ -f "$pluggable_path" ]]; then + echo "$pluggable_path" + fi +} + +# Function to tail one or all worker log files (-L|--logs) +tail_logs() { + local requested=$1 # may be empty → tail all + local log_files=() + + if [[ -z "$requested" || "$requested" == "all" ]]; then + for d in $(list_core_worker_dirs) $(list_pluggable_worker_dirs); do + local f + f=$(resolve_log_file "$d") + [[ -n "$f" ]] && log_files+=("$f") + done + else + # Resolve alias (e.g. "api" → "api-deployment") via WORKERS / PLUGGABLE_WORKERS + local canonical="${WORKERS[$requested]:-${PLUGGABLE_WORKERS[$requested]:-}}" + if [[ -z "$canonical" || "$canonical" == "all" ]]; then + print_status $RED "Error: Unknown worker type for logs: $requested" + print_status $BLUE "Tip: omit the worker type to tail all logs" + exit 1 fi + local f + f=$(resolve_log_file "$canonical") + if [[ -z "$f" ]]; then + print_status $YELLOW "No log file found for $canonical. Did you start it with -d?" + exit 0 + fi + log_files+=("$f") + fi + + if [[ ${#log_files[@]} -eq 0 ]]; then + print_status $YELLOW "No log files found. Workers must be started in detached mode (-d or 'all') for logs to be written to files." + exit 0 + fi + + print_status $BLUE "Tailing ${#log_files[@]} log file(s) — Ctrl+C to stop" + for f in "${log_files[@]}"; do + print_status $GREEN " $f" done + exec tail -F "${log_files[@]}" +} +# Function to delete worker log files (-C|--clear-logs) +clear_logs() { + local count=0 + for d in $(list_core_worker_dirs) $(list_pluggable_worker_dirs); do + local f + f=$(resolve_log_file "$d") + if [[ -n "$f" ]]; then + rm -f "$f" + print_status $GREEN "Removed: $f" + ((count++)) || true + fi + done + if [[ $count -eq 0 ]]; then + print_status $YELLOW "No worker log files to clear" + else + print_status $BLUE "Cleared $count log file(s)" + fi +} + +# Function to kill a single worker by its canonical dir name. +# Uses the anchored matcher from get_worker_pids so callback vs ide_callback +# don't bleed into each other. +kill_one_worker() { + local worker_dir=$1 + local pids + pids=$(get_worker_pids "$worker_dir" | tr '\n' ' ' | sed 's/ $//') + if [[ -z "$pids" ]]; then + print_status $YELLOW " $worker_dir: not running" + return + fi + print_status $YELLOW " $worker_dir: killing $pids" + echo "$pids" | xargs kill -TERM 2>/dev/null || true + sleep 1 + echo "$pids" | xargs kill -KILL 2>/dev/null || true +} + +# Function to kill workers +kill_workers() { + print_status $YELLOW "Killing all running workers..." + local pids + pids=$(pgrep -f "uv run celery.*worker" || true) + if [[ -z "$pids" ]]; then + print_status $GREEN "No workers running" + return + fi + print_status $YELLOW "Killing worker processes: $pids" + echo "$pids" | xargs kill -TERM 2>/dev/null || true + sleep 2 + echo "$pids" | xargs kill -KILL 2>/dev/null || true print_status $GREEN "All workers stopped" } @@ -309,38 +437,21 @@ show_status() { print_status $BLUE "Worker Status:" echo "==============" - local workers_to_check="api-deployment general file_processing callback log_consumer notification scheduler executor" - - # Add discovered pluggable workers - if [[ ${#PLUGGABLE_WORKERS[@]} -gt 0 ]]; then - for pluggable_name in "${!PLUGGABLE_WORKERS[@]}"; do - local canonical_name="${PLUGGABLE_WORKERS[$pluggable_name]}" - # Only add canonical names (skip aliases) - if [[ "$pluggable_name" == "$canonical_name" ]]; then - workers_to_check="$workers_to_check $canonical_name" - fi - done - fi + # Derive list from WORKERS (skips aliases & "all") so adding a new + # worker in one place keeps status in sync. Pluggable workers included. + local workers_to_check + workers_to_check="$(list_core_worker_dirs) $(list_pluggable_worker_dirs)" for worker in $workers_to_check; do - local worker_dir="$WORKERS_DIR/$worker" - local health_port="${WORKER_HEALTH_PORTS[$worker]}" - local pids=$(get_worker_pids "$worker") + local pids + pids=$(get_worker_pids "$worker" | tr '\n' ' ' | sed 's/ $//') - echo -n " $worker: " + printf ' %-22s ' "$worker:" if [[ -n "$pids" ]]; then - print_status $GREEN "RUNNING (PID: $pids)" - - # Check health endpoint if possible - if command -v curl >/dev/null 2>&1; then - local health_url="http://localhost:$health_port/health" - if curl -s --max-time 2 "$health_url" >/dev/null 2>&1; then - echo " Health: http://localhost:$health_port/health - OK" - else - echo " Health: http://localhost:$health_port/health - UNREACHABLE" - fi - fi + local count + count=$(echo "$pids" | wc -w) + print_status $GREEN "RUNNING ($count proc, PID: $pids)" else print_status $RED "STOPPED" fi @@ -416,6 +527,9 @@ run_worker() { "${EXECUTOR_WORKER_TYPE}") export EXECUTOR_HEALTH_PORT="$health_port" ;; + "ide_callback") + export IDE_CALLBACK_HEALTH_PORT="$health_port" + ;; *) # Handle pluggable workers dynamically if [[ -n "${PLUGGABLE_WORKERS[$worker_type]:-}" ]]; then @@ -492,6 +606,9 @@ run_worker() { "${EXECUTOR_WORKER_TYPE}") cmd_args+=("--concurrency=2") ;; + "ide_callback") + cmd_args+=("--concurrency=2") + ;; *) # Default for pluggable and other workers if [[ -n "${PLUGGABLE_WORKERS[$worker_type]:-}" ]]; then @@ -539,7 +656,7 @@ run_all_workers() { print_status $GREEN "Starting all workers..." # Define core workers - local core_workers="api-deployment general file_processing callback log_consumer notification scheduler executor" + local core_workers="api-deployment general file_processing callback log_consumer notification scheduler executor ide_callback" # Add discovered pluggable workers if [[ ${#PLUGGABLE_WORKERS[@]} -gt 0 ]]; then @@ -594,6 +711,9 @@ POOL_TYPE="" CUSTOM_HOSTNAME="" KILL_WORKERS=false SHOW_STATUS=false +LOGS_MODE=false +CLEAR_LOGS_MODE=false +RESTART_MODE=false while [[ $# -gt 0 ]]; do case $1 in @@ -605,6 +725,14 @@ while [[ $# -gt 0 ]]; do DETACH=true shift ;; + -L|--logs) + LOGS_MODE=true + shift + ;; + -C|--clear-logs) + CLEAR_LOGS_MODE=true + shift + ;; -l|--log-level) LOG_LEVEL="$2" shift 2 @@ -633,6 +761,10 @@ while [[ $# -gt 0 ]]; do KILL_WORKERS=true shift ;; + -r|--restart) + RESTART_MODE=true + shift + ;; -s|--status) SHOW_STATUS=true shift @@ -660,10 +792,43 @@ if [[ "$KILL_WORKERS" == "true" ]]; then fi if [[ "$SHOW_STATUS" == "true" ]]; then + discover_pluggable_workers show_status exit 0 fi +if [[ "$LOGS_MODE" == "true" ]]; then + discover_pluggable_workers + tail_logs "$WORKER_TYPE" + exit 0 +fi + +if [[ "$CLEAR_LOGS_MODE" == "true" ]]; then + discover_pluggable_workers + clear_logs + exit 0 +fi + +# Restart = kill matching + fall through to the normal launch path. +# - If WORKER_TYPE is set (and not "all"), kill only that worker. +# - Otherwise kill everything; the launch path will treat the unset +# WORKER_TYPE as "all" once we set it below. +if [[ "$RESTART_MODE" == "true" ]]; then + discover_pluggable_workers + if [[ -n "$WORKER_TYPE" && "$WORKER_TYPE" != "all" ]]; then + restart_target_dir="${WORKERS[$WORKER_TYPE]:-${PLUGGABLE_WORKERS[$WORKER_TYPE]:-}}" + if [[ -z "$restart_target_dir" ]]; then + print_status $RED "Error: Unknown worker type for restart: $WORKER_TYPE" + exit 1 + fi + print_status $BLUE "Restarting $restart_target_dir..." + kill_one_worker "$restart_target_dir" + else + kill_workers + WORKER_TYPE="all" + fi +fi + # Validate worker type if [[ -z "$WORKER_TYPE" ]]; then print_status $RED "Error: Worker type is required" diff --git a/workers/scheduler/tasks.py b/workers/scheduler/tasks.py index 2e22946a58..3215ecc104 100644 --- a/workers/scheduler/tasks.py +++ b/workers/scheduler/tasks.py @@ -7,7 +7,7 @@ import traceback from typing import Any -from celery import shared_task +from queue_backend import dispatch, worker_task from shared.enums.status_enums import PipelineStatus from shared.enums.worker_enums import QueueName from shared.infrastructure.config import WorkerConfig @@ -145,16 +145,13 @@ def _execute_scheduled_workflow( f"[exec:{execution_id}] [pipeline:{context.pipeline_id}] Triggering async execution for workflow {context.workflow_id}" ) - # Use Celery to dispatch async execution task directly (like backend scheduler does) - from celery import current_app - logger.info( f"[exec:{execution_id}] [pipeline:{context.pipeline_id}] Dispatching async_execute_bin task for scheduled execution" ) try: - # Dispatch the Celery task directly to the general queue - async_result = current_app.send_task( + # Dispatch through the queue_backend seam (Celery underneath today). + async_result = dispatch( "async_execute_bin", args=[ context.organization_id, # schema_name (organization_id) @@ -211,7 +208,7 @@ def _execute_scheduled_workflow( ) -@shared_task(name="scheduler.tasks.execute_pipeline_task", bind=True) +@worker_task(name="scheduler.tasks.execute_pipeline_task", bind=True) def execute_pipeline_task( self, workflow_id: Any, @@ -234,7 +231,7 @@ def execute_pipeline_task( ) -@shared_task(name="execute_pipeline_task_v2", bind=True) +@worker_task(name="execute_pipeline_task_v2", bind=True) def execute_pipeline_task_v2( self, organization_id: Any, @@ -415,7 +412,7 @@ def execute_pipeline_task_v2( # Health check task for monitoring -@shared_task(name="scheduler_health_check") +@worker_task(name="scheduler_health_check") def health_check() -> dict[str, Any]: """Health check task for scheduler worker. diff --git a/workers/scheduler/worker.py b/workers/scheduler/worker.py index 1a316d1790..a8e59f71fc 100644 --- a/workers/scheduler/worker.py +++ b/workers/scheduler/worker.py @@ -3,6 +3,7 @@ Celery worker for scheduled tasks and background processing. """ +from queue_backend import worker_task from shared.enums.worker_enums import WorkerType from shared.infrastructure.config.builder import WorkerBuilder from shared.infrastructure.config.registry import WorkerRegistry @@ -58,7 +59,7 @@ def check_scheduler_health(): ) -@app.task(bind=True) +@worker_task(bind=True) def healthcheck(self): """Health check task for monitoring systems.""" return { diff --git a/workers/shared/patterns/notification/helper.py b/workers/shared/patterns/notification/helper.py index 977f5a875f..cd79d6bd24 100644 --- a/workers/shared/patterns/notification/helper.py +++ b/workers/shared/patterns/notification/helper.py @@ -6,7 +6,7 @@ import logging -from celery import current_app +from queue_backend import dispatch # Import shared data models from @unstract/core from unstract.core.data_models import ( @@ -73,7 +73,7 @@ def send_notification_to_worker( payload_dict = payload.to_webhook_payload() # Send task to notification worker - current_app.send_task( + dispatch( "send_webhook_notification", args=[ url, diff --git a/workers/tests/test_dispatch_sites_characterisation.py b/workers/tests/test_dispatch_sites_characterisation.py index 95278a81b6..42de9a61ef 100644 --- a/workers/tests/test_dispatch_sites_characterisation.py +++ b/workers/tests/test_dispatch_sites_characterisation.py @@ -1,30 +1,29 @@ -"""Characterisation tests for the two raw `current_app.send_task` dispatch sites. +"""Behavioural tests for the two dispatch sites — now routed through the +``queue_backend.dispatch()`` seam. -These are the only two places in workers/ that bypass `@shared_task`-based -dispatch and call `current_app.send_task(...)` directly with a string task -name. PR #8 (mass `@shared_task` -> `@worker_task` migration) will replace -both with a unified `dispatch()` helper. +Both sites used to call ``current_app.send_task(...)`` directly with a +string task name. They now go through ``queue_backend.dispatch(...)`` +instead. The dispatch contract — task name, positional args, keyword +args, and target queue — is unchanged by the migration; ``dispatch()`` +is a transparent pass-through to ``current_app.send_task`` today. -This test suite locks down the **current** dispatch contract — task name, -positional args, keyword args, and target queue — so the migration can -be proved equivalent. It does NOT exercise the receiving tasks; it only -captures what is dispatched. +This suite locks down the contract at each call site so a future change +that quietly drops a kwarg or rewires a queue will fail loudly. -Sites characterised: -1. ``shared/patterns/notification/helper.py:76`` — webhook notification -2. ``scheduler/tasks.py:157`` — scheduled workflow async dispatch +Sites: +1. ``shared/patterns/notification/helper.py:send_notification_to_worker`` +2. ``scheduler/tasks.py:_execute_scheduled_workflow`` """ from unittest.mock import MagicMock, patch import pytest - # --- Site 1: shared/patterns/notification/helper.py --- class TestNotificationDispatchSite: - """Characterise ``send_notification_to_worker`` -> ``current_app.send_task``.""" + """Pin ``send_notification_to_worker`` -> ``queue_backend.dispatch``.""" def _make_payload(self): """Build a minimal NotificationPayload-shaped mock. @@ -41,9 +40,7 @@ def _make_payload(self): def test_dispatch_task_name_and_queue(self): from shared.patterns.notification.helper import send_notification_to_worker - with patch( - "shared.patterns.notification.helper.current_app" - ) as mock_app: + with patch("shared.patterns.notification.helper.dispatch") as mock_dispatch: send_notification_to_worker( url="https://example.com/hook", payload=self._make_payload(), @@ -52,8 +49,8 @@ def test_dispatch_task_name_and_queue(self): auth_header=None, ) - mock_app.send_task.assert_called_once() - call = mock_app.send_task.call_args + mock_dispatch.assert_called_once() + call = mock_dispatch.call_args assert call.args[0] == "send_webhook_notification" assert call.kwargs["queue"] == "notifications" @@ -61,9 +58,7 @@ def test_dispatch_positional_args_layout(self): """Positional args MUST be [url, payload_dict, headers, timeout].""" from shared.patterns.notification.helper import send_notification_to_worker - with patch( - "shared.patterns.notification.helper.current_app" - ) as mock_app: + with patch("shared.patterns.notification.helper.dispatch") as mock_dispatch: send_notification_to_worker( url="https://example.com/hook", payload=self._make_payload(), @@ -72,7 +67,7 @@ def test_dispatch_positional_args_layout(self): auth_header=None, ) - args = mock_app.send_task.call_args.kwargs["args"] + args = mock_dispatch.call_args.kwargs["args"] assert len(args) == 4 assert args[0] == "https://example.com/hook" # url assert args[1] == {"event": "test_event", "id": 42} # payload_dict @@ -83,9 +78,7 @@ def test_dispatch_kwargs_layout(self): """Kwargs MUST contain max_retries, retry_delay, platform.""" from shared.patterns.notification.helper import send_notification_to_worker - with patch( - "shared.patterns.notification.helper.current_app" - ) as mock_app: + with patch("shared.patterns.notification.helper.dispatch") as mock_dispatch: send_notification_to_worker( url="https://example.com/hook", payload=self._make_payload(), @@ -96,7 +89,7 @@ def test_dispatch_kwargs_layout(self): platform="SLACK", ) - kwargs = mock_app.send_task.call_args.kwargs["kwargs"] + kwargs = mock_dispatch.call_args.kwargs["kwargs"] assert kwargs == { "max_retries": 5, "retry_delay": 10, @@ -106,7 +99,7 @@ def test_dispatch_kwargs_layout(self): def test_dispatch_returns_true_on_success(self): from shared.patterns.notification.helper import send_notification_to_worker - with patch("shared.patterns.notification.helper.current_app"): + with patch("shared.patterns.notification.helper.dispatch"): result = send_notification_to_worker( url="https://example.com/hook", payload=self._make_payload(), @@ -120,10 +113,8 @@ def test_dispatch_returns_true_on_success(self): def test_dispatch_returns_false_on_send_task_failure(self): from shared.patterns.notification.helper import send_notification_to_worker - with patch( - "shared.patterns.notification.helper.current_app" - ) as mock_app: - mock_app.send_task.side_effect = RuntimeError("broker down") + with patch("shared.patterns.notification.helper.dispatch") as mock_dispatch: + mock_dispatch.side_effect = RuntimeError("broker down") result = send_notification_to_worker( url="https://example.com/hook", payload=self._make_payload(), @@ -139,7 +130,7 @@ def test_dispatch_returns_false_on_send_task_failure(self): class TestSchedulerDispatchSite: - """Characterise ``_execute_scheduled_workflow`` -> ``current_app.send_task``.""" + """Pin ``_execute_scheduled_workflow`` -> ``queue_backend.dispatch``.""" def _make_context(self): """Minimal ScheduledPipelineContext-shaped mock. @@ -166,32 +157,32 @@ def _make_api_client(self, execution_id="exec-123"): def test_dispatch_task_name(self): from scheduler.tasks import _execute_scheduled_workflow - with patch("celery.current_app") as mock_app: + with patch("scheduler.tasks.dispatch") as mock_dispatch: _execute_scheduled_workflow(self._make_api_client(), self._make_context()) - mock_app.send_task.assert_called_once() - assert mock_app.send_task.call_args.args[0] == "async_execute_bin" + mock_dispatch.assert_called_once() + assert mock_dispatch.call_args.args[0] == "async_execute_bin" def test_dispatch_routes_to_general_queue(self): from scheduler.tasks import _execute_scheduled_workflow from shared.enums.worker_enums import QueueName - with patch("celery.current_app") as mock_app: + with patch("scheduler.tasks.dispatch") as mock_dispatch: _execute_scheduled_workflow(self._make_api_client(), self._make_context()) - assert mock_app.send_task.call_args.kwargs["queue"] == QueueName.GENERAL + assert mock_dispatch.call_args.kwargs["queue"] == QueueName.GENERAL def test_dispatch_positional_args_layout(self): """Positional args MUST be [org_id, workflow_id, execution_id, {}, True].""" from scheduler.tasks import _execute_scheduled_workflow - with patch("celery.current_app") as mock_app: + with patch("scheduler.tasks.dispatch") as mock_dispatch: _execute_scheduled_workflow( self._make_api_client(execution_id="exec-xyz"), self._make_context(), ) - args = mock_app.send_task.call_args.kwargs["args"] + args = mock_dispatch.call_args.kwargs["args"] assert len(args) == 5 assert args[0] == "org-test" # organization_id (schema_name) assert args[1] == "wf-001" # workflow_id @@ -206,10 +197,10 @@ def test_dispatch_kwargs_layout(self): ctx = self._make_context() ctx.use_file_history = True - with patch("celery.current_app") as mock_app: + with patch("scheduler.tasks.dispatch") as mock_dispatch: _execute_scheduled_workflow(self._make_api_client(), ctx) - kwargs = mock_app.send_task.call_args.kwargs["kwargs"] + kwargs = mock_dispatch.call_args.kwargs["kwargs"] assert kwargs == { "use_file_history": True, "pipeline_id": "pipe-007", @@ -224,28 +215,26 @@ def test_no_dispatch_when_execution_creation_fails(self): api = MagicMock() api.create_workflow_execution.return_value = {} # no execution_id - with patch("celery.current_app") as mock_app: + with patch("scheduler.tasks.dispatch") as mock_dispatch: result = _execute_scheduled_workflow(api, self._make_context()) # The dispatch contract: nothing is sent when execution creation fails. - mock_app.send_task.assert_not_called() + mock_dispatch.assert_not_called() # And the function returns an error result (not raised exception). assert result.status == SchedulerExecutionStatus.ERROR - def test_dispatch_returns_error_result_when_send_task_raises(self): - """If current_app.send_task raises (broker down, queue gone, etc.), - _execute_scheduled_workflow MUST catch it, log, and return a - SchedulerExecutionResult.error(...) — NOT propagate the exception. - - PR #8 (dispatch migration) must preserve this semantic. If the new - dispatch() helper either swallows the error silently or re-raises a - different exception type, this test fails. + def test_dispatch_returns_error_result_when_dispatch_raises(self): + """If ``dispatch`` raises (broker down, queue gone, etc.), + ``_execute_scheduled_workflow`` MUST catch it, log, and return a + ``SchedulerExecutionResult.error(...)`` — NOT propagate the + exception. This guard makes sure a future routing change can't + silently flip the failure mode. """ from scheduler.tasks import _execute_scheduled_workflow from shared.models.scheduler_models import SchedulerExecutionStatus - with patch("celery.current_app") as mock_app: - mock_app.send_task.side_effect = RuntimeError("broker down") + with patch("scheduler.tasks.dispatch") as mock_dispatch: + mock_dispatch.side_effect = RuntimeError("broker down") result = _execute_scheduled_workflow( self._make_api_client(), self._make_context() ) @@ -253,20 +242,31 @@ def test_dispatch_returns_error_result_when_send_task_raises(self): # Exception was caught (not propagated). assert result.status == SchedulerExecutionStatus.ERROR # Execution itself was created before the dispatch failed, so the - # error result is about the dispatch, not the creation. - assert result.execution_id == "exec-123" or result.execution_id is None + # error result is about the dispatch and must carry the + # execution_id created moments earlier. + assert result.execution_id == "exec-123" # --- Cross-site invariant --- class TestDispatchSiteInventory: - """If a third raw current_app.send_task site appears, this test breaks - so PR #8's migration can't silently miss it.""" + """All task dispatch must flow through the queue_backend seam. + + Raw ``current_app.send_task`` calls outside ``queue_backend/`` are now + forbidden — they bypass the substrate-routing chokepoint. If one + reappears, this test fails so the regression can't ship silently. + """ - def test_only_two_known_dispatch_sites_in_workers(self): - """Verify the count of raw current_app.send_task references in - workers/ source matches the two we have characterised.""" + def test_no_raw_dispatch_sites_outside_seam(self): + """Zero raw ``current_app.send_task`` references should exist in + workers/ source outside the queue_backend/ seam. + + ``queue_backend/dispatch.py`` is the canonical home: it's the one + file that calls ``current_app.send_task`` (today, while the + substrate is still Celery). Everything else must call + ``queue_backend.dispatch(...)`` instead. + """ import pathlib import re @@ -274,7 +274,7 @@ def test_only_two_known_dispatch_sites_in_workers(self): # Anchor skip to the top-level directory relative to workers_root so # we don't accidentally exclude legitimately-named subdirectories # (e.g. workers/shared/tests_helpers/). - skip_top_dirs = {"tests", "__pycache__", "htmlcov", ".venv"} + skip_top_dirs = {"tests", "__pycache__", "htmlcov", ".venv", "queue_backend"} pattern = re.compile(r"current_app\.send_task\b") hits = [] @@ -287,16 +287,24 @@ def test_only_two_known_dispatch_sites_in_workers(self): if pattern.search(line): hits.append(f"{py.relative_to(workers_root)}:{line_no}") - # Expected exactly two — in helper.py and scheduler/tasks.py. - # If a third appears, this test fails so PR #8 doesn't miss it. - assert len(hits) == 2, ( - f"Expected exactly 2 raw current_app.send_task sites in workers/, found " - f"{len(hits)}:\n " + "\n ".join(hits) + assert hits == [], ( + "Raw current_app.send_task call(s) found outside queue_backend/. " + "All task dispatch must go through queue_backend.dispatch(). " + "Found:\n " + "\n ".join(hits) ) - # Sanity: the two we know about - joined = " ".join(hits) - assert "shared/patterns/notification/helper.py" in joined - assert "scheduler/tasks.py" in joined + + def test_canonical_seam_exists(self): + """The queue_backend/dispatch.py seam must exist — it's where every + dispatch site routes through. Behavioural equivalence to the + legacy ``current_app.send_task`` calls is covered by the + ``TestDispatchEquivalence`` suite in test_queue_backend_seam.py; + this test only guards the file's existence so a future refactor + that moves the module by mistake fails loudly. + """ + import pathlib + + seam = pathlib.Path(__file__).parent.parent / "queue_backend" / "dispatch.py" + assert seam.exists(), f"Canonical dispatch seam missing: {seam}" if __name__ == "__main__": diff --git a/workers/tests/test_queue_backend_seam.py b/workers/tests/test_queue_backend_seam.py new file mode 100644 index 0000000000..bc1333cb26 --- /dev/null +++ b/workers/tests/test_queue_backend_seam.py @@ -0,0 +1,229 @@ +"""Equivalence tests for the queue-backend seam. + +Today the seam is a transparent pass-through to Celery. These tests +lock that in so future PRs adding per-task routing (PG Queue) can be +proved to preserve the Celery default path. + +Two layers: + +1. **dispatch()** must produce the same ``current_app.send_task`` call + as the raw idiom used at the two existing dispatch sites + (``notification/helper.py:76``, ``scheduler/tasks.py:157``). + +2. **@worker_task** must be the same object as ``@shared_task`` so that + functions decorated with it register identically with the Celery app. +""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +# --- dispatch() equivalence --- + + +class TestDispatchEquivalence: + """dispatch() forwards to current_app.send_task with the same shape.""" + + def test_dispatches_task_name(self): + from queue_backend import dispatch + + with patch("queue_backend.dispatch.current_app") as mock_app: + dispatch("send_webhook_notification") + + mock_app.send_task.assert_called_once() + assert mock_app.send_task.call_args.args[0] == "send_webhook_notification" + + def test_forwards_args(self): + from queue_backend import dispatch + + with patch("queue_backend.dispatch.current_app") as mock_app: + dispatch("any_task", args=["a", "b", 42]) + + assert mock_app.send_task.call_args.kwargs["args"] == ["a", "b", 42] + + def test_forwards_kwargs(self): + from queue_backend import dispatch + + with patch("queue_backend.dispatch.current_app") as mock_app: + dispatch("any_task", kwargs={"max_retries": 5, "platform": "SLACK"}) + + assert mock_app.send_task.call_args.kwargs["kwargs"] == { + "max_retries": 5, + "platform": "SLACK", + } + + def test_forwards_queue(self): + from queue_backend import dispatch + + with patch("queue_backend.dispatch.current_app") as mock_app: + dispatch("any_task", queue="notifications") + + assert mock_app.send_task.call_args.kwargs["queue"] == "notifications" + + def test_defaults_args_to_empty_list(self): + """Omitting args must produce [], not None — Celery treats them differently.""" + from queue_backend import dispatch + + with patch("queue_backend.dispatch.current_app") as mock_app: + dispatch("any_task") + + assert mock_app.send_task.call_args.kwargs["args"] == [] + + def test_defaults_kwargs_to_empty_dict(self): + """Omitting kwargs must produce {}, not None.""" + from queue_backend import dispatch + + with patch("queue_backend.dispatch.current_app") as mock_app: + dispatch("any_task") + + assert mock_app.send_task.call_args.kwargs["kwargs"] == {} + + def test_returns_underlying_handle(self): + """Whatever send_task returns is what dispatch returns.""" + from queue_backend import dispatch + + sentinel = object() + with patch("queue_backend.dispatch.current_app") as mock_app: + mock_app.send_task.return_value = sentinel + result = dispatch("any_task") + + assert result is sentinel + + def test_matches_notification_helper_shape(self): + """dispatch() output is identical to the raw send_task at notification/helper.py:76. + + Reference call from helper.py (paraphrased): + + current_app.send_task( + "send_webhook_notification", + args=[url, payload_dict, headers, timeout], + kwargs={"max_retries": ..., "retry_delay": ..., "platform": ...}, + queue="notifications", + ) + + The seam must produce the same call. + """ + from queue_backend import dispatch + + with patch("queue_backend.dispatch.current_app") as mock_app: + dispatch( + "send_webhook_notification", + args=["https://example.com/hook", {"event": "x"}, {}, 10], + kwargs={"max_retries": 3, "retry_delay": 10, "platform": "API"}, + queue="notifications", + ) + + call = mock_app.send_task.call_args + assert call.args[0] == "send_webhook_notification" + assert call.kwargs["args"] == ["https://example.com/hook", {"event": "x"}, {}, 10] + assert call.kwargs["kwargs"] == { + "max_retries": 3, + "retry_delay": 10, + "platform": "API", + } + assert call.kwargs["queue"] == "notifications" + + def test_matches_scheduler_dispatch_shape(self): + """dispatch() output is identical to the raw send_task at scheduler/tasks.py:157.""" + from queue_backend import dispatch + + with patch("queue_backend.dispatch.current_app") as mock_app: + dispatch( + "async_execute_bin", + args=["org-1", "wf-1", "exec-1", {}, True], + kwargs={"use_file_history": False, "pipeline_id": "pipe-1"}, + queue="general", + ) + + call = mock_app.send_task.call_args + assert call.args[0] == "async_execute_bin" + assert call.kwargs["args"] == ["org-1", "wf-1", "exec-1", {}, True] + assert call.kwargs["kwargs"] == { + "use_file_history": False, + "pipeline_id": "pipe-1", + } + assert call.kwargs["queue"] == "general" + + +# --- @worker_task equivalence --- + + +class TestWorkerTaskEquivalence: + """@worker_task produces Celery-task-equivalent objects today. + + The seam is a thin function wrapper (not an identity alias) so PR #15 + can grow consumer-registration logic without restructuring callers. + """ + + def test_worker_task_is_callable_passthrough(self): + """worker_task delegates to shared_task — verified by calling it with + the same args and asserting the produced task carries the same name. + """ + from celery import shared_task + + from queue_backend import worker_task + + @worker_task(name="queue_backend_test.passthrough") + def via_seam(): + return "ok" + + @shared_task(name="queue_backend_test.passthrough_native") + def via_native(): + return "ok" + + # Both should produce something Celery considers a registered task. + for t in (via_seam, via_native): + assert hasattr(t, "apply") or hasattr(t, "delay") + + def test_bare_decorator_form(self): + """@worker_task on a function registers a Celery task.""" + from queue_backend import worker_task + + @worker_task + def some_function(x): + return x * 2 + + # shared_task returns a Task instance (or a PromiseProxy on import). + # We don't care about the concrete type — just that the decorator + # produced something call-able with .delay/.apply. + assert hasattr(some_function, "apply") or hasattr(some_function, "delay") + + def test_parameterised_decorator_form(self): + """@worker_task(name=..., queue=...) accepts kwargs like @shared_task.""" + from queue_backend import worker_task + + @worker_task(name="queue_backend_test.parameterised", queue="general") + def some_function(): + return "ok" + + assert hasattr(some_function, "apply") or hasattr(some_function, "delay") + + +# --- Module surface --- + + +class TestPublicSurface: + """Pin the public API — guards future PRs against accidental signature changes.""" + + def test_exports_dispatch(self): + import queue_backend + + assert hasattr(queue_backend, "dispatch") + assert callable(queue_backend.dispatch) + + def test_exports_worker_task(self): + import queue_backend + + assert hasattr(queue_backend, "worker_task") + assert callable(queue_backend.worker_task) + + def test_all_exports(self): + import queue_backend + + assert set(queue_backend.__all__) == {"dispatch", "worker_task"} + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 658ea89e452a2b1a45269651739190de6f0440af Mon Sep 17 00:00:00 2001 From: ali Date: Mon, 1 Jun 2026 12:57:20 +0530 Subject: [PATCH 2/3] UN-3497 [FIX] Address PR review feedback on queue_backend seam MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Critical: * run-worker.sh:300 pgrep regression — anchored matcher required a trailing `-` after `-worker`, but default hostname is `${type}-worker@%h` (no dash unless WORKER_INSTANCE_ID set). Status, restart, and kill_one_worker silently reported STOPPED for any worker not setting that env var. Pattern now matches `(@|-)`. * test_queue_backend_seam.py — replaced `hasattr(t, "apply")` tautology with real Celery registration assertions (force PromiseProxy resolution, verify `current_app.tasks` membership, round-trip via `.apply().get()`). * test_queue_backend_seam.py — added explicit decorator-kwarg passthrough test pinning `name`, `bind`, `autoretry_for`, `max_retries`, `default_retry_delay`. Guards against a refactor like `return shared_task(*args)` silently stripping every retry policy. Important: * run-worker.sh:423 kill_workers — bulk path used loose `pgrep -f "uv run celery.*worker"` which would kill unrelated celery procs. Iterate canonical dir list and delegate to kill_one_worker so the anchored matcher is the single source of truth. * run-worker.sh:414 silent EPERM — dropped `2>/dev/null || true` from kill invocations; re-check PIDs after sleep; warn on survivors. * test_dispatch_sites_characterisation.py — replaced regex-based inventory canary with AST walker. Now catches aliased imports (`from celery import current_app as app`), locally-constructed apps (`Celery(...).send_task(...)`), and `apply_async`. Excludes `plugins/manual_review` (pre-existing `.apply_async` baggage, out of this PR's scope). * test_dispatch_sites_characterisation.py — added `test_substrate_failure_surfaces_as_false` patching at the substrate boundary (`queue_backend.dispatch.current_app.send_task`) with `ConnectionError`. Complements the helper-side test so the full helper → seam → broker failure path is exercised without either side being trivially mocked. Suggestions: * dispatch.py — dropped redundant `args is not None else []` coercion; Celery normalises `None` internally. Avoids subtle list-vs-tuple default behaviour change for third-party routers checking `isinstance(args, tuple)`. * dispatch.py — return type Any → `DispatchHandle` Protocol exposing `.id: str` (already in use at scheduler/tasks.py). Documents the real invariant the PG Queue handle will also need. Relaxed `list[Any]` → `Sequence[Any]`, `dict[str, Any]` → `Mapping[str, Any]`. * Removed `PR #15` placeholder + env var name pinning (`WORKER_PG_QUEUE_ENABLED_TASKS`) from docstrings — describe behaviour, not unstable identifiers. * test_queue_backend_seam.py:11 — replaced stale `scheduler/tasks.py:157` line ref with symbol-based citation (`_execute_scheduled_workflow` / `send_notification_to_worker`). SonarCloud (S1192): * run-worker.sh — defined `readonly IDE_CALLBACK_WORKER_TYPE` constant and replaced the 8 literal occurrences, mirroring the existing `EXECUTOR_WORKER_TYPE` pattern. Test count: 28 → 31 (seam 15→17, characterisation 13→14). Co-Authored-By: Claude Opus 4.7 (1M context) --- workers/queue_backend/__init__.py | 6 +- workers/queue_backend/decorator.py | 5 +- workers/queue_backend/dispatch.py | 48 ++++--- workers/run-worker.sh | 67 ++++++--- .../test_dispatch_sites_characterisation.py | 94 +++++++++--- workers/tests/test_queue_backend_seam.py | 135 ++++++++++++------ 6 files changed, 251 insertions(+), 104 deletions(-) diff --git a/workers/queue_backend/__init__.py b/workers/queue_backend/__init__.py index 16ba181f6d..5fe4ed4510 100644 --- a/workers/queue_backend/__init__.py +++ b/workers/queue_backend/__init__.py @@ -8,9 +8,9 @@ * ``dispatch(task_name, args, kwargs, queue)`` -> ``current_app.send_task(...)`` * ``@worker_task`` -> ``@shared_task`` -In a later PR these will gain per-task routing: when a task name appears -in ``WORKER_PG_QUEUE_ENABLED_TASKS``, dispatch routes through the PG Queue -backend instead. Default (env var empty) keeps 100% of traffic on Celery. +A later phase will route specific tasks through a non-Celery substrate +(PG Queue) based on configuration; until then everything goes to Celery. +The exact routing mechanism is intentionally not pinned here. Call sites should migrate to this module so the eventual substrate switch is a single-flag operation rather than a codebase-wide rewrite. diff --git a/workers/queue_backend/decorator.py b/workers/queue_backend/decorator.py index f62232b714..9f8b666fe0 100644 --- a/workers/queue_backend/decorator.py +++ b/workers/queue_backend/decorator.py @@ -26,8 +26,9 @@ def worker_task(*args: Any, **kwargs: Any) -> Any: """Register a function as a worker task via the queue_backend seam. Today this is a one-line passthrough to ``celery.shared_task``. The - indirection is the seam: when PR #15 adds PG Queue routing, the - consumer-registration logic lands here without touching call sites. + indirection is the seam: when a later phase adds PG Queue routing, + the consumer-registration logic lands here without touching call + sites. The return type is ``Any`` because ``shared_task`` returns different objects depending on call form — a ``PromiseProxy`` for the bare diff --git a/workers/queue_backend/dispatch.py b/workers/queue_backend/dispatch.py index 6edcb3b95e..f1c21ca110 100644 --- a/workers/queue_backend/dispatch.py +++ b/workers/queue_backend/dispatch.py @@ -1,46 +1,58 @@ """Transport-agnostic task dispatch. Today: thin pass-through to ``celery.current_app.send_task``. -Future: per-task routing between Celery and PG Queue based on -``WORKER_PG_QUEUE_ENABLED_TASKS``. +A later phase will introduce per-task routing through a non-Celery +substrate (PG Queue); call sites stay untouched. -The signature intentionally exposes only what the current two raw -``current_app.send_task`` sites actually use (args, kwargs, queue). -More Celery options can be added when a real call site needs them — -not before. +The signature intentionally exposes only what the current call sites +actually use (args, kwargs, queue). More Celery options can be added +when a real call site needs them — not before. """ from __future__ import annotations -from typing import Any +from collections.abc import Mapping, Sequence +from typing import Any, Protocol from celery import current_app +class DispatchHandle(Protocol): + """The minimum contract every dispatch substrate must satisfy. + + Today this is satisfied by Celery's ``AsyncResult`` (which exposes + ``.id``). A future PG Queue handle will need to expose the same + attribute so existing callers — e.g. ``scheduler/tasks.py`` — keep + working unchanged. + """ + + id: str + + def dispatch( task_name: str, *, - args: list[Any] | None = None, - kwargs: dict[str, Any] | None = None, + args: Sequence[Any] | None = None, + kwargs: Mapping[str, Any] | None = None, queue: str | None = None, -) -> Any: +) -> DispatchHandle: """Enqueue a task by name. Args: task_name: Registered task name (e.g. "send_webhook_notification"). - args: Positional task args. Defaults to ``[]``. - kwargs: Keyword task args. Defaults to ``{}``. + args: Positional task args. Forwarded verbatim; Celery normalises + ``None`` internally. + kwargs: Keyword task args. Forwarded verbatim; Celery normalises + ``None`` internally. queue: Target queue name. Defaults to the task's bound queue. Returns: - Whatever the underlying transport returns (today: a Celery - ``AsyncResult``). Callers should not assume a particular type - beyond truthiness — this is the seam point that lets PG Queue - return a different handle later. + A handle to the enqueued task. ``.id`` is guaranteed; everything + else is substrate-specific and callers must not rely on it. """ return current_app.send_task( task_name, - args=args if args is not None else [], - kwargs=kwargs if kwargs is not None else {}, + args=args, + kwargs=kwargs, queue=queue, ) diff --git a/workers/run-worker.sh b/workers/run-worker.sh index 4a16d11247..5913b0ec60 100755 --- a/workers/run-worker.sh +++ b/workers/run-worker.sh @@ -23,6 +23,7 @@ ENV_FILE="$WORKERS_DIR/.env" # Worker type constant for the executor worker readonly EXECUTOR_WORKER_TYPE="executor" +readonly IDE_CALLBACK_WORKER_TYPE="ide_callback" # Available workers declare -A WORKERS=( @@ -41,8 +42,8 @@ declare -A WORKERS=( ["scheduler"]="scheduler" ["schedule"]="scheduler" ["${EXECUTOR_WORKER_TYPE}"]="${EXECUTOR_WORKER_TYPE}" - ["ide-callback"]="ide_callback" - ["ide_callback"]="ide_callback" + ["ide-callback"]="${IDE_CALLBACK_WORKER_TYPE}" + ["${IDE_CALLBACK_WORKER_TYPE}"]="${IDE_CALLBACK_WORKER_TYPE}" ["all"]="all" ) @@ -59,7 +60,7 @@ declare -A WORKER_QUEUES=( ["notification"]="notifications,notifications_webhook,notifications_email,notifications_sms,notifications_priority" ["scheduler"]="scheduler" ["${EXECUTOR_WORKER_TYPE}"]="celery_executor_legacy" - ["ide_callback"]="ide_callback" + ["${IDE_CALLBACK_WORKER_TYPE}"]="${IDE_CALLBACK_WORKER_TYPE}" ) # Worker health ports @@ -72,7 +73,7 @@ declare -A WORKER_HEALTH_PORTS=( ["notification"]="8085" ["scheduler"]="8087" ["${EXECUTOR_WORKER_TYPE}"]="8088" - ["ide_callback"]="8089" + ["${IDE_CALLBACK_WORKER_TYPE}"]="8089" ) # Function to display usage @@ -294,10 +295,13 @@ validate_env() { # Function to get worker PIDs # Anchors the match on the hostname value's start (preceded by a non-word -# character) so e.g. "callback" doesn't also match "ide_callback". +# character) so e.g. "callback" doesn't also match "ide_callback". The +# trailing `(@|-)` covers both Celery hostname forms emitted by run_worker: +# --hostname=callback-worker@%h (default, no WORKER_INSTANCE_ID) +# --hostname=callback-worker-${id}@%h (when WORKER_INSTANCE_ID is set) get_worker_pids() { local worker_type=$1 - pgrep -f -- "[^[:alnum:]_]${worker_type}-worker-" 2>/dev/null || true + pgrep -f -- "[^[:alnum:]_]${worker_type}-worker(@|-)" || true } # Function to resolve canonical worker dir names from the WORKERS map @@ -401,35 +405,52 @@ clear_logs() { # Function to kill a single worker by its canonical dir name. # Uses the anchored matcher from get_worker_pids so callback vs ide_callback -# don't bleed into each other. +# don't bleed into each other. Surfaces kill failures (EPERM, ESRCH) and +# verifies the processes are gone before reporting success. kill_one_worker() { local worker_dir=$1 local pids pids=$(get_worker_pids "$worker_dir" | tr '\n' ' ' | sed 's/ $//') if [[ -z "$pids" ]]; then print_status $YELLOW " $worker_dir: not running" - return + return 0 fi print_status $YELLOW " $worker_dir: killing $pids" - echo "$pids" | xargs kill -TERM 2>/dev/null || true + # shellcheck disable=SC2086 + kill -TERM $pids || print_status $RED " $worker_dir: kill -TERM failed" sleep 1 - echo "$pids" | xargs kill -KILL 2>/dev/null || true + local survivors + survivors=$(get_worker_pids "$worker_dir" | tr '\n' ' ' | sed 's/ $//') + if [[ -n "$survivors" ]]; then + # shellcheck disable=SC2086 + kill -KILL $survivors || print_status $RED " $worker_dir: kill -KILL failed" + sleep 1 + survivors=$(get_worker_pids "$worker_dir" | tr '\n' ' ' | sed 's/ $//') + if [[ -n "$survivors" ]]; then + print_status $RED " $worker_dir: survivors after SIGKILL: $survivors" + return 1 + fi + fi + return 0 } -# Function to kill workers +# Function to kill all known workers. Iterates the canonical dir list and +# delegates to kill_one_worker so the anchored matcher is the single source +# of truth — avoids killing unrelated `celery worker` processes on the host. kill_workers() { print_status $YELLOW "Killing all running workers..." - local pids - pids=$(pgrep -f "uv run celery.*worker" || true) - if [[ -z "$pids" ]]; then - print_status $GREEN "No workers running" - return + local any_failed=0 + local workers_to_kill + workers_to_kill="$(list_core_worker_dirs) $(list_pluggable_worker_dirs)" + for worker in $workers_to_kill; do + kill_one_worker "$worker" || any_failed=1 + done + if [[ $any_failed -eq 0 ]]; then + print_status $GREEN "All workers stopped" + else + print_status $RED "One or more workers could not be stopped" + return 1 fi - print_status $YELLOW "Killing worker processes: $pids" - echo "$pids" | xargs kill -TERM 2>/dev/null || true - sleep 2 - echo "$pids" | xargs kill -KILL 2>/dev/null || true - print_status $GREEN "All workers stopped" } # Function to show worker status @@ -527,7 +548,7 @@ run_worker() { "${EXECUTOR_WORKER_TYPE}") export EXECUTOR_HEALTH_PORT="$health_port" ;; - "ide_callback") + "${IDE_CALLBACK_WORKER_TYPE}") export IDE_CALLBACK_HEALTH_PORT="$health_port" ;; *) @@ -606,7 +627,7 @@ run_worker() { "${EXECUTOR_WORKER_TYPE}") cmd_args+=("--concurrency=2") ;; - "ide_callback") + "${IDE_CALLBACK_WORKER_TYPE}") cmd_args+=("--concurrency=2") ;; *) diff --git a/workers/tests/test_dispatch_sites_characterisation.py b/workers/tests/test_dispatch_sites_characterisation.py index 42de9a61ef..d07721bdb2 100644 --- a/workers/tests/test_dispatch_sites_characterisation.py +++ b/workers/tests/test_dispatch_sites_characterisation.py @@ -110,7 +110,10 @@ def test_dispatch_returns_true_on_success(self): assert result is True - def test_dispatch_returns_false_on_send_task_failure(self): + def test_helper_swallows_dispatch_exception(self): + """The helper's try/except contract: any exception from ``dispatch`` + is caught and surfaced as ``False``. Patches the helper-side seam, + so this is purely a Python-level error-handling assertion.""" from shared.patterns.notification.helper import send_notification_to_worker with patch("shared.patterns.notification.helper.dispatch") as mock_dispatch: @@ -125,6 +128,31 @@ def test_dispatch_returns_false_on_send_task_failure(self): assert result is False + def test_substrate_failure_surfaces_as_false(self): + """End-to-end seam → substrate failure path. + + Leaves ``dispatch`` itself unpatched and patches the substrate + boundary (``queue_backend.dispatch.current_app``). Models a real + broker outage: ``send_task`` raises ``ConnectionError``, + ``dispatch`` propagates, the helper's ``except`` catches it, + return value is ``False``. Complements the helper-side test + above — together they exercise the full failure path without + either side being trivially mocked.""" + from shared.patterns.notification.helper import send_notification_to_worker + + with patch("queue_backend.dispatch.current_app") as mock_app: + mock_app.send_task.side_effect = ConnectionError("broker down") + result = send_notification_to_worker( + url="https://example.com/hook", + payload=self._make_payload(), + auth_type="NONE", + auth_key=None, + auth_header=None, + ) + + assert result is False + mock_app.send_task.assert_called_once() + # --- Site 2: scheduler/tasks.py --- @@ -259,38 +287,70 @@ class TestDispatchSiteInventory: """ def test_no_raw_dispatch_sites_outside_seam(self): - """Zero raw ``current_app.send_task`` references should exist in - workers/ source outside the queue_backend/ seam. + """Zero raw Celery dispatch calls should exist in workers/ source + outside the queue_backend/ seam. - ``queue_backend/dispatch.py`` is the canonical home: it's the one - file that calls ``current_app.send_task`` (today, while the + ``queue_backend/dispatch.py`` is the canonical home: it's the + one file that calls ``current_app.send_task`` (today, while the substrate is still Celery). Everything else must call ``queue_backend.dispatch(...)`` instead. + + Uses an AST walker rather than a regex so the canary covers: + + * aliased imports (``from celery import current_app as app``; + ``app.send_task(...)``), + * locally-constructed apps (``Celery(...).send_task(...)``), + * ``.apply_async`` (the other half of how Celery sends tasks, + including ``signature(...).apply_async()``), + * multi-line dotted access after autoformatter wrap. """ + import ast import pathlib - import re workers_root = pathlib.Path(__file__).parent.parent # Anchor skip to the top-level directory relative to workers_root so # we don't accidentally exclude legitimately-named subdirectories # (e.g. workers/shared/tests_helpers/). skip_top_dirs = {"tests", "__pycache__", "htmlcov", ".venv", "queue_backend"} - - pattern = re.compile(r"current_app\.send_task\b") - hits = [] + # ``plugins/manual_review`` predates this seam and uses + # ``task.apply_async`` lambda overrides to inject queues — out of + # scope for the queue_backend transport seam PR. Tracked as a + # follow-up so this canary is exact for everything we own. + skip_subpaths = {"plugins/manual_review"} + forbidden_attrs = {"send_task", "apply_async"} + + class DispatchCallFinder(ast.NodeVisitor): + def __init__(self) -> None: + self.hits: list[int] = [] + + def visit_Call(self, node: ast.Call) -> None: + func = node.func + if isinstance(func, ast.Attribute) and func.attr in forbidden_attrs: + self.hits.append(node.lineno) + self.generic_visit(node) + + hits: list[str] = [] for py in workers_root.rglob("*.py"): - rel_parts = py.relative_to(workers_root).parts + rel = py.relative_to(workers_root) + rel_parts = rel.parts if rel_parts and rel_parts[0] in skip_top_dirs: continue - text = py.read_text() - for line_no, line in enumerate(text.splitlines(), start=1): - if pattern.search(line): - hits.append(f"{py.relative_to(workers_root)}:{line_no}") + rel_str = rel.as_posix() + if any(rel_str.startswith(skip + "/") for skip in skip_subpaths): + continue + try: + tree = ast.parse(py.read_text(), filename=str(py)) + except SyntaxError: + continue + finder = DispatchCallFinder() + finder.visit(tree) + for line_no in finder.hits: + hits.append(f"{py.relative_to(workers_root)}:{line_no}") assert hits == [], ( - "Raw current_app.send_task call(s) found outside queue_backend/. " - "All task dispatch must go through queue_backend.dispatch(). " - "Found:\n " + "\n ".join(hits) + "Raw Celery dispatch call(s) (send_task / apply_async) found " + "outside queue_backend/. All task dispatch must go through " + "queue_backend.dispatch(). Found:\n " + "\n ".join(hits) ) def test_canonical_seam_exists(self): diff --git a/workers/tests/test_queue_backend_seam.py b/workers/tests/test_queue_backend_seam.py index bc1333cb26..bf42aaf7ab 100644 --- a/workers/tests/test_queue_backend_seam.py +++ b/workers/tests/test_queue_backend_seam.py @@ -7,11 +7,13 @@ Two layers: 1. **dispatch()** must produce the same ``current_app.send_task`` call - as the raw idiom used at the two existing dispatch sites - (``notification/helper.py:76``, ``scheduler/tasks.py:157``). + as the raw idiom used at the two existing dispatch sites: + ``shared/patterns/notification/helper.py::send_notification_to_worker`` + and ``scheduler/tasks.py::_execute_scheduled_workflow``. -2. **@worker_task** must be the same object as ``@shared_task`` so that - functions decorated with it register identically with the Celery app. +2. **@worker_task** must register a task with the Celery app so that + functions decorated with it behave identically to ones decorated + with ``@shared_task``. """ from __future__ import annotations @@ -19,6 +21,7 @@ from unittest.mock import patch import pytest +from celery import current_app # --- dispatch() equivalence --- @@ -62,23 +65,28 @@ def test_forwards_queue(self): assert mock_app.send_task.call_args.kwargs["queue"] == "notifications" - def test_defaults_args_to_empty_list(self): - """Omitting args must produce [], not None — Celery treats them differently.""" + def test_omitted_args_forwarded_as_none(self): + """Omitted ``args`` is forwarded verbatim as ``None``. + + Celery's own ``send_task`` normalises ``None`` to its native default + (a tuple) internally — the seam doesn't coerce, so any third-party + router that checks ``isinstance(args, tuple)`` sees Celery's default. + """ from queue_backend import dispatch with patch("queue_backend.dispatch.current_app") as mock_app: dispatch("any_task") - assert mock_app.send_task.call_args.kwargs["args"] == [] + assert mock_app.send_task.call_args.kwargs["args"] is None - def test_defaults_kwargs_to_empty_dict(self): - """Omitting kwargs must produce {}, not None.""" + def test_omitted_kwargs_forwarded_as_none(self): + """Same as args — omitted kwargs reaches Celery as ``None``.""" from queue_backend import dispatch with patch("queue_backend.dispatch.current_app") as mock_app: dispatch("any_task") - assert mock_app.send_task.call_args.kwargs["kwargs"] == {} + assert mock_app.send_task.call_args.kwargs["kwargs"] is None def test_returns_underlying_handle(self): """Whatever send_task returns is what dispatch returns.""" @@ -92,9 +100,9 @@ def test_returns_underlying_handle(self): assert result is sentinel def test_matches_notification_helper_shape(self): - """dispatch() output is identical to the raw send_task at notification/helper.py:76. + """Output mirrors the raw send_task at ``send_notification_to_worker``. - Reference call from helper.py (paraphrased): + Reference call from helper (paraphrased): current_app.send_task( "send_webhook_notification", @@ -102,8 +110,6 @@ def test_matches_notification_helper_shape(self): kwargs={"max_retries": ..., "retry_delay": ..., "platform": ...}, queue="notifications", ) - - The seam must produce the same call. """ from queue_backend import dispatch @@ -126,7 +132,7 @@ def test_matches_notification_helper_shape(self): assert call.kwargs["queue"] == "notifications" def test_matches_scheduler_dispatch_shape(self): - """dispatch() output is identical to the raw send_task at scheduler/tasks.py:157.""" + """Output mirrors the raw send_task at ``_execute_scheduled_workflow``.""" from queue_backend import dispatch with patch("queue_backend.dispatch.current_app") as mock_app: @@ -151,54 +157,101 @@ def test_matches_scheduler_dispatch_shape(self): class TestWorkerTaskEquivalence: - """@worker_task produces Celery-task-equivalent objects today. + """@worker_task registers tasks indistinguishably from @shared_task today. - The seam is a thin function wrapper (not an identity alias) so PR #15 - can grow consumer-registration logic without restructuring callers. + The seam is a thin function wrapper (not an identity alias) so a later + phase can grow consumer-registration logic without restructuring callers. + Assertions go through Celery's task registry so they fail loudly if the + decorator stops producing real Celery tasks. """ - def test_worker_task_is_callable_passthrough(self): - """worker_task delegates to shared_task — verified by calling it with - the same args and asserting the produced task carries the same name. - """ + def test_bare_decorator_registers_with_celery(self): + """@worker_task on a function registers the task by its module-qualified name.""" + from queue_backend import worker_task + + @worker_task + def queue_backend_test_bare(x): + return x * 2 + + # Force PromiseProxy resolution — MagicMock.name wouldn't survive this. + resolved_name = queue_backend_test_bare.name + assert resolved_name in current_app.tasks + # Round-trip the registered task to confirm it actually runs. + assert current_app.tasks[resolved_name].apply(args=(3,)).get() == 6 + + def test_parameterised_decorator_uses_explicit_name(self): + """@worker_task(name=..., queue=...) registers under the explicit name.""" + from queue_backend import worker_task + + @worker_task(name="queue_backend_test.parameterised", queue="general") + def some_function(): + return "ok" + + assert some_function.name == "queue_backend_test.parameterised" + assert "queue_backend_test.parameterised" in current_app.tasks + + def test_worker_task_matches_shared_task_registration(self): + """A function decorated with @worker_task is the same kind of object as + one decorated with @shared_task — same registration semantics, same + invocation interface.""" from celery import shared_task from queue_backend import worker_task - @worker_task(name="queue_backend_test.passthrough") + @worker_task(name="queue_backend_test.via_seam") def via_seam(): return "ok" - @shared_task(name="queue_backend_test.passthrough_native") + @shared_task(name="queue_backend_test.via_native") def via_native(): return "ok" - # Both should produce something Celery considers a registered task. - for t in (via_seam, via_native): - assert hasattr(t, "apply") or hasattr(t, "delay") + for task, expected_name in ( + (via_seam, "queue_backend_test.via_seam"), + (via_native, "queue_backend_test.via_native"), + ): + assert task.name == expected_name + assert expected_name in current_app.tasks + assert current_app.tasks[expected_name].apply().get() == "ok" + + def test_forwards_decorator_kwargs(self): + """All Celery decorator kwargs reach @shared_task. - def test_bare_decorator_form(self): - """@worker_task on a function registers a Celery task.""" + Guards against a refactor like ``return shared_task(*args)`` that + would silently drop every retry policy, name override, and bind + flag in the codebase. + """ from queue_backend import worker_task - @worker_task - def some_function(x): - return x * 2 + @worker_task( + name="queue_backend_test.kwargs", + bind=True, + autoretry_for=(ValueError,), + max_retries=7, + default_retry_delay=42, + ) + def with_policy(self): + return "ok" - # shared_task returns a Task instance (or a PromiseProxy on import). - # We don't care about the concrete type — just that the decorator - # produced something call-able with .delay/.apply. - assert hasattr(some_function, "apply") or hasattr(some_function, "delay") + assert with_policy.name == "queue_backend_test.kwargs" + registered = current_app.tasks["queue_backend_test.kwargs"] + assert ValueError in (registered.autoretry_for or ()) + assert registered.max_retries == 7 + assert registered.default_retry_delay == 42 - def test_parameterised_decorator_form(self): - """@worker_task(name=..., queue=...) accepts kwargs like @shared_task.""" + def test_bare_decorator_form_uses_module_qualified_name(self): + """``@worker_task`` (no parens) gives Celery's auto-generated name.""" from queue_backend import worker_task - @worker_task(name="queue_backend_test.parameterised", queue="general") - def some_function(): + @worker_task + def queue_backend_test_auto_named(): return "ok" - assert hasattr(some_function, "apply") or hasattr(some_function, "delay") + # Default name is ``.``. + assert queue_backend_test_auto_named.name.endswith( + ".queue_backend_test_auto_named" + ) + assert queue_backend_test_auto_named.name in current_app.tasks # --- Module surface --- From c6af203b00605b860aded91617f0abe5eefabab4 Mon Sep 17 00:00:00 2001 From: ali Date: Mon, 1 Jun 2026 13:20:22 +0530 Subject: [PATCH 3/3] UN-3497 [FIX] Extract get_worker_pids_oneline helper Address SonarCloud S1192: ``get_worker_pids ... | tr '\n' ' ' | sed 's/ $//'`` was repeated at 4 call sites in run-worker.sh. Wrap the formatting pipeline in a single helper so the literal lives in one place. Co-Authored-By: Claude Opus 4.7 (1M context) --- workers/run-worker.sh | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/workers/run-worker.sh b/workers/run-worker.sh index 5913b0ec60..6597907a1a 100755 --- a/workers/run-worker.sh +++ b/workers/run-worker.sh @@ -304,6 +304,14 @@ get_worker_pids() { pgrep -f -- "[^[:alnum:]_]${worker_type}-worker(@|-)" || true } +# Returns get_worker_pids output as a single space-separated string with +# the trailing space stripped — the form xargs and `kill` want. Wrapping +# this avoids repeating the `tr | sed` pipeline at every caller. +get_worker_pids_oneline() { + local worker_type=$1 + get_worker_pids "$worker_type" | tr '\n' ' ' | sed 's/ $//' +} + # Function to resolve canonical worker dir names from the WORKERS map # (skips aliases and "all"). Includes discovered pluggable workers. list_core_worker_dirs() { @@ -410,7 +418,7 @@ clear_logs() { kill_one_worker() { local worker_dir=$1 local pids - pids=$(get_worker_pids "$worker_dir" | tr '\n' ' ' | sed 's/ $//') + pids=$(get_worker_pids_oneline "$worker_dir") if [[ -z "$pids" ]]; then print_status $YELLOW " $worker_dir: not running" return 0 @@ -420,12 +428,12 @@ kill_one_worker() { kill -TERM $pids || print_status $RED " $worker_dir: kill -TERM failed" sleep 1 local survivors - survivors=$(get_worker_pids "$worker_dir" | tr '\n' ' ' | sed 's/ $//') + survivors=$(get_worker_pids_oneline "$worker_dir") if [[ -n "$survivors" ]]; then # shellcheck disable=SC2086 kill -KILL $survivors || print_status $RED " $worker_dir: kill -KILL failed" sleep 1 - survivors=$(get_worker_pids "$worker_dir" | tr '\n' ' ' | sed 's/ $//') + survivors=$(get_worker_pids_oneline "$worker_dir") if [[ -n "$survivors" ]]; then print_status $RED " $worker_dir: survivors after SIGKILL: $survivors" return 1 @@ -465,7 +473,7 @@ show_status() { for worker in $workers_to_check; do local pids - pids=$(get_worker_pids "$worker" | tr '\n' ' ' | sed 's/ $//') + pids=$(get_worker_pids_oneline "$worker") printf ' %-22s ' "$worker:"