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."""

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

@@ -0,0 +1,86 @@
+import logging
+import time
+from typing import Any, List, Tuple
+
+from pydantic import BaseModel, field_validator
+
+from .base import CallbackBase
+
+logger = logging.getLogger(__name__)
+
+
+class CloudDownloaderConfig(BaseModel):
+    """Model for validating CloudDownloader configuration."""
+
+    paths: List[Tuple[str, str]]
+
+    @field_validator("paths")
+    @classmethod
+    def validate_paths(cls, v: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+        # Supported cloud storage URI schemes
+        valid_schemes = ("s3://", "gs://", "abfss://", "azure://")
+
+        for i, (cloud_uri, _) in enumerate(v):
+            if not any(cloud_uri.startswith(scheme) for scheme in valid_schemes):
+                raise ValueError(
+                    f"paths[{i}][0] (cloud_uri) must start with one of {valid_schemes}, "
+                    f"got '{cloud_uri}'"
+                )
+        return v
+
+
+class CloudDownloader(CallbackBase):
+    """Callback that downloads files from cloud storage before model files are downloaded.
+
+    This callback expects self.kwargs to contain a 'paths' field which should be
+    a list of tuples, where each tuple contains (cloud_uri, local_path) strings.
+
+    Supported cloud storage URIs: s3://, gs://, abfss://, azure://
+
+    Example:
+        ```
+        from ray.llm._internal.common.callbacks.cloud_downloader import CloudDownloader
+        from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+        config = LLMConfig(
+            ...
+            callback_config={
+                "callback_class": CloudDownloader,
+                "callback_kwargs": {
+                    "paths": [
+                        ("s3://bucket/path/to/file.txt", "/local/path/to/file.txt"),
+                        ("gs://bucket/path/to/file.txt", "/local/path/to/file.txt"),
+                    ]
+                }
+            }
+            ...
+        )
+        ```
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the CloudDownloader callback.
+
+        Args:
+            **kwargs: Keyword arguments passed to the callback as a dictionary.
+                Must contain a 'paths' field with a list of (cloud_uri, local_path) tuples.
+        """
+        super().__init__(**kwargs)
+
+        # Validate configuration using Pydantic
+        if "paths" not in self.kwargs:
+            raise ValueError("CloudDownloader requires 'paths' field in kwargs")
+
+        CloudDownloaderConfig.model_validate(self.kwargs)
+
+    def on_before_download_model_files_distributed(self) -> None:
+        """Download files from cloud storage to local paths before model files are downloaded."""
+        from ray.llm._internal.common.utils.cloud_utils import CloudFileSystem
+
+        paths = self.kwargs["paths"]
+        start_time = time.monotonic()
+        for cloud_uri, local_path in paths:
+            CloudFileSystem.download_files(path=local_path, bucket_uri=cloud_uri)
+        end_time = time.monotonic()
+        logger.info(
+            f"CloudDownloader: Files downloaded in {end_time - start_time} seconds"
+        )