diff --git a/.github/workflows/pytest.yaml b/.github/workflows/pytest.yaml index 8cb994a89..7e91db0fa 100644 --- a/.github/workflows/pytest.yaml +++ b/.github/workflows/pytest.yaml @@ -47,13 +47,4 @@ jobs: id: tests run: | set -ex - pytest --splits 10 --group 1 --cov-fail-under=10 - pytest --splits 10 --group 2 --cov-fail-under=10 - pytest --splits 10 --group 3 --cov-fail-under=10 - pytest --splits 10 --group 4 --cov-fail-under=10 - pytest --splits 10 --group 5 --cov-fail-under=10 - pytest --splits 10 --group 6 --cov-fail-under=10 - pytest --splits 10 --group 7 --cov-fail-under=10 - pytest --splits 10 --group 8 --cov-fail-under=10 - pytest --splits 10 --group 9 --cov-fail-under=10 - pytest --splits 10 --group 10 --cov-fail-under=10 + pytest --cov-fail-under 50 diff --git a/docs/source/conf.py b/docs/source/conf.py index dbd0f1b83..47d197ca9 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -365,14 +365,13 @@ def _modules_to_rst() -> List[types.ModuleType]: document_modules: List[types.Module] = [ streaming, streaming.base.compression, + streaming.base.coord, streaming.base.format, streaming.base.hashing, streaming.base.partition, - streaming.base.shared, streaming.base.shuffle, streaming.base.storage, streaming.base.util, - streaming.base.world, ] exclude_modules: List[types.Module] = [streaming.base, streaming._version] for name in streaming.__dict__: diff --git a/setup.py b/setup.py index c70a84143..74430d7d7 100644 --- a/setup.py +++ b/setup.py @@ -58,6 +58,7 @@ 'azure-storage-blob>=12.0.0,<13', 'azure-storage-file-datalake>=12.11.0,<13', 'azure-identity>=1.13.0', + 'psutil>=5.9.4', ] extra_deps = {} diff --git a/simulation/core/sim_dataset.py b/simulation/core/sim_dataset.py index f9dc758a0..bf9de42eb 100644 --- a/simulation/core/sim_dataset.py +++ b/simulation/core/sim_dataset.py @@ -7,7 +7,6 @@ import os import shutil import time -import warnings from math import ceil from typing import Optional, Sequence, Union @@ -18,6 +17,7 @@ from streaming.base import Stream, StreamingDataset from streaming.base.batching import generate_work +from streaming.base.coord.world import World from streaming.base.format import get_index_basename from streaming.base.spanner import Spanner from streaming.base.util import bytes_to_int, number_abbrev_to_int @@ -33,30 +33,36 @@ class SimulationDataset(StreamingDataset): nodes (int): Number of nodes. devices (int): Number of devices. workers (int): Number of workers. - streams (Optional[Sequence[Stream]]): One or more streams to stream/cache samples from, + epoch_size (Union[int, str], optional): Number of samples to draw per epoch balanced + across all streams. If ``None``, takes its value from the total number of underlying + samples. Provide this field if you are weighting streams relatively to target a larger + or smaller epoch size. Defaults to ``None``. Can also take in human-readable number + abbreviations (e.g., ``"100k"``, ``"64M"``, ``"77b"``, etc). Defaults to ``None``. + streams (Sequence[Stream], optional): One or more streams to stream/cache samples from, which may be upsampled or downsampled. StreamingDataset uses either ``streams`` or ``remote``/``local``. Defaults to ``None``. - remote (Optional[str]): Remote path or directory to download the dataset from. If ``None``, + remote (str, optional): Remote path or directory to download the dataset from. If ``None``, its data must exist locally. StreamingDataset uses either ``streams`` or ``remote``/``local``. Defaults to ``None``. - local (Optional[str]): Local working directory to download shards to. This is where shards + local (str, optional): Local working directory to download shards to. This is where shards are cached while they are being used. Uses a temp directory if not set. StreamingDataset uses either ``streams`` or ``remote``/``local``. Defaults to ``None``. - split (Optional[str]): Which dataset split to use, if any. If provided, we stream from/to + split (str, optional): Which dataset split to use, if any. If provided, we stream from/to the ``split`` subdirs of ``remote`` and ``local``. Defaults to ``None``. download_retry (int): Number of download re-attempts before giving up. Defaults to ``2``. download_timeout (float): Number of seconds to wait for a shard to download before raising an exception. Defaults to ``60``. - validate_hash (Optional[str]): Optional hash or checksum algorithm to use to validate + validate_hash (str, optional): Optional hash or checksum algorithm to use to validate shards. Defaults to ``None``. keep_zip (bool): Whether to keep or delete the compressed form when decompressing downloaded shards. If ``False``, keep iff remote is local or no remote. Defaults to ``False``. - epoch_size (Union[int, str], optional): Number of samples to draw per epoch balanced across all - streams. If ``None``, takes its value from the total number of underlying samples. - Provide this field if you are weighting streams relatively to target a larger or - smaller epoch size. Defaults to ``None``. Can also take in human-readable number - abbreviations (e.g., ``"100k"``, ``"64M"``, ``"77b"``, and so on). Defaults to ``None``. + allow_unsafe_types (bool): If a shard contains Pickle, which allows arbitrary code + execution during deserialization, whether to keep going if ``True`` or raise an error + if ``False``. Defaults to ``False``. + config_root (str, optional): Streaming configuration root directory, used for collision + detection, filelock paths, etc. If ``None``, uses a ``/streaming/`` subdir under your + system's temp root. Defaults to ``None``. predownload (int, optional): Target number of samples to download per worker in advance of current sample. Workers will attempt to download ahead by this many samples during, but not before, training. Recommendation is to provide a value greater than per device @@ -68,6 +74,12 @@ class SimulationDataset(StreamingDataset): Set to ``None`` to disable shard eviction. Supports integer bytes as well as string human-readable bytes (e.g., ``100b``, ``64kb``, ``77mb``, and so on). Defaults to ``None``. + sampling_method (str): Which sampling method to use, either ``balanced`` or ``fixed``. + Defaults to ``balanced``. + sampling_granularity (int): When picking samples for a stream's final partial repeat, + how many samples to pick from the same shard at a time (``1`` for evenly balanced + across shards, ``1000`` to pick 1000 samples from the same shard at a time, etc). + Defaults to ``1``. partition_algo (str): Which partitioning algorithm to use. Defaults to ``relaxed``. num_canonical_nodes (int, optional): Canonical number of nodes for shuffling with resumption. The sample space is divided evenly according to the number of canonical @@ -86,51 +98,45 @@ class SimulationDataset(StreamingDataset): shuffle (bool): Whether to iterate over the samples in randomized order. Defaults to ``False``. shuffle_algo (str): Which shuffling algorithm to use. Defaults to ``py1e``. - shuffle_seed (int): Seed for Deterministic data shuffling. Defaults to ``9176``. + shuffle_seed (int): Seed for deterministic data shuffling. Defaults to ``9176``. shuffle_block_size (int, optional): Unit of shuffle. A canonical node's samples are split into blocks of this size, and samples within each block are shuffled. If ``None``, its value is calculated as ``max(4_000_000 // num_canonical_nodes), 1 << 18)``. Defaults to ``None``. - sampling_method (str): Which sampling method to use, either ``balanced`` or ``fixed``. - Defaults to ``balanced``. - sampling_granularity (int): When picking samples for a stream's final partial repeat, - how many samples to pick from the same shard at a time (``1`` for evenly balanced - across shards, ``1000`` to pick 1000 samples from the same shard at a time, etc). - Defaults to ``1``. batching_method (str): Which batching method to use, either ``random``, ``stratified``, or ``per_stream``. Defaults to ``random``. - allow_unsafe_types (bool): If a shard contains Pickle, which allows arbitrary code - execution during deserialization, whether to keep going if ``True`` or raise an error - if ``False``. Defaults to ``False``. """ - def __init__(self, - nodes: int, - devices: int, - workers: int, - streams: Optional[Sequence[Stream]] = None, - remote: Optional[str] = None, - local: Optional[str] = None, - split: Optional[str] = None, - download_retry: int = 2, - download_timeout: float = 60, - validate_hash: Optional[str] = None, - keep_zip: bool = False, - epoch_size: Optional[Union[int, str]] = None, - predownload: Optional[int] = None, - cache_limit: Optional[Union[int, str]] = None, - partition_algo: str = 'relaxed', - num_canonical_nodes: Optional[int] = None, - batch_size: Optional[int] = None, - shuffle: bool = False, - shuffle_algo: str = 'py1e', - shuffle_seed: int = 9176, - shuffle_block_size: Optional[int] = None, - sampling_method: str = 'balanced', - sampling_granularity: int = 1, - batching_method: str = 'random', - allow_unsafe_types: bool = False) -> None: - + def __init__( + self, + *, + nodes: int, + devices: int, + workers: int, + epoch_size: Optional[Union[int, str]] = None, + streams: Optional[Sequence[Stream]] = None, + remote: Optional[str] = None, + local: Optional[str] = None, + split: Optional[str] = None, + download_retry: int = 2, + download_timeout: float = 60, + validate_hash: Optional[str] = None, + keep_zip: bool = False, + allow_unsafe_types: bool = False, + config_root: Optional[str] = None, + predownload: Optional[int] = None, + cache_limit: Optional[Union[int, str]] = None, + sampling_method: str = 'balanced', + sampling_granularity: int = 1, + partition_algo: str = 'relaxed', + num_canonical_nodes: Optional[int] = None, + batch_size: Optional[int] = None, + shuffle: bool = False, + shuffle_algo: str = 'py1e', + shuffle_seed: int = 9176, + shuffle_block_size: Optional[int] = None, + batching_method: str = 'random', + ) -> None: # Time how long it takes for StreamingDataset instantiation t0 = time.time() @@ -138,59 +144,32 @@ def __init__(self, self.nodes = nodes self.devices = devices self.workers = workers - self.cache_limit = cache_limit - self.partition_algo = partition_algo - self.predownload = predownload + + # Purely StreamingDataset arguments (which do not live in Streams). + self.config_root = self._get_config_root(config_root) + self.predownload = self._get_predownload(predownload, batch_size) + self.cache_limit = self._get_cache_limit(cache_limit) + self.sampling_method = self._get_sampling_method(sampling_method) + self.sampling_granularity = self._get_sampling_granularity(sampling_granularity) + self.partition_algo = self._get_partition_algo(partition_algo) + self.num_canonical_nodes: int self.batch_size = batch_size self.shuffle = shuffle - self.shuffle_algo = shuffle_algo - self.shuffle_seed = shuffle_seed - self.shuffle_block_size = shuffle_block_size - self.sampling_method = sampling_method - self.sampling_granularity = sampling_granularity - self.batching_method = batching_method - self.num_canonical_nodes = num_canonical_nodes - self.allow_unsafe_types = allow_unsafe_types + self.shuffle_algo = self._get_shuffle_algo(shuffle_algo) + self.shuffle_seed = self._get_shuffle_seed(shuffle_seed) + self.input_shuffle_block_size = shuffle_block_size + self.shuffle_block_size: int # Set below. + self.batching_method = self._get_batching_method(batching_method) + + # StreamingDataset arguments which depend on other such arguments. + world = World() + self.num_canonical_nodes = self._get_num_canonical_nodes(num_canonical_nodes, + self.shuffle_algo, world) + self.shuffle_block_size = self._get_shuffle_block_size(shuffle_block_size, + self.num_canonical_nodes, world) self.initial_physical_nodes = nodes - # Set num_canonical_nodes based on the shuffling algorithm chosen. - if self.num_canonical_nodes is None: - if self.shuffle_algo in ['py1s', 'py2s']: - self.num_canonical_nodes = 64 * self.nodes - else: - self.num_canonical_nodes = self.nodes - - # Set shuffle_block_size if not provided, based on num_canonical_nodes. - if self.shuffle_block_size is None: - self.shuffle_block_size = max(4_000_000 // self.num_canonical_nodes, 1 << 18) - - # Check streams vs remote/local. - if bool(streams) == (bool(remote) or bool(local)): - raise ValueError( - 'You must provide either `streams` or `remote`/`local`, but not both.') - - # Check sampling method is one of "balanced" or "fixed". - if self.sampling_method not in ['balanced', 'fixed']: - raise ValueError( - f'Invalid sampling method: {sampling_method}. Must be one of `balanced` or `fixed`.' - ) - - # Check sampling method is one of "balanced" or "fixed". - if self.batching_method not in ['random', 'per_stream', 'stratified']: - raise ValueError( - f'Invalid batching method: {batching_method}. Must be one of `random`, \ - `per_stream`, or `stratified`.') - - # Check that predownload is at least per device batch size, and set it if currently `None`. - if self.predownload is not None and self.batch_size is not None and \ - self.predownload < self.batch_size: - warnings.warn(f'predownload < batch_size ({self.predownload} < {self.batch_size}).' + - f'This may result in slower batch time. Recommendation is to set ' + - f'predownload to at-least batch_size.') - elif self.predownload is None: - self.predownload = 8 * self.batch_size if self.batch_size is not None else 64 - self.batch_size = batch_size or 1 # Convert epoch size from string to int, if needed. Cannot be negative. @@ -202,26 +181,22 @@ def __init__(self, # Initialize the Stream defaults and normalize to a list of Streams. if streams: - default = { - 'remote': remote, - 'local': local, - 'split': split, - 'download_retry': download_retry, - 'download_timeout': download_timeout, - 'validate_hash': validate_hash, - 'keep_zip': keep_zip, - } for stream in streams: - stream.apply_default(default) + stream.apply_defaults(split=split, + download_retry=download_retry, + download_timeout=download_timeout, + validate_hash=validate_hash, + keep_zip=keep_zip, + allow_unsafe_types=allow_unsafe_types) else: - default = Stream(remote=remote, + streams = Stream(remote=remote, local=local, split=split, download_retry=download_retry, download_timeout=download_timeout, validate_hash=validate_hash, - keep_zip=keep_zip) - streams = [default] + keep_zip=keep_zip, + allow_unsafe_types=allow_unsafe_types), # Validate the stream weighting scheme (relative or absolute) to catch errors before we go # to the trouble of loading them. @@ -270,7 +245,7 @@ def __init__(self, local_foldernames = [] for stream_id, stream in enumerate(self.streams): logger.info(f' Processing index file for stream {stream_id + 1}') - stream_shards = stream.get_shards(self.world, self.allow_unsafe_types) + stream_shards = stream.get_shards(self.world) num_stream_samples = sum(map(len, stream_shards)) index_filename = os.path.join(stream.local, stream.split, get_index_basename()) index_filenames.append(index_filename) @@ -421,9 +396,6 @@ def get_num_canonical_nodes(self) -> int: Returns: int: The dataset's number of canonical nodes. """ - if not isinstance(self.num_canonical_nodes, int): - raise TypeError(f'`self.num_canonical_nodes` must be an int. ' + - f'Got {type(self.num_canonical_nodes)} instead.') return self.num_canonical_nodes def get_batch_size(self) -> int: @@ -459,9 +431,6 @@ def get_predownload(self) -> int: Returns: int: The dataset's predownload. """ - if not isinstance(self.predownload, int): - raise TypeError(f'`self.predownload` must be an int. ' + - f'Got {type(self.predownload)} instead.') return self.predownload def get_cache_limit(self) -> Optional[int]: @@ -531,9 +500,6 @@ def get_shuffle_block_size(self) -> int: Returns: int: The dataset's shuffle block size. """ - if not isinstance(self.shuffle_block_size, int): - raise TypeError(f'`self.shuffle_block_size` must be an int. ' + - f'Got {type(self.shuffle_block_size)} instead.') return self.shuffle_block_size def get_epoch_size(self) -> int: diff --git a/simulation/core/sim_world.py b/simulation/core/sim_world.py index c36364530..0e9060460 100644 --- a/simulation/core/sim_world.py +++ b/simulation/core/sim_world.py @@ -3,7 +3,7 @@ """Contains info about the nodes, ranks, and workers of the run for simulation purposes.""" -from streaming.base.world import World +from streaming.base.coord.world import World class SimulationWorld(World): diff --git a/simulation/core/yaml_processing.py b/simulation/core/yaml_processing.py index 07ebfee2a..18942cb14 100644 --- a/simulation/core/yaml_processing.py +++ b/simulation/core/yaml_processing.py @@ -197,11 +197,29 @@ def create_simulation_dataset(nodes: int, devices: int, workers: int, global_bat sampling_granularity = train_dataset.get('sampling_granularity', 1) batching_method = train_dataset.get('batching_method', 'random') - dataset = SimulationDataset(nodes, devices, workers, streams, remote, local, split, - download_retry, download_timeout, validate_hash, keep_zip, - epoch_size, predownload, cache_limit, partition_algo, - num_canonical_nodes, batch_size, shuffle, shuffle_algo, - shuffle_seed, shuffle_block_size, sampling_method, - sampling_granularity, batching_method) + dataset = SimulationDataset(nodes=nodes, + devices=devices, + workers=workers, + streams=streams, + remote=remote, + local=local, + split=split, + download_retry=download_retry, + download_timeout=download_timeout, + validate_hash=validate_hash, + keep_zip=keep_zip, + epoch_size=epoch_size, + predownload=predownload, + cache_limit=cache_limit, + partition_algo=partition_algo, + num_canonical_nodes=num_canonical_nodes, + batch_size=batch_size, + shuffle=shuffle, + shuffle_algo=shuffle_algo, + shuffle_seed=shuffle_seed, + shuffle_block_size=shuffle_block_size, + sampling_method=sampling_method, + sampling_granularity=sampling_granularity, + batching_method=batching_method) return dataset diff --git a/streaming/base/batching/__init__.py b/streaming/base/batching/__init__.py index 4d4ebe3ae..774d54b1a 100644 --- a/streaming/base/batching/__init__.py +++ b/streaming/base/batching/__init__.py @@ -12,7 +12,7 @@ from streaming.base.batching.per_stream import generate_work_per_stream_batching from streaming.base.batching.random import generate_work_random_batching from streaming.base.batching.stratified import generate_work_stratified_batching -from streaming.base.world import World +from streaming.base.coord.world import World if TYPE_CHECKING: from streaming.base.dataset import StreamingDataset diff --git a/streaming/base/batching/per_stream.py b/streaming/base/batching/per_stream.py index f2f8bd6dc..2b69c2add 100644 --- a/streaming/base/batching/per_stream.py +++ b/streaming/base/batching/per_stream.py @@ -10,9 +10,9 @@ import numpy as np from numpy.typing import NDArray +from streaming.base.coord.world import World from streaming.base.partition import get_partitions from streaming.base.shuffle import get_shuffle -from streaming.base.world import World if TYPE_CHECKING: from streaming.base.dataset import StreamingDataset @@ -63,9 +63,6 @@ def generate_work_per_stream_batching(dataset: StreamingDataset, world: World, e # same as the ratio of the stream's samples to overall samples. # This ensures that the overall training shuffle block size is still approximately # equal to what is set by the user, and allows for reasoning about cache_limit as well. - if not isinstance(dataset.shuffle_block_size, int): - raise TypeError(f'Dataset `shuffle_block_size` must be an integer. ' + - f'Got {type(dataset.shuffle_block_size)} instead.') shuffle_block_portion = int(dataset.shuffle_block_size * stream.proportion) stream_shuffle = get_shuffle(dataset.shuffle_algo, shuffle_units, dataset.num_canonical_nodes, dataset.shuffle_seed, epoch, diff --git a/streaming/base/batching/random.py b/streaming/base/batching/random.py index ddd0fe050..035219bd6 100644 --- a/streaming/base/batching/random.py +++ b/streaming/base/batching/random.py @@ -10,9 +10,9 @@ import numpy as np from numpy.typing import NDArray +from streaming.base.coord.world import World from streaming.base.partition import get_partitions from streaming.base.shuffle import get_shuffle -from streaming.base.world import World if TYPE_CHECKING: from streaming.base.dataset import StreamingDataset @@ -58,9 +58,6 @@ def generate_work_random_batching(dataset: StreamingDataset, world: World, epoch # If we need to shuffle, shuffle in a node-aware and *underlying* shard-aware way. if dataset.shuffle: - if not isinstance(dataset.shuffle_block_size, int): - raise TypeError(f'Dataset `shuffle_block_size` must be an integer. ' + - f'Got {type(dataset.shuffle_block_size)} instead.') shuffle = get_shuffle(dataset.shuffle_algo, shuffle_units, dataset.num_canonical_nodes, dataset.shuffle_seed, epoch, dataset.shuffle_block_size) big_ids = np.where(big_ids != -1, shuffle[big_ids], -1) diff --git a/streaming/base/batching/stratified.py b/streaming/base/batching/stratified.py index 40bcea209..da7bbdf5f 100644 --- a/streaming/base/batching/stratified.py +++ b/streaming/base/batching/stratified.py @@ -11,9 +11,9 @@ import numpy as np from numpy.typing import NDArray +from streaming.base.coord.world import World from streaming.base.partition import get_partitions from streaming.base.shuffle import get_shuffle -from streaming.base.world import World if TYPE_CHECKING: from streaming.base.dataset import StreamingDataset @@ -75,9 +75,6 @@ def generate_work_stratified_batching(dataset: StreamingDataset, world: World, e # same as the ratio of the stream's samples to overall samples. # This ensures that the overall training shuffle block size is still approximately # equal to what is set by the user, and allows for reasoning about cache_limit as well. - if not isinstance(dataset.shuffle_block_size, int): - raise TypeError(f'Dataset `shuffle_block_size` must be an integer. ' + - f'Got {type(dataset.shuffle_block_size)} instead.') shuffle_block_portion = int(dataset.shuffle_block_size * stream.proportion) stream_shuffle = get_shuffle(dataset.shuffle_algo, shuffle_units, dataset.num_canonical_nodes, dataset.shuffle_seed, epoch, diff --git a/streaming/base/coord/__init__.py b/streaming/base/coord/__init__.py new file mode 100644 index 000000000..430dfc136 --- /dev/null +++ b/streaming/base/coord/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Coordination among ranks and workers.""" + +from streaming.base.coord.job import JobDirectory, JobRegistry +from streaming.base.coord.shmem import (SharedArray, SharedBarrier, SharedMemory, SharedScalar, + get_shm_prefix) +from streaming.base.coord.world import World + +__all__ = [ + 'JobDirectory', 'JobRegistry', 'SharedArray', 'SharedBarrier', 'SharedMemory', + 'get_shm_prefix', 'SharedScalar', 'World' +] diff --git a/streaming/base/coord/file/__init__.py b/streaming/base/coord/file/__init__.py new file mode 100644 index 000000000..90f2459d8 --- /dev/null +++ b/streaming/base/coord/file/__init__.py @@ -0,0 +1,9 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Coordinating processes using files.""" + +from streaming.base.coord.file.lock import SoftFileLock +from streaming.base.coord.file.waiting import create_file, wait_for_creation, wait_for_deletion + +__all__ = ['create_file', 'wait_for_creation', 'wait_for_deletion', 'SoftFileLock'] diff --git a/streaming/base/coord/file/lock.py b/streaming/base/coord/file/lock.py new file mode 100644 index 000000000..0c8c59689 --- /dev/null +++ b/streaming/base/coord/file/lock.py @@ -0,0 +1,184 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Soft file locking via file open mode 'x'.""" + +import os +from types import TracebackType +from typing import Optional, Type, Union + +from typing_extensions import Self + +from streaming.base.coord.process import get_live_processes +from streaming.base.coord.waiting import wait + +__all__ = ['SoftFileLock'] + + +class SoftFileLock: + """Soft file locking via file open mode 'x'. + + Args: + filename (str): Path to lock. + timeout (float, optional): How long to wait in seconds before raising an exception. + Set to ``None`` to never time out. Defaults to ``30``. + tick (float): Check interval in seconds. Defaults to ``0.007``. + """ + + def __init__( + self, + filename: str, + timeout: Optional[float] = 30, + tick: float = 0.007, + ) -> None: + if not filename: + raise ValueError('Path to file lock is empty.') + + if timeout is not None: + if timeout <= 0: + raise ValueError( + f'Timeout must be positive float seconds, but got: {timeout} sec.') + + if tick <= 0: + raise ValueError(f'Tick must be positive float seconds, but got: {tick} sec.') + + self.filename = filename + self.timeout = timeout + self.tick = tick + + self._normalize(filename) + + @classmethod + def _write(cls, filename: str, pid: int) -> None: + """Write the locking process's pid. + + Args: + filename (str): Path to lock. + """ + with open(filename, 'x') as file: + file.write(str(pid)) + + @classmethod + def _read(cls, filename: str) -> int: + """Read the locking process's pid. + + Args: + filename (str): Path to lock. + """ + with open(filename, 'r') as file: + return int(file.read()) + + @classmethod + def _normalize(cls, filename: str) -> None: + """Ensure parent dirs exist and lock files held by dead processes do not exist. + + Args: + filename (str): Path to lock. + """ + # Ensure the file's parent directory exists so we can write it in one shot. + dirname = os.path.dirname(filename) + if dirname: + os.makedirs(dirname, exist_ok=True) + + # If no file, we don't need to do anything. + if not os.path.exists(filename): + return + + # If we fail to open the file and parse the pid, bail out while deleting it. + try: + pid = cls._read(filename) + except: + os.remove(filename) + return + + # If the pid is not among the living, delete the file. + if pid not in get_live_processes(): + os.remove(filename) + + @classmethod + def _get_timeout( + cls, + init_timeout: Optional[float], + timeout: Optional[Union[str, float]] = 'auto', + ) -> Optional[float]: + """Determine the timeout for a given acquire(). + + Args: + init_timeout (float, optional): Default timeout provided to init. + timeout (str | float, optional): Override timeout for just this method call. + + Returns: + float, optional: Normalized timeout as positive float seconds or ``None`` to disable. + """ + if timeout is None: + # No timeout. + ret = timeout + elif isinstance(timeout, float): + # Override timeout. + if timeout <= 0: + raise ValueError( + f'Timeout must be positive float seconds, but got: {timeout} sec.') + ret = timeout + elif timeout == 'auto': + # Default timeout. + ret = init_timeout + else: + raise ValueError(f'Timeout must either be positive float seconds, ``None`` to ' + + f'disable timing out, or ``auto`` to use the default passed to ' + + f'init, but got: {timeout}.') + return ret + + def acquire( + self, + timeout: Optional[Union[str, float]] = 'auto', + ) -> None: + """Acquire this lock. + + Args: + timeout (str | float, optional): Override timeout for just this method call. + """ + + def stop() -> bool: + try: + with open(self.filename, 'x') as out: + text = str(os.getpid()) + out.write(text) + return True + except: + return False + + norm_timeout = self._get_timeout(self.timeout, timeout) + wait(stop, norm_timeout, self.tick) + + def release(self) -> None: + """Release this lock.""" + if os.path.isfile(self.filename): + os.remove(self.filename) + elif os.path.exists(self.filename): + raise ValueError(f'Path exists, but is not a file: {self.filename}.') + else: + raise ValueError(f'Path does not exist: {self.filename}.') + + def __enter__(self) -> Self: + """Enter context manager. + + Returns: + Self: This lock. + """ + self.acquire() + return self + + def __exit__( + self, + err_type: Optional[Type[BaseException]] = None, + err: Optional[BaseException] = None, + trace: Optional[TracebackType] = None, + ) -> None: + """Exit context manager. + + Args: + err_type (Type[BaseException], optional): Exc type. + err (BaseException, optional): Exc. + trace (TracebackType, optional): Traceback. + """ + self.release() diff --git a/streaming/base/coord/file/waiting.py b/streaming/base/coord/file/waiting.py new file mode 100644 index 000000000..daf4a8996 --- /dev/null +++ b/streaming/base/coord/file/waiting.py @@ -0,0 +1,71 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Waiting on files.""" + +import os +from typing import Any, Optional + +from streaming.base.coord.waiting import wait + +__all__ = ['wait_for_creation', 'wait_for_deletion', 'create_file'] + + +def wait_for_creation( + path: str, + timeout: Optional[float] = 30, + tick: float = 0.007, + lock: Optional[Any] = None, +) -> None: + """Wait for the creation of a path on the local filesystem. + + Args: + path (str): Local path to wait on the creation of. + timeout (float, optional): How long to wait before raising an exception, in seconds. + Defaults to ``30``. + tick (float): Check interval, in seconds. Defaults to ``0.007``. + lock (Any, optional): Context manager (this is intended for locks) to be held when + checking the predicate. Defaults to ``None``. + """ + + def stop(): + return os.path.exists(path) + + wait(stop, timeout, tick, lock) + + +def wait_for_deletion( + path: str, + timeout: Optional[float] = 30, + tick: float = 0.007, + lock: Optional[Any] = None, +) -> None: + """Wait for the deletion of a path on the local filesystem. + + Args: + path (str): Local path to wait on the deletion of. + timeout (float, optional): How long to wait before raising an exception, in seconds. + Defaults to ``30``. + tick (float): Check interval, in seconds. Defaults to ``0.007``. + lock (Any, optional): Context manager (this is intended for locks) to be held when + checking the predicate. Defaults to ``None``. + """ + + def stop(): + return not os.path.exists(path) + + wait(stop, timeout, tick, lock) + + +def create_file(filename: str) -> None: + """Create a file at the given path on the local filesystem. + + Raises an exception if the path already exists. + + Args: + filename (str): Filename to create. + """ + dirname = os.path.dirname(filename) + os.makedirs(dirname, exist_ok=True) + with open(filename, 'x'): + pass diff --git a/streaming/base/coord/job/__init__.py b/streaming/base/coord/job/__init__.py new file mode 100644 index 000000000..82e765f4b --- /dev/null +++ b/streaming/base/coord/job/__init__.py @@ -0,0 +1,9 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Handling for jobs, which are collections of StreamingDataset replicas with the same config.""" + +from streaming.base.coord.job.directory import JobDirectory +from streaming.base.coord.job.registry import JobRegistry + +__all__ = ['JobDirectory', 'JobRegistry'] diff --git a/streaming/base/coord/job/directory.py b/streaming/base/coord/job/directory.py new file mode 100644 index 000000000..03adb2a6a --- /dev/null +++ b/streaming/base/coord/job/directory.py @@ -0,0 +1,49 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""A directory containing all dataset-wide filesystem state for a Streaming job.""" + +import os +from typing import Sequence + +from streaming.base.coord.job.registry import JobRegistry +from streaming.base.coord.world import World +from streaming.base.stream import Stream + +__all__ = ['JobDirectory'] + + +class JobDirectory: + """Represents a Streaming job lease. On ``__del__``, cleans up after itself. + + When it goes out of scope naturally, this Job will delete its config dir and its hold on all + the local dirs it is streaming to. + + If this process dies badly and the destructor is not reached, the same cleanup will be done by + some future process incidentally as it registers or unregisters a Streaming job. It can tell it + died by a combination of pid and process create time. + + Args: + registry (JobRegistry): Stremaing job registry. + """ + + def __init__(self, registry: JobRegistry, streams: Sequence[Stream], world: World) -> None: + self.registry = registry + self.streams = streams + self.world = world + self.job_hash = registry.register(streams, world) + + def get_filename(self, path: str) -> str: + """Get a filename by relative path under its job dir. + + Args: + path (str): Path relative to job dir. + + Returns: + str: Filename. + """ + return os.path.join(self.registry.config_root, self.job_hash, path) + + def __del__(self) -> None: + """Destructor.""" + self.registry.unregister(self.job_hash, self.world) diff --git a/streaming/base/coord/job/entry.py b/streaming/base/coord/job/entry.py new file mode 100644 index 000000000..6ebf88a6f --- /dev/null +++ b/streaming/base/coord/job/entry.py @@ -0,0 +1,65 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""An entry in a Streaming job registry file.""" + +from typing import Any, Dict, List, Optional + +from typing_extensions import Self + +__all__ = ['JobEntry'] + + +class JobEntry: + """Info about a Streaming job for local dir reuse detection purposes. + + Args: + index (int, optional): The job's index in the total list. + job_hash (str): Job hash. + stream_hashes (List[str]): Stream hashes. + stream_locals (List[str], optional): Stream locals, if available. + process_id (int): PID of local rank zero of the Streaming job. + register_time (int): Process registration time. + """ + + def __init__( + self, + *, + index: Optional[int] = None, + job_hash: str, + stream_hashes: List[str], + stream_locals: Optional[List[str]] = None, + process_id: int, + register_time: int, + ) -> None: + self.index = index + self.job_hash = job_hash + self.stream_hashes = stream_hashes + self.stream_locals = stream_locals + self.process_id = process_id + self.register_time = register_time + + @classmethod + def from_json(cls, obj: Dict[str, Any]) -> Self: + """Load from JSON. + + Args: + obj (Dict[str, Any]): Source JSON object. + + Returns: + Self: Loaded JobEntry. + """ + return cls(job_hash=obj['job_hash'], + stream_hashes=obj['stream_hashes'], + stream_locals=obj.get('stream_locals'), + process_id=obj['process_id'], + register_time=obj['register_time']) + + def to_json(self) -> Dict[str, Any]: + return { + 'job_hash': self.job_hash, + 'stream_hashes': self.stream_hashes, + # stream_locals is not saved, only their hashes. + 'process_id': self.process_id, + 'register_time': self.register_time, + } diff --git a/streaming/base/coord/job/file.py b/streaming/base/coord/job/file.py new file mode 100644 index 000000000..213394eac --- /dev/null +++ b/streaming/base/coord/job/file.py @@ -0,0 +1,130 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""A Streaming job registry file.""" + +import json +import os +from typing import Dict, List + +from typing_extensions import Self + +from streaming.base.coord.job.entry import JobEntry + +__all__ = ['RegistryFile'] + + +class RegistryFile: + """StreamingDataset job registry, which is backed by a JSON file. + + Args: + jobs (List[JobEntry]): List of StreamingDataset jobs. + """ + + def __init__(self, jobs: List[JobEntry]) -> None: + self.jobs = [] + self.job_hash2job = {} + self.stream_hash2job = {} + self.num_jobs = 0 + for job in jobs: + self.add(job) + + @classmethod + def read(cls, filename: str) -> Self: + if os.path.exists(filename): + obj = json.load(open(filename)) + else: + obj = {} + jobs = obj.get('jobs') or [] + jobs = [JobEntry.from_json(job) for job in jobs] + return cls(jobs) + + def write(self, filename: str) -> None: + jobs = [job.to_json() for job in filter(bool, self.jobs)] + obj = {'jobs': jobs} + with open(filename, 'w') as out: + json.dump(obj, out) + + def __len__(self) -> int: + """Get the number of jobs registered. + + Returns: + int: Number of registered jobs. + """ + return self.num_jobs + + def add(self, job: JobEntry) -> None: + """Register a Stremaing job. + + Args: + job (Job): The job. + """ + # Check that stream locals line up. + if job.stream_locals: + if len(job.stream_hashes) != len(job.stream_locals): + raise ValueError(f'If locals are provided, must have one local per stream hash, ' + + f'but got: {len(job.stream_hashes)} hashes vs ' + + f'{len(job.stream_locals)} locals.') + norm_stream_locals = job.stream_locals + else: + norm_stream_locals = [None] * len(job.stream_hashes) + + # Check dataset hash for reuse. + if job.job_hash in self.job_hash2job: + if job.stream_locals: + raise ValueError(f'Reused dataset local path(s): {job.stream_locals}.') + else: + raise ValueError(f'Reused dataset local path(s): stream hashes = ' + + f'{job.stream_hashes}, dataset hash = {job.job_hash}.') + + # Check each stream hash for reuse. + for stream_hash, norm_stream_local in zip(job.stream_hashes, norm_stream_locals): + if stream_hash in self.stream_hash2job: + if norm_stream_local: + raise ValueError('Reused stream local path: {norm_stream_local}.') + else: + raise ValueError('Reused stream local path: stream hash = {stream_hash}.') + + # Do the insertion. + job.index = len(self.jobs) + self.jobs.append(job) + self.job_hash2job[job.job_hash] = job + for stream_hash in job.stream_hashes: + self.stream_hash2job[stream_hash] = job + self.num_jobs += 1 + + def remove(self, job_hash: str) -> None: + """Deregister a Streaming job. + + Args: + job_hash (str): Job hash. + """ + job = self.job_hash2job.get(job_hash) + if not job: + raise ValueError(f'Job hash not found: {job_hash}.') + + if job.index is None: + raise ValueError('Internal error in job registration: job index is missing.') + + self.jobs[job.index] = None + del self.job_hash2job[job.job_hash] + for stream_hash in job.stream_hashes: + del self.stream_hash2job[stream_hash] + self.num_jobs -= 1 + + def filter(self, pid2create_time: Dict[int, int]) -> List[str]: + """Filter our collection of Streaming jobs. + + Args: + pid2create_time (Dict[int, int]): Mapping of pid to creation time. + + Returns: + List[str]: List of hashes of removed datasets. + """ + del_job_hashes = [] + for job in filter(bool, self.jobs): + create_time = pid2create_time.get(job.process_id) + if not create_time or job.register_time < create_time: + self.remove(job.job_hash) + del_job_hashes.append(job.job_hash) + return del_job_hashes diff --git a/streaming/base/coord/job/registry.py b/streaming/base/coord/job/registry.py new file mode 100644 index 000000000..56b46b0a2 --- /dev/null +++ b/streaming/base/coord/job/registry.py @@ -0,0 +1,257 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""A directory containing all Streaming-wide filesystem state. + +Useful for detecting collisions between different jobs' local dirs. +""" + +import os +from hashlib import sha3_224 +from shutil import rmtree +from time import time_ns +from typing import Dict, List, Optional, Sequence, Tuple + +from psutil import process_iter + +from streaming.base.coord.file.lock import SoftFileLock +from streaming.base.coord.file.waiting import wait_for_creation, wait_for_deletion +from streaming.base.coord.job.entry import JobEntry +from streaming.base.coord.job.file import RegistryFile +from streaming.base.coord.world import World +from streaming.base.stream import Stream + +__all__ = ['JobRegistry'] + + +class JobRegistry: + """StreamingDataset job registry, for the purpose of detecting local dir reuse. + + This class is safe for concurrent access via a filelock. + + Args: + config_root (str): Streaming configuration root directory, used for collision detection, + filelock paths, etc. Defaults to ``/tmp/streaming``, using the equivalent temp root on + your system. + timeout (float, optional): How long to wait before raising an exception, in seconds. + Defaults to ``30``. + tick (float): Check interval, in seconds. Defaults to ``0.007``. + """ + + def __init__( + self, + config_root: str, + timeout: Optional[float] = 30, + tick: float = 0.007, + ) -> None: + self.config_root = config_root + self.timeout = timeout + self.tick = tick + + self.lock_filename = os.path.join(config_root, 'registry.lock') + self.lock = SoftFileLock(self.lock_filename) + + self.registry_filename = os.path.join(config_root, 'registry.json') + + def _get_live_procs(self) -> Dict[int, int]: + """List the pids and creation times of every live process in the system. + + The creation times protect us from PID reuse. + + Returns: + Dict[int, int]: Mapping of pid to integer creation time. + """ + ret = {} + for obj in process_iter(['pid', 'create_time']): + ret[obj.pid] = int(obj.create_time() * 1e9) + return ret + + def _hash(self, data: bytes) -> str: + """Get a short, deterministic, fixed-length code for the given data. + + Args: + data (bytes): The data to hash. + + Returns: + str: Truncated hex digest. + """ + return sha3_224(data).hexdigest()[:8] + + def _hash_streams(self, streams: Sequence[Stream]) -> Tuple[List[str], List[str], str]: + """Get a short, opaque str key for a StreamingDataset and each of its Streams. + + This is useful for collision detection. + + Args: + streams (Sequence[Stream]): List of this StreamingDataset's Streams, which in + combination with process IDs and creation times lets us uniquely identify a + Streaming job. + + Returns: + Tuple[str, List[str], List[str]]: Triple of (normalized stream locals, stream hashes, + and dataset hash). + """ + # Get a list of the normalized locals of each Stream. + stream_locals = [] + for stream in streams: + local = os.path.join(stream.local, stream.split) + local = os.path.normpath(local) + local = os.path.abspath(local) + stream_locals.append(local) + + # Collect the locals into a deduped set. + stream_locals_set = set() + for local in stream_locals: + if local in stream_locals_set: + raise ValueError(f'Reused local path: {local}.') + stream_locals_set.add(local) + + # Verify that no local is contained within another local. + for local in stream_locals: + parts = local.split(os.path.sep)[1:] + for num_parts in range(1, len(parts) - 1): # Leftmost is '' because they start with /. + parent = os.path.sep.join(parts[:num_parts]) + if parent in stream_locals_set: + raise ValueError(f'One local path contains another local path: {parent} vs ' + + f'{local}.') + + # Hash each local. + stream_hashes = [] + for local in sorted(stream_locals): + data = local.encode('utf-8') + stream_hash = self._hash(data) + stream_hashes.append(stream_hash) + + # Hash the dataset. + text = ','.join(stream_hashes) + data = text.encode('utf-8') + job_hash = self._hash(data) + + return stream_locals, stream_hashes, job_hash + + def _make_dir(self, job_hash: str) -> None: + """Create a Streaming job config dir. + + Args: + job_hash: Streaming config subdir for this job. + """ + dirname = os.path.join(self.config_root, job_hash) + os.makedirs(dirname) + + def _remove_dir(self, job_hash: str) -> None: + """Delete a Streaming job config dir. + + Args: + job_hash: Streaming config subdir for this job. + """ + dirname = os.path.join(self.config_root, job_hash) + rmtree(dirname) + + def _register(self, streams: Sequence[Stream]) -> str: + """Register this collection of StreamingDataset replicas. + + Called by the local leader. + + Args: + streams (Sequence[Stream]): List of this StreamingDataset's Streams, which in + combination with process IDs and creation times lets us uniquely identify a + Streaming job. + + Returns: + str: Streaming config subdir for this job. + """ + register_time = time_ns() + pid2create_time = self._get_live_procs() + pid = os.getpid() + create_time = pid2create_time.get(pid) + if create_time is None: + raise RuntimeError('`psutil` thinks we are dead, and yet here we are: pid = {pid}.') + + stream_locals, stream_hashes, job_hash = self._hash_streams(streams) + + entry = JobEntry(job_hash=job_hash, + stream_hashes=stream_hashes, + stream_locals=stream_locals, + process_id=pid, + register_time=register_time) + + with self.lock: + conf = RegistryFile.read(self.registry_filename) + conf.add(entry) + del_job_hashes = conf.filter(pid2create_time) + conf.write(self.registry_filename) + map(self._remove_dir, del_job_hashes) + self._make_dir(job_hash) + + return job_hash + + def _lookup(self, streams: Sequence[Stream]) -> str: + """Look up this collection of StreamingDataset replicas. + + Called by the local leader. + + Args: + streams (Sequence[Stream]): List of this StreamingDataset's Streams, which in + combination with process IDs and creation times lets us uniquely identify a + Streaming job. + + Returns: + str: Streaming config subdir for this job. + """ + _, _, job_hash = self._hash_streams(streams) + return job_hash + + def register(self, streams: Sequence[Stream], world: World) -> str: + """Register or look up this collection of StreamingDataset replicas. + + Called by all ranks. + + Args: + streams (Sequence[Stream]): List of this StreamingDataset's Streams, which in + combination with process IDs and creation times lets us uniquely identify a + Streaming job. + world (World): Rank-wise world state. + + Returns: + str: Subdir for this collection of StreamingDataset replicas. + """ + if world.is_local_leader: + job_hash = self._register(streams) + else: + job_hash = self._lookup(streams) + dirname = os.path.join(self.config_root, job_hash) + wait_for_creation(dirname, self.timeout, self.tick, self.lock) + return job_hash + + def _unregister(self, job_hash: str) -> None: + """Unregister this collection of StreamingDataset replicas. + + Called by the local leader. + + Args: + job_hash (str): Subdir identifying this Streaming job. + """ + pid2create_time = self._get_live_procs() + + with self.lock: + conf = RegistryFile.read(self.registry_filename) + conf.remove(job_hash) + del_job_hashes = conf.filter(pid2create_time) + conf.write(self.registry_filename) + map(self._remove_dir, del_job_hashes) + self._remove_dir(job_hash) + + def unregister(self, job_hash: str, world: World) -> None: + """Unregister this collection of StreamingDataset replicas. + + Called by all ranks. + + Args: + job_hash (str): Subdir identifying this Streaming job. + world (World): Rank-wise world state. + """ + if world.is_local_leader: + self._unregister(job_hash) + else: + dirname = os.path.join(self.config_root, job_hash) + wait_for_deletion(dirname, self.timeout, self.tick, self.lock) diff --git a/streaming/base/coord/mmap/__init__.py b/streaming/base/coord/mmap/__init__.py new file mode 100644 index 000000000..c208a57a0 --- /dev/null +++ b/streaming/base/coord/mmap/__init__.py @@ -0,0 +1,26 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Share data across processes with mmap().""" + +from streaming.base.coord.mmap.barrier import MemMapBarrier, barrier +from streaming.base.coord.mmap.base import MemMap +from streaming.base.coord.mmap.buffer import MemMapBuffer, buffer +from streaming.base.coord.mmap.ndarray import MemMapNDArray, ndarray + +# isort: off +from streaming.base.coord.mmap.number import ( + MemMapFloat16, MemMapFloat32, MemMapFloat64, MemMapFloating, MemMapInexact, MemMapInt8, + MemMapInt16, MemMapInt32, MemMapInt64, MemMapInteger, MemMapNumber, MemMapSignedInteger, + MemMapUInt8, MemMapUInt16, MemMapUInt32, MemMapUInt64, MemMapUnsignedInteger, float16, float32, + float64, int8, int16, int32, int64, uint8, uint16, uint32, uint64) +# isort: on + +__all__ = [ + 'MemMapBarrier', 'barrier', 'MemMap', 'MemMapBuffer', 'buffer', 'MemMapNumber', + 'MemMapInteger', 'MemMapSignedInteger', 'MemMapInt8', 'MemMapInt16', 'MemMapInt32', + 'MemMapInt64', 'MemMapUnsignedInteger', 'MemMapUInt8', 'MemMapUInt16', 'MemMapUInt32', + 'MemMapUInt64', 'MemMapInexact', 'MemMapFloating', 'MemMapFloat16', 'MemMapFloat32', + 'MemMapFloat64', 'int8', 'int16', 'int32', 'int64', 'uint8', 'uint16', 'uint32', 'uint64', + 'float16', 'float32', 'float64', 'MemMapNDArray', 'ndarray' +] diff --git a/streaming/base/coord/mmap/barrier.py b/streaming/base/coord/mmap/barrier.py new file mode 100644 index 000000000..f7eae5f61 --- /dev/null +++ b/streaming/base/coord/mmap/barrier.py @@ -0,0 +1,164 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Share a barrier across processes using mmap().""" + +from enum import IntEnum +from typing import Optional + +import numpy as np + +from streaming.base.coord.file import SoftFileLock +from streaming.base.coord.mmap.ndarray import MemMapNDArray +from streaming.base.coord.waiting import wait + +__all__ = ['MemMapBarrier', 'barrier'] + + +class BarrierFlag(IntEnum): + """Whether to stop or go, used internally by MemMapBarrier.""" + + STOP = 0 + GO = 1 + + +class MemMapBarrier: + """A barrier backed by a memory-mapped file and a file lock. + + Args: + create (bool): If ``True``, create. If ``False``, attach. + mmap_filename (str): Path to memory-mapped file. + lock_filename (str): Path to SoftFileLock file. + timeout (float, optional): How long to wait before raising an exception, in seconds. + Defaults to ``30``. + tick (float): Check interval, in seconds. Defaults to ``0.007``. + """ + + def __init__( + self, + create: bool, + mmap_filename: str, + lock_filename: str, + timeout: Optional[float] = 30, + tick: float = 0.007, + ) -> None: + value = 0 if create else None + self._arr = MemMapNDArray(mmap_filename, 3, np.int32, value) + self._num_enter = 0 + self._num_exit = -1 + self._flag = BarrierFlag.GO + self._lock = SoftFileLock(lock_filename) + self._timeout = timeout + self._tick = tick + + @property + def _num_enter(self) -> int: + """Getter for _num_enter. + + Returns: + int: Entered process count. + """ + return int(self._arr[0]) + + @_num_enter.setter + def _num_enter(self, num_enter: int) -> None: + """Setter for _num_enter. + + Args: + num_enter (int): Entered process count. + """ + self._arr[0] = num_enter + + @property + def _num_exit(self) -> int: + """Getter for _num_exit. + + Returns: + int: Exited process count. + """ + return int(self._arr[1]) + + @_num_exit.setter + def _num_exit(self, num_exit: int) -> None: + """Setter for _num_exit. + + Args: + num_exit (int): Exited process count. + """ + self._arr[1] = num_exit + + @property + def _flag(self) -> BarrierFlag: + """Getter for _flag. + + Returns: + BarrierFlat: Flag value. + """ + return BarrierFlag(self._arr[2]) + + @_flag.setter + def _flag(self, flag: BarrierFlag) -> None: + """Setter for _flag. + + Args: + flag (BarrierFlag): Flag value. + """ + self._arr[2] = flag + + def __call__(self, total: int) -> None: + """A set number of processes enter, wait, and exit the barrier. + + Args: + num_procs (int): How many processes are sharing this barrier. + """ + # Initialize `_num_exit` to the number of processes. + with self._lock: + if self._num_exit == -1: + self._num_exit = total + + # If we are the first to arrive, wait for everyone to exit, then set `_flag` to `STOP`. + self._lock.acquire() + if not self._num_enter: + self._lock.release() + wait(lambda: self._num_exit == total, self._timeout, self._tick, self._lock) + self._lock.acquire() + self._flag = BarrierFlag.STOP + + # Note that we entered.z + self._num_enter += 1 + + # If we are the last to arrive, reset `_enter` and `_exit`, and set `_flag` to `GO`. + if self._num_enter == total: + self._num_enter = 0 + self._num_exit = 0 + self._flag = BarrierFlag.GO + self._lock.release() + + # Everybody waits until `_flag` is set to `GO`. + wait(lambda: self._flag == BarrierFlag.GO, self._timeout, self._tick, self._lock) + + # Note that we exited. + with self._lock: + self._num_exit += 1 + if self._num_exit == total: + self._num_exit = -1 + + +def barrier( + create: bool, + mmap_filename: str, + lock_filename: str, + timeout: Optional[float] = 30, + tick: float = 0.007, +) -> MemMapBarrier: + """Get a barrier backed by a memory-mapped file and a file lock. + + Args: + create (bool): If ``True``, create. If ``False``, attach. + mmap_filename (str): Path to memory-mapped file. + lock_filename (str): Path to SoftFileLock file. + timeout (float, optional): How long to wait before raising an exception, in seconds. + Defaults to ``30``. + tick (float): Check interval, in seconds. Defaults to ``0.007``. + """ + return MemMapBarrier(create, mmap_filename, lock_filename, timeout, tick) diff --git a/streaming/base/coord/mmap/base.py b/streaming/base/coord/mmap/base.py new file mode 100644 index 000000000..5f21b8822 --- /dev/null +++ b/streaming/base/coord/mmap/base.py @@ -0,0 +1,372 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Base functionality for sharing data across processes using mmap().""" + +import os +from mmap import mmap +from typing import Generic, Optional, Tuple, TypeVar, Union + +import numpy as np +from numpy.typing import DTypeLike, NDArray + +from streaming.base.coord.file.waiting import wait_for_creation +from streaming.base.coord.mmap.file import (get_file_header_size, get_file_size, read_file_header, + write_file) + +__all__ = ['T', 'MemMap', 'Number'] + +Number = Union[int, float, complex, np.number] + + +def _get_wildcarded_shape(shape: Optional[Union[int, Tuple[int]]]) -> Optional[Tuple[int]]: + """Normalize and validate a shape argument. + + Args: + shape (int | Tuple[int], optional): Input shape. + + Returns: + Tuple[int], optional: Normalized shape containing at least one dimension, and at most one + wildcard, or else ``None`` to mean accept any shape. + """ + if shape is None: + return None + + if shape == (): + wild_shape = 1, + elif isinstance(shape, int): + wild_shape = shape, + else: + wild_shape = shape + + num_wild = 0 + for dim in wild_shape: + if dim == -1: + num_wild += 1 + elif dim < 1: + raise ValueError(f'Each dimension must be a positive integer, with at most one ' + + f'wildcard, but got shape: {shape}.') + + if 1 < num_wild: + raise ValueError(f'Shape contains multiple ({num_wild}) wildcards: {shape}.') + + return wild_shape + + +def _get_exact_shape(shape: Optional[Union[int, Tuple[int]]]) -> Tuple[int]: + """Normalize and validate a shape argument. + + Args: + shape (int | Tuple[int], optional): Input shape. + + Returns: + Tuple[int]: Normalized shape containing at least one dimension, and no wildcards. + """ + if shape is None: + exact_shape = 1, + elif isinstance(shape, int): + exact_shape = shape, + else: + exact_shape = shape + + for dim in exact_shape: + if dim < 1: + raise ValueError(f'Each dimension must be a positive integer, but got shape: {shape}.') + + return exact_shape + + +def _normalize_ndarray(value: NDArray[np.number]) -> NDArray[np.number]: + """Normalize an input ndarray to rule out the zero ndim case. + + Args: + value (NDArray[np.number]): Input ndarray. + + Returns: + NDArray[np.number]: Normalized ndarray. + """ + if not value.ndim: + return np.expand_dims(value, 0) + + return value + + +def _normalize_number(value: Number, dtype: Optional[np.dtype]) -> NDArray[np.number]: + """Normalize an input number (np.number, int, float, etc). to ndarray. + + Args: + value (Number): Input number. + dtype (np.dtype, optional): Explicitly provided dtype, if any. + + Returns: + NDArray[np.number]: Normalized number. + """ + if isinstance(value, np.number): + arr = np.array([value]) + elif dtype is not None: + arr = np.array([value], dtype) + else: + arr = np.array([value]) + return arr + + +def _is_shape_nonempty(shape: Tuple[int]) -> bool: + """Tell whether the shape is empty on any axes. + + Args: + shape (Tuple[int]): The shape. + + Returns: + bool: Whether the shape is valid. + """ + for dim in shape: + if dim < 1: + return False + + return True + + +def _accepts_shape(shape: Tuple[int], wildcarded_shape: Optional[Tuple[int]]) -> bool: + """Tell whether the shape restriction accepts the observed shape. + + Args: + shape (Tuple[int]): Observed shape. + wildcarded_shape (Tuple[int], optional): Shape restriction, if provided. + + Returns: + bool: Whether acceptable. + """ + if wildcarded_shape is None: + return True + + if len(shape) != len(wildcarded_shape): + return False + + for dim, wild_dim in zip(shape, wildcarded_shape): + if wild_dim != dim and wild_dim != -1: + return False + + return True + + +def _accepts_dtype(have: np.dtype, want: Optional[np.dtype]) -> bool: + """Tell whether the dtype restriction accepts the observed dtype. + + Args: + have (np.dtype): Observed dtype. + want (np.dtype, optional): Dtype restriction. + + Returns: + bool: Whether acceptable. + """ + if not want: + return True + + return have == want + + +def _create_file( + filename: str, + value: Union[Number, NDArray[np.number]], + shape: Optional[Union[int, Tuple[int]]] = None, + dtype: Optional[np.dtype] = None, +) -> Tuple[Tuple[int], np.dtype]: + """Create the file backing the memory mapping given optional explicit shape and dtype. + + Args: + filename (str): Path to file to memory map. + value (Number | NDArray[np.number], optional): Creates as this value. + shape (int | Tuple[int], optional): Exact required number of elements along each axis, if + known in advance. At most one wildcard ``-1`` may be used. Defaults to ``None``. + dtype (np.dtype, optional): Exact required data type, if known in advance. Defaults to + ``None``. + + Returns: + Tuple[Tuple[int], np.dtype, int]: The file's exact shape and dtype. + """ + # The file must not already exist. + if os.path.exists(filename): + raise ValueError(f'`value` was provided, so we are initializing the file to memory map ' + + f'ourselves. Out of caution, we require that the file not already ' + + f'exist. But it does: {filename}.') + + # What are we initializing it with? + if isinstance(value, np.ndarray): + # Normalize the ndarray to have at least one axis. + arr = _normalize_ndarray(value) + + # Its normalized shape is the target shape. + exact_shape = arr.shape + + # It must not have any zero-length axes. + if not _is_shape_nonempty(arr.shape): + raise ValueError(f'``value` must be non-empty on each axis, but got shape: {shape}.') + + # If shape is provided explicitly as well, they must match. + wildcarded_shape = _get_wildcarded_shape(shape) + if not _accepts_shape(arr.shape, wildcarded_shape): + raise ValueError(f'`value` and `shape` were both provided, but they do not match: ' + + f'`value` shape {value.shape} vs explicit `shape` {shape}.') + else: # Number. + # Normalize the Number to be an ndarray of shape (1,). + arr = _normalize_number(value, dtype) + + # If shape is provided, broadcast the value to all elements, else one element. + exact_shape = _get_exact_shape(shape) + + # If dtype is provided explicitly as well, they must match. + if not _accepts_dtype(arr.dtype, dtype): + raise ValueError(f'`value` and `dtype` were both provided, but they do not match: ' + + f'`value` dtype {arr.dtype} vs explicit `dtype` {dtype}.') + + # Create the file (dense or sparse). + write_file(arr, exact_shape, filename) + + return exact_shape, arr.dtype + + +def _check_file( + filename: str, + shape: Optional[Union[int, Tuple[int]]] = None, + dtype: Optional[np.dtype] = None, + timeout: Optional[float] = 30, + tick: float = 0.007, +) -> Tuple[Tuple[int], np.dtype]: + """Validate the file backing the memory mapping given optional explicit shape and dtype. + + Args: + filename (str): Path to file to memory map. + shape (int | Tuple[int], optional): Exact required number of elements along each axis, if + known in advance. At most one wildcard ``-1`` may be used. Defaults to ``None``. + dtype (np.dtype, optional): Exact required data type, if known in advance. Defaults to + ``None``. + timeout (float, optional): How long to wait before raising an exception, in seconds. + Defaults to ``30``. + tick (float): Check interval, in seconds. Defaults to ``0.007``. + + Returns: + Tuple[Tuple[int], np.dtype]: The file's exact shape and dtype. + """ + # Wait for the file to exist. + wait_for_creation(filename, timeout, tick) + + # Read the shape and dtpye in the file header. + with open(filename, 'rb') as file: + got_shape, got_dtype = read_file_header(file) + + # From those, derive the expected file size, and compare to actual. + want_size = get_file_size(got_shape, got_dtype) + got_size = os.stat(filename).st_size + if got_size != want_size: + raise ValueError(f'`File size did not match what we expected given shape and dtype as ' + + f'stated in its header: file {filename}, shape {got_shape}, dtype ' + + f'{got_dtype}, expected size {want_size}, actual size {got_size}.') + + # If shape is provided explicitly as well, they must match. + wildcarded_shape = _get_wildcarded_shape(shape) + if not _accepts_shape(got_shape, wildcarded_shape): + raise ValueError(f'`shape` was explicitly provided, but does not match what is in the ' + + f'file: `filename` {filename}, `shape` {shape}, actual shape {got_size}.') + + # If dtype is provided explicitly as well, they must match. + if not _accepts_dtype(got_dtype, dtype): + dtype_name = dtype.name if dtype else None + raise ValueError(f'`dtype` was explicitly provided, but does not match what is in the ' + + f'file: `filename` {filename}, `dtype` {dtype_name}, actual dtype ' + + f'{got_dtype.name}.') + + return got_shape, got_dtype + + +def _ensure_file( + filename: str, + shape: Optional[Union[int, Tuple[int]]] = None, + dtype: Optional[np.dtype] = None, + value: Optional[Union[Number, NDArray[np.number]]] = None, + timeout: Optional[float] = 30, + tick: float = 0.007, +) -> Tuple[Tuple[int], np.dtype]: + """Create the file backing the memory mapping given optional explicit shape and dtype. + + Args: + filename (str): Path to file to memory map. + shape (int | Tuple[int], optional): Exact required number of elements along each axis, if + known in advance. At most one wildcard ``-1`` may be used. Defaults to ``None``. + dtype (np.dtype, optional): Exact required data type, if known in advance. Defaults to + ``None``. + value (Number | NDArray[np.number], optional): If a number, creates as this value. If + ``None``, attaches. Defaults to ``None``. + timeout (float, optional): How long to wait before raising an exception, in seconds. + Defaults to ``30``. + tick (float): Check interval, in seconds. Defaults to ``0.007``. + + Returns: + Tuple[Tuple[int], np.dtype, int]: The file's exact shape and dtype. + """ + if value is not None: + return _create_file(filename, value, shape, dtype) + else: + return _check_file(filename, shape, dtype, timeout, tick) + + +T = TypeVar('T', bound=np.dtype) + + +class MemMap(Generic[T]): + """An ndarray backed by a memory-mapped file. + + Format: + + [dtype: str, padded with nulls to uint64().itemsize] + [ndim: uint64] + [shape: ndim x uint64] + [data: prod(shape) * dtype] + + Args: + filename (str): Path to file to memory map. + shape (int | Tuple[int], optional): Exact required number of elements along each axis, if + known in advance. At most one wildcard ``-1`` may be used. Defaults to ``None``. + dtype (DTypeLike, optional): Exact required data type, if known in advance. Defaults to + ``None``. + value (Number | NDArray[np.number], optional): If a number, creates as this value. If + ``None``, attaches. Defaults to ``None``. + timeout (float, optional): How long to wait before raising an exception, in seconds. + Defaults to ``30``. + tick (float): Check interval, in seconds. Defaults to ``0.007``. + """ + + def __init__( + self, + filename: str, + shape: Optional[Union[int, Tuple[int]]] = None, + dtype: Optional[DTypeLike] = None, + value: Optional[Union[Number, NDArray[np.number]]] = None, + timeout: Optional[float] = 30, + tick: float = 0.007, + ) -> None: + norm_dtype = np.dtype(dtype) if dtype is not None else None + self.shape, self.dtype = _ensure_file(filename, shape, norm_dtype, value, timeout, tick) + self.offset = get_file_header_size(self.shape) + self.filename = filename + self.file = open(filename, 'r+b', 0) + self.mmap = mmap(self.file.fileno(), 0) + + def flush(self) -> None: + """Flush the mmap.""" + self.mmap.flush() + + def close(self) -> None: + """Close the mmap and file handle.""" + if not self.mmap.closed: + self.mmap.close() + if not self.file.closed: + self.file.close() + + def delete(self) -> None: + """Ensure everything is closed, then delete the file.""" + self.close() + os.remove(self.filename) + + def __del__(self) -> None: + """Destructor.""" + self.close() diff --git a/streaming/base/coord/mmap/buffer.py b/streaming/base/coord/mmap/buffer.py new file mode 100644 index 000000000..d0fe113d8 --- /dev/null +++ b/streaming/base/coord/mmap/buffer.py @@ -0,0 +1,81 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Share a buffer across processes using mmap().""" + +from typing import Optional + +import numpy as np + +from streaming.base.coord.mmap.base import MemMap + +__all__ = ['MemMapBuffer', 'buffer'] + + +class MemMapBuffer(MemMap[np.dtype]): + """A buffer backed by a memory-mapped file. + + Args: + filename (str): Path to file to memory map. + size (int, optional): Exact required size in bytes, if known in advance. Defaults to + ``None``. + value (bytes, optional): If provided, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__( + self, + filename: str, + size: Optional[int] = None, + value: Optional[bytes] = None, + ) -> None: + norm_value = np.frombuffer(value, np.uint8) if value is not None else None + super().__init__(filename, size, np.uint8, norm_value) + + def __len__(self) -> int: + """Get the number of bytes in the buffer. + + Returns: + int: Number of bytes in the buffer. + """ + return self.shape[0] + + def __getitem__(self, index: slice) -> bytes: + """Get the data at the index(es). + + Args: + index (slice): The index(es). + + Returns: + bytes: The data. + """ + return self.mmap[index] + + def __setitem__(self, index: slice, data: bytes) -> None: + """Set the data at the index(es). + + Args: + index (int | slice): The index(es). + data (bytes): The data. + """ + self.mmap[index] = data + + +def buffer( + filename: str, + size: Optional[int] = None, + value: Optional[bytes] = None, +) -> MemMapBuffer: + """Get a buffer backed by a memory-mapped file. + + Args: + filename (str): Path to file to memory map. + size (int, optional): Exact required size in bytes, if known in advance. Defaults to + ``None``. + value (bytes, optional): If provided, creates as this value. If ``None``, attaches. + Defaults to ``None``. + + Returns: + MemMapBuffer: A buffer backed by a memory-mapped file. + """ + return MemMapBuffer(filename, size, value) diff --git a/streaming/base/coord/mmap/file.py b/streaming/base/coord/mmap/file.py new file mode 100644 index 000000000..07591a9fe --- /dev/null +++ b/streaming/base/coord/mmap/file.py @@ -0,0 +1,201 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Base functionality for sharing data across processes using mmap().""" + +import os +from typing import IO, Tuple + +import numpy as np +from numpy.typing import NDArray + +__all__ = [ + 'get_file_header_size', 'encode_file_header', 'get_file_body_size', 'encode_file_body', + 'get_file_size', 'encode_file', 'read_file_header', 'write_sparse_file', 'write_dense_file', + 'write_file' +] + +# For memory alignment purposes. +_biggest_uint = np.uint64 + + +def get_file_header_size(shape: Tuple[int]) -> int: + """Get the expected size of the header in bytes. + + Args: + shape (Tuple[int]): Normalized ndarray shape. + """ + return (1 + 1 + len(shape)) * _biggest_uint().itemsize + + +def encode_file_header(shape: Tuple[int], dtype: np.dtype) -> bytes: + """Serialize shape/dtype metadata. + + Args: + shape (Tuple[int]): Normalized ndarray shape. + dtype (np.dtype): Normalized ndarray dtype. + + Returns: + bytes: Header data. + """ + max_itemsize = _biggest_uint().itemsize + dtype_bytes = np.dtype(dtype).name.encode('utf-8')[:max_itemsize] + pad_bytes = b'\0' * (max_itemsize - len(dtype_bytes)) + ndim_bytes = np.uint64(len(shape)).tobytes() + shape_bytes = np.array(shape, np.uint64).tobytes() + return dtype_bytes + pad_bytes + ndim_bytes + shape_bytes + + +def get_file_body_size(shape: Tuple[int], dtype: np.dtype) -> int: + """Get the expected size of the body in bytes. + + Args: + shape (Tuple[int]): Normalized ndarray shape. + dtype (np.dtype): Normalized ndarray dtype. + """ + numel = int(np.prod(shape)) if shape else 1 + norm_dtype = np.dtype(dtype) + return numel * norm_dtype.itemsize + + +def encode_file_body(arr: NDArray[np.number]) -> bytes: + """Serialize data. + + Args: + arr (NDArray[np.number]): Array to write. + + Returns: + bytes: Body data. + """ + return arr.tobytes() + + +def get_file_size(shape: Tuple[int], dtype: np.dtype) -> int: + """Get the expected size of the file in bytes. + + Args: + shape (Tuple[int]): Normalized ndarray shape. + dtype (np.dtype): Normalized ndarray dtype. + """ + header_size = get_file_header_size(shape) + body_size = get_file_body_size(shape, dtype) + return header_size + body_size + + +def encode_file(arr: NDArray[np.number]) -> bytes: + """Serialize metadata and data. + + Args: + arr (NDArray[np.number]): Array to write. + + Returns: + bytes: Body data. + """ + header_data = encode_file_header(arr.shape, arr.dtype) + body_data = encode_file_body(arr) + return header_data + body_data + + +def read_file_header(file: IO[bytes]) -> Tuple[Tuple[int], np.dtype]: + """Deserialize shape/dtype info. + + Args: + data (bytes): Header data. + + Returns: + Tuple[Tuple[int], Tuple[np.number]]: Shape and type. + """ + max_itemsize = _biggest_uint().itemsize + + # Check against minimum possible length. + to_read = 2 * max_itemsize + data = file.read(to_read) + if len(data) < to_read: + raise ValueError(f'Header data is too short: read {len(data)} bytes, but need at the ' + + f'very least {to_read} bytes.') + + # Get dtype. + part = data[:max_itemsize] + part = part[:part.index(b'\0')] + dtype_name = part.decode('utf-8') + dtype = np.dtype(dtype_name) + + # Get ndim. + part = data[max_itemsize:] + arr = np.frombuffer(part, np.uint64) + ndim = int(arr[0]) + if ndim < 0: + raise ValueError(f'Header ndim is negative: {ndim}.') + + # Check against required length in our case. + to_read = ndim * max_itemsize + data = file.read(to_read) + if len(data) < to_read: + offset = 2 * max_itemsize + raise ValueError(f'Header data is too short: got {offset + len(data)} total bytes, but ' + + f'need {offset + to_read} total bytes.') + + # Get shape. + arr = np.frombuffer(data, np.int64) + shape = tuple(arr.tolist()) + + return shape, dtype + + +def write_sparse_file(shape: Tuple[int], dtype: np.dtype, filename: str) -> None: + """Write a sparse file of the given size, which reads as zeros unless otherwise written. + + Args: + shape (Tuple[int]): Normalized ndarray shape. + dtype (np.dtype): Normalized ndarray dtype. + filename (str): Path to file. + """ + # Write header to a temp file. + tmp_filename = filename + '.tmp' + header = encode_file_header(shape, dtype) + with open(tmp_filename, 'wb') as out: + out.write(header) + + # Truncate to desired sparse size. + file_size = get_file_size(shape, dtype) + os.truncate(tmp_filename, file_size) + + # Rename to final name. + os.rename(tmp_filename, filename) + + +def write_dense_file(arr: NDArray[np.number], filename: str) -> None: + """Write a regular file with the given ndarray. + + Args: + arr (NDArray[np.number]): Array to write. + filename (str): Path to file. + """ + # Encode the file. + data = encode_file(arr) + + # Write header and body to a temp file. + tmp_filename = filename + '.tmp' + with open(tmp_filename, 'wb') as out: + out.write(data) + + # Rename to final name. + os.rename(tmp_filename, filename) + + +def write_file(arr: NDArray[np.number], shape: Tuple[int], filename: str) -> None: + """Write the ndarray as either a regular or sparse file. + + Args: + arr (NDArray[np.number]): Array to write or value to broadcast. + shape (Tuple[int]): Normalized ndarray shape. + filename (str): Path to file. + """ + if (arr == 0).all(): + write_sparse_file(shape, arr.dtype, filename) + else: + if arr.size == 1: + arr = arr.flatten() + arr = arr.repeat(np.prod(shape)) + arr = arr.reshape(shape) + write_dense_file(arr, filename) diff --git a/streaming/base/coord/mmap/ndarray.py b/streaming/base/coord/mmap/ndarray.py new file mode 100644 index 000000000..0994ee4ce --- /dev/null +++ b/streaming/base/coord/mmap/ndarray.py @@ -0,0 +1,89 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Share an ndarray across processes using mmap().""" + +from typing import Any, Optional, Tuple, Union + +import numpy as np +from numpy.typing import DTypeLike, NDArray + +from streaming.base.coord.mmap.base import MemMap, Number, T + +__all__ = ['MemMapNDArray', 'ndarray'] + +IndexType = Union[int, slice, np.integer, NDArray[np.integer]] + + +class MemMapNDArray(MemMap[T]): + """An ndarray backed by a memory-mapped file. + + Args: + filename (str): Path to file to memory map. + shape (int | Tuple[int], optional): Exact required number of elements along each axis, if + known in advance. At most one wildcard ``-1`` may be used. Defaults to ``None``. + dtype (DTypeLike, optional): The value's data type. Defaults to ``None``. + value (Number | NDArray[np.number], optional): If a number, creates as this value. If + ``None``, attaches. Defaults to ``None``. + """ + + def __init__( + self, + filename: str, + shape: Optional[Union[int, Tuple[int]]] = None, + dtype: Optional[DTypeLike] = None, + value: Optional[Union[Number, NDArray[np.number]]] = None, + ) -> None: + super().__init__(filename, shape, dtype, value) + + def numpy(self) -> NDArray: + """Get a numpy array backed by our internal memory mapped buffer. + + This is a method instead of being cached due to adventures in fork/spawn issues. + + Returns: + NDArray[T]: Our internal buffer as an ndarray. + """ + return np.ndarray(self.shape, self.dtype, self.mmap, self.offset) + + def __getitem__(self, index: IndexType) -> Union[T, NDArray]: + """Get the item at the index. + + Args: + index (IndexType): The index(es). + + Returns: + T | NDArray[T]: The item(s). + """ + return self.numpy()[index] + + def __setitem__(self, index: IndexType, item: Any) -> None: + """Set the item at the index. + + Args: + index (IndexType): The index(es). + item (Number): The item(s). + """ + self.numpy()[index] = item + + +def ndarray( + filename: str, + shape: Optional[Union[int, Tuple[int]]] = None, + dtype: Optional[DTypeLike] = None, + value: Optional[Union[Number, NDArray[np.number]]] = None, +) -> MemMapNDArray[np.dtype]: + """Get an ndarray backed by a memory-mapped file. + + Args: + filename (str): Path to file to memory map. + shape (int | Tuple[int], optional): Exact required number of elements along each axis, if + known in advance. At most one wildcard ``-1`` may be used. Defaults to ``None``. + dtype (DTypeLike, optional): The value's data type. Defaults to ``None``. + value (Number | NDArray[np.number], optional): If a number, creates as this value. If + ``None``, attaches. Defaults to ``None``. + + Returns: + MemMapNDArray[T]: An ndarray backed by a memory-mapped file. + """ + return MemMapNDArray(filename, shape, dtype, value) diff --git a/streaming/base/coord/mmap/number.py b/streaming/base/coord/mmap/number.py new file mode 100644 index 000000000..dd35acb80 --- /dev/null +++ b/streaming/base/coord/mmap/number.py @@ -0,0 +1,406 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Share a single number across processes using mmap(). + +Class hierarchy: + + MemMap + ├── ... + ├── MemMapNumber + │   ├── MemMapInexact + │   │   └── MemMapFloating + │   │   ├── MemMapFloat16 + │   │   ├── MemMapFloat32 + │   │   └── MemMapFloat64 + │   └── MemMapInteger + │   ├── MemMapSignedInteger + │   │   ├── MemMapInt16 + │   │   ├── MemMapInt32 + │   │   ├── MemMapInt64 + │   │   └── MemMapInt8 + │   └── MemMapUnsignedInteger + │   ├── MemMapUInt16 + │   ├── MemMapUInt32 + │   ├── MemMapUInt64 + │   └── MemMapUInt8 + └── ... +""" + +from typing import Optional, Union + +import numpy as np +from numpy.typing import DTypeLike, NDArray + +from streaming.base.coord.mmap.base import MemMap, Number, T + +__all__ = [ + 'MemMapNumber', 'MemMapInteger', 'MemMapSignedInteger', 'MemMapInt8', 'MemMapInt16', + 'MemMapInt32', 'MemMapInt64', 'MemMapUnsignedInteger', 'MemMapUInt8', 'MemMapUInt16', + 'MemMapUInt32', 'MemMapUInt64', 'MemMapInexact', 'MemMapFloating', 'MemMapFloat16', + 'MemMapFloat32', 'MemMapFloat64', 'int8', 'int16', 'int32', 'int64', 'uint8', 'uint16', + 'uint32', 'uint64', 'float16', 'float32', 'float64' +] + + +class MemMapNumber(MemMap[T]): + """A number backed by a memory-mapped file. + + Args: + filename (str): Path to file to memory map. + dtype (DTypeLike, optional): The value's data type. Defaults to ``None``. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__( + self, + filename: str, + dtype: Optional[DTypeLike] = None, + value: Optional[Union[Number, NDArray[np.number]]] = None, + ) -> None: + super().__init__(filename, 1, dtype, value) + + def get(self) -> T: + """Get value. + + Returns: + np.number: The value. + """ + return np.frombuffer(self.mmap, self.dtype, 1, self.offset)[0] + + def set(self, value: Number) -> None: + """Set value. + + Args: + value (Number): The value. + """ + cls = getattr(np, self.dtype.name) + self.mmap[self.offset:] = cls(value).tobytes() + + +class MemMapInteger(MemMapNumber): + """An integer backed by a mempry-mapped file.""" + + pass + + +class MemMapSignedInteger(MemMapInteger): + """A signed integer backecd by a memory-mapped file.""" + + pass + + +class MemMapInt8(MemMapSignedInteger): + """An int8 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__(self, filename: str, value: Optional[Number] = None) -> None: + super().__init__(filename, np.int8, value) + + +class MemMapInt16(MemMapSignedInteger): + """An int16 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__(self, filename: str, value: Optional[Number] = None) -> None: + super().__init__(filename, np.int16, value) + + +class MemMapInt32(MemMapSignedInteger): + """An int32 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__(self, filename: str, value: Optional[Number] = None) -> None: + super().__init__(filename, np.int32, value) + + +class MemMapInt64(MemMapSignedInteger): + """An int64 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__(self, filename: str, value: Optional[Number] = None) -> None: + super().__init__(filename, np.int64, value) + + +class MemMapUnsignedInteger(MemMapInteger): + """An unsigned integer backecd by a memory-mapped file.""" + + pass + + +class MemMapUInt8(MemMapUnsignedInteger): + """An int8 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__(self, filename: str, value: Optional[Number] = None) -> None: + super().__init__(filename, np.uint8, value) + + +class MemMapUInt16(MemMapUnsignedInteger): + """An int16 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__(self, filename: str, value: Optional[Number] = None) -> None: + super().__init__(filename, np.uint16, value) + + +class MemMapUInt32(MemMapUnsignedInteger): + """An int32 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__(self, filename: str, value: Optional[Number] = None) -> None: + super().__init__(filename, np.int32, value) + + +class MemMapUInt64(MemMapUnsignedInteger): + """An int64 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__(self, filename: str, value: Optional[Number] = None) -> None: + super().__init__(filename, np.uint64, value) + + +class MemMapInexact(MemMapNumber): + """An inexact number backed by a mempry-mapped file.""" + + pass + + +class MemMapFloating(MemMapInexact): + """A floating point number backecd by a memory-mapped file.""" + + pass + + +class MemMapFloat16(MemMapFloating): + """A float16 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__(self, filename: str, value: Optional[Number] = None) -> None: + super().__init__(filename, np.float16, value) + + +class MemMapFloat32(MemMapFloating): + """A float32 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__(self, filename: str, value: Optional[Number] = None) -> None: + super().__init__(filename, np.float32, value) + + +class MemMapFloat64(MemMapFloating): + """A float64 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + """ + + def __init__(self, filename: str, value: Optional[Number] = None) -> None: + super().__init__(filename, np.float64, value) + + +def int8(filename: str, value: Optional[Number] = None) -> MemMapInt8: + """Get a int8 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + + Returns: + MemMapInt8: A int8 backed by a memory-mapped file. + """ + return MemMapInt8(filename, value) + + +def int16(filename: str, value: Optional[Number] = None) -> MemMapInt16: + """Get a int16 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + + Returns: + MemMapInt16: A int16 backed by a memory-mapped file. + """ + return MemMapInt16(filename, value) + + +def int32(filename: str, value: Optional[Number] = None) -> MemMapInt32: + """Get a int32 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + + Returns: + MemMapInt32: A int32 backed by a memory-mapped file. + """ + return MemMapInt32(filename, value) + + +def int64(filename: str, value: Optional[Number] = None) -> MemMapInt64: + """Get a int64 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + + Returns: + MemMapInt64: A int64 backed by a memory-mapped file. + """ + return MemMapInt64(filename, value) + + +def uint8(filename: str, value: Optional[Number] = None) -> MemMapUInt8: + """Get a uint8 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + + Returns: + MemMapUInt8: A uint8 backed by a memory-mapped file. + """ + return MemMapUInt8(filename, value) + + +def uint16(filename: str, value: Optional[Number] = None) -> MemMapUInt16: + """Get a uint16 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + + Returns: + MemMapUInt16: A uint16 backed by a memory-mapped file. + """ + return MemMapUInt16(filename, value) + + +def uint32(filename: str, value: Optional[Number] = None) -> MemMapUInt32: + """Get a uint32 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + + Returns: + MemMapUInt32: A uint32 backed by a memory-mapped file. + """ + return MemMapUInt32(filename, value) + + +def uint64(filename: str, value: Optional[Number] = None) -> MemMapUInt64: + """Get a uint64 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + + Returns: + MemMapUInt64: A uint64 backed by a memory-mapped file. + """ + return MemMapUInt64(filename, value) + + +def float16(filename: str, value: Optional[Number] = None) -> MemMapFloat16: + """Get a float16 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + + Returns: + MemMapFloat16: A float16 backed by a memory-mapped file. + """ + return MemMapFloat16(filename, value) + + +def float32(filename: str, value: Optional[Number] = None) -> MemMapFloat32: + """Get a float32 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + + Returns: + MemMapFloat32: A float32 backed by a memory-mapped file. + """ + return MemMapFloat32(filename, value) + + +def float64(filename: str, value: Optional[Number] = None) -> MemMapFloat64: + """Get a float64 backed by a memory-mapped file. + + Args: + filename (str): File to memory map. + value (Number, optional): If a number, creates as this value. If ``None``, attaches. + Defaults to ``None``. + + Returns: + MemMapFloat64: A float64 backed by a memory-mapped file. + """ + return MemMapFloat64(filename, value) diff --git a/streaming/base/coord/process.py b/streaming/base/coord/process.py new file mode 100644 index 000000000..facedb420 --- /dev/null +++ b/streaming/base/coord/process.py @@ -0,0 +1,20 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Utility methods for coordination related to processes.""" + +from typing import Dict + +from psutil import process_iter + + +def get_live_processes() -> Dict[int, int]: + """Get the live processes by pid and creation time. + + Returns: + Dict[int, int]: Mapping of pid to creation time in integer nanoseconds. + """ + ret = {} + for obj in process_iter(['pid', 'create_time']): + ret[obj.pid] = int(obj.create_time() * 1e9) + return ret diff --git a/streaming/base/coord/shmem/__init__.py b/streaming/base/coord/shmem/__init__.py new file mode 100644 index 000000000..66a487542 --- /dev/null +++ b/streaming/base/coord/shmem/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Objects that live in shared memory. + +For when using `threading` or `multiprocessing` from the python standard library won't do, because +we are coordinating separately instantiated pytorch worker processes. +""" + +from streaming.base.coord.shmem.array import SharedArray as SharedArray +from streaming.base.coord.shmem.barrier import SharedBarrier as SharedBarrier +from streaming.base.coord.shmem.memory import SharedMemory as SharedMemory +from streaming.base.coord.shmem.prefix import _get_path as _get_path +from streaming.base.coord.shmem.prefix import get_shm_prefix as get_shm_prefix +from streaming.base.coord.shmem.scalar import SharedScalar as SharedScalar + +__all__ = ['SharedArray', 'SharedBarrier', 'SharedMemory', 'get_shm_prefix', 'SharedScalar'] diff --git a/streaming/base/shared/array.py b/streaming/base/coord/shmem/array.py similarity index 97% rename from streaming/base/shared/array.py rename to streaming/base/coord/shmem/array.py index 9b4d8b0c8..191b3f29d 100644 --- a/streaming/base/shared/array.py +++ b/streaming/base/coord/shmem/array.py @@ -8,7 +8,7 @@ import numpy as np from numpy.typing import NDArray -from streaming.base.shared.memory import SharedMemory +from streaming.base.coord.shmem.memory import SharedMemory class SharedArray: diff --git a/streaming/base/shared/barrier.py b/streaming/base/coord/shmem/barrier.py similarity index 58% rename from streaming/base/shared/barrier.py rename to streaming/base/coord/shmem/barrier.py index 7af279232..221b5a70b 100644 --- a/streaming/base/shared/barrier.py +++ b/streaming/base/coord/shmem/barrier.py @@ -9,10 +9,10 @@ from time import sleep import numpy as np -from filelock import FileLock from streaming.base.constant import TICK -from streaming.base.shared.array import SharedArray +from streaming.base.coord.file import SoftFileLock +from streaming.base.coord.shmem.array import SharedArray # Time out to wait before raising exception TIMEOUT = 60 @@ -26,33 +26,32 @@ class SharedBarrier: processes. Args: - filelock_path (str): Path to lock file on local filesystem. + lock_filename (str): Path to lock file on local filesystem. shm_name (str): Shared memory object name in /dev/shm. """ - def __init__(self, filelock_path: str, shm_name: str) -> None: + def __init__(self, lock_filename: str, shm_name: str) -> None: # Create lock. - self.filelock_path = filelock_path - self.lock = FileLock(self.filelock_path) + self._lock = SoftFileLock(lock_filename) # Create three int32 fields in shared memory: num_enter, num_exit, flag. self._arr = SharedArray(3, np.int32, shm_name) - self.num_enter = 0 - self.num_exit = -1 - self.flag = True + self._num_enter = 0 + self._num_exit = -1 + self._flag = True @property - def num_enter(self) -> int: - """Get property num_enter. + def _num_enter(self) -> int: + """Get property _num_enter. Returns: int: Number of processes that have entered the barrier. """ return self._arr[0] - @num_enter.setter - def num_enter(self, num_enter: int) -> None: - """Set property num_enter. + @_num_enter.setter + def _num_enter(self, num_enter: int) -> None: + """Set property _num_enter. Args: num_enter (int): Number of processes that have entered the barrier. @@ -60,17 +59,17 @@ def num_enter(self, num_enter: int) -> None: self._arr[0] = num_enter @property - def num_exit(self) -> int: - """Get property num_exit. + def _num_exit(self) -> int: + """Get property _num_exit. Returns: int: Number of processes that have exited the barrier. """ return self._arr[1] - @num_exit.setter - def num_exit(self, num_exit: int) -> None: - """Set property num_exit. + @_num_exit.setter + def _num_exit(self, num_exit: int) -> None: + """Set property _num_exit. Args: num_exit (int): Number of processes that have exited the barrier. @@ -78,17 +77,17 @@ def num_exit(self, num_exit: int) -> None: self._arr[1] = num_exit @property - def flag(self) -> bool: - """Get property flag. + def _flag(self) -> bool: + """Get property _flag. Returns: bool: The flag value. """ return bool(self._arr[2]) - @flag.setter - def flag(self, flag: bool) -> None: - """Set property flag. + @_flag.setter + def _flag(self, flag: bool) -> None: + """Set property _flag. Args: flag (bool): The flag value. @@ -101,36 +100,36 @@ def __call__(self, num_procs: int) -> None: Args: num_procs (int): How many processes are sharing this barrier. """ - # Initialize num_exit to the number of processes. - with self.lock: - if self.num_exit == -1: - self.num_exit = num_procs + # Initialize _num_exit to the number of processes. + with self._lock: + if self._num_exit == -1: + self._num_exit = num_procs # If we are the first to arrive, wait for everyone to exit, then set flag to "don't go". - self.lock.acquire() - if not self.num_enter: - self.lock.release() - while self.num_exit != num_procs: + self._lock.acquire() + if not self._num_enter: + self._lock.release() + while self._num_exit != num_procs: sleep(TICK) - self.lock.acquire() - self.flag = False + self._lock.acquire() + self._flag = False # Note that we entered. - self.num_enter += 1 + self._num_enter += 1 # If we are the last to arrive, reset `enter` and `exit`, and set flag to "go". - if self.num_enter == num_procs: - self.num_enter = 0 - self.num_exit = 0 - self.flag = True - self.lock.release() + if self._num_enter == num_procs: + self._num_enter = 0 + self._num_exit = 0 + self._flag = True + self._lock.release() # Everybody waits until the flag is set to "go". - while not self.flag: + while not self._flag: sleep(TICK) # Note that we exited. - with self.lock: - self.num_exit += 1 - if self.num_exit == num_procs: - self.num_exit = -1 + with self._lock: + self._num_exit += 1 + if self._num_exit == num_procs: + self._num_exit = -1 diff --git a/streaming/base/shared/memory.py b/streaming/base/coord/shmem/memory.py similarity index 100% rename from streaming/base/shared/memory.py rename to streaming/base/coord/shmem/memory.py diff --git a/streaming/base/shared/prefix.py b/streaming/base/coord/shmem/prefix.py similarity index 98% rename from streaming/base/shared/prefix.py rename to streaming/base/coord/shmem/prefix.py index 7e8936086..9887def0f 100644 --- a/streaming/base/shared/prefix.py +++ b/streaming/base/coord/shmem/prefix.py @@ -15,8 +15,8 @@ from torch import distributed as dist from streaming.base.constant import LOCALS, TICK -from streaming.base.shared import SharedMemory -from streaming.base.world import World +from streaming.base.coord.shmem import SharedMemory +from streaming.base.coord.world import World def _each_prefix_int() -> Iterator[int]: diff --git a/streaming/base/shared/scalar.py b/streaming/base/coord/shmem/scalar.py similarity index 93% rename from streaming/base/shared/scalar.py rename to streaming/base/coord/shmem/scalar.py index e1dbfcaf8..d18fae607 100644 --- a/streaming/base/shared/scalar.py +++ b/streaming/base/coord/shmem/scalar.py @@ -5,7 +5,7 @@ from typing import Any -from streaming.base.shared.array import SharedArray +from streaming.base.coord.shmem.array import SharedArray class SharedScalar: diff --git a/streaming/base/coord/waiting.py b/streaming/base/coord/waiting.py new file mode 100644 index 000000000..92a640630 --- /dev/null +++ b/streaming/base/coord/waiting.py @@ -0,0 +1,72 @@ +# Copyright 2022-2024 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Waiting on predicates.""" + +from contextlib import nullcontext +from time import sleep, time +from typing import Any, Callable, Optional + +__all__ = ['wait'] + + +def _say_duration(duration: float) -> str: + """Pretty-print a duration. + + Args: + duration (float): The duration as a float. + + Returns: + str: The duration as a str. + """ + return f'{duration:.3f}'.rstrip('0').rstrip('.') + + +def wait( + stop: Callable[[], bool], + timeout: Optional[float] = 30, + tick: float = 0.007, + lock: Optional[Any] = None, +) -> None: + """Wait for the predicate to succeed. + + Args: + stop (Callable[[], bool]): When this check returns True, you break out of the retry loop. + timeout (float, optional): How long to wait before raising an exception, in seconds. + Defaults to ``30``. + tick (float): Check interval, in seconds. Defaults to ``0.007``. + lock (Any, optional): Context manager (this is intended for locks) to be held when + checking the predicate. Defaults to ``None``. + """ + start = time() + + if timeout is not None and timeout <= 0: + raise ValueError(f'Timeout must be positive if provided, but got: ' + + f'{_say_duration(timeout)} sec.') + + if tick <= 0: + raise ValueError(f'Tick must be positive if provided, but got: {_say_duration(tick)} sec.') + + if lock is not None: + if not hasattr(lock, '__enter__'): + raise ValueError(f'Lock must support `__enter__`, but got: {lock}.') + + if not hasattr(lock, '__exit__'): + raise ValueError(f'Lock must support `__exit__`, but got: {lock}.') + + norm_lock = lock + else: + norm_lock = nullcontext() + + while True: + with norm_lock: + if stop(): + break + + if timeout is not None: + now = time() + if timeout <= now - start: + raise RuntimeError(f'Wait timed out: timeout {_say_duration(timeout)} sec vs ' + + f'elapsed {_say_duration(now - start)} sec.') + + sleep(tick) diff --git a/streaming/base/world.py b/streaming/base/coord/world.py similarity index 100% rename from streaming/base/world.py rename to streaming/base/coord/world.py diff --git a/streaming/base/dataloader.py b/streaming/base/dataloader.py index bb3ac6c3c..9068c481e 100644 --- a/streaming/base/dataloader.py +++ b/streaming/base/dataloader.py @@ -9,8 +9,8 @@ from torch.utils.data import DataLoader from transformers import BatchEncoding, BatchFeature +from streaming.base.coord.world import World from streaming.base.dataset import StreamingDataset -from streaming.base.world import World class StreamingDataLoader(DataLoader): diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index e9f5b96d6..89070def9 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -7,35 +7,32 @@ import logging import os import sys -import warnings from concurrent.futures import ThreadPoolExecutor, wait from concurrent.futures._base import Future from enum import IntEnum from math import ceil +from tempfile import gettempdir from threading import Event, Lock from time import sleep, time_ns from typing import Any, Dict, Iterator, Optional, Sequence, Tuple, Union +from warnings import warn import numpy as np -from filelock import FileLock from numpy.typing import NDArray -from torch import distributed as dist from torch.utils.data import IterableDataset from streaming.base.array import Array from streaming.base.batching import generate_work -from streaming.base.constant import (BARRIER, BARRIER_FILELOCK, CACHE_FILELOCK, CACHE_USAGE, - EPOCH_DATA, EPOCH_SHAPE, NEXT_EPOCH, RESUME, - SHARD_ACCESS_TIMES, SHARD_STATES, TICK) -from streaming.base.distributed import maybe_init_dist +from streaming.base.constant import TICK +from streaming.base.coord import mmap as mm +from streaming.base.coord.file import SoftFileLock +from streaming.base.coord.job import JobDirectory, JobRegistry +from streaming.base.coord.world import World from streaming.base.format import get_index_basename from streaming.base.sampling import get_sampling -from streaming.base.shared import (SharedArray, SharedBarrier, SharedMemory, SharedScalar, - _get_path, get_shm_prefix) from streaming.base.spanner import Spanner from streaming.base.stream import Stream from streaming.base.util import bytes_to_int, number_abbrev_to_int -from streaming.base.world import World # An arbitrary time in the future, used for cold shard eviction. NEVER = np.iinfo(np.uint64).max @@ -183,30 +180,34 @@ class StreamingDataset(Array, IterableDataset): "num_canonical_nodes": "int" } - StreamingDataset init takes two kinds of arguments: + StreamingDataset init takes two categories of arguments: - * What to iterate: + * What to iterate (the Stream arguments): - * One or more streams (you must provide either ``streams`` or ``remote``/``local``): + * Stream paths. To provide your own Streams, set ``streams`` and optionally ``epoch_size``. + To have StreamingDataset implicitly create one for you instead, set ``remote`` and/or + ``local``. + * ``epoch_size`` * ``streams`` * ``remote`` * ``local`` - * Knobs to control streaming behavior, which, if multiple streams are provided, - become defaults applied to each of them: + * Stream settings. These fields are all either set in Stream init, or else set by default + here in StreamingDataset init. * ``split`` * ``download_retry`` * ``download_timeout`` * ``validate_hash`` * ``keep_zip`` + * ``allow_unsafe_types`` - * Absolute dataset size, if streams were weighted relatively: + * How to iterate (the StreamingDataset arguments): - * ``epoch_size`` + * Configuration: - * How to iterate: + * ``config_root`` * Shard lifecycle: @@ -235,8 +236,12 @@ class StreamingDataset(Array, IterableDataset): * ``batching_method`` - Args: + epoch_size (Union[int, str], optional): Number of samples to draw per epoch balanced + across all streams. If ``None``, takes its value from the total number of underlying + samples. Provide this field if you are weighting streams relatively to target a larger + or smaller epoch size. Defaults to ``None``. Can also take in human-readable number + abbreviations (e.g., ``"100k"``, ``"64M"``, ``"77b"``, etc). Defaults to ``None``. streams (Sequence[Stream], optional): One or more streams to stream/cache samples from, which may be upsampled or downsampled. StreamingDataset uses either ``streams`` or ``remote``/``local``. Defaults to ``None``. @@ -256,11 +261,12 @@ class StreamingDataset(Array, IterableDataset): keep_zip (bool): Whether to keep or delete the compressed form when decompressing downloaded shards. If ``False``, keep iff remote is local or no remote. Defaults to ``False``. - epoch_size (Union[int, str], optional): Number of samples to draw per epoch balanced - across all streams. If ``None``, takes its value from the total number of underlying - samples. Provide this field if you are weighting streams relatively to target a larger - or smaller epoch size. Defaults to ``None``. Can also take in human-readable number - abbreviations (e.g., ``"100k"``, ``"64M"``, ``"77b"``, etc). Defaults to ``None``. + allow_unsafe_types (bool): If a shard contains Pickle, which allows arbitrary code + execution during deserialization, whether to keep going if ``True`` or raise an error + if ``False``. Defaults to ``False``. + config_root (str, optional): Streaming configuration root directory, used for collision + detection, filelock paths, etc. If ``None``, uses a ``/streaming/`` subdir under your + system's temp root. Defaults to ``None``. predownload (int, optional): Target number of samples to download per worker in advance of current sample. Workers will attempt to download ahead by this many samples during, but not before, training. Recommendation is to provide a value greater than per device @@ -303,49 +309,64 @@ class StreamingDataset(Array, IterableDataset): ``None``. batching_method (str): Which batching method to use, either ``random``, ``stratified``, or ``per_stream``. Defaults to ``random``. - allow_unsafe_types (bool): If a shard contains Pickle, which allows arbitrary code - execution during deserialization, whether to keep going if ``True`` or raise an error - if ``False``. Defaults to ``False``. """ - def __init__(self, - *, - streams: Optional[Sequence[Stream]] = None, - remote: Optional[str] = None, - local: Optional[str] = None, - split: Optional[str] = None, - download_retry: int = 2, - download_timeout: float = 60, - validate_hash: Optional[str] = None, - keep_zip: bool = False, - epoch_size: Optional[Union[int, str]] = None, - predownload: Optional[int] = None, - cache_limit: Optional[Union[int, str]] = None, - sampling_method: str = 'balanced', - sampling_granularity: int = 1, - partition_algo: str = 'relaxed', - num_canonical_nodes: Optional[int] = None, - batch_size: Optional[int] = None, - shuffle: bool = False, - shuffle_algo: str = 'py1e', - shuffle_seed: int = 9176, - shuffle_block_size: Optional[int] = None, - batching_method: str = 'random', - allow_unsafe_types: bool = False) -> None: - # Global arguments (which do not live in Streams). - self.predownload = predownload - self.cache_limit = cache_limit - self.sampling_method = sampling_method - self.sampling_granularity = sampling_granularity - self.partition_algo = partition_algo - self.num_canonical_nodes = num_canonical_nodes + def __init__( + self, + *, + epoch_size: Optional[Union[int, str]] = None, + streams: Optional[Sequence[Stream]] = None, + remote: Optional[str] = None, + local: Optional[str] = None, + split: Optional[str] = None, + download_retry: int = 2, + download_timeout: float = 60, + validate_hash: Optional[str] = None, + keep_zip: bool = False, + allow_unsafe_types: bool = False, + config_root: Optional[str] = None, + predownload: Optional[int] = None, + cache_limit: Optional[Union[int, str]] = None, + sampling_method: str = 'balanced', + sampling_granularity: int = 1, + partition_algo: str = 'relaxed', + num_canonical_nodes: Optional[int] = None, + batch_size: Optional[int] = None, + shuffle: bool = False, + shuffle_algo: str = 'py1e', + shuffle_seed: int = 9176, + shuffle_block_size: Optional[int] = None, + batching_method: str = 'random', + ) -> None: + # Initialize the World context. + # + # Beware: This information is for the per-rank process. DataLoader worker processes may see + # different values for these fields. We are saving the rank World here because we cannot + # instantiate a World inside the StreamingDataset destructor. + self._rank_world = world = World() + + # Purely StreamingDataset arguments (which do not live in Streams). + self.config_root = self._get_config_root(config_root) + self._test_config_root(self.config_root) + self.predownload = self._get_predownload(predownload, batch_size) + self.cache_limit = self._get_cache_limit(cache_limit) + self.sampling_method = self._get_sampling_method(sampling_method) + self.sampling_granularity = self._get_sampling_granularity(sampling_granularity) + self.partition_algo = self._get_partition_algo(partition_algo) + self.num_canonical_nodes: int self.batch_size = batch_size self.shuffle = shuffle - self.shuffle_algo = shuffle_algo - self.shuffle_seed = shuffle_seed - self.shuffle_block_size = shuffle_block_size - self.batching_method = batching_method - self.allow_unsafe_types = allow_unsafe_types + self.shuffle_algo = self._get_shuffle_algo(shuffle_algo) + self.shuffle_seed = self._get_shuffle_seed(shuffle_seed) + self.input_shuffle_block_size = shuffle_block_size + self.shuffle_block_size: int + self.batching_method = self._get_batching_method(batching_method) + + # Purely StreamingDataset arguments which depend on other such arguments. + self.num_canonical_nodes = self._get_num_canonical_nodes(num_canonical_nodes, + self.shuffle_algo, world) + self.shuffle_block_size = self._get_shuffle_block_size(self.input_shuffle_block_size, + self.num_canonical_nodes, world) # Initialize initial_physical_nodes to None. If we are resuming, then we will set it to the # number of physical nodes of the initial run in the _resume function. @@ -356,50 +377,6 @@ def __init__(self, raise ValueError( 'You must provide either `streams` or `remote`/`local`, but not both.') - # Check sampling method is one of "balanced" or "fixed". - if self.sampling_method not in ['balanced', 'fixed']: - raise ValueError( - f'Invalid sampling method: {sampling_method}. ' + \ - f'Must be one of `balanced` or `fixed`.' - ) - - # Check sampling granularity. - if self.sampling_granularity <= 0: - raise ValueError(f'`sampling_granularity` must be a positive integer, but got: ' + - f'{self.sampling_granularity}.') - - # Check batching method is one of "random", "stratified", or "per_stream". - if self.batching_method not in ['random', 'stratified', 'per_stream']: - raise ValueError( - f'Invalid batching method: {batching_method}. ' + \ - f'Must be one of `random`, `stratified`, or `per_stream.' - ) - - # issue deprecation warning for py1b shuffle algorithm. - if self.shuffle_algo == 'py1b': - warnings.warn('The \'py1b\' shuffle algorithm will soon be deprecated. \ - Please use the more performant \'py1br\' algorithm instead.', - DeprecationWarning, - stacklevel=2) - - # Check shuffle seed. - if self.shuffle_seed < 0: - raise ValueError(f'`shuffle_seed` must be a non-negative integer, but got: ' + - f'{self.shuffle_seed}.') - - # Check that predownload is at least per device batch size, and set it if currently `None`. - if self.predownload is not None and self.batch_size is not None and \ - self.predownload < self.batch_size: - warnings.warn(f'predownload < batch_size ({self.predownload} < {self.batch_size}).' + - f'This may result in slower batch time. Recommendation is to set ' + - f'predownload to at-least batch_size.') - elif self.predownload is None: - logger.warning(f'Because `predownload` was not specified, it will default to ' + - f'8*batch_size if batch_size is not None, otherwise 64. Prior to ' + - f'Streaming v0.7.0, `predownload` defaulted to ' + - f'max(batch_size, 256 * batch_size // num_canonical_nodes).') - self.predownload = 8 * self.batch_size if self.batch_size is not None else 64 - # Convert epoch size from string to int, if needed. Cannot be negative. epoch_size_value = None if epoch_size: @@ -407,48 +384,32 @@ def __init__(self, if epoch_size_value < 0: raise ValueError(f'Epoch size cannot be negative. Received {epoch_size_value}.') - # Initialize torch dist ourselves, if necessary. - destroy_dist = maybe_init_dist() - # Initialize the Stream defaults and normalize to a list of Streams. if streams: - default = { - 'remote': remote, - 'local': local, - 'split': split, - 'download_retry': download_retry, - 'download_timeout': download_timeout, - 'validate_hash': validate_hash, - 'keep_zip': keep_zip, - } for stream in streams: - stream.apply_default(default) + stream.apply_defaults(split=split, + download_retry=download_retry, + download_timeout=download_timeout, + validate_hash=validate_hash, + keep_zip=keep_zip, + allow_unsafe_types=allow_unsafe_types) else: - default = Stream(remote=remote, + streams = Stream(remote=remote, local=local, split=split, download_retry=download_retry, download_timeout=download_timeout, validate_hash=validate_hash, - keep_zip=keep_zip) - streams = [default] + keep_zip=keep_zip, + allow_unsafe_types=allow_unsafe_types), # Validate the stream weighting scheme (relative or absolute) to catch errors before we go # to the trouble of loading them. Stream.validate_weights(streams) - # Set streams. + # Download each stream's index, load their shards, and map streams <-> shards <-> samples. self.streams = streams self.num_streams = len(streams) - - # Initialize the World context. - # - # Beware: This information is for the per-rank process. DataLoader worker processes may see - # different values for these fields. We are saving the rank World here because we cannot - # instantiate a World inside the StreamingDataset destructor. - self._rank_world = world = World() - - # Download each stream's index, load their shards, and map streams <-> shards. self.num_samples = 0 self.shards = [] stream_per_shard = [] @@ -457,7 +418,7 @@ def __init__(self, self.sample_offset_per_stream = np.zeros(self.num_streams, np.int64) self.samples_per_stream = np.zeros(self.num_streams, np.int64) for stream_id, stream in enumerate(self.streams): - stream_shards = stream.get_shards(world, self.allow_unsafe_types) + stream_shards = stream.get_shards(world) num_stream_samples = sum(map(len, stream_shards)) if not num_stream_samples: index_filename = os.path.join(stream.local, stream.split, get_index_basename()) @@ -474,8 +435,6 @@ def __init__(self, # Check that cache limit is possible. if self.cache_limit: - if isinstance(self.cache_limit, str): - self.cache_limit = bytes_to_int(self.cache_limit) min_cache_usage = sum((stream.get_index_size() for stream in streams)) if self.cache_limit <= min_cache_usage: raise ValueError(f'Minimum cache usage ({min_cache_usage} bytes) is larger than ' + @@ -505,47 +464,46 @@ def __init__(self, # Length (__len__) is the resampled epoch size divided over the number of devices. self.length = ceil(self.epoch_size / world.num_ranks) - # Register/lookup our shared memory prefix and filelock root directory. - streams_local = [os.path.abspath(os.path.join(x.local, x.split)) for x in streams] - streams_remote = [ - os.path.join(x.remote, x.split) if x.remote is not None else None for x in streams - ] - self._shm_prefix_int, self._locals_shm = get_shm_prefix(streams_local, streams_remote, - world) - self._filelock_root = os.path.join(os.path.sep, 'tmp', 'streaming') - os.makedirs(self._filelock_root, exist_ok=True) - - # Create the shared memory-backed barrier, without its lock, which is unpickleable. - self._shared_barrier = SharedBarrier( - os.path.join(self._filelock_root, _get_path(self._shm_prefix_int, BARRIER_FILELOCK)), - _get_path(self._shm_prefix_int, BARRIER)) + # Init registry, then register this Streaming job. + self.registry = JobRegistry(self.config_root) + self.job = JobDirectory(self.registry, streams, world) + job_file = self.job.get_filename + + # Rank barrier. + self._rank_barrier = mm.barrier(world.is_local_leader, job_file('rank_barrier.npy'), + job_file('rank_barrier.lock')) + + # Worker barrier. + self._worker_barrier = mm.barrier(world.is_local_leader, job_file('worker_barrier.npy'), + job_file('worker_barrier.lock')) # Epoch counter. # # Note: we do not assume that the end of __iter__() will ever be reached, so we need to # increment the epoch counter at the start of __iter__() instead of at the end, so we need # to track what the next epoch is, not the current epoch. - self._next_epoch = SharedScalar(np.int64, _get_path(self._shm_prefix_int, NEXT_EPOCH)) + value = 0 if world.is_local_leader else None + self._next_epoch = mm.int64(job_file('epoch.npy'), value) - # Cache filelock. Protects downloading and evicting shards. - self._cache_filelock_path = os.path.join(self._filelock_root, - _get_path(self._shm_prefix_int, CACHE_FILELOCK)) - self._cache_filelock: FileLock + # Cache filelock. + # + # Protects downloading and evicting shards. + self._cache_lock = SoftFileLock(job_file('cache.lock')) # Cache usage in bytes. - self._cache_usage = SharedScalar(np.int64, _get_path(self._shm_prefix_int, CACHE_USAGE)) + self._cache_usage = mm.int64(job_file('cache_usage.npy'), value) # Shard states array. Tells if a shard is missing, downloading, or present (eviction # happens under the lock). - self._shard_states = SharedArray(self.num_shards, np.uint8, - _get_path(self._shm_prefix_int, SHARD_STATES)) + self._shard_states = mm.ndarray(job_file('shard_states.npy'), self.num_shards, np.uint8, + value) # Time of last access per shard. This is used to decide which shard(s) to evict when we run # out of space. - self._shard_access_times = SharedArray(self.num_shards, np.uint64, - _get_path(self._shm_prefix_int, SHARD_ACCESS_TIMES)) + self._shard_access_times = mm.ndarray(job_file('shard_access_times.npy'), self.num_shards, + np.uint64, value) - # Initialize shared memory objects. + # Initialize interprocess state. if world.is_local_leader: # Set initial epoch (before any resumption). self.next_epoch = 0 @@ -569,32 +527,257 @@ def __init__(self, self._shard_states[shard_id] = _ShardState.LOCAL if size else _ShardState.REMOTE self._shard_access_times[shard_id] = time_ns() - if dist.is_available() and dist.is_initialized(): - dist.barrier() + # These fields are set each __iter__(). + self._iterator: _Iterator # Tracks thread positions. + self._executor: ThreadPoolExecutor # Multi-threading. + self._event: Event # Exception handling. - if destroy_dist: - dist.destroy_process_group() + # Init is not done for anyone until all interprocess state is populated by local leader. + self._rank_barrier(world.ranks_per_node) - # Placeholder for a shared memory object where load_state_dict() saves its data to be - # picked up by __iter__(). - self._resume_shm: SharedMemory + @classmethod + def _test_config_root(cls, config_root: str) -> None: + """Validate that the provided config root is usable. - # Placeholder for an _Iterator which tracks state during __iter__(). - self._iterator: _Iterator + If you are unable to get root or 777 perms, you may encounter problems in registering your + Streaming jobs for collision detection, getting unique interprocess filelock paths, etc. + You can sort of get around this by changing config root to a directory you control, but + this may negatively impact collision detection. - # For exception handling in __iter__ threads. - self._executor: ThreadPoolExecutor - self._event: Event + Args: + config_root (str): Streaming configuration root directory. + """ + os.makedirs(config_root, exist_ok=True) + filename = os.path.join(config_root, 'test.txt') + try: + with open(filename, 'wb') as out: + out.write(b'') + except: + raise ValueError('Please provide a `config_root` dir that is writeable and readable.') - del self._shared_barrier.lock # Remote the lock that makes it unpickleable. + @classmethod + def _get_config_root(cls, config_root: Optional[str]) -> str: + """Get the default Streaming configuration root directory. - def __del__(self) -> None: - """Destructor, which releases its local working directories.""" - if hasattr(self, '_locals_shm'): - try: - self._locals_shm.buf[:4] = np.int32(0).tobytes() - except: - pass + Args: + config_root (str, optional): Config root, if explicitly provided. + + Returns: + str: Streaming configuration root directory. + """ + return os.path.join(gettempdir(), 'streaming') + + @classmethod + def _get_predownload(cls, predownload: Optional[int], batch_size: Optional[int]) -> int: + if predownload is not None: + if batch_size is not None and predownload < batch_size: + warn(f'`predownload` < `batch_size` ({predownload} < {batch_size}). This may ' + + f'result in slower batch time. The recommendation is to set `predownload` ' + + f'to at least `batch_size`.') + norm_predownload = predownload + else: + logger.warning(f'Because `predownload` was not specified, it will default to ' + + f'`8 * batch_size` if batch_size is not None, otherwise 64. Prior to ' + + f'Streaming v0.7.0, `predownload` defaulted to ' + + f'`max(batch_size, 256 * batch_size // num_canonical_nodes)`.') + if batch_size is None: + norm_predownload = 64 + else: + norm_predownload = 8 * batch_size + return norm_predownload + + @classmethod + def _get_cache_limit(cls, cache_limit: Optional[Union[int, str]]) -> Optional[int]: + """Get cache limit. + + Args: + cache_limit (int | str, optional): Input cache limit. + + Returns: + int, optional: Normalized cache limit. + """ + if cache_limit is not None: + if isinstance(cache_limit, str): + norm_cache_limit = bytes_to_int(cache_limit) + else: + norm_cache_limit = cache_limit + if norm_cache_limit <= 0: + raise ValueError(f'Cache limit, if set, must be positive, but got: ' + + f'{cache_limit} -> {norm_cache_limit}.') + else: + norm_cache_limit = cache_limit + return norm_cache_limit + + @classmethod + def _get_sampling_method(cls, sampling_method: str) -> str: + """Get sampling method. + + Args: + sampling_method (str): Input sampling method. + + Returns: + str: Normalized sampling method, + """ + methods = 'balanced', 'fixed' + + if sampling_method not in methods: + raise ValueError(f'`sampling_method` must be one of {sorted(methods)}, but got: ' + + f'{sampling_method}.') + + return sampling_method + + @classmethod + def _get_sampling_granularity(cls, sampling_granularity: int) -> int: + """Get sampling granularity. + + Args: + samping_granularity (int): Input sampling granularity. + + Returns: + int: Normalized sampling granularity. + """ + # Check sampling granularity. + if sampling_granularity < 1: + raise ValueError(f'`sampling_granularity` must be a positive integer, but got: ' + + f'{sampling_granularity}.') + + return sampling_granularity + + @classmethod + def _get_partition_algo(cls, partition_algo: str) -> str: + """Get partition algo. + + Args: + partition_algo (str): Input parittion algo. + + Returns: + str: Normalized partition algo. + """ + from streaming.base.partition import algos + + if partition_algo not in algos: + raise ValueError(f'`partition_algo` must be one of {sorted(algos)}, but got: ' + + f'{partition_algo}.') + + return partition_algo + + @classmethod + def _get_num_canonical_nodes(cls, num_canonical_nodes: Optional[int], shuffle_algo: str, + world: World) -> int: + """Get num canonical nodes. + + This method is called upon resume() (from iter) -- not init -- by some 2 of 3 code paths, + while the last one sets num canonical nodes directly from checkpoint state. + + Args: + num_canonical_nodes (int, optional): Input num canonical nodes. + shuffle_algo (str): Shuffle algo. + world (World): Our place in the world. + + Returns: + int: Normalized num canonical nodes. + """ + if num_canonical_nodes is not None: + if num_canonical_nodes < 1: + raise ValueError('`num_canonical_nodes`, if provided, must be a positive integer.') + norm_num_canonical_nodes = num_canonical_nodes + else: + if shuffle_algo in {'py1s', 'py2s'}: + norm_num_canonical_nodes = 64 * world.num_nodes + else: + if world.is_local_leader: + logger.warning( + f'Because `num_canonical_nodes` was not specified, and `shuffle_algo` ' + + f'is {shuffle_algo}, it will default to be equal to the number of ' + + f'physical nodes. Prior to Streaming v0.7.0, `num_canonical_nodes` ' + + f'defaulted to `64 * physical nodes`.') + norm_num_canonical_nodes = world.num_nodes + return norm_num_canonical_nodes + + @classmethod + def _get_shuffle_algo(cls, shuffle_algo: str) -> str: + """Get shuffle algo. + + Args: + shuffle_algo (str): Input shuffle algo. + + Returns: + str: Normalized shuffle algo. + """ + from streaming.base.shuffle import algos + + if shuffle_algo not in algos: + raise ValueError(f'`shuffle_algo` must be one of {sorted(algos)}, but got: ' + + f'{shuffle_algo}.') + elif shuffle_algo == 'py1b': + logger.warning('The `py1b` shuffle algorithm will soon be deprecated. Please use ' + + 'the more performant `py1br` algorithm instead.', + DeprecationWarning, + stacklevel=2) + + return shuffle_algo + + @classmethod + def _get_shuffle_seed(cls, shuffle_seed: int) -> int: + """Get shuffle seed. + + Args: + shuffle_seed (int): Input shuffle seed. + + Returns: + int: Normalized shuffle seed. + """ + # Check shuffle seed. + if not (0 <= shuffle_seed < 2**32): + raise ValueError(f'`shuffle_seed` must be in `0 <= x < 2**32`, but got: ' + + f'{shuffle_seed}.') + + return shuffle_seed + + @classmethod + def _get_shuffle_block_size(cls, shuffle_block_size: Optional[int], num_canonical_nodes: int, + world: World) -> int: + """Get shuffle block size. + + This method is called upon resume() (from iter) -- not init -- because resuming sets the + official number of canonical nodes, which we depend on. + + Args: + shuffle_block_size (int, optional): Input shuffle block size. + num_canonical_nodes (int): Number of canonical nodes. + world (World): Our place in the world. + + Returns: + int: Normalized shuffle block size. + """ + if shuffle_block_size is not None: + norm_shuffle_block_size = shuffle_block_size + else: + if world.is_local_leader: + logger.warning(f'Because `shuffle_block_size` was not specified, it will ' + + f'default to `max(4_000_000 // num_canonical_nodes, 1 << 18)` if ' + + f'`num_canonical_nodes` is not None, otherwise 262144. Prior to ' + + f'Streaming v0.7.0, `shuffle_block_size` defaulted to 262144.') + norm_shuffle_block_size = max(4_000_000 // num_canonical_nodes, 1 << 18) + return norm_shuffle_block_size + + @classmethod + def _get_batching_method(cls, batching_method: str) -> str: + """Get batching method. + + Args: + batching_method (str): Input batching method. + + Returns: + str: Normalized batching method. + """ + from streaming.base.batching import batching_methods + + if batching_method not in batching_methods: + raise ValueError(f'`batching_method` must be one of {sorted(batching_methods)}, but ' + + f'got: {batching_method}.') + + return batching_method @property def size(self) -> int: @@ -649,17 +832,6 @@ def __len__(self) -> int: """ return self.length - def _set_shuffle_block_size(self, world: World): - """Set the shuffle block size value.""" - if self.shuffle_block_size is None: - if not world.worker_of_rank: - logger.warning(f'Because `shuffle_block_size` was not specified, it will ' + - f'default to max(4_000_000 // num_canonical_nodes, 1 << 18) if ' + - f'num_canonical_nodes is not None, otherwise 262144. Prior to ' + - f'Streaming v0.7.0, `shuffle_block_size` defaulted to 262144.') - self.shuffle_block_size = max(4_000_000 // self.num_canonical_nodes, 1 << 18) \ - if self.num_canonical_nodes is not None else 1 << 18 - def _resume(self, world: World, epoch: int) -> Tuple[int, int]: """Either resume from checkpoint or start at the beginning. @@ -670,60 +842,25 @@ def _resume(self, world: World, epoch: int) -> Tuple[int, int]: Returns: Tuple[int, int]: What epoch this is, and sample offset in that epoch. """ - # Get the resume state, if it exists. - name = _get_path(self._shm_prefix_int, RESUME) - try: - shm = SharedMemory(name=name, create=False) - except FileNotFoundError: - # There is nothing to resume. - if not self.num_canonical_nodes: - if self.shuffle_algo in ['py1s', 'py2s']: - self.num_canonical_nodes = 64 * world.num_nodes - else: - if not world.worker_of_rank: - logger.warning( - f'Because `num_canonical_nodes` was not specified, and ' + - f'`shuffle_algo` is {self.shuffle_algo}, it will default to ' + - f'be equal to physical nodes. Prior to Streaming ' + - f'v0.7.0, `num_canonical_nodes` defaulted to 64 * physical ' + - f'nodes.') - self.num_canonical_nodes = world.num_nodes - self._set_shuffle_block_size(world) + # If there is no checkpoint, bail. + filename = self.job.get_filename('checkpoint.json') + if not os.path.exists(filename): return epoch, 0 - # SharedMemory buffers may contain additional null bytes at the end. - buf = bytes(shm.buf) - index = buf.find(b'\0') - buf = buf[:index] if index != -1 else buf - obj = json.loads(buf.decode('utf-8')) - - # Check if the resume state is stale. + # If the checkpoint is stale, bail. + obj = json.load(open(filename)) if obj['epoch'] < epoch: - if not self.num_canonical_nodes: - if self.shuffle_algo in ['py1s', 'py2s']: - self.num_canonical_nodes = 64 * world.num_nodes - else: - if not world.worker_of_rank: - logger.warning( - f'Because `num_canonical_nodes` was not specified, and ' + - f'`shuffle_algo` is {self.shuffle_algo}, it will default to ' + - f'be equal to physical nodes. Prior to Streaming ' + - f'v0.7.0, `num_canonical_nodes` defaulted to 64 * physical ' + - f'nodes.') - self.num_canonical_nodes = world.num_nodes - self._set_shuffle_block_size(world) return epoch, 0 - # Load the correct resumption meta data. + # Load from checkpoint. epoch = obj['epoch'] sample_in_epoch = obj['sample_in_epoch'] - self.num_canonical_nodes = obj['num_canonical_nodes'] self.shuffle_seed = obj['shuffle_seed'] - # Ensure that we are backwards compatible with old checkpoint dataset state, since the - # 'initial_physical_nodes' key may not be present. - self.initial_physical_nodes = obj.get('initial_physical_nodes', None) - self._set_shuffle_block_size(world) - + # Using get() for backwards compatibility as older versions of Streaming did not have this. + self.initial_physical_nodes = obj.get('initial_physical_nodes') + self.num_canonical_nodes = obj['num_canonical_nodes'] + self.shuffle_block_size = self._get_shuffle_block_size(self.input_shuffle_block_size, + self.num_canonical_nodes, world) return epoch, sample_in_epoch def _resume_incr_epoch(self, world: World) -> Tuple[int, int]: @@ -735,17 +872,12 @@ def _resume_incr_epoch(self, world: World) -> Tuple[int, int]: Returns: Tuple[int, int]: What epoch this is, and sample offset in that epoch. """ - # Lazily create the shared barrier's FileLock, which contains a threading Lock, which is - # unpickleable. - if not hasattr(self._shared_barrier, 'lock'): - self._shared_barrier.lock = FileLock(self._shared_barrier.filelock_path) - # Either resume from checkpoint, or start from scratch. presumed_epoch = self.next_epoch epoch, sample_in_epoch = self._resume(world, presumed_epoch) # Wait for everyone to get the epoch above. - self._shared_barrier(world.workers_per_node) + self._worker_barrier(world.workers_per_node) # Set the new next epoch. if world.is_local_leader: @@ -803,13 +935,9 @@ def load_state_dict(self, obj: Dict[str, Any]) -> None: Args: obj (Dict[str, Any]): The state. """ - name = _get_path(self._shm_prefix_int, RESUME) - data = json.dumps(obj, sort_keys=True).encode('utf-8') - # Some platforms choose to allocate chunks of memory based upon that platform's memory page - # size, hence the exact size of the shared memory block that was returned may be larger - # than what was requested. - self._resume_shm = SharedMemory(name=name, size=len(data)) - self._resume_shm.buf[:len(data)] = data + filename = self.job.get_filename('checkpoint.json') + with open(filename, 'w') as out: + json.dump(obj, out, sort_keys=True) def resample_streams( self, @@ -888,59 +1016,6 @@ def resample_streams( sample_ids = np.concatenate(sample_ids).astype(np.int64) return shuffle_units, sample_ids - def _share_work(self, sample_ids: NDArray[np.int64]) -> Tuple[SharedMemory, SharedMemory]: - """Put an epoch's sample ordering into shared memory. - - Args: - sample_ids (NDArray[np.int64]): Sample IDs. - - Returns: - Tuple[SharedMemory, SharedMemory]: Shared memory arrays containing shape and data. - """ - ndim = 5 - - # Validate shape. - if sample_ids.ndim != ndim: - raise ValueError(f'Sample IDs must be of {ndim}D shape (num physical nodes, ' + - f'ranks per node, workers per rank, batches per worker, ' + - f'batch size). Instead, found as {sample_ids.ndim}D shape.') - - # Save the generated epoch shape to shared memory. - name = _get_path(self._shm_prefix_int, EPOCH_SHAPE) - size = ndim * np.int64().nbytes - shape_shm = SharedMemory(name=name, create=True, size=size, auto_cleanup=False) - shape_shm.buf[:size] = np.array(sample_ids.shape, np.int64).tobytes() - - # Save the generated epoch data to shared memory. - name = _get_path(self._shm_prefix_int, EPOCH_DATA) - size = sample_ids.size * np.int64().nbytes - data_shm = SharedMemory(name=name, create=True, size=size, auto_cleanup=False) - data_shm.buf[:size] = sample_ids.tobytes() - - return shape_shm, data_shm - - def _attach_work(self) -> Tuple[NDArray[np.int64], SharedMemory, SharedMemory]: - """Get an epoch's sample ordering from shared memory. - - Returns: - NDArray[np.int64]: Sample IDs. - """ - ndim = 5 - - # Load the generated epoch shape from shared memory. - name = _get_path(self._shm_prefix_int, EPOCH_SHAPE) - size = ndim * np.int64().nbytes - shape_shm = SharedMemory(name=name, create=False, size=size, auto_cleanup=False) - shape = tuple(np.ndarray(5, buffer=shape_shm.buf, dtype=np.int64)) - - # Attach to the generated epoch data in shared memory. - name = _get_path(self._shm_prefix_int, EPOCH_DATA) - size = int(np.prod(shape)) * np.int64().nbytes - data_shm = SharedMemory(name=name, create=False, size=size, auto_cleanup=False) - sample_ids = np.ndarray(shape, buffer=data_shm.buf, dtype=np.int64) - - return sample_ids, shape_shm, data_shm - def _get_work(self, world: World, epoch: int, sample_in_epoch: int) -> NDArray[np.int64]: """Get this worker's partition of this epoch's sample space. @@ -952,37 +1027,37 @@ def _get_work(self, world: World, epoch: int, sample_in_epoch: int) -> NDArray[n Returns: Optional[NDArray[np.int64]]: Our partition of the epoch. """ - # Lazily create the shared barrier's FileLock, which contains a threading Lock, which is - # unpickleable. - if not hasattr(self._shared_barrier, 'lock'): - self._shared_barrier.lock = FileLock(self._shared_barrier.filelock_path) + filename = self.job.get_filename('epoch_sample_ids.npy') - # Do expensive work that may use a lot of cores/memory just once, in the local leader. if world.is_local_leader: epoch_sample_ids = generate_work(self.batching_method, self, world, epoch, sample_in_epoch) - shape_shm, data_shm = self._share_work(epoch_sample_ids) - self._shared_barrier(world.workers_per_node) + io = mm.ndarray(filename, value=epoch_sample_ids) + self._worker_barrier(world.workers_per_node) else: - self._shared_barrier(world.workers_per_node) - epoch_sample_ids, shape_shm, data_shm = self._attach_work() + self._worker_barrier(world.workers_per_node) + io = mm.ndarray(filename) + epoch_sample_ids = io.numpy() - # Each worker gets their portion of the work. worker_sample_ids = epoch_sample_ids[world.node, world.rank_of_node, world.worker_of_rank].flatten() - self._shared_barrier(world.workers_per_node) + self._worker_barrier(world.workers_per_node) + + if not world.is_local_leader: + io.close() - # Now clean up after ourselves. - shape_shm.cleanup() - data_shm.cleanup() + self._worker_barrier(world.workers_per_node) + + if world.is_local_leader: + io.delete() return worker_sample_ids def _evict_shard(self, shard_id: int) -> None: """Evict the given shard. - Assumes you hold ``_cache_filelock``, preventing anyone else from modifying the cache. We + Assumes you hold ``_cache_lock``, preventing anyone else from modifying the cache. We expect that shard deletions are very fast. This method is called internally by ``prepare_shard`` to clear space for more downloads. @@ -1047,12 +1122,7 @@ def evict_shard(self, shard_id: int) -> None: Args: shard_id (int): Shard to evict. """ - # Lock the cache. FileLocks contain threading Locks, which are not pickleable, which is - # incompatible with spawn, so must be created lazily. - if not hasattr(self, CACHE_FILELOCK): - self._cache_filelock = FileLock(self._cache_filelock_path) - - with self._cache_filelock: + with self._cache_lock: self._evict_shard(shard_id) def evict_coldest_shard(self) -> None: @@ -1060,12 +1130,7 @@ def evict_coldest_shard(self) -> None: This method is multithread/multiprocess-safe. """ - # Lock the cache. FileLocks contain threading Locks, which are not pickleable, which is - # incompatible with spawn, so must be created lazily. - if not hasattr(self, CACHE_FILELOCK): - self._cache_filelock = FileLock(self._cache_filelock_path) - - with self._cache_filelock: + with self._cache_lock: self._evict_coldest_shard() def prepare_shard(self, shard_id: int, blocking: bool = True) -> None: @@ -1081,12 +1146,7 @@ def prepare_shard(self, shard_id: int, blocking: bool = True) -> None: blocking (bool): Whether to wait or skip if the shard is currently being downloaded by someone else. """ - # Lock the cache. FileLocks contain threading Locks, which are not pickleable, which is - # incompatible with spawn, so must be created lazily. - if not hasattr(self, CACHE_FILELOCK): - self._cache_filelock = FileLock(self._cache_filelock_path) - lock = self._cache_filelock - lock.acquire() + self._cache_lock.acquire() # Get the state of the shard to download. state = self._shard_states[shard_id] @@ -1110,21 +1170,20 @@ def prepare_shard(self, shard_id: int, blocking: bool = True) -> None: self._evict_coldest_shard() # With the above preamble done, we can release the cache lock. - lock.release() + self._cache_lock.release() # Perform the download (shard will not be modified by others in PREPARING state). delta = stream.prepare_shard(shard) # Download completed, so note the time and transition shard state to LOCAL. - lock.acquire() - self.cache_usage += delta - self._shard_access_times[shard_id] = time_ns() - self._shard_states[shard_id] = _ShardState.LOCAL - lock.release() + with self._cache_lock: + self.cache_usage += delta + self._shard_access_times[shard_id] = time_ns() + self._shard_states[shard_id] = _ShardState.LOCAL elif state == _ShardState.PREPARING: # Someone else is currently downloading the shard. Release the lock for others to make # progress. - lock.release() + self._cache_lock.release() # Do we wait on them? if blocking: @@ -1146,16 +1205,16 @@ def prepare_shard(self, shard_id: int, blocking: bool = True) -> None: raw_filename = os.path.join(stream.local, stream.split, raw_info.basename) # Find raw. if not os.path.isfile(raw_filename): # Is raw missing? self._shard_states[shard_id] = _ShardState.PREPARING # Lock the shard. - lock.release() # Unblock other workers. + self._cache_lock.release() # Unblock other workers. delta = stream.prepare_shard(shard) # Decompress and remove zip. - lock.acquire() # Briefly take the lock back. + self._cache_lock.acquire() # Briefly take the lock back. self._shard_states[shard_id] = _ShardState.LOCAL # Restore shard state. self.cache_usage += delta # Update accounting. self._shard_access_times[shard_id] = time_ns() # Touch the shard. - lock.release() + self._cache_lock.release() else: # Unknown state. - lock.release() + self._cache_lock.release() raise RuntimeError(f'Invalid shard state: {state}') def get_item(self, sample_id: int, retry: int = 7) -> Any: diff --git a/streaming/base/shared/__init__.py b/streaming/base/shared/__init__.py deleted file mode 100644 index 0bc6c0e2e..000000000 --- a/streaming/base/shared/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright 2022-2024 MosaicML Streaming authors -# SPDX-License-Identifier: Apache-2.0 - -"""Objects that live in shared memory. - -For when using `threading` or `multiprocessing` from the python standard library won't do, because -we are coordinating separately instantiated pytorch worker processes. -""" - -from streaming.base.shared.array import SharedArray as SharedArray -from streaming.base.shared.barrier import SharedBarrier as SharedBarrier -from streaming.base.shared.memory import SharedMemory as SharedMemory -from streaming.base.shared.prefix import _get_path as _get_path -from streaming.base.shared.prefix import get_shm_prefix as get_shm_prefix -from streaming.base.shared.scalar import SharedScalar as SharedScalar - -__all__ = ['SharedArray', 'SharedBarrier', 'SharedMemory', 'get_shm_prefix', 'SharedScalar'] diff --git a/streaming/base/stream.py b/streaming/base/stream.py index 2eb806c17..f5bba8744 100644 --- a/streaming/base/stream.py +++ b/streaming/base/stream.py @@ -15,12 +15,12 @@ from streaming.base.compression import decompress from streaming.base.constant import TICK +from streaming.base.coord.world import World from streaming.base.distributed import barrier, get_local_rank from streaming.base.format import FileInfo, Reader, get_index_basename, reader_from_json from streaming.base.hashing import get_hash from streaming.base.storage import download_file from streaming.base.util import retry, wait_for_file_to_exist -from streaming.base.world import World class Stream: @@ -86,6 +86,10 @@ class Stream: keep_zip (bool, optional): Whether to keep or delete the compressed form when decompressing downloaded shards. If ``False``, keep if and only if remote is local or no remote. Defaults to ``None``. + allow_unsafe_types (bool, optional): If a shard contains Pickle, which allows arbitrary + code execution during deserialization, whether to keep going if ``True`` or raise an + error if ``False``. Inherits from its owning StreamingDataset if ``None``. Defaults to + ``None``. """ def __init__(self, @@ -99,7 +103,8 @@ def __init__(self, download_retry: Optional[int] = None, download_timeout: Optional[float] = None, validate_hash: Optional[str] = None, - keep_zip: Optional[bool] = None) -> None: + keep_zip: Optional[bool] = None, + allow_unsafe_types: Optional[bool] = None) -> None: self.remote = remote self._local = local self.split = split or '' @@ -161,6 +166,10 @@ def __init__(self, self.keep_zip = keep_zip self.safe_keep_zip = self.keep_zip or self.remote in {None, self.local} + self._allow_unsafe_types = allow_unsafe_types + if allow_unsafe_types is not None: + self.allow_unsafe_types = allow_unsafe_types + def _get_temporary_directory(self) -> str: """Construct a path to a temporary directory based on remote and split.""" root = tempfile.gettempdir() @@ -169,28 +178,43 @@ def _get_temporary_directory(self) -> str: hash = hashlib.blake2s(self.remote.encode('utf-8'), digest_size=16).hexdigest() return os.path.join(root, hash, self.split) - def apply_default(self, default: dict) -> None: + def apply_defaults(self, *, split: Optional[str], download_retry: int, download_timeout: float, + validate_hash: Optional[str], keep_zip: bool, + allow_unsafe_types: bool) -> None: """Apply defaults, setting any unset fields. We use pairs of (name, _name) in order to make type checking happy. Args: - default (Self): Stream containing default values for all optional fields. + split (str, optional): Which dataset split to use, if any. If provided, we stream + from/to the ``split`` subdirs of ``remote`` and ``local``. + download_retry (int): Number of download re-attempts before giving up. + download_timeout (float): Number of seconds to wait for a shard to download before + raising an exception. + validate_hash (str, optional): Optional hash or checksum algorithm to use to validate + shards. + keep_zip (bool): Whether to keep or delete the compressed form when decompressing + downloaded shards. If ``False``, keep iff remote is local or no remote. + allow_unsafe_types (bool): If a shard contains Pickle, which allows arbitrary code + execution during deserialization, whether to keep going if ``True`` or raise an + error if ``False``. """ if not (self.remote or self._local): raise ValueError('`remote` and/or `local` path must be provided') if not self.split: - self.split = default['split'] or '' + self.split = split or '' if self._download_retry is None: - self.download_retry = default['download_retry'] + self.download_retry = download_retry if self._download_timeout is None: - self.download_timeout = default['download_timeout'] + self.download_timeout = download_timeout if self.validate_hash is None: - self.validate_hash = default['validate_hash'] or None + self.validate_hash = validate_hash if self._keep_zip is None: - self.keep_zip = default['keep_zip'] - self.safe_keep_zip = default['keep_zip'] or self.remote in {None, self.local} + self.keep_zip = keep_zip + self.safe_keep_zip = keep_zip or self.remote in {None, self.local} + if self._allow_unsafe_types is None: + self.allow_unsafe_types = allow_unsafe_types @classmethod def validate_weights(cls, streams: Sequence[Self]) -> Tuple[bool, bool]: @@ -421,18 +445,18 @@ def prepare_shard(self, shard: Reader) -> int: delta += self._prepare_shard_part(raw_info, zip_info, shard.compression) return delta - def get_shards(self, world: World, allow_unsafe_types: bool) -> List[Reader]: + def get_shards(self, world: World) -> List[Reader]: """Load this Stream's index, retrieving its shard readers. Args: world (World): Distributed context. - allow_unsafe_types (bool): If a shard contains Pickle, which allows arbitrary code - execution during deserialization, whether to keep going if ``True`` or raise an - error. Returns: `List[Reader]: Shard readers. """ + if self.allow_unsafe_types is None: + raise RuntimeError('`allow_unsafe_types` was not provided.') + # Download the index file if it does not exist locally. basename = get_index_basename() filename = os.path.join(self.local, self.split, basename) # pyright: ignore @@ -472,7 +496,7 @@ def get_shards(self, world: World, allow_unsafe_types: bool) -> List[Reader]: shards = [] for info in obj['shards']: shard = reader_from_json(self.local, self.split, info) - shard.validate(allow_unsafe_types) + shard.validate(self.allow_unsafe_types) shards.append(shard) return shards diff --git a/streaming/base/util.py b/streaming/base/util.py index 7cdd15628..e89b901ce 100644 --- a/streaming/base/util.py +++ b/streaming/base/util.py @@ -21,9 +21,9 @@ import torch.distributed as dist from streaming.base.constant import SHM_TO_CLEAN +from streaming.base.coord.shmem.prefix import _get_path from streaming.base.distributed import get_local_rank, maybe_init_dist from streaming.base.format.index import get_index_basename -from streaming.base.shared.prefix import _get_path logger = logging.getLogger(__name__) diff --git a/tests/test_barrier.py b/tests/test_barrier.py index 6c5873d7f..1abc44899 100644 --- a/tests/test_barrier.py +++ b/tests/test_barrier.py @@ -11,7 +11,7 @@ import pytest -from streaming.base.shared import SharedArray, SharedBarrier +from streaming.base.coord.shmem import SharedArray, SharedBarrier class TestSharedBarrier: @@ -22,21 +22,21 @@ def test_params(self, filelock_path: str, shm_name: str): barrier = SharedBarrier(filelock_path, shm_name) assert isinstance(barrier._arr, SharedArray) assert barrier._arr.shape == (3,) - assert barrier.num_enter == 0 - assert barrier.num_exit == -1 - assert barrier.flag is True + assert barrier._num_enter == 0 + assert barrier._num_exit == -1 + assert barrier._flag is True @pytest.mark.parametrize('num_enter', [3, 10]) @pytest.mark.parametrize('num_exit', [4, 9]) @pytest.mark.parametrize('flag', [True, False]) def test_setter_getter(self, num_enter: int, num_exit: int, flag: bool): barrier = SharedBarrier('/tmp/dir/filelock_path', 'barrier_shm_name') - barrier.num_enter = num_enter - assert barrier.num_enter == num_enter - barrier.num_exit = num_exit - assert barrier.num_exit == num_exit - barrier.flag = flag - assert barrier.flag == flag + barrier._num_enter = num_enter + assert barrier._num_enter == num_enter + barrier._num_exit = num_exit + assert barrier._num_exit == num_exit + barrier._flag = flag + assert barrier._flag == flag def run(self, num_process: int, barrier: Any, shared_list: ListProxy): sleep(random()) diff --git a/tests/test_reader.py b/tests/test_reader.py index 48957da9f..083aee848 100644 --- a/tests/test_reader.py +++ b/tests/test_reader.py @@ -337,6 +337,5 @@ def test_predownload_batch_size_warning(local_remote_dir: Any): num_samples=117, size_limit=1 << 8) with pytest.warns(UserWarning, - match='predownload < batch_size.*This may result in slower ' + - 'batch time. Recommendation is to set'): + match='This may result in slower batch time. The recommendation is to set'): _ = StreamingDataset(local=local_dir, remote=remote_dir, predownload=4, batch_size=8) diff --git a/tests/test_shared.py b/tests/test_shared.py index 44d376d93..2aa15a63e 100644 --- a/tests/test_shared.py +++ b/tests/test_shared.py @@ -5,8 +5,8 @@ import pytest -from streaming.base.shared import get_shm_prefix -from streaming.base.world import World +from streaming.base.coord.shmem import get_shm_prefix +from streaming.base.coord.world import World @pytest.mark.usefixtures('local_remote_dir') diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 7ef98dfec..494f92b68 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -318,7 +318,7 @@ def test_dataloader_stratified_batching_user_set(local_remote_dir: Tuple[str, @pytest.mark.parametrize('stream_2_size', list(range(1, 65, 10))) @pytest.mark.usefixtures('local_remote_dir') -def test_stratified_batching_Exception(local_remote_dir: Tuple[str, str], stream_2_size: int): +def test_stratified_batching_exception(local_remote_dir: Tuple[str, str], stream_2_size: int): local, remote = local_remote_dir local1 = os.path.join(local, 'stream1') @@ -631,7 +631,7 @@ def test_dataloader_single_device(local_remote_dir: Tuple[str, str], batch_size: @pytest.mark.parametrize('shuffle', [True]) @pytest.mark.parametrize('sampling_method', ['balanfixed', 'fixedd', '', 'random', 'ayo']) @pytest.mark.usefixtures('local_remote_dir') -def test_sampling_method_invalid_Exception(local_remote_dir: Any, batch_size: int, seed: int, +def test_sampling_method_invalid_exception(local_remote_dir: Any, batch_size: int, seed: int, shuffle: bool, sampling_method: str): remote_dir, local_dir = local_remote_dir convert_to_mds(out_root=remote_dir, @@ -639,7 +639,7 @@ def test_sampling_method_invalid_Exception(local_remote_dir: Any, batch_size: in num_samples=117, size_limit=1 << 8) - with pytest.raises(ValueError, match=f'Invalid sampling method:*'): + with pytest.raises(ValueError): _ = StreamingDataset(local=local_dir, remote=remote_dir, shuffle=shuffle, @@ -782,6 +782,7 @@ def test_streamingdataloader_mid_epoch_resumption(local_remote_dir: Any, batch_s sample_order.extend(batch['id'][:]) del dataloader + del dataset.job del dataset clean_stale_shared_memory() @@ -861,6 +862,10 @@ def test_multiple_dataset_instantiation(local_remote_dir: Any, shuffle_seed: tup assert len(set(train_sample_order)) == len(set(val_sample_order)), 'Duplicate samples' +@pytest.mark.skip('Even though a streaming dataset is local (has no remote), we cannot draw ' + + 'conclusions about what exact phases of its files are present and would ' + + 'require prepare work (e.g., unzipping) for use, which would have to be ' + + 'managed in one place, so this test is sadly invalid.') def test_same_local_no_remote(local_remote_dir: Tuple[str, str]): local_0, _ = local_remote_dir convert_to_mds(out_root=local_0, @@ -893,5 +898,5 @@ def test_same_local_diff_remote(local_remote_dir: Tuple[str, str]): # Build StreamingDataset _ = StreamingDataset(local=local_0, remote=remote_0, batch_size=4, num_canonical_nodes=1) # Build StreamingDataset - with pytest.raises(ValueError, match='Reused local directory.*vs.*Provide a different one.'): + with pytest.raises(ValueError): _ = StreamingDataset(local=local_0, remote=remote_1, batch_size=2, num_canonical_nodes=1) diff --git a/tests/test_util.py b/tests/test_util.py index cbc171191..64cd9989e 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -12,7 +12,7 @@ import pytest from streaming.base.constant import RESUME -from streaming.base.shared.prefix import _get_path +from streaming.base.coord.shmem.prefix import _get_path from streaming.base.storage.download import download_file from streaming.base.storage.upload import CloudUploader from streaming.base.util import (bytes_to_int, clean_stale_shared_memory, get_list_arg,