chore: import upstream snapshot with attribution

python/ray/llm/_internal/common/utils/cloud_filesystem/__init__.py

@@ -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",
+]

python/ray/llm/_internal/common/utils/cloud_filesystem/azure_filesystem.py

@@ -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)

python/ray/llm/_internal/common/utils/cloud_filesystem/base.py

@@ -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

python/ray/llm/_internal/common/utils/cloud_filesystem/gcs_filesystem.py

@@ -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)

python/ray/llm/_internal/common/utils/cloud_filesystem/pyarrow_filesystem.py

@@ -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

python/ray/llm/_internal/common/utils/cloud_filesystem/s3_filesystem.py

@@ -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