chore: import upstream snapshot with attribution

python/ray/serve/task_processor.py

@@ -0,0 +1,359 @@
+import asyncio
+import logging
+import threading
+import time
+from typing import Any, Dict, List, Optional
+
+from celery import Celery
+from celery.signals import task_failure, task_unknown
+
+from ray.serve import get_replica_context
+from ray.serve._private.constants import (
+    DEFAULT_CONSUMER_CONCURRENCY,
+    SERVE_LOGGER_NAME,
+)
+from ray.serve.schema import (
+    CeleryAdapterConfig,
+    TaskProcessorAdapter,
+    TaskProcessorConfig,
+    TaskResult,
+)
+from ray.util.annotations import PublicAPI
+
+logger = logging.getLogger(SERVE_LOGGER_NAME)
+
+
+CELERY_WORKER_POOL = "worker_pool"
+CELERY_WORKER_CONCURRENCY = "worker_concurrency"
+CELERY_TASK_IGNORE_RESULT = "task_ignore_result"
+CELERY_TASK_ACKS_LATE = "task_acks_late"
+CELERY_TASK_REJECT_ON_WORKER_LOST = "task_reject_on_worker_lost"
+
+CELERY_DEFAULT_APP_CONFIG = [
+    CELERY_WORKER_POOL,
+    CELERY_WORKER_CONCURRENCY,
+    CELERY_TASK_IGNORE_RESULT,
+    CELERY_TASK_ACKS_LATE,
+    CELERY_TASK_REJECT_ON_WORKER_LOST,
+]
+
+
+@PublicAPI(stability="alpha")
+class CeleryTaskProcessorAdapter(TaskProcessorAdapter):
+    """
+    Celery-based task processor adapter.
+    This adapter does NOT support any async operations.
+    All operations must be performed synchronously.
+    """
+
+    _app: Celery
+    _config: TaskProcessorConfig
+    _worker_thread: Optional[threading.Thread] = None
+    _worker_hostname: Optional[str] = None
+    _worker_concurrency: int = DEFAULT_CONSUMER_CONCURRENCY
+
+    def __init__(self, config: TaskProcessorConfig, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if not isinstance(config.adapter_config, CeleryAdapterConfig):
+            raise TypeError(
+                "TaskProcessorConfig.adapter_config must be an instance of CeleryAdapterConfig"
+            )
+
+        # Check if any app_custom_config keys conflict with default Celery app config
+        if config.adapter_config.app_custom_config:
+            conflicting_keys = set(
+                config.adapter_config.app_custom_config.keys()
+            ) & set(CELERY_DEFAULT_APP_CONFIG)
+            if conflicting_keys:
+                raise ValueError(
+                    f"The following configuration keys cannot be changed via app_custom_config: {sorted(conflicting_keys)}. "
+                    f"These are managed internally by the CeleryTaskProcessorAdapter."
+                )
+
+        self._config = config
+
+        # Celery adapter does not support any async capabilities
+        # self._async_capabilities is already an empty set from parent class
+
+    def initialize(self, consumer_concurrency: int = DEFAULT_CONSUMER_CONCURRENCY):
+        self._app = Celery(
+            self._config.queue_name,
+            backend=self._config.adapter_config.backend_url,
+            broker=self._config.adapter_config.broker_url,
+        )
+
+        app_configuration = {
+            CELERY_WORKER_POOL: "threads",
+            CELERY_WORKER_CONCURRENCY: consumer_concurrency,
+            CELERY_TASK_IGNORE_RESULT: False,  # Store task results so they can be retrieved after completion
+            CELERY_TASK_ACKS_LATE: True,  # Acknowledge tasks only after completion (not when received) for better reliability
+            CELERY_TASK_REJECT_ON_WORKER_LOST: True,  # Reject and requeue tasks when worker is lost to prevent data loss
+        }
+
+        if self._config.adapter_config.app_custom_config:
+            app_configuration.update(self._config.adapter_config.app_custom_config)
+
+        self._app.conf.update(app_configuration)
+
+        queue_config = {
+            self._config.queue_name: {
+                "exchange": self._config.queue_name,
+                "exchange_type": "direct",
+                "routing_key": self._config.queue_name,
+            },
+        }
+
+        if self._config.failed_task_queue_name:
+            queue_config[self._config.failed_task_queue_name] = {
+                "exchange": self._config.failed_task_queue_name,
+                "exchange_type": "direct",
+                "routing_key": self._config.failed_task_queue_name,
+            }
+
+        if self._config.unprocessable_task_queue_name:
+            queue_config[self._config.unprocessable_task_queue_name] = {
+                "exchange": self._config.unprocessable_task_queue_name,
+                "exchange_type": "direct",
+                "routing_key": self._config.unprocessable_task_queue_name,
+            }
+
+        self._app.conf.update(
+            task_queues=queue_config,
+            task_routes={
+                # Default tasks go to main queue
+                "*": {"queue": self._config.queue_name},
+            },
+        )
+
+        if self._config.adapter_config.broker_transport_options is not None:
+            self._app.conf.update(
+                broker_transport_options=self._config.adapter_config.broker_transport_options,
+            )
+
+        if self._config.failed_task_queue_name:
+            task_failure.connect(self._handle_task_failure)
+
+        if self._config.unprocessable_task_queue_name:
+            task_unknown.connect(self._handle_unknown_task)
+
+    def register_task_handle(self, func, name=None):
+        # Celery does not support async task handlers
+        if asyncio.iscoroutinefunction(func):
+            raise NotImplementedError(
+                "Async task handlers are not supported yet. "
+                "Please use synchronous functions for task handlers."
+            )
+
+        task_options = {
+            "autoretry_for": (Exception,),
+            "retry_kwargs": {"max_retries": self._config.max_retries},
+            "retry_backoff": True,
+            "retry_backoff_max": 60,  # Max backoff of 60 seconds
+            "retry_jitter": False,  # Disable jitter for predictable testing
+        }
+        if self._config.adapter_config.task_custom_config:
+            task_options.update(self._config.adapter_config.task_custom_config)
+
+        if name:
+            self._app.task(name=name, **task_options)(func)
+        else:
+            self._app.task(**task_options)(func)
+
+    def enqueue_task_sync(
+        self, task_name, args=None, kwargs=None, **options
+    ) -> TaskResult:
+        task_response = self._app.send_task(
+            task_name,
+            args=args,
+            kwargs=kwargs,
+            queue=self._config.queue_name,
+            **options,
+        )
+
+        return TaskResult(
+            id=task_response.id,
+            status=task_response.status,
+            created_at=time.time(),
+            result=task_response.result,
+        )
+
+    def get_task_status_sync(self, task_id) -> TaskResult:
+        task_details = self._app.AsyncResult(task_id)
+        return TaskResult(
+            id=task_details.id,
+            result=task_details.result,
+            status=task_details.status,
+        )
+
+    def start_consumer(self, **kwargs):
+        """Starts the Celery worker thread."""
+        if self._worker_thread is not None and self._worker_thread.is_alive():
+            logger.info("Celery worker thread is already running.")
+            return
+
+        unique_id = get_replica_context().replica_tag
+        self._worker_hostname = f"{self._app.main}_{unique_id}"
+
+        worker_args = [
+            "worker",
+            f"--hostname={self._worker_hostname}",
+            "-Q",
+            self._config.queue_name,
+        ]
+
+        self._worker_thread = threading.Thread(
+            target=self._app.worker_main,
+            args=(worker_args,),
+        )
+        self._worker_thread.start()
+
+        logger.info(
+            f"Celery worker thread started with hostname: {self._worker_hostname}"
+        )
+
+    def stop_consumer(self, timeout: float = 10.0):
+        """Signals the Celery worker to shut down and waits for it to terminate."""
+        if self._worker_thread is None or not self._worker_thread.is_alive():
+            logger.info("Celery worker thread is not running.")
+            return
+
+        logger.info("Sending shutdown signal to Celery worker...")
+
+        # Use the worker's hostname for targeted shutdown
+        self._app.control.broadcast(
+            "shutdown", destination=[f"celery@{self._worker_hostname}"]
+        )
+        self._worker_thread.join(timeout=timeout)
+
+        if self._worker_thread.is_alive():
+            logger.warning(f"Worker thread did not terminate after {timeout} seconds.")
+        else:
+            logger.info("Celery worker thread has stopped.")
+
+        self._worker_thread = None
+
+    def cancel_task_sync(self, task_id):
+        """
+        Cancels a task synchronously. Only supported for Redis and RabbitMQ brokers by Celery.
+        More details can be found here: https://docs.celeryq.dev/en/stable/userguide/workers.html#revoke-revoking-tasks
+        """
+        self._app.control.revoke(task_id)
+
+    def get_metrics_sync(self) -> Dict[str, Any]:
+        """
+        Returns the metrics of the Celery worker synchronously.
+        More details can be found here: https://docs.celeryq.dev/en/stable/reference/celery.app.control.html#celery.app.control.Inspect.stats
+        """
+        return self._app.control.inspect().stats()
+
+    def health_check_sync(self) -> List[Dict]:
+        """
+        Checks the health of the Celery worker synchronously.
+        Returns a list of dictionaries, each containing the worker name and a dictionary with the health status.
+        Example: [{'celery@192.168.1.100': {'ok': 'pong'}}]
+        More details can be found here: https://docs.celeryq.dev/en/stable/reference/celery.app.control.html#celery.app.control.Control.ping
+        """
+        return self._app.control.ping()
+
+    def _handle_task_failure(
+        self,
+        sender: Any = None,
+        task_id: str = None,
+        args: Any = None,
+        kwargs: Any = None,
+        einfo: Any = None,
+        **kw,
+    ):
+        """Handle task failures and route them to appropriate dead letter queues.
+
+        This method is called when a task fails after all retry attempts have been
+        exhausted. It logs the failure and moves the task to failed_task_queue
+
+        Args:
+            sender: The task object that failed
+            task_id: Unique identifier of the failed task
+            args: Positional arguments passed to the task
+            kwargs: Keyword arguments passed to the task
+            einfo: Exception info object containing exception details and traceback
+            **kw: Additional keyword arguments passed by Celery
+        """
+        logger.info(
+            f"Task failure detected for task_id: {task_id}, einfo: {str(einfo)}"
+        )
+
+        dlq_args = [
+            task_id,
+            str(einfo.exception),
+            str(args),
+            str(kwargs),
+            str(einfo),
+        ]
+
+        if self._config.failed_task_queue_name:
+            self._move_task_to_queue(
+                self._config.failed_task_queue_name,
+                sender.name,
+                dlq_args,
+            )
+
+            logger.error(
+                f"Task {task_id} failed after max retries. Exception: {einfo}. Moved it to the {self._config.failed_task_queue_name} queue."
+            )
+
+    def _handle_unknown_task(
+        self,
+        sender: Any = None,
+        name: str = None,
+        id: str = None,
+        message: Any = None,
+        exc: Any = None,
+        **kwargs,
+    ):
+        """Handle unknown or unregistered tasks received by Celery.
+
+        This method is called when Celery receives a task that it doesn't recognize
+        (i.e., a task that hasn't been registered with the Celery app). These tasks
+        are moved to the unprocessable task queue if configured.
+
+        Args:
+            sender: The Celery app or worker that detected the unknown task
+            name: Name of the unknown task
+            id: Task ID of the unknown task
+            message: The raw message received for the unknown task
+            exc: The exception raised when trying to process the unknown task
+            **kwargs: Additional context information from Celery
+        """
+        logger.info(
+            f"Unknown task detected by Celery. Name: {name}, ID: {id}, Exc: {str(exc)}"
+        )
+
+        if self._config.unprocessable_task_queue_name:
+            self._move_task_to_queue(
+                self._config.unprocessable_task_queue_name,
+                name,
+                [
+                    name,
+                    id,
+                    str(message),
+                    str(exc),
+                    str(kwargs),
+                ],
+            )
+
+    def _move_task_to_queue(self, queue_name: str, task_name: str, args: list):
+        """Helper function to move a task to a specified queue."""
+        try:
+            logger.info(
+                f"Moving task: {task_name} to queue: {queue_name}, args: {args}"
+            )
+            self._app.send_task(
+                name=task_name,
+                queue=queue_name,
+                args=args,
+            )
+        except Exception as e:
+            logger.error(
+                f"Failed to move task: {task_name} to queue: {queue_name}, error: {e}"
+            )
+            raise e