chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
@@ -0,0 +1,29 @@
"""Cloud filesystem module for provider-specific implementations.
This module provides a unified interface for cloud storage operations across
different providers (S3, GCS, Azure) while allowing provider-specific optimizations.
"""
from ray.llm._internal.common.utils.cloud_filesystem.azure_filesystem import (
AzureFileSystem,
)
from ray.llm._internal.common.utils.cloud_filesystem.base import (
BaseCloudFileSystem,
)
from ray.llm._internal.common.utils.cloud_filesystem.gcs_filesystem import (
GCSFileSystem,
)
from ray.llm._internal.common.utils.cloud_filesystem.pyarrow_filesystem import (
PyArrowFileSystem,
)
from ray.llm._internal.common.utils.cloud_filesystem.s3_filesystem import (
S3FileSystem,
)
__all__ = [
"BaseCloudFileSystem",
"PyArrowFileSystem",
"GCSFileSystem",
"AzureFileSystem",
"S3FileSystem",
]
@@ -0,0 +1,82 @@
"""Azure-specific filesystem implementation.
This module provides an Azure-specific implementation that delegates to PyArrowFileSystem.
This maintains backward compatibility while allowing for future optimizations using
native Azure tools (azcopy, azure-storage-blob SDK).
"""
from typing import List, Optional, Union
from ray.llm._internal.common.utils.cloud_filesystem.base import BaseCloudFileSystem
from ray.llm._internal.common.utils.cloud_filesystem.pyarrow_filesystem import (
PyArrowFileSystem,
)
class AzureFileSystem(BaseCloudFileSystem):
"""Azure-specific implementation of cloud filesystem operations.
**Note**: This implementation currently delegates to PyArrowFileSystem to maintain
stability. Optimized implementation using azure-storage-blob SDK and azcopy
will be added in a future PR.
"""
@staticmethod
def get_file(
object_uri: str, decode_as_utf_8: bool = True
) -> Optional[Union[str, bytes]]:
"""Download a file from cloud storage into memory.
Args:
object_uri: URI of the file (abfss:// or azure://)
decode_as_utf_8: If True, decode the file as UTF-8
Returns:
File contents as string or bytes, or None if file doesn't exist
"""
return PyArrowFileSystem.get_file(object_uri, decode_as_utf_8)
@staticmethod
def list_subfolders(folder_uri: str) -> List[str]:
"""List the immediate subfolders in a cloud directory.
Args:
folder_uri: URI of the directory (abfss:// or azure://)
Returns:
List of subfolder names (without trailing slashes)
"""
return PyArrowFileSystem.list_subfolders(folder_uri)
@staticmethod
def download_files(
path: str,
bucket_uri: str,
substrings_to_include: Optional[List[str]] = None,
suffixes_to_exclude: Optional[List[str]] = None,
) -> None:
"""Download files from cloud storage to a local directory.
Args:
path: Local directory where files will be downloaded
bucket_uri: URI of cloud directory
substrings_to_include: Only include files containing these substrings
suffixes_to_exclude: Exclude certain files from download (e.g .safetensors)
"""
PyArrowFileSystem.download_files(
path, bucket_uri, substrings_to_include, suffixes_to_exclude
)
@staticmethod
def upload_files(
local_path: str,
bucket_uri: str,
) -> None:
"""Upload files to cloud storage.
Args:
local_path: The local path of the files to upload.
bucket_uri: The bucket uri to upload the files to, must start with
`abfss://` or `azure://`.
"""
PyArrowFileSystem.upload_files(local_path, bucket_uri)
@@ -0,0 +1,81 @@
"""Abstract base class for cloud filesystem implementations.
This module defines the interface that all cloud storage provider implementations
must follow, ensuring consistency across different providers while allowing
provider-specific optimizations.
"""
from abc import ABC, abstractmethod
from typing import List, Optional, Union
class BaseCloudFileSystem(ABC):
"""Abstract base class for cloud filesystem implementations.
This class defines the interface that all cloud storage provider implementations
must implement. Provider-specific classes (S3FileSystem, GCSFileSystem, etc.)
will inherit from this base class and provide optimized implementations for
their respective cloud storage platforms.
"""
@staticmethod
@abstractmethod
def get_file(
object_uri: str, decode_as_utf_8: bool = True
) -> Optional[Union[str, bytes]]:
"""Download a file from cloud storage into memory.
Args:
object_uri: URI of the file (s3://, gs://, abfss://, or azure://)
decode_as_utf_8: If True, decode the file as UTF-8
Returns:
File contents as string or bytes, or None if file doesn't exist
"""
pass
@staticmethod
@abstractmethod
def list_subfolders(folder_uri: str) -> List[str]:
"""List the immediate subfolders in a cloud directory.
Args:
folder_uri: URI of the directory (s3://, gs://, abfss://, or azure://)
Returns:
List of subfolder names (without trailing slashes)
"""
pass
@staticmethod
@abstractmethod
def download_files(
path: str,
bucket_uri: str,
substrings_to_include: Optional[List[str]] = None,
suffixes_to_exclude: Optional[List[str]] = None,
) -> None:
"""Download files from cloud storage to a local directory.
Args:
path: Local directory where files will be downloaded
bucket_uri: URI of cloud directory
substrings_to_include: Only include files containing these substrings
suffixes_to_exclude: Exclude certain files from download (e.g .safetensors)
"""
pass
@staticmethod
@abstractmethod
def upload_files(
local_path: str,
bucket_uri: str,
) -> None:
"""Upload files to cloud storage.
Args:
local_path: The local path of the files to upload.
bucket_uri: The bucket uri to upload the files to, must start with
`s3://`, `gs://`, `abfss://`, or `azure://`.
"""
pass
@@ -0,0 +1,81 @@
"""GCS-specific filesystem implementation.
This module provides a GCS-specific implementation.
This maintains backward compatibility while allowing for future optimizations using
native GCS tools (gsutil, google-cloud-storage SDK).
"""
from typing import List, Optional, Union
from ray.llm._internal.common.utils.cloud_filesystem.base import BaseCloudFileSystem
from ray.llm._internal.common.utils.cloud_filesystem.pyarrow_filesystem import (
PyArrowFileSystem,
)
class GCSFileSystem(BaseCloudFileSystem):
"""GCS-specific implementation of cloud filesystem operations.
**Note**: This implementation currently delegates to PyArrowFileSystem to maintain
stability. Optimized implementation using google-cloud-storage SDK and gsutil
will be added in a future PR.
"""
@staticmethod
def get_file(
object_uri: str, decode_as_utf_8: bool = True
) -> Optional[Union[str, bytes]]:
"""Download a file from cloud storage into memory.
Args:
object_uri: URI of the file (gs://)
decode_as_utf_8: If True, decode the file as UTF-8
Returns:
File contents as string or bytes, or None if file doesn't exist
"""
return PyArrowFileSystem.get_file(object_uri, decode_as_utf_8)
@staticmethod
def list_subfolders(folder_uri: str) -> List[str]:
"""List the immediate subfolders in a cloud directory.
Args:
folder_uri: URI of the directory (gs://)
Returns:
List of subfolder names (without trailing slashes)
"""
return PyArrowFileSystem.list_subfolders(folder_uri)
@staticmethod
def download_files(
path: str,
bucket_uri: str,
substrings_to_include: Optional[List[str]] = None,
suffixes_to_exclude: Optional[List[str]] = None,
) -> None:
"""Download files from cloud storage to a local directory.
Args:
path: Local directory where files will be downloaded
bucket_uri: URI of cloud directory
substrings_to_include: Only include files containing these substrings
suffixes_to_exclude: Exclude certain files from download (e.g .safetensors)
"""
PyArrowFileSystem.download_files(
path, bucket_uri, substrings_to_include, suffixes_to_exclude
)
@staticmethod
def upload_files(
local_path: str,
bucket_uri: str,
) -> None:
"""Upload files to cloud storage.
Args:
local_path: The local path of the files to upload.
bucket_uri: The bucket uri to upload the files to, must start with `gs://`.
"""
PyArrowFileSystem.upload_files(local_path, bucket_uri)
@@ -0,0 +1,389 @@
"""PyArrow-based filesystem implementation for cloud storage.
This module provides a PyArrow-based implementation of the cloud filesystem
interface, supporting S3, GCS, and Azure storage providers.
"""
import os
from concurrent.futures import ThreadPoolExecutor
from typing import List, Optional, Tuple, Union
from urllib.parse import urlparse
import pyarrow.fs as pa_fs
from ray.llm._internal.common.observability.logging import get_logger
from ray.llm._internal.common.utils.cloud_filesystem.base import BaseCloudFileSystem
logger = get_logger(__name__)
class PyArrowFileSystem(BaseCloudFileSystem):
"""PyArrow-based implementation of cloud filesystem operations.
This class provides a unified interface for cloud storage operations using
PyArrow's filesystem abstraction. It supports S3, GCS, and Azure storage
providers.
"""
@staticmethod
def get_fs_and_path(object_uri: str) -> Tuple[pa_fs.FileSystem, str]:
"""Get the appropriate filesystem and path from a URI.
Args:
object_uri: URI of the file (s3://, gs://, abfss://, or azure://)
If URI contains 'anonymous@', anonymous access is used.
Example: s3://anonymous@bucket/path
Returns:
Tuple of (filesystem, path)
"""
if object_uri.startswith("pyarrow-"):
object_uri = object_uri[8:]
anonymous = False
# Check for anonymous access pattern (only for S3/GCS)
# e.g. s3://anonymous@bucket/path
if "@" in object_uri and not (
object_uri.startswith("abfss://") or object_uri.startswith("azure://")
):
parts = object_uri.split("@", 1)
# Check if the first part ends with "anonymous"
if parts[0].endswith("anonymous"):
anonymous = True
# Remove the anonymous@ part, keeping the scheme
scheme = parts[0].split("://")[0]
object_uri = f"{scheme}://{parts[1]}"
if object_uri.startswith("s3://"):
endpoint = os.getenv("AWS_ENDPOINT_URL_S3", None)
virtual_hosted_style = os.getenv("AWS_S3_ADDRESSING_STYLE", None)
fs = pa_fs.S3FileSystem(
anonymous=anonymous,
endpoint_override=endpoint,
force_virtual_addressing=(virtual_hosted_style == "virtual"),
)
path = object_uri[5:] # Remove "s3://"
elif object_uri.startswith("gs://"):
fs = pa_fs.GcsFileSystem(anonymous=anonymous)
path = object_uri[5:] # Remove "gs://"
elif object_uri.startswith("abfss://"):
fs, path = PyArrowFileSystem._create_abfss_filesystem(object_uri)
elif object_uri.startswith("azure://"):
fs, path = PyArrowFileSystem._create_azure_filesystem(object_uri)
else:
raise ValueError(f"Unsupported URI scheme: {object_uri}")
return fs, path
@staticmethod
def _create_azure_filesystem(object_uri: str) -> Tuple[pa_fs.FileSystem, str]:
"""Create an Azure filesystem for Azure Blob Storage or ABFSS.
Args:
object_uri: Azure URI (azure://container@account.blob.core.windows.net/path or
abfss://container@account.dfs.core.windows.net/path)
Returns:
Tuple of (PyArrow FileSystem, path without scheme prefix)
Raises:
ImportError: If required dependencies are not installed.
ValueError: If the Azure URI format is invalid.
"""
try:
import adlfs
from azure.identity import DefaultAzureCredential
except ImportError:
raise ImportError(
"You must `pip install adlfs azure-identity` "
"to use Azure/ABFSS URIs. "
"Note that these must be preinstalled on all nodes in the Ray cluster."
)
# Parse and validate the Azure URI
parsed = urlparse(object_uri)
scheme = parsed.scheme.lower()
# Validate URI format: scheme://container@account.domain/path
if not parsed.netloc or "@" not in parsed.netloc:
raise ValueError(
f"Invalid {scheme.upper()} URI format - missing container@account: {object_uri}"
)
container_part, hostname_part = parsed.netloc.split("@", 1)
# Validate container name (must be non-empty)
if not container_part:
raise ValueError(
f"Invalid {scheme.upper()} URI format - empty container name: {object_uri}"
)
# Validate hostname format based on scheme
valid_hostname = False
if scheme == "abfss":
valid_hostname = hostname_part.endswith(".dfs.core.windows.net")
expected_domains = ".dfs.core.windows.net"
elif scheme == "azure":
valid_hostname = hostname_part.endswith(
".blob.core.windows.net"
) or hostname_part.endswith(".dfs.core.windows.net")
expected_domains = ".blob.core.windows.net or .dfs.core.windows.net"
if not hostname_part or not valid_hostname:
raise ValueError(
f"Invalid {scheme.upper()} URI format - invalid hostname (must end with {expected_domains}): {object_uri}"
)
# Extract and validate account name
azure_storage_account_name = hostname_part.split(".")[0]
if not azure_storage_account_name:
raise ValueError(
f"Invalid {scheme.upper()} URI format - empty account name: {object_uri}"
)
# Create the adlfs filesystem
adlfs_fs = adlfs.AzureBlobFileSystem(
account_name=azure_storage_account_name,
credential=DefaultAzureCredential(),
)
# Wrap with PyArrow's PyFileSystem for compatibility
fs = pa_fs.PyFileSystem(pa_fs.FSSpecHandler(adlfs_fs))
# Return the path without the scheme prefix
path = f"{container_part}{parsed.path}"
return fs, path
@staticmethod
def _create_abfss_filesystem(object_uri: str) -> Tuple[pa_fs.FileSystem, str]:
"""Create an ABFSS filesystem for Azure Data Lake Storage Gen2.
This is a wrapper around _create_azure_filesystem for backward compatibility.
Args:
object_uri: ABFSS URI (abfss://container@account.dfs.core.windows.net/path)
Returns:
Tuple of (PyArrow FileSystem, path without abfss:// prefix)
"""
return PyArrowFileSystem._create_azure_filesystem(object_uri)
@staticmethod
def _filter_files(
fs: pa_fs.FileSystem,
source_path: str,
destination_path: str,
substrings_to_include: Optional[List[str]] = None,
suffixes_to_exclude: Optional[List[str]] = None,
) -> List[Tuple[str, str]]:
"""Filter files from cloud storage based on inclusion and exclusion criteria.
Args:
fs: PyArrow filesystem instance
source_path: Source path in cloud storage
destination_path: Local destination path
substrings_to_include: Only include files containing these substrings
suffixes_to_exclude: Exclude files ending with these suffixes
Returns:
List of tuples containing (source_file_path, destination_file_path)
"""
file_selector = pa_fs.FileSelector(source_path, recursive=True)
file_infos = fs.get_file_info(file_selector)
path_pairs = []
for file_info in file_infos:
if file_info.type != pa_fs.FileType.File:
continue
rel_path = file_info.path[len(source_path) :].lstrip("/")
# Apply filters
if substrings_to_include:
if not any(
substring in rel_path for substring in substrings_to_include
):
continue
if suffixes_to_exclude:
if any(rel_path.endswith(suffix) for suffix in suffixes_to_exclude):
continue
path_pairs.append(
(file_info.path, os.path.join(destination_path, rel_path))
)
return path_pairs
@staticmethod
def get_file(
object_uri: str, decode_as_utf_8: bool = True
) -> Optional[Union[str, bytes]]:
"""Download a file from cloud storage into memory.
Args:
object_uri: URI of the file (s3://, gs://, abfss://, or azure://)
decode_as_utf_8: If True, decode the file as UTF-8
Returns:
File contents as string or bytes, or None if file doesn't exist
"""
try:
fs, path = PyArrowFileSystem.get_fs_and_path(object_uri)
# Check if file exists
if not fs.get_file_info(path).type == pa_fs.FileType.File:
logger.info(f"URI {object_uri} does not exist.")
return None
# Read file
with fs.open_input_file(path) as f:
body = f.read()
if decode_as_utf_8:
body = body.decode("utf-8")
return body
except Exception as e:
logger.warning(f"Error reading {object_uri}: {e}")
return None
@staticmethod
def list_subfolders(folder_uri: str) -> List[str]:
"""List the immediate subfolders in a cloud directory.
Args:
folder_uri: URI of the directory (s3://, gs://, abfss://, or azure://)
Returns:
List of subfolder names (without trailing slashes)
"""
# Ensure that the folder_uri has a trailing slash.
folder_uri = f"{folder_uri.rstrip('/')}/"
try:
fs, path = PyArrowFileSystem.get_fs_and_path(folder_uri)
# List directory contents
file_infos = fs.get_file_info(pa_fs.FileSelector(path, recursive=False))
# Filter for directories and extract subfolder names
subfolders = []
for file_info in file_infos:
if file_info.type == pa_fs.FileType.Directory:
# Extract just the subfolder name without the full path
subfolder = os.path.basename(file_info.path.rstrip("/"))
subfolders.append(subfolder)
return subfolders
except Exception as e:
logger.error(f"Error listing subfolders in {folder_uri}: {e}")
return []
@staticmethod
def download_files(
path: str,
bucket_uri: str,
substrings_to_include: Optional[List[str]] = None,
suffixes_to_exclude: Optional[List[str]] = None,
max_concurrency: int = 10,
chunk_size: int = 64 * 1024 * 1024,
) -> None:
"""Download files from cloud storage to a local directory.
Args:
path: Local directory where files will be downloaded
bucket_uri: URI of cloud directory
substrings_to_include: Only include files containing these substrings
suffixes_to_exclude: Exclude certain files from download
max_concurrency: Maximum number of concurrent files to download (default: 10)
chunk_size: Size of transfer chunks (default: 64MB)
"""
try:
fs, source_path = PyArrowFileSystem.get_fs_and_path(bucket_uri)
# Ensure destination exists
os.makedirs(path, exist_ok=True)
# If no filters, use direct copy_files
if not substrings_to_include and not suffixes_to_exclude:
pa_fs.copy_files(
source=source_path,
destination=path,
source_filesystem=fs,
destination_filesystem=pa_fs.LocalFileSystem(),
use_threads=True,
chunk_size=chunk_size,
)
return
# List and filter files
files_to_download = PyArrowFileSystem._filter_files(
fs, source_path, path, substrings_to_include, suffixes_to_exclude
)
if not files_to_download:
logger.info("Filters do not match any of the files, skipping download")
return
def download_single_file(file_paths):
source_file_path, dest_file_path = file_paths
# Create destination directory if needed
dest_dir = os.path.dirname(dest_file_path)
if dest_dir:
os.makedirs(dest_dir, exist_ok=True)
# Use PyArrow's copy_files for individual files,
pa_fs.copy_files(
source=source_file_path,
destination=dest_file_path,
source_filesystem=fs,
destination_filesystem=pa_fs.LocalFileSystem(),
use_threads=True,
chunk_size=chunk_size,
)
return dest_file_path
max_workers = min(max_concurrency, len(files_to_download))
with ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = [
executor.submit(download_single_file, file_paths)
for file_paths in files_to_download
]
for future in futures:
try:
future.result()
except Exception as e:
logger.error(f"Failed to download file: {e}")
raise
except Exception as e:
logger.exception(f"Error downloading files from {bucket_uri}: {e}")
raise
@staticmethod
def upload_files(
local_path: str,
bucket_uri: str,
) -> None:
"""Upload files to cloud storage.
Args:
local_path: The local path of the files to upload.
bucket_uri: The bucket uri to upload the files to, must start with
`s3://`, `gs://`, `abfss://`, or `azure://`.
"""
try:
fs, dest_path = PyArrowFileSystem.get_fs_and_path(bucket_uri)
pa_fs.copy_files(
source=local_path,
destination=dest_path,
source_filesystem=pa_fs.LocalFileSystem(),
destination_filesystem=fs,
)
except Exception as e:
logger.exception(f"Error uploading files to {bucket_uri}: {e}")
raise
@@ -0,0 +1,397 @@
"""S3-specific filesystem implementation using boto3.
This module provides an S3-specific implementation that uses boto3 (AWS SDK for Python)
for reliable and efficient S3 operations.
"""
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, List, Optional, Union
try:
import boto3
from botocore import UNSIGNED
from botocore.config import Config
except ImportError:
boto3 = None # type: ignore[assignment]
UNSIGNED = None # type: ignore[assignment]
Config = None # type: ignore[assignment]
from ray.llm._internal.common.observability.logging import get_logger
from ray.llm._internal.common.utils.cloud_filesystem.base import BaseCloudFileSystem
logger = get_logger(__name__)
def _check_boto3() -> None:
"""Raise a clear error if boto3/botocore are not installed."""
if boto3 is None:
raise ImportError(
"boto3 and botocore are required for S3 operations but are not installed. "
"Install them with: pip install boto3"
)
class S3FileSystem(BaseCloudFileSystem):
"""S3-specific implementation of cloud filesystem operations using boto3.
This implementation uses boto3 (AWS SDK for Python) for reliable and efficient
operations with S3 storage.
"""
@staticmethod
def _parse_s3_uri(uri: str) -> tuple[str, str, bool]:
"""Parse S3 URI into bucket and key.
Args:
uri: S3 URI (e.g., s3://bucket/path/to/object or s3://anonymous@bucket/path/to/object)
Returns:
Tuple of (bucket_name, key, is_anonymous)
Raises:
ValueError: If URI is not a valid S3 URI
"""
# Check if anonymous@ prefix exists
is_anonymous = False
if uri.startswith("s3://anonymous@"):
is_anonymous = True
uri = uri.replace("s3://anonymous@", "s3://", 1)
if not uri.startswith("s3://"):
raise ValueError(f"Invalid S3 URI: {uri}")
# Remove s3:// prefix and split into bucket and key
path = uri[5:] # Remove "s3://"
parts = path.split("/", 1)
bucket = parts[0]
key = parts[1] if len(parts) > 1 else ""
return bucket, key, is_anonymous
@staticmethod
def _get_s3_client(max_pool_connections: int = 50, anonymous: bool = False):
"""Create a new S3 client instance with connection pooling.
Args:
max_pool_connections: Maximum number of connections in the pool.
Should be >= max_workers for optimal performance.
anonymous: Whether to use anonymous access to S3.
Returns:
boto3 S3 client with connection pooling configured
"""
_check_boto3()
# Configure connection pooling for better concurrent performance
config = Config(
max_pool_connections=max_pool_connections,
# Retry configuration for transient failures
retries={
"max_attempts": 3,
"mode": "adaptive", # Adapts retry behavior based on error type
},
# TCP keepalive helps with long-running connections
tcp_keepalive=True,
signature_version=UNSIGNED if anonymous else None,
)
return boto3.client("s3", config=config)
@staticmethod
def get_file(
object_uri: str, decode_as_utf_8: bool = True
) -> Optional[Union[str, bytes]]:
"""Download a file from cloud storage into memory.
Args:
object_uri: URI of the file (s3://)
decode_as_utf_8: If True, decode the file as UTF-8
Returns:
File contents as string or bytes, or None if file doesn't exist
"""
_check_boto3()
try:
bucket, key, is_anonymous = S3FileSystem._parse_s3_uri(object_uri)
s3_client = S3FileSystem._get_s3_client(anonymous=is_anonymous)
# Download file directly into memory
response = s3_client.get_object(Bucket=bucket, Key=key)
body = response["Body"].read()
if decode_as_utf_8:
return body.decode("utf-8")
return body
except Exception as e:
logger.error(f"Error reading {object_uri}: {e}")
@staticmethod
def list_subfolders(folder_uri: str) -> List[str]:
"""List the immediate subfolders in a cloud directory.
Args:
folder_uri: URI of the directory (s3://)
Returns:
List of subfolder names (without trailing slashes)
"""
_check_boto3()
try:
bucket, prefix, is_anonymous = S3FileSystem._parse_s3_uri(folder_uri)
# Ensure that the prefix has a trailing slash
if prefix and not prefix.endswith("/"):
prefix = f"{prefix}/"
s3_client = S3FileSystem._get_s3_client(anonymous=is_anonymous)
# List objects with delimiter to get only immediate subfolders
response = s3_client.list_objects_v2(
Bucket=bucket, Prefix=prefix, Delimiter="/"
)
subfolders = []
# CommonPrefixes contains the subdirectories
for common_prefix in response.get("CommonPrefixes", []):
folder_path = common_prefix["Prefix"]
# Extract the folder name from the full prefix
# Remove the parent prefix and trailing slash
folder_name = folder_path[len(prefix) :].rstrip("/")
if folder_name:
subfolders.append(folder_name)
return subfolders
except Exception as e:
logger.error(f"Error listing subfolders in {folder_uri}: {e}")
return []
@staticmethod
def _calculate_optimal_workers(
num_files: int, total_size: int, default_max: int = 100, default_min: int = 10
) -> int:
"""Calculate optimal number of workers based on file characteristics.
Args:
num_files: Number of files to download
total_size: Total size of all files in bytes
default_max: Maximum workers to cap at
default_min: Minimum workers to use
Returns:
Optimal number of workers between default_min and default_max
"""
if num_files == 0:
return default_min
avg_file_size = total_size / num_files if total_size > 0 else 0
# Strategy: More workers for smaller files, fewer for larger files
if avg_file_size < 1024 * 1024: # < 1MB (small files)
# Use more workers for many small files
workers = min(num_files, default_max)
elif avg_file_size < 10 * 1024 * 1024: # 1-10MB (medium files)
# Use moderate workers
workers = min(num_files // 2, default_max // 2)
else: # > 10MB (large files)
# Use fewer workers since each download is bandwidth-intensive
workers = min(20, num_files)
# Ensure workers is between min and max
return max(default_min, min(workers, default_max))
@staticmethod
def _download_single_file(
s3_client: Any, bucket: str, key: str, local_file_path: str
) -> tuple[str, bool]:
"""Download a single file from S3.
Args:
s3_client: Shared boto3 S3 client
bucket: S3 bucket name
key: S3 object key
local_file_path: Local path where file will be saved
Returns:
Tuple of (key, success)
"""
try:
# Create parent directories if needed
os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
s3_client.download_file(bucket, key, local_file_path)
logger.debug(f"Downloaded {key} to {local_file_path}")
return key, True
except Exception as e:
logger.error(f"Failed to download {key}: {e}")
return key, False
@staticmethod
def download_files(
path: str,
bucket_uri: str,
substrings_to_include: Optional[List[str]] = None,
suffixes_to_exclude: Optional[List[str]] = None,
max_workers: Optional[int] = None,
) -> None:
"""Download files from cloud storage to a local directory concurrently.
Args:
path: Local directory where files will be downloaded
bucket_uri: URI of cloud directory
substrings_to_include: Only include files containing these substrings
suffixes_to_exclude: Exclude certain files from download (e.g .safetensors)
max_workers: Maximum number of concurrent downloads. If None, automatically
calculated based on file count and sizes (min: 10, max: 100)
"""
try:
bucket, prefix, is_anonymous = S3FileSystem._parse_s3_uri(bucket_uri)
# Ensure the destination directory exists
os.makedirs(path, exist_ok=True)
# Ensure prefix has trailing slash for directory listing
if prefix and not prefix.endswith("/"):
prefix = f"{prefix}/"
# Create initial client for listing (will recreate with proper pool size later)
s3_client = S3FileSystem._get_s3_client(anonymous=is_anonymous)
# List all objects in the bucket with the given prefix
paginator = s3_client.get_paginator("list_objects_v2")
pages = paginator.paginate(Bucket=bucket, Prefix=prefix)
# Collect all files to download and track total size
files_to_download = []
total_size = 0
for page in pages:
for obj in page.get("Contents", []):
key = obj["Key"]
size = obj.get("Size", 0)
# Skip if it's a directory marker
if key.endswith("/"):
continue
# Get the relative path (remove the prefix)
relative_path = key[len(prefix) :]
# Apply include filters
if substrings_to_include:
if not any(
substr in relative_path for substr in substrings_to_include
):
continue
# Apply exclude filters
if suffixes_to_exclude:
if any(
relative_path.endswith(suffix.lstrip("*"))
for suffix in suffixes_to_exclude
):
continue
# Construct local file path
local_file_path = os.path.join(path, relative_path)
files_to_download.append((bucket, key, local_file_path))
total_size += size
# Download files concurrently
if not files_to_download:
logger.info(f"No files matching filters to download from {bucket_uri}")
return
# Dynamically calculate workers if not provided
if max_workers is None:
max_workers = S3FileSystem._calculate_optimal_workers(
num_files=len(files_to_download),
total_size=total_size,
default_max=100,
default_min=10,
)
# Create shared client with proper connection pool size for downloads
s3_client = S3FileSystem._get_s3_client(
max_pool_connections=max_workers + 10, anonymous=is_anonymous
)
failed_downloads = []
with ThreadPoolExecutor(max_workers=max_workers) as executor:
# Submit all download tasks with shared client
future_to_key = {
executor.submit(
S3FileSystem._download_single_file,
s3_client, # Pass shared client to each worker
bucket,
key,
local_path,
): key
for bucket, key, local_path in files_to_download
}
# Process completed downloads
for future in as_completed(future_to_key):
key, success = future.result()
if not success:
failed_downloads.append(key)
# Report any failures
if failed_downloads:
logger.error(
f"Failed to download {len(failed_downloads)} files: {failed_downloads[:5]}..."
)
except Exception as e:
logger.exception(f"Error downloading files from {bucket_uri}: {e}")
raise
@staticmethod
def upload_files(
local_path: str,
bucket_uri: str,
) -> None:
"""Upload files to cloud storage.
Args:
local_path: The local path of the files to upload.
bucket_uri: The bucket uri to upload the files to, must start with `s3://`.
"""
try:
bucket, prefix, is_anonymous = S3FileSystem._parse_s3_uri(bucket_uri)
# Ensure prefix has trailing slash for directory upload
if prefix and not prefix.endswith("/"):
prefix = f"{prefix}/"
s3_client = S3FileSystem._get_s3_client(anonymous=is_anonymous)
local_path_obj = Path(local_path)
# Walk through the local directory and upload each file
if local_path_obj.is_file():
# Upload a single file
file_name = local_path_obj.name
s3_key = f"{prefix}{file_name}" if prefix else file_name
s3_client.upload_file(str(local_path_obj), bucket, s3_key)
logger.debug(f"Uploaded {local_path_obj} to s3://{bucket}/{s3_key}")
elif local_path_obj.is_dir():
# Upload directory recursively
for file_path in local_path_obj.rglob("*"):
if file_path.is_file():
# Calculate relative path from local_path
relative_path = file_path.relative_to(local_path_obj)
# Construct S3 key
s3_key = f"{prefix}{relative_path.as_posix()}"
# Upload file
s3_client.upload_file(str(file_path), bucket, s3_key)
logger.debug(f"Uploaded {file_path} to s3://{bucket}/{s3_key}")
else:
raise ValueError(
f"Path {local_path} does not exist or is not a file/directory"
)
except Exception as e:
logger.exception(f"Error uploading files to {bucket_uri}: {e}")
raise
@@ -0,0 +1,546 @@
import asyncio
import inspect
import os
import time
from pathlib import Path
from typing import (
Any,
Awaitable,
Callable,
Dict,
List,
NamedTuple,
Optional,
TypeVar,
Union,
)
from pydantic import Field, field_validator
from ray.llm._internal.common.base_pydantic import BaseModelExtended
from ray.llm._internal.common.observability.logging import get_logger
from ray.llm._internal.common.utils.cloud_filesystem import (
AzureFileSystem,
GCSFileSystem,
PyArrowFileSystem,
S3FileSystem,
)
T = TypeVar("T")
logger = get_logger(__name__)
def is_remote_path(path: str) -> bool:
"""Check if the path is a remote path.
Args:
path: The path to check.
Returns:
True if the path is a remote path, False otherwise.
"""
return (
path.startswith("s3://")
or path.startswith("gs://")
or path.startswith("abfss://")
or path.startswith("azure://")
or path.startswith("pyarrow-")
)
class ExtraFiles(BaseModelExtended):
bucket_uri: str
destination_path: str
class CloudMirrorConfig(BaseModelExtended):
"""Unified mirror config for cloud storage (S3, GCS, or Azure).
Args:
bucket_uri: URI of the bucket (s3://, gs://, abfss://, or azure://)
extra_files: Additional files to download
"""
bucket_uri: Optional[str] = None
extra_files: List[ExtraFiles] = Field(default_factory=list)
@field_validator("bucket_uri")
@classmethod
def check_uri_format(cls, value):
if value is None:
return value
if not is_remote_path(value):
raise ValueError(
f'Got invalid value "{value}" for bucket_uri. '
'Expected a URI that starts with "s3://", "gs://", "abfss://", or "azure://".'
)
return value
@property
def storage_type(self) -> str:
"""Returns the storage type ('s3', 'gcs', 'abfss', or 'azure') based on the URI prefix."""
if self.bucket_uri is None:
return None
elif self.bucket_uri.startswith("s3://"):
return "s3"
elif self.bucket_uri.startswith("gs://"):
return "gcs"
elif self.bucket_uri.startswith("abfss://"):
return "abfss"
elif self.bucket_uri.startswith("azure://"):
return "azure"
return None
class LoraMirrorConfig(BaseModelExtended):
lora_model_id: str
bucket_uri: str
max_total_tokens: Optional[int]
sync_args: Optional[List[str]] = None
@field_validator("bucket_uri")
@classmethod
def check_uri_format(cls, value):
if value is None:
return value
if not is_remote_path(value):
raise ValueError(
f'Got invalid value "{value}" for bucket_uri. '
'Expected a URI that starts with "s3://", "gs://", "abfss://", or "azure://".'
)
return value
@property
def _bucket_name_and_path(self) -> str:
for prefix in ["s3://", "gs://", "abfss://", "azure://"]:
if self.bucket_uri.startswith(prefix):
return self.bucket_uri[len(prefix) :]
return self.bucket_uri
@property
def bucket_name(self) -> str:
bucket_part = self._bucket_name_and_path.split("/")[0]
# For ABFSS and Azure URIs, extract container name from container@account format
if self.bucket_uri.startswith(("abfss://", "azure://")) and "@" in bucket_part:
return bucket_part.split("@")[0]
return bucket_part
@property
def bucket_path(self) -> str:
return "/".join(self._bucket_name_and_path.split("/")[1:])
class CloudFileSystem:
"""A unified interface for cloud file system operations.
This class provides a simple interface for common operations on cloud storage
systems (S3, GCS, Azure) by delegating to provider-specific implementations
for optimal performance.
"""
@staticmethod
def _get_provider_fs(bucket_uri: str):
"""Get the appropriate provider-specific filesystem class based on URI.
Args:
bucket_uri: URI of the cloud storage (s3://, gs://, abfss://, or azure://)
Returns:
The appropriate filesystem class (S3FileSystem, GCSFileSystem, or AzureFileSystem)
Raises:
ValueError: If the URI scheme is not supported
"""
if bucket_uri.startswith("pyarrow-"):
return PyArrowFileSystem
elif bucket_uri.startswith("s3://"):
return S3FileSystem
elif bucket_uri.startswith("gs://"):
return GCSFileSystem
elif bucket_uri.startswith(("abfss://", "azure://")):
return AzureFileSystem
else:
raise ValueError(f"Unsupported URI scheme: {bucket_uri}")
@staticmethod
def get_file(
object_uri: str, decode_as_utf_8: bool = True
) -> Optional[Union[str, bytes]]:
"""Download a file from cloud storage into memory.
Args:
object_uri: URI of the file (s3://, gs://, abfss://, or azure://)
decode_as_utf_8: If True, decode the file as UTF-8
Returns:
File contents as string or bytes, or None if file doesn't exist
"""
fs_class = CloudFileSystem._get_provider_fs(object_uri)
return fs_class.get_file(object_uri, decode_as_utf_8)
@staticmethod
def list_subfolders(folder_uri: str) -> List[str]:
"""List the immediate subfolders in a cloud directory.
Args:
folder_uri: URI of the directory (s3://, gs://, abfss://, or azure://)
Returns:
List of subfolder names (without trailing slashes)
"""
fs_class = CloudFileSystem._get_provider_fs(folder_uri)
return fs_class.list_subfolders(folder_uri)
@staticmethod
def download_files(
path: str,
bucket_uri: str,
substrings_to_include: Optional[List[str]] = None,
suffixes_to_exclude: Optional[List[str]] = None,
) -> None:
"""Download files from cloud storage to a local directory.
Args:
path: Local directory where files will be downloaded
bucket_uri: URI of cloud directory
substrings_to_include: Only include files containing these substrings
suffixes_to_exclude: Exclude certain files from download (e.g .safetensors)
"""
fs_class = CloudFileSystem._get_provider_fs(bucket_uri)
fs_class.download_files(
path, bucket_uri, substrings_to_include, suffixes_to_exclude
)
@staticmethod
def download_model(
destination_path: str,
bucket_uri: str,
tokenizer_only: bool,
exclude_safetensors: bool = False,
) -> None:
"""Download a model from cloud storage.
This downloads a model in the format expected by the HuggingFace transformers
library.
Args:
destination_path: Path where the model will be stored
bucket_uri: URI of the cloud directory containing the model
tokenizer_only: If True, only download tokenizer-related files
exclude_safetensors: If True, skip download of safetensor files
"""
try:
# Get the provider-specific filesystem
fs_class = CloudFileSystem._get_provider_fs(bucket_uri)
# Construct hash file URI
hash_uri = bucket_uri.rstrip("/") + "/hash"
# Try to download and read hash file
hash_content = fs_class.get_file(hash_uri, decode_as_utf_8=True)
if hash_content is not None:
f_hash = hash_content.strip()
logger.info(
f"Detected hash file in bucket {bucket_uri}. "
f"Using {f_hash} as the hash."
)
else:
f_hash = "0000000000000000000000000000000000000000"
logger.info(
f"Hash file does not exist in bucket {bucket_uri}. "
f"Using default hash {f_hash} - expected behavior - a hash file is optional. "
)
# Write hash to refs/main
main_dir = os.path.join(destination_path, "refs")
os.makedirs(main_dir, exist_ok=True)
with open(os.path.join(main_dir, "main"), "w") as f:
f.write(f_hash)
# Create destination directory
destination_dir = os.path.join(destination_path, "snapshots", f_hash)
os.makedirs(destination_dir, exist_ok=True)
logger.info(f'Downloading model files to directory "{destination_dir}".')
# Download files
tokenizer_file_substrings = (
["tokenizer", "config.json", "chat_template"] if tokenizer_only else []
)
safetensors_to_exclude = [".safetensors"] if exclude_safetensors else None
CloudFileSystem.download_files(
path=destination_dir,
bucket_uri=bucket_uri,
substrings_to_include=tokenizer_file_substrings,
suffixes_to_exclude=safetensors_to_exclude,
)
except Exception as e:
logger.exception(f"Error downloading model from {bucket_uri}: {e}")
raise
@staticmethod
def upload_files(
local_path: str,
bucket_uri: str,
) -> None:
"""Upload files to cloud storage.
Args:
local_path: The local path of the files to upload.
bucket_uri: The bucket uri to upload the files to, must start with
`s3://`, `gs://`, `abfss://`, or `azure://`.
"""
fs_class = CloudFileSystem._get_provider_fs(bucket_uri)
fs_class.upload_files(local_path, bucket_uri)
@staticmethod
def upload_model(
local_path: str,
bucket_uri: str,
) -> None:
"""Upload a model to cloud storage.
Args:
local_path: The local path of the model.
bucket_uri: The bucket uri to upload the model to, must start with `s3://` or `gs://`.
"""
try:
# If refs/main exists, upload as hash, and treat snapshots/<hash> as the model.
# Otherwise, this is a custom model, we do not assume folder hierarchy.
refs_main = Path(local_path, "refs", "main")
if refs_main.exists():
model_path = os.path.join(
local_path, "snapshots", refs_main.read_text().strip()
)
CloudFileSystem.upload_files(
local_path=model_path, bucket_uri=bucket_uri
)
CloudFileSystem.upload_files(
local_path=str(refs_main),
bucket_uri=os.path.join(bucket_uri, "hash"),
)
else:
CloudFileSystem.upload_files(
local_path=local_path, bucket_uri=bucket_uri
)
logger.info(f"Uploaded model files to {bucket_uri}.")
except Exception as e:
logger.exception(f"Error uploading model to {bucket_uri}: {e}")
raise
class _CacheEntry(NamedTuple):
value: Any
expire_time: Optional[float]
class CloudObjectCache:
"""A cache that works with both sync and async fetch functions.
The purpose of this data structure is to cache the result of a function call
usually used to fetch a value from a cloud object store.
The idea is this:
- Cloud operations are expensive
- In LoRA specifically, we would fetch remote storage to download the model weights
at each request.
- If the same model is requested many times, we don't want to inflate the time to first token.
- We control the cache via not only the least recently used eviction policy, but also
by expiring cache entries after a certain time.
- If the object is missing, we cache the missing status for a small duration while if
the object exists, we cache the object for a longer duration.
"""
def __init__(
self,
max_size: int,
fetch_fn: Union[Callable[[str], Any], Callable[[str], Awaitable[Any]]],
missing_expire_seconds: Optional[int] = None,
exists_expire_seconds: Optional[int] = None,
missing_object_value: Any = object(),
):
"""Initialize the cache.
Args:
max_size: Maximum number of items to store in cache
fetch_fn: Function to fetch values (can be sync or async)
missing_expire_seconds: How long to cache missing objects (None for no expiration)
exists_expire_seconds: How long to cache existing objects (None for no expiration)
missing_object_value: Sentinel value used to represent a missing object in the cache.
"""
self._cache: Dict[str, _CacheEntry] = {}
self._max_size = max_size
self._fetch_fn = fetch_fn
self._missing_expire_seconds = missing_expire_seconds
self._exists_expire_seconds = exists_expire_seconds
self._is_async = inspect.iscoroutinefunction(fetch_fn) or (
callable(fetch_fn) and inspect.iscoroutinefunction(fetch_fn.__call__)
)
self._missing_object_value = missing_object_value
# Lock for thread-safe cache access
self._lock = asyncio.Lock()
async def aget(self, key: str) -> Any:
"""Async get value from cache or fetch it if needed."""
if not self._is_async:
raise ValueError("Cannot use async get() with sync fetch function")
async with self._lock:
value, should_fetch = self._check_cache(key)
if not should_fetch:
return value
# Fetch new value
value = await self._fetch_fn(key)
self._update_cache(key, value)
return value
def get(self, key: str) -> Any:
"""Sync get value from cache or fetch it if needed."""
if self._is_async:
raise ValueError("Cannot use sync get() with async fetch function")
# For sync access, we use a simple check-then-act pattern
# This is safe because sync functions are not used in async context
value, should_fetch = self._check_cache(key)
if not should_fetch:
return value
# Fetch new value
value = self._fetch_fn(key)
self._update_cache(key, value)
return value
def _check_cache(self, key: str) -> tuple[Any, bool]:
"""Check if key exists in cache and is valid.
Args:
key: The cache key to check.
Returns:
Tuple of (value, should_fetch)
where should_fetch is True if we need to fetch a new value
"""
now = time.monotonic()
if key in self._cache:
value, expire_time = self._cache[key]
if expire_time is None or now < expire_time:
return value, False
return None, True
def _update_cache(self, key: str, value: Any) -> None:
"""Update cache with new value."""
now = time.monotonic()
# Calculate expiration
expire_time = None
if (
self._missing_expire_seconds is not None
or self._exists_expire_seconds is not None
):
if value is self._missing_object_value:
expire_time = (
now + self._missing_expire_seconds
if self._missing_expire_seconds
else None
)
else:
expire_time = (
now + self._exists_expire_seconds
if self._exists_expire_seconds
else None
)
# Enforce size limit by removing oldest entry if needed
# This is an O(n) operation but it's fine since the cache size is usually small.
if len(self._cache) >= self._max_size:
oldest_key = min(
self._cache, key=lambda k: self._cache[k].expire_time or float("inf")
)
del self._cache[oldest_key]
self._cache[key] = _CacheEntry(value, expire_time)
def __len__(self) -> int:
return len(self._cache)
class CloudModelAccessor:
"""Unified accessor for models stored in cloud storage (S3 or GCS).
Args:
model_id: The model id to download or upload.
mirror_config: The mirror config for the model.
"""
def __init__(self, model_id: str, mirror_config: CloudMirrorConfig):
self.model_id = model_id
self.mirror_config = mirror_config
def _get_lock_path(self, suffix: str = "") -> Path:
return Path(
"~", f"{self.model_id.replace('/', '--')}{suffix}.lock"
).expanduser()
def _get_model_path(self) -> Path:
if Path(self.model_id).exists():
return Path(self.model_id)
# Delayed import to avoid circular dependencies
from huggingface_hub.constants import HF_HUB_CACHE
return Path(
HF_HUB_CACHE, f"models--{self.model_id.replace('/', '--')}"
).expanduser()
def remote_object_cache(
max_size: int,
missing_expire_seconds: Optional[int] = None,
exists_expire_seconds: Optional[int] = None,
missing_object_value: Any = None,
) -> Callable[[Callable[..., T]], Callable[..., T]]:
"""A decorator that provides async caching using CloudObjectCache.
This is a direct replacement for the remote_object_cache/cachetools combination,
using CloudObjectCache internally to maintain cache state.
Args:
max_size: Maximum number of items to store in cache
missing_expire_seconds: How long to cache missing objects
exists_expire_seconds: How long to cache existing objects
missing_object_value: Value to use for missing objects
Returns:
A decorator that wraps an async function with cache lookup.
"""
def decorator(func: Callable[..., T]) -> Callable[..., T]:
# Create a single cache instance for this function
cache = CloudObjectCache(
max_size=max_size,
fetch_fn=func,
missing_expire_seconds=missing_expire_seconds,
exists_expire_seconds=exists_expire_seconds,
missing_object_value=missing_object_value,
)
async def wrapper(*args, **kwargs):
# Extract the key from either first positional arg or object_uri kwarg
key = args[0] if args else kwargs.get("object_uri")
return await cache.aget(key)
return wrapper
return decorator
@@ -0,0 +1,324 @@
import enum
import os
from pathlib import Path
from typing import List, Optional
from filelock import FileLock
from ray.llm._internal.common.callbacks.base import CallbackBase
from ray.llm._internal.common.observability.logging import get_logger
from ray.llm._internal.common.utils.cloud_utils import (
CloudFileSystem,
CloudMirrorConfig,
CloudModelAccessor,
is_remote_path,
)
from ray.llm._internal.common.utils.import_utils import try_import
torch = try_import("torch")
logger = get_logger(__name__)
STREAMING_LOAD_FORMATS = ["runai_streamer", "runai_streamer_sharded", "tensorizer"]
class NodeModelDownloadable(enum.Enum):
"""Defines which files to download from cloud storage."""
MODEL_AND_TOKENIZER = enum.auto()
TOKENIZER_ONLY = enum.auto()
EXCLUDE_SAFETENSORS = enum.auto()
NONE = enum.auto()
def __bool__(self):
return self != NodeModelDownloadable.NONE
def union(self, other: "NodeModelDownloadable") -> "NodeModelDownloadable":
"""Return a NodeModelDownloadable that is a union of this and the other."""
if (
self == NodeModelDownloadable.MODEL_AND_TOKENIZER
or other == NodeModelDownloadable.MODEL_AND_TOKENIZER
):
return NodeModelDownloadable.MODEL_AND_TOKENIZER
if (
self == NodeModelDownloadable.EXCLUDE_SAFETENSORS
or other == NodeModelDownloadable.EXCLUDE_SAFETENSORS
):
return NodeModelDownloadable.EXCLUDE_SAFETENSORS
if (
self == NodeModelDownloadable.TOKENIZER_ONLY
or other == NodeModelDownloadable.TOKENIZER_ONLY
):
return NodeModelDownloadable.TOKENIZER_ONLY
return NodeModelDownloadable.NONE
def get_model_entrypoint(model_id: str) -> str:
"""Get the path to entrypoint of the model on disk if it exists, otherwise return the model id as is.
Entrypoint is typically <HF_HUB_CACHE>/models--<model_id>/
Args:
model_id: Hugging Face model ID.
Returns:
The path to the entrypoint of the model on disk if it exists, otherwise the model id as is.
"""
from huggingface_hub.constants import HF_HUB_CACHE
model_dir = Path(
HF_HUB_CACHE, f"models--{model_id.replace('/', '--')}"
).expanduser()
if not model_dir.exists():
return model_id
return str(model_dir.absolute())
def get_model_location_on_disk(model_id: str) -> str:
"""Get the location of the model on disk if exists, otherwise return the model id as is.
Args:
model_id: Hugging Face model ID.
Returns:
The path to the model on disk if it exists, otherwise the model id as is.
"""
model_dir = Path(get_model_entrypoint(model_id))
model_id_or_path = model_id
model_dir_refs_main = Path(model_dir, "refs", "main")
if model_dir.exists():
if model_dir_refs_main.exists():
# If refs/main exists, use the snapshot hash to find the model
# and check if *config.json (could be config.json for general models
# or adapter_config.json for LoRA adapters) exists to make sure it
# follows HF model repo structure.
with open(model_dir_refs_main, "r") as f:
snapshot_hash = f.read().strip()
snapshot_hash_path = Path(model_dir, "snapshots", snapshot_hash)
if snapshot_hash_path.exists() and list(
Path(snapshot_hash_path).glob("*config.json")
):
model_id_or_path = str(snapshot_hash_path.absolute())
else:
# If it doesn't have refs/main, it is a custom model repo
# and we can just return the model_dir.
model_id_or_path = str(model_dir.absolute())
return model_id_or_path
class CloudModelDownloader(CloudModelAccessor):
"""Unified downloader for models stored in cloud storage (S3 or GCS).
Args:
model_id: The model id to download.
mirror_config: The mirror config for the model.
"""
def get_model(
self,
tokenizer_only: bool,
exclude_safetensors: bool = False,
) -> str:
"""Gets a model from cloud storage and stores it locally.
Args:
tokenizer_only: whether to download only the tokenizer files.
exclude_safetensors: whether to download safetensors files to disk.
Returns:
File path of model if downloaded, else the model id.
"""
bucket_uri = self.mirror_config.bucket_uri
if bucket_uri is None:
return self.model_id
# Use different lock paths for different download types to avoid race conditions
# where a tokenizer-only download completes and subsequent full model downloads
# incorrectly assume the model weights are already cached.
if tokenizer_only:
lock_suffix = "-tokenizer"
elif exclude_safetensors:
lock_suffix = "-exclude-safetensors"
else:
lock_suffix = "-full"
lock_path = self._get_lock_path(suffix=lock_suffix)
path = self._get_model_path()
storage_type = self.mirror_config.storage_type
try:
# Timeout 0 means there will be only one attempt to acquire
# the file lock. If it cannot be acquired, a TimeoutError
# will be thrown.
# This ensures that subsequent processes don't duplicate work.
with FileLock(lock_path, timeout=0):
try:
if exclude_safetensors:
logger.info("Skipping download of safetensors files.")
CloudFileSystem.download_model(
destination_path=path,
bucket_uri=bucket_uri,
tokenizer_only=tokenizer_only,
exclude_safetensors=exclude_safetensors,
)
logger.info(
"Finished downloading %s for %s from %s storage",
"tokenizer" if tokenizer_only else "model and tokenizer",
self.model_id,
storage_type.upper() if storage_type else "cloud",
)
except RuntimeError:
logger.exception(
"Failed to download files for model %s from %s storage",
self.model_id,
storage_type.upper() if storage_type else "cloud",
)
except TimeoutError:
# If the directory is already locked, then wait but do not do anything.
with FileLock(lock_path, timeout=-1):
pass
return get_model_location_on_disk(self.model_id)
def get_extra_files(self) -> List[str]:
"""Gets user-specified extra files from cloud storage and stores them in
provided paths.
Returns: list of file paths of extra files if downloaded.
"""
paths = []
extra_files = self.mirror_config.extra_files or []
if not extra_files:
return paths
lock_path = self._get_lock_path(suffix="-extra_files")
storage_type = self.mirror_config.storage_type
logger.info(
f"Downloading extra files for {self.model_id} from {storage_type} storage"
)
try:
# Timeout 0 means there will be only one attempt to acquire
# the file lock. If it cannot be acquired, a TimeoutError
# will be thrown.
# This ensures that subsequent processes don't duplicate work.
with FileLock(lock_path, timeout=0):
for extra_file in extra_files:
path = Path(
os.path.expandvars(extra_file.destination_path)
).expanduser()
paths.append(path)
CloudFileSystem.download_files(
path=path,
bucket_uri=extra_file.bucket_uri,
)
except TimeoutError:
# If the directory is already locked, then wait but do not do anything.
with FileLock(lock_path, timeout=-1):
pass
return paths
def _log_download_info(
*, source: str, download_model: NodeModelDownloadable, download_extra_files: bool
):
if download_model == NodeModelDownloadable.NONE:
if download_extra_files:
logger.info("Downloading extra files from %s", source)
else:
logger.info("Not downloading anything from %s", source)
elif download_model == NodeModelDownloadable.TOKENIZER_ONLY:
if download_extra_files:
logger.info("Downloading tokenizer and extra files from %s", source)
else:
logger.info("Downloading tokenizer from %s", source)
elif download_model == NodeModelDownloadable.MODEL_AND_TOKENIZER:
if download_extra_files:
logger.info("Downloading model, tokenizer, and extra files from %s", source)
else:
logger.info("Downloading model and tokenizer from %s", source)
def download_model_files(
model_id: Optional[str] = None,
mirror_config: Optional[CloudMirrorConfig] = None,
download_model: NodeModelDownloadable = NodeModelDownloadable.MODEL_AND_TOKENIZER,
download_extra_files: bool = True,
callback: Optional[CallbackBase] = None,
) -> Optional[str]:
"""
Download the model files from the cloud storage. We support two ways to specify
the remote model path in the cloud storage:
Approach 1:
- model_id: The vanilla model id such as "meta-llama/Llama-3.1-8B-Instruct".
- mirror_config: Config for downloading model from cloud storage.
Approach 2:
- model_id: The remote path (s3:// or gs://) in the cloud storage.
- mirror_config: None.
In this approach, we will create a CloudMirrorConfig from the model_id and use that
to download the model.
Args:
model_id: The model id.
mirror_config: Config for downloading model from cloud storage.
download_model: What parts of the model to download.
download_extra_files: Whether to download extra files specified in the mirror config.
callback: Callback to run before downloading model files.
Returns:
The local path to the downloaded model, or the original model ID
if no cloud storage mirror is configured or if the model is not downloaded.
"""
# Create the torch cache kernels directory if it doesn't exist.
# This is a workaround for a torch issue, where the kernels directory
# cannot be created by torch if the parent directory doesn't exist.
torch_cache_home = torch.hub._get_torch_home()
os.makedirs(os.path.join(torch_cache_home, "kernels"), exist_ok=True)
model_path_or_id = model_id
if callback is not None:
callback.run_callback_sync("on_before_download_model_files_distributed")
if model_id is None:
return None
if mirror_config is None:
if is_remote_path(model_id):
logger.info(
"Creating a CloudMirrorConfig from remote model path %s", model_id
)
mirror_config = CloudMirrorConfig(bucket_uri=model_id)
else:
logger.info("No cloud storage mirror configured")
return model_id
storage_type = mirror_config.storage_type
source = (
f"{storage_type.upper()} mirror" if storage_type else "Cloud storage mirror"
)
_log_download_info(
source=source,
download_model=download_model,
download_extra_files=download_extra_files,
)
downloader = CloudModelDownloader(model_id, mirror_config)
if download_model != NodeModelDownloadable.NONE:
model_path_or_id = downloader.get_model(
tokenizer_only=download_model == NodeModelDownloadable.TOKENIZER_ONLY,
exclude_safetensors=download_model
== NodeModelDownloadable.EXCLUDE_SAFETENSORS,
)
if download_extra_files:
downloader.get_extra_files()
return model_path_or_id
@@ -0,0 +1,90 @@
"""Utility functions for importing modules in the LLM module."""
import importlib
import logging
from types import ModuleType
from typing import Any, NoReturn, Optional, Type
logger = logging.getLogger(__name__)
def try_import(name: str, error: bool = False) -> Optional[ModuleType]:
"""Try importing the module and returns the module (or None).
Args:
name: The name of the module to import.
error: Whether to raise an error if the module cannot be imported.
Returns:
The module, or None if it cannot be imported.
Raises:
ImportError: If error=True and the module is not installed.
"""
try:
return importlib.import_module(name)
except ImportError:
if error:
raise ImportError(f"Could not import {name}")
else:
logger.warning("Could not import %s", name)
return None
def raise_llm_engine_import_error(
vllm_error: ImportError,
sglang_error: ImportError,
) -> NoReturn:
"""Raise a descriptive ImportError when both vLLM and SGLang fail to import.
Distinguishes between a package not being installed (ModuleNotFoundError
whose .name matches the top-level package) and a broken installation
(any other ImportError, e.g. a missing .so or a missing transitive dep).
Args:
vllm_error: The ImportError raised when importing vLLM.
sglang_error: The ImportError raised when importing SGLang.
"""
vllm_not_installed = (
isinstance(vllm_error, ModuleNotFoundError) and vllm_error.name == "vllm"
)
sglang_not_installed = (
isinstance(sglang_error, ModuleNotFoundError) and sglang_error.name == "sglang"
)
if vllm_not_installed and sglang_not_installed:
raise ImportError(
"Neither vLLM nor SGLang is installed. At least one is required "
"for Ray Serve LLM protocol models. Install with: "
"`pip install ray[llm]` or `pip install sglang[all]`"
)
messages = []
if not vllm_not_installed:
messages.append(
"vLLM is installed but failed to import. This may indicate a "
"CUDA version mismatch or a missing vLLM dependency. "
f"Original error: {vllm_error}"
)
if not sglang_not_installed:
messages.append(
"SGLang is installed but failed to import. This may indicate a "
"missing SGLang dependency. "
f"Original error: {sglang_error}"
)
# Chain to the error that is actually relevant: vLLM's if it is broken,
# otherwise sglang's (i.e. vLLM was simply not installed).
cause = vllm_error if not vllm_not_installed else sglang_error
raise ImportError("\n".join(messages)) from cause
def load_class(path: str) -> Type[Any]:
"""Load class from string path."""
if ":" in path:
module_path, class_name = path.rsplit(":", 1)
else:
module_path, class_name = path.rsplit(".", 1)
module = try_import(module_path, error=True)
callback_class = getattr(module, class_name)
return callback_class
@@ -0,0 +1,233 @@
"""
Generic LoRA utilities and abstractions.
This module provides canonical LoRA utility functions for both serve and batch components.
It serves as the single source of truth for LoRA operations and builds on the generic
download primitives from download_utils.py.
"""
import json
import os
import subprocess
import time
from functools import wraps
from typing import Any, Callable, List, Optional, TypeVar, Union
from ray.llm._internal.common.constants import (
CLOUD_OBJECT_EXISTS_EXPIRE_S,
CLOUD_OBJECT_MISSING_EXPIRE_S,
LORA_ADAPTER_CONFIG_NAME,
)
# Import the global ID manager from common models
from ray.llm._internal.common.models import make_async
from ray.llm._internal.common.observability.logging import get_logger
from ray.llm._internal.common.utils.cloud_utils import (
CloudFileSystem,
is_remote_path,
remote_object_cache,
)
from ray.llm._internal.common.utils.download_utils import (
CloudMirrorConfig,
CloudModelDownloader,
)
logger = get_logger(__name__)
# Sentinel object for missing cloud objects
CLOUD_OBJECT_MISSING = object()
DEFAULT_LORA_MAX_TOTAL_TOKENS = 4096
T = TypeVar("T")
def get_base_model_id(model_id: str) -> str:
"""Get base model id for a given model id."""
return model_id.split(":")[0]
def get_lora_id(lora_model_id: str) -> str:
"""Get lora id for a given lora model id."""
return ":".join(lora_model_id.split(":")[1:])
def clean_model_id(model_id: str) -> str:
"""Clean model ID for filesystem usage by replacing slashes with dashes."""
return model_id.replace("/", "--")
def clear_directory(dir: str) -> None:
"""Clear a directory recursively, ignoring missing directories."""
try:
subprocess.run(f"rm -r {dir}", shell=True, check=False)
except FileNotFoundError:
pass
def retry_with_exponential_backoff(
max_tries: int,
exception_to_check: type[Exception],
base_delay: float = 1,
max_delay: float = 32,
exponential_base: float = 2,
) -> Callable[[Callable[..., T]], Callable[..., T]]:
"""Retry decorator with exponential backoff."""
def decorator(func: Callable[..., T]) -> Callable[..., T]:
@wraps(func)
def wrapper(*args: Any, **kwargs: Any) -> T:
delay = base_delay
last_exception = None
for attempt in range(max_tries):
try:
return func(*args, **kwargs)
except exception_to_check as e:
last_exception = e
if attempt == max_tries - 1: # Last attempt
raise last_exception
# Log the failure and retry
logger.warning(
f"Attempt {attempt + 1}/{max_tries} failed: {str(e)}. "
f"Retrying in {delay} seconds..."
)
time.sleep(delay)
# Calculate next delay with exponential backoff
delay = min(delay * exponential_base, max_delay)
# This should never be reached due to the raise in the loop
raise last_exception if last_exception else RuntimeError(
"Unexpected error in retry logic"
)
return wrapper
return decorator
def sync_files_with_lock(
bucket_uri: str,
local_path: str,
timeout: Optional[float] = None,
substrings_to_include: Optional[List[str]] = None,
) -> None:
"""Sync files from bucket_uri to local_path with file locking."""
from filelock import FileLock
logger.info("Downloading %s to %s", bucket_uri, local_path)
with FileLock(local_path + ".lock", timeout=timeout or -1):
try:
CloudFileSystem.download_files(
path=local_path,
bucket_uri=bucket_uri,
substrings_to_include=substrings_to_include,
)
except Exception as e:
logger.error(
"Failed to sync files from %s to %s: %s",
bucket_uri,
local_path,
str(e),
)
raise
@make_async
def _get_object_from_cloud(object_uri: str) -> Union[str, object]:
"""Gets an object from the cloud."""
if object_uri.endswith("/"):
raise ValueError(f'object_uri {object_uri} must not end with a "/".')
body_str = CloudFileSystem.get_file(object_uri)
if body_str is None:
logger.info(f"{object_uri} does not exist.")
return CLOUD_OBJECT_MISSING
else:
return body_str
@remote_object_cache(
max_size=4096,
missing_expire_seconds=CLOUD_OBJECT_MISSING_EXPIRE_S,
exists_expire_seconds=CLOUD_OBJECT_EXISTS_EXPIRE_S,
missing_object_value=CLOUD_OBJECT_MISSING,
)
async def get_object_from_cloud(object_uri: str) -> Union[str, object]:
"""Gets an object from the cloud with caching."""
return await _get_object_from_cloud(object_uri)
async def get_lora_finetuned_context_length(bucket_uri: str) -> Optional[int]:
"""Gets the sequence length used to tune the LoRA adapter."""
if bucket_uri.endswith("/"):
bucket_uri = bucket_uri.rstrip("/")
object_uri = f"{bucket_uri}/{LORA_ADAPTER_CONFIG_NAME}"
object_str_or_missing_message = await get_object_from_cloud(object_uri)
if object_str_or_missing_message is CLOUD_OBJECT_MISSING:
logger.debug(f"LoRA adapter config file not found at {object_uri}")
return None
try:
adapter_config_str = object_str_or_missing_message
adapter_config = json.loads(adapter_config_str)
return adapter_config.get("max_length")
except (json.JSONDecodeError, AttributeError) as e:
logger.warning(f"Failed to parse LoRA adapter config at {object_uri}: {e}")
return None
def get_lora_model_ids(
dynamic_lora_loading_path: str,
base_model_id: str,
) -> List[str]:
"""Get the model IDs of all the LoRA models.
The dynamic_lora_loading_path is expected to hold subfolders each for
a different lora checkpoint. Each subfolder name will correspond to
the unique identifier for the lora checkpoint. The lora model is
accessible via <base_model_id>:<lora_id>. Therefore, we prepend
the base_model_id to each subfolder name.
Args:
dynamic_lora_loading_path: the cloud folder that contains all the LoRA
weights.
base_model_id: model ID of the base model.
Returns:
List of LoRA fine-tuned model IDs. Does not include the base model
itself.
"""
lora_subfolders = CloudFileSystem.list_subfolders(dynamic_lora_loading_path)
lora_model_ids = []
for subfolder in lora_subfolders:
lora_model_ids.append(f"{base_model_id}:{subfolder}")
return lora_model_ids
def download_lora_adapter(
lora_name: str,
remote_path: Optional[str] = None,
) -> str:
"""Download a LoRA adapter from remote storage.
This maintains backward compatibility with existing code.
"""
assert not is_remote_path(
lora_name
), "lora_name cannot be a remote path (s3:// or gs://)"
if remote_path is None:
return lora_name
lora_path = os.path.join(remote_path, lora_name)
mirror_config = CloudMirrorConfig(bucket_uri=lora_path)
downloader = CloudModelDownloader(lora_name, mirror_config)
return downloader.get_model(tokenizer_only=False)
@@ -0,0 +1,126 @@
from pathlib import Path
import typer
from filelock import FileLock
from typing_extensions import Annotated
from ray.llm._internal.common.observability.logging import get_logger
from ray.llm._internal.common.utils.cloud_utils import (
CloudFileSystem,
CloudMirrorConfig,
CloudModelAccessor,
is_remote_path,
)
from ray.llm._internal.common.utils.download_utils import (
get_model_entrypoint,
)
logger = get_logger(__name__)
class CloudModelUploader(CloudModelAccessor):
"""Unified uploader to upload models to cloud storage (S3 or GCS).
Args:
model_id: The model id to upload.
mirror_config: The mirror config for the model.
"""
def upload_model(self) -> str:
"""Upload the model to cloud storage (s3 or gcs).
Returns:
The remote path of the uploaded model.
"""
bucket_uri = self.mirror_config.bucket_uri
lock_path = self._get_lock_path()
path = self._get_model_path()
storage_type = self.mirror_config.storage_type
try:
# Timeout 0 means there will be only one attempt to acquire
# the file lock. If it cannot be acquired, a TimeoutError
# will be thrown.
# This ensures that subsequent processes don't duplicate work.
with FileLock(lock_path, timeout=0):
try:
CloudFileSystem.upload_model(
local_path=path,
bucket_uri=bucket_uri,
)
logger.info(
"Finished uploading %s to %s storage",
self.model_id,
storage_type.upper() if storage_type else "cloud",
)
except RuntimeError:
logger.exception(
"Failed to upload model %s to %s storage",
self.model_id,
storage_type.upper() if storage_type else "cloud",
)
except TimeoutError:
# If the directory is already locked, then wait but do not do anything.
with FileLock(lock_path, timeout=-1):
pass
return bucket_uri
def upload_model_files(model_id: str, bucket_uri: str) -> str:
"""Upload the model files to cloud storage (s3 or gcs).
If `model_id` is a local path, the files will be uploaded to the cloud storage.
If `model_id` is a huggingface model id, the model will be downloaded from huggingface
and then uploaded to the cloud storage.
Args:
model_id: The huggingface model id, or local model path to upload.
bucket_uri: The bucket uri to upload the model to, must start with `s3://` or `gs://`.
Returns:
The remote path of the uploaded model.
"""
assert not is_remote_path(
model_id
), f"model_id must NOT be a remote path: {model_id}"
assert is_remote_path(bucket_uri), f"bucket_uri must be a remote path: {bucket_uri}"
if not Path(model_id).exists():
maybe_downloaded_model_path = get_model_entrypoint(model_id)
if not Path(maybe_downloaded_model_path).exists():
logger.info(
"Assuming %s is huggingface model id, and downloading it.", model_id
)
import huggingface_hub
huggingface_hub.snapshot_download(repo_id=model_id)
# Try to get the model path again after downloading.
maybe_downloaded_model_path = get_model_entrypoint(model_id)
assert Path(
maybe_downloaded_model_path
).exists(), f"Failed to download the model {model_id} to {maybe_downloaded_model_path}"
return upload_model_files(maybe_downloaded_model_path, bucket_uri)
else:
return upload_model_files(maybe_downloaded_model_path, bucket_uri)
uploader = CloudModelUploader(model_id, CloudMirrorConfig(bucket_uri=bucket_uri))
return uploader.upload_model()
def upload_model_cli(
model_source: Annotated[
str,
typer.Option(
help="HuggingFace model ID to download, or local model path to upload",
),
],
bucket_uri: Annotated[
str,
typer.Option(
help="The bucket uri to upload the model to, must start with `s3://` or `gs://`",
),
],
):
"""Upload the model files to cloud storage (s3 or gcs)."""
upload_model_files(model_source, bucket_uri)