chore: import upstream snapshot with attribution

python/ray/llm/_internal/common/callbacks/base.py

@@ -0,0 +1,144 @@
+import asyncio
+import inspect
+import logging
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Type, Union
+
+if TYPE_CHECKING:
+    from ray.llm._internal.common.utils.download_utils import NodeModelDownloadable
+    from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CallbackCtx:
+    """
+    Context object passed to all callback hooks.
+    Callbacks can read and modify fields as needed.
+    """
+
+    worker_node_download_model: Optional["NodeModelDownloadable"] = None
+    """Model download configuration for worker nodes. Used to specify how
+    models should be downloaded and cached on worker nodes in distributed
+    deployments."""
+    placement_group: Optional[Any] = None
+    """Ray placement group for resource allocation and scheduling. Controls
+    where and how resources are allocated across the cluster."""
+    runtime_env: Optional[Dict[str, Any]] = None
+    """Runtime environment configuration for the Ray workers. Includes
+    dependencies, environment variables, and other runtime settings."""
+    custom_data: Dict[str, Any] = field(default_factory=dict)
+    """Flexible dictionary for callback-specific state and data. Allows
+    callbacks to store and share custom information during initialization."""
+    run_init_node: bool = True
+    """Whether to run model downloads during initialization. Set to False
+    to skip downloading models."""
+
+
+class CallbackBase:
+    """Base class for custom initialization implementations.
+
+    This class defines the interface for custom initialization logic
+    for LLMEngine to be called in node_initialization.
+    """
+
+    def __init__(
+        self,
+        llm_config: "LLMConfig",
+        raise_error_on_callback: bool = True,
+        ctx_kwargs: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        self.raise_error_on_callback = raise_error_on_callback
+        self.kwargs = kwargs
+        self.llm_config = llm_config
+
+        # Create and store CallbackCtx internally using ctx_kwargs
+        ctx_kwargs = ctx_kwargs or {}
+        self.ctx = CallbackCtx(**ctx_kwargs)
+
+    async def on_before_node_init(self) -> None:
+        """Called before node initialization begins."""
+        pass
+
+    async def on_after_node_init(self) -> None:
+        """Called after node initialization completes."""
+        pass
+
+    def on_before_download_model_files_distributed(self) -> None:
+        """Called before model files are downloaded on each node."""
+        pass
+
+    def _get_method(self, method_name: str) -> Tuple[Callable, bool]:
+        """Get a callback method."""
+        if not hasattr(self, method_name):
+            raise AttributeError(
+                f"Callback {type(self).__name__} does not have method '{method_name}'"
+            )
+        return getattr(self, method_name), inspect.iscoroutinefunction(
+            getattr(self, method_name)
+        )
+
+    def _handle_callback_error(self, method_name: str, e: Exception) -> None:
+        if self.raise_error_on_callback:
+            raise Exception(
+                f"Error running callback method '{method_name}' on {type(self).__name__}: {str(e)}"
+            ) from e
+        else:
+            logger.error(
+                f"Error running callback method '{method_name}' on {type(self).__name__}: {str(e)}"
+            )
+
+    async def run_callback(self, method_name: str) -> None:
+        """Run a callback method either synchronously or asynchronously.
+
+        Args:
+            method_name: The name of the method to call on the callback
+
+        Returns:
+            None
+        """
+        method, is_async = self._get_method(method_name)
+
+        try:
+            if is_async:
+                await method()
+            else:
+                method()
+        except Exception as e:
+            self._handle_callback_error(method_name, e)
+
+    def run_callback_sync(self, method_name: str) -> None:
+        """Run a callback method synchronously
+
+        Args:
+            method_name: The name of the method to call on the callback
+        Returns:
+            None
+        """
+        method, is_async = self._get_method(method_name)
+
+        try:
+            if is_async:
+                try:
+                    loop = asyncio.get_running_loop()
+                    loop.run_until_complete(method())
+                except RuntimeError:
+                    asyncio.run(method())
+            else:
+                method()
+        except Exception as e:
+            self._handle_callback_error(method_name, e)
+
+
+@dataclass
+class CallbackConfig:
+    """Configuration for the callback to be used in LLMConfig"""
+
+    callback_class: Union[str, Type[CallbackBase]] = CallbackBase
+    """Class to use for the callback. Can be custom user defined class"""
+    callback_kwargs: Dict[str, Any] = field(default_factory=dict)
+    """Keyword arguments to pass to the Callback class at construction."""
+    raise_error_on_callback: bool = True
+    """Whether to raise an error if a callback method fails."""