megatron/core/datasets/object_storage_utils.py

# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
import os
from dataclasses import dataclass
from typing import Any, Dict, Protocol, Tuple

import torch

try:
    import boto3
    import botocore.exceptions as exceptions
except ModuleNotFoundError:
    pass

from megatron.core.msc_utils import MultiStorageClientFeature

S3_PREFIX = "s3://"
MSC_PREFIX = "msc://"


@dataclass
class ObjectStorageConfig:
    """Config when the data (.bin) file and the index (.idx) file are in object storage

    Attributes:

        path_to_idx_cache (str): The local directory where we will store the index (.idx) file

        bin_chunk_nbytes (int): If the number of bytes is too small, then we send a request to S3
        at each call of the `read` method in _S3BinReader, which is slow, because each request
        has a fixed cost independent of the size of the byte range requested. If the number of
        bytes is too large, then we only rarely have to send requests to S3, but it takes a lot
        of time to complete the request when we do, which can block training. We've found that
        256 * 1024 * 1024 (i.e., 256 MiB) has worked well (though we have not put that much
        effort into tuning it), so we default to it.
    """

    path_to_idx_cache: str

    bin_chunk_nbytes: int = 256 * 1024 * 1024


# S3Config is deprecated, use ObjectStorageConfig instead
S3Config = ObjectStorageConfig


class S3Client(Protocol):
    """The protocol which all s3 clients should abide by"""

    def download_file(self, Bucket: str, Key: str, Filename: str) -> None:
        """Download the file from S3 to the local file system"""
        ...

    def upload_file(self, Filename: str, Bucket: str, Key: str) -> None:
        """Upload the file to S3"""
        ...

    def head_object(self, Bucket: str, Key: str) -> Dict[str, Any]:
        """Get the metadata of the file in S3"""
        ...

    def get_object(self, Bucket: str, Key: str, Range: str) -> Dict[str, Any]:
        """Get the file from S3"""
        ...

    def close(self) -> None:
        """Close the S3 client"""
        ...


def _remove_s3_prefix(path: str) -> str:
    """Remove the S3 prefix from a path

    Args:
        path (str): The path

    Returns:
        str: The path without the S3 prefix
    """
    return path.removeprefix(S3_PREFIX)


def _is_s3_path(path: str) -> bool:
    """Ascertain whether a path is in S3

    Args:
        path (str): The path

    Returns:
        bool: True if the path is in S3, False otherwise
    """
    return path.startswith(S3_PREFIX)


def _remove_msc_prefix(path: str) -> str:
    """
    Remove the MSC prefix from a path

    Args:
        path (str): The path

    Returns:
        str: The path without the MSC prefix
    """
    return path.removeprefix(MSC_PREFIX)


def _is_msc_path(path: str) -> bool:
    """Checks whether a path is in MSC path (msc://profile/path/to/file)

    Args:
        path (str): The path

    Returns:
        bool: True if the path is in MSC path, False otherwise
    """
    return path.startswith(MSC_PREFIX)


def _s3_download_file(client: S3Client, s3_path: str, local_path: str) -> None:
    """Download the object at the given S3 path to the given local file system path

    Args:
        client (S3Client): The S3 client

        s3_path (str): The S3 source path

        local_path (str): The local destination path
    """
    dirname = os.path.dirname(local_path)
    os.makedirs(dirname, exist_ok=True)
    parsed_s3_path = parse_s3_path(s3_path)
    client.download_file(parsed_s3_path[0], parsed_s3_path[1], local_path)


def _s3_object_exists(client: S3Client, path: str) -> bool:
    """Ascertain whether the object at the given S3 path exists in S3

    Args:
        client (S3Client): The S3 client

        path (str): The S3 path

    Raises:
        botocore.exceptions.ClientError: The error code is 404

    Returns:
        bool: True if the object exists in S3, False otherwise
    """
    parsed_s3_path = parse_s3_path(path)
    try:
        _ = client.head_object(bucket=parsed_s3_path[0], key=parsed_s3_path[1])
    except exceptions.ClientError as e:
        if e.response["Error"]["Code"] != "404":
            raise e
    return True


