chore: import upstream snapshot with attribution
This commit is contained in:
@@ -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)
|
||||
Reference in New Issue
Block a user