wehub-resource-sync
87 lines
3.0 KiB
Python
87 lines
3.0 KiB
Python
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"
|
|
)
|