def is_object_storage_path(path: str) -> bool:
    """Ascertain whether a path is in object storage

    Args:
        path (str): The path

    Returns:
        bool: True if the path is in object storage (s3:// or msc://), False otherwise
    """
    return _is_s3_path(path) or _is_msc_path(path)


def get_index_cache_path(idx_path: str, object_storage_config: ObjectStorageConfig) -> str:
    """Get the index cache path for the given path

    Args:
        idx_path (str): The path to the index file

        object_storage_config (ObjectStorageConfig): The object storage config

    Returns:
        str: The index cache path
    """
    if _is_s3_path(idx_path):
        cache_idx_path = os.path.join(
            object_storage_config.path_to_idx_cache, _remove_s3_prefix(idx_path)
        )
    elif _is_msc_path(idx_path):
        cache_idx_path = os.path.join(
            object_storage_config.path_to_idx_cache, _remove_msc_prefix(idx_path)
        )
    else:
        raise ValueError(f"Invalid path: {idx_path}")

    return cache_idx_path


def parse_s3_path(path: str) -> Tuple[str, str]:
    """Parses the given S3 path returning correspsonding bucket and key.

    Args:
        path (str): The S3 path

    Returns:
        Tuple[str, str]: A (bucket, key) tuple
    """
    assert _is_s3_path(path)
    parts = path.replace(S3_PREFIX, "").split("/")
    bucket = parts[0]
    if len(parts) > 1:
        key = "/".join(parts[1:])
        assert S3_PREFIX + bucket + "/" + key == path
    else:
        key = ""
    return bucket, key


def get_object_storage_access(path: str) -> str:
    """Get the object storage access"""
    return "s3" if _is_s3_path(path) else "msc"


def dataset_exists(path_prefix: str, idx_path: str, bin_path: str) -> bool:
    """Check if the dataset exists on object storage

    Args:
        path_prefix (str): The prefix to the index (.idx) and data (.bin) files

        idx_path (str): The path to the index file

        bin_path (str): The path to the data file

    Returns:
        bool: True if the dataset exists on object storage, False otherwise
    """
    if _is_s3_path(path_prefix):
        s3_client = boto3.client("s3")
        return _s3_object_exists(s3_client, idx_path) and _s3_object_exists(s3_client, bin_path)
    elif _is_msc_path(path_prefix):
        msc = MultiStorageClientFeature.import_package()
        return msc.exists(idx_path) and msc.exists(bin_path)
    else:
        raise ValueError(f"Invalid path: {path_prefix}")


def cache_index_file(remote_path: str, local_path: str) -> None:
    """Download a file from object storage to a local path with distributed training support.
    The download only happens on Rank 0, and other ranks will wait for the file to be available.

    Note that this function does not include any barrier synchronization. The caller (typically
    in blended_megatron_dataset_builder.py) is responsible for ensuring proper synchronization
    between ranks using torch.distributed.barrier() after this function returns.

    Args:
        remote_path (str): The URL of the file to download (e.g., s3://bucket/path/file.idx
            or msc://profile/path/file.idx)
        local_path (str): The local destination path where the file should be saved

    Raises:
        ValueError: If the remote_path is not a valid S3 or MSC path
    """
    torch_dist_enabled = torch.distributed.is_initialized()

    if torch_dist_enabled:
        rank = torch.distributed.get_rank()
    else:
        rank = 0

    if _is_s3_path(remote_path):
        s3_client = boto3.client("s3")

        if not torch_dist_enabled or rank == 0:
            _s3_download_file(s3_client, remote_path, local_path)

        assert os.path.exists(local_path)
    elif _is_msc_path(remote_path):
        msc = MultiStorageClientFeature.import_package()

        if not torch_dist_enabled or rank == 0:
            msc.download_file(remote_path, local_path)

        assert os.path.exists(local_path)
    else:
        raise ValueError(f"Invalid path: {remote_path}")