chore: import upstream snapshot with attribution

python/ray/serve/task_consumer.py

@@ -0,0 +1,229 @@
+import inspect
+import logging
+from functools import wraps
+from typing import Callable, Optional
+
+from ray._common.utils import import_attr
+from ray.serve._private.constants import (
+    DEFAULT_CONSUMER_CONCURRENCY,
+    SERVE_LOGGER_NAME,
+)
+from ray.serve._private.task_consumer import TaskConsumerWrapper
+from ray.serve._private.utils import copy_class_metadata
+from ray.serve.schema import (
+    TaskProcessorAdapter,
+    TaskProcessorConfig,
+)
+from ray.util.annotations import PublicAPI
+
+logger = logging.getLogger(SERVE_LOGGER_NAME)
+
+
+def _instantiate_adapter(
+    task_processor_config: TaskProcessorConfig,
+    consumer_concurrency: int = DEFAULT_CONSUMER_CONCURRENCY,
+) -> TaskProcessorAdapter:
+    adapter = task_processor_config.adapter
+
+    # Handle string-based adapter specification (module path)
+    if isinstance(adapter, str):
+        adapter_class = import_attr(adapter)
+
+    elif callable(adapter):
+        adapter_class = adapter
+
+    else:
+        raise TypeError(
+            f"Adapter must be either a string path or a callable class, got {type(adapter).__name__}: {adapter}"
+        )
+
+    try:
+        adapter_instance = adapter_class(task_processor_config)
+    except Exception as e:
+        raise RuntimeError(f"Failed to instantiate {adapter_class.__name__}: {e}")
+
+    if not isinstance(adapter_instance, TaskProcessorAdapter):
+        raise TypeError(
+            f"{adapter_class.__name__} must inherit from TaskProcessorAdapter, got {type(adapter_instance).__name__}"
+        )
+
+    try:
+        adapter_instance.initialize(consumer_concurrency)
+    except Exception as e:
+        raise RuntimeError(f"Failed to initialize {adapter_class.__name__}: {e}")
+
+    return adapter_instance
+
+
+@PublicAPI(stability="alpha")
+def instantiate_adapter_from_config(
+    task_processor_config: TaskProcessorConfig,
+) -> TaskProcessorAdapter:
+    """
+    Create a TaskProcessorAdapter instance from the provided configuration and call .initialize(). This function supports two ways to specify an adapter:
+
+    1. String path: A fully qualified module path to an adapter class
+       Example: "ray.serve.task_processor.CeleryTaskProcessorAdapter"
+
+    2. Class reference: A direct reference to an adapter class
+       Example: CeleryTaskProcessorAdapter
+
+    Args:
+        task_processor_config: Configuration object containing adapter specification.
+    Returns:
+        An initialized TaskProcessorAdapter instance ready for use.
+
+    Raises:
+        ValueError: If the adapter string path is malformed or cannot be imported.
+        TypeError: If the adapter is not a string or callable class.
+
+    Example:
+        .. code-block:: python
+
+            config = TaskProcessorConfig(
+                adapter="my.module.CustomAdapter",
+                adapter_config={"param": "value"},
+                queue_name="my_queue"
+            )
+            adapter = instantiate_adapter_from_config(config)
+    """
+
+    return _instantiate_adapter(task_processor_config)
+
+
+@PublicAPI(stability="alpha")
+def task_consumer(*, task_processor_config: TaskProcessorConfig):
+    """
+    Decorator to mark a class as a TaskConsumer.
+
+    Args:
+        task_processor_config: Configuration for the task processor (required)
+
+    Note:
+        This decorator must be used with parentheses:
+        @task_consumer(task_processor_config=config)
+
+    Returns:
+        A wrapper class that inherits from the target class and implements the task consumer functionality.
+
+    Example:
+        .. code-block:: python
+
+            from ray import serve
+            from ray.serve.task_consumer import task_consumer, task_handler
+
+            @serve.deployment
+            @task_consumer(task_processor_config=config)
+            class MyTaskConsumer:
+
+                @task_handler(name="my_task")
+                def my_task(self, *args, **kwargs):
+                    pass
+
+    """
+
+    def decorator(target_cls):
+        class _TaskConsumerWrapper(target_cls, TaskConsumerWrapper):
+            _adapter: TaskProcessorAdapter
+
+            def __init__(self, *args, **kwargs):
+                target_cls.__init__(self, *args, **kwargs)
+
+            def initialize_callable(self, consumer_concurrency: int):
+                self._adapter = _instantiate_adapter(
+                    task_processor_config, consumer_concurrency
+                )
+
+                for name, method in inspect.getmembers(
+                    target_cls, predicate=inspect.isfunction
+                ):
+                    if getattr(method, "_is_task_handler", False):
+                        task_name = getattr(method, "_task_name", name)
+
+                        # Create a callable that properly binds the method to this instance
+                        bound_method = getattr(self, name)
+
+                        self._adapter.register_task_handle(bound_method, task_name)
+
+                try:
+                    self._adapter.start_consumer()
+                    logger.info("task consumer started successfully")
+                except Exception as e:
+                    logger.error(f"Failed to start task consumer: {e}")
+                    raise
+
+            def __del__(self):
+                self._adapter.stop_consumer()
+
+                if hasattr(target_cls, "__del__"):
+                    target_cls.__del__(self)
+
+        copy_class_metadata(_TaskConsumerWrapper, target_cls)
+
+        return _TaskConsumerWrapper
+
+    return decorator
+
+
+@PublicAPI(stability="alpha")
+def task_handler(
+    _func: Optional[Callable] = None, *, name: Optional[str] = None
+) -> Callable:
+    """
+    Decorator to mark a method as a task handler.
+    Optionally specify a task name. Default is the method name.
+
+    Arguments:
+        _func: The function to decorate.
+        name: The name of the task. Default is the method name.
+
+    Returns:
+        A wrapper function that is marked as a task handler.
+
+    Example:
+        .. code-block:: python
+
+            from ray import serve
+            from ray.serve.task_consumer import task_consumer, task_handler
+
+            @serve.deployment
+            @task_consumer(task_processor_config=config)
+            class MyTaskConsumer:
+
+                @task_handler(name="my_task")
+                def my_task(self, *args, **kwargs):
+                    pass
+
+    """
+
+    # Validate name parameter if provided
+    if name is not None and (not isinstance(name, str) or not name.strip()):
+        raise ValueError(f"Task name must be a non-empty string, got {name}")
+
+    def decorator(f):
+        # Handle both sync and async functions
+        if inspect.iscoroutinefunction(f):
+
+            @wraps(f)
+            async def async_wrapper(*args, **kwargs):
+                return await f(*args, **kwargs)
+
+            async_wrapper._is_task_handler = True  # type: ignore
+            async_wrapper._task_name = name or f.__name__  # type: ignore
+            return async_wrapper
+        else:
+
+            @wraps(f)
+            def wrapper(*args, **kwargs):
+                return f(*args, **kwargs)
+
+            wrapper._is_task_handler = True  # type: ignore
+            wrapper._task_name = name or f.__name__  # type: ignore
+            return wrapper
+
+    if _func is not None:
+        # Used without arguments: @task_handler
+        return decorator(_func)
+    else:
+        # Used with arguments: @task_handler(name="...")
+        return decorator