From 7867b107b774ea568a085293f9654f559d508843 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Mon, 25 Dec 2023 23:46:48 -0800 Subject: [PATCH 01/54] Move epoch_size arg. --- streaming/base/dataset.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index a3453c570..b22aeb434 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -237,6 +237,11 @@ class StreamingDataset(Array, IterableDataset): 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,6 @@ 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``. 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 @@ -310,6 +310,7 @@ class StreamingDataset(Array, IterableDataset): def __init__(self, *, + epoch_size: Optional[Union[int, str]] = None, streams: Optional[Sequence[Stream]] = None, remote: Optional[str] = None, local: Optional[str] = None, @@ -318,7 +319,6 @@ def __init__(self, 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', From afe835a7632c8c0a7f2074f6d2c74351970d8230 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Mon, 25 Dec 2023 23:50:02 -0800 Subject: [PATCH 02/54] Move allow_unsafe_types arg. --- streaming/base/dataset.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index b22aeb434..a74b29d48 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -201,6 +201,7 @@ class StreamingDataset(Array, IterableDataset): * ``download_timeout`` * ``validate_hash`` * ``keep_zip`` + * ``allow_unsafe_types`` * Absolute dataset size, if streams were weighted relatively: @@ -261,6 +262,9 @@ 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``. + 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``. 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,9 +307,6 @@ 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, @@ -319,6 +320,7 @@ def __init__(self, download_timeout: float = 60, validate_hash: Optional[str] = None, keep_zip: bool = False, + allow_unsafe_types: bool = False, predownload: Optional[int] = None, cache_limit: Optional[Union[int, str]] = None, sampling_method: str = 'balanced', @@ -330,8 +332,7 @@ def __init__(self, shuffle_algo: str = 'py1e', shuffle_seed: int = 9176, shuffle_block_size: Optional[int] = None, - batching_method: str = 'random', - allow_unsafe_types: bool = False) -> None: + batching_method: str = 'random') -> None: # Global arguments (which do not live in Streams). self.predownload = predownload self.cache_limit = cache_limit @@ -345,7 +346,6 @@ def __init__(self, self.shuffle_seed = shuffle_seed self.shuffle_block_size = shuffle_block_size self.batching_method = batching_method - self.allow_unsafe_types = allow_unsafe_types # 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. @@ -457,7 +457,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, allow_unsafe_types) 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()) From 010c613a4dc2a4a98412fb74c61f332f1cb9e08b Mon Sep 17 00:00:00 2001 From: James Knighton Date: Mon, 25 Dec 2023 23:52:32 -0800 Subject: [PATCH 03/54] Fix usage. --- streaming/base/dataset.py | 2 -- streaming/base/stream.py | 4 ++-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index a74b29d48..aae33f64d 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -413,8 +413,6 @@ 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, diff --git a/streaming/base/stream.py b/streaming/base/stream.py index 0f3187592..0feebd11a 100644 --- a/streaming/base/stream.py +++ b/streaming/base/stream.py @@ -7,7 +7,7 @@ import json import os import tempfile -from typing import List, Optional, Sequence, Tuple +from typing import Any, Dict, List, Optional, Sequence, Tuple import numpy as np from numpy.typing import NDArray @@ -169,7 +169,7 @@ 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_default(self, default: Dict[str, Any]) -> None: """Apply defaults, setting any unset fields. We use pairs of (name, _name) in order to make type checking happy. From 2f4875b97836091eb9d623c1f482aa08908c5ec0 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 00:01:05 -0800 Subject: [PATCH 04/54] Propagate allow_unsafe_types as a normal Stream argument. --- simulation/core/sim_dataset.py | 22 +++++++++++----------- streaming/base/dataset.py | 20 +++++++++++--------- streaming/base/stream.py | 23 +++++++++++++++++------ 3 files changed, 39 insertions(+), 26 deletions(-) diff --git a/simulation/core/sim_dataset.py b/simulation/core/sim_dataset.py index 8dbc5a83d..cc3df2ad6 100644 --- a/simulation/core/sim_dataset.py +++ b/simulation/core/sim_dataset.py @@ -203,25 +203,25 @@ 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, + 'allow_unsafe_types': allow_unsafe_types, } for stream in streams: stream.apply_default(default) else: - default = 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] + stream = Stream(remote=remote, + local=local, + split=split, + download_retry=download_retry, + download_timeout=download_timeout, + validate_hash=validate_hash, + keep_zip=keep_zip, + allow_unsafe_types=allow_unsafe_types) + streams = [stream] # Validate the stream weighting scheme (relative or absolute) to catch errors before we go # to the trouble of loading them. @@ -270,7 +270,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) diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index aae33f64d..c8f8a91e0 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -418,18 +418,20 @@ def __init__(self, 'download_timeout': download_timeout, 'validate_hash': validate_hash, 'keep_zip': keep_zip, + 'allow_unsafe_types': allow_unsafe_types, } for stream in streams: stream.apply_default(default) else: - default = 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] + stream = Stream(remote=remote, + local=local, + split=split, + download_retry=download_retry, + download_timeout=download_timeout, + validate_hash=validate_hash, + keep_zip=keep_zip, + allow_unsafe_types=allow_unsafe_types) + streams = [stream] # Validate the stream weighting scheme (relative or absolute) to catch errors before we go # to the trouble of loading them. @@ -455,7 +457,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, 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()) diff --git a/streaming/base/stream.py b/streaming/base/stream.py index 0feebd11a..de42a76c7 100644 --- a/streaming/base/stream.py +++ b/streaming/base/stream.py @@ -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() @@ -191,6 +200,8 @@ def apply_default(self, default: Dict[str, Any]) -> None: 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} + if self._allow_unsafe_types is None: + self.allow_unsafe_types = default['allow_unsafe_types'] @classmethod def validate_weights(cls, streams: Sequence[Self]) -> Tuple[bool, bool]: @@ -421,18 +432,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 +483,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 From 333605eaf31e00db2586a0991ec730cda62f3239 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 00:08:31 -0800 Subject: [PATCH 05/54] Explicit list the kwargs for Stream.apply_defaults(). --- simulation/core/sim_dataset.py | 15 ++++++--------- streaming/base/dataset.py | 15 ++++++--------- streaming/base/stream.py | 33 +++++++++++++++++++++++---------- 3 files changed, 35 insertions(+), 28 deletions(-) diff --git a/simulation/core/sim_dataset.py b/simulation/core/sim_dataset.py index cc3df2ad6..bc3646b2c 100644 --- a/simulation/core/sim_dataset.py +++ b/simulation/core/sim_dataset.py @@ -202,16 +202,13 @@ def __init__(self, # Initialize the Stream defaults and normalize to a list of Streams. if streams: - default = { - 'split': split, - 'download_retry': download_retry, - 'download_timeout': download_timeout, - 'validate_hash': validate_hash, - 'keep_zip': keep_zip, - 'allow_unsafe_types': allow_unsafe_types, - } 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: stream = Stream(remote=remote, local=local, diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index c8f8a91e0..668ac70ef 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -412,16 +412,13 @@ def __init__(self, # Initialize the Stream defaults and normalize to a list of Streams. if streams: - default = { - 'split': split, - 'download_retry': download_retry, - 'download_timeout': download_timeout, - 'validate_hash': validate_hash, - 'keep_zip': keep_zip, - 'allow_unsafe_types': allow_unsafe_types, - } 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: stream = Stream(remote=remote, local=local, diff --git a/streaming/base/stream.py b/streaming/base/stream.py index de42a76c7..bbd7f1fbb 100644 --- a/streaming/base/stream.py +++ b/streaming/base/stream.py @@ -7,7 +7,7 @@ import json import os import tempfile -from typing import Any, Dict, List, Optional, Sequence, Tuple +from typing import List, Optional, Sequence, Tuple import numpy as np from numpy.typing import NDArray @@ -178,30 +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[str, Any]) -> 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 = default['allow_unsafe_types'] + self.allow_unsafe_types = allow_unsafe_types @classmethod def validate_weights(cls, streams: Sequence[Self]) -> Tuple[bool, bool]: From 2d8a9052d8bf71c097b99e51c7284cb2d7d9b786 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 00:19:05 -0800 Subject: [PATCH 06/54] Tweak docstrings. --- streaming/base/dataset.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index 668ac70ef..294e14c64 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -183,18 +183,21 @@ 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`` @@ -203,11 +206,7 @@ class StreamingDataset(Array, IterableDataset): * ``keep_zip`` * ``allow_unsafe_types`` - * Absolute dataset size, if streams were weighted relatively: - - * ``epoch_size`` - - * How to iterate: + * How to iterate (the StreamingDataset arguments): * Shard lifecycle: @@ -236,7 +235,6 @@ 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 From d298479871775b0daa6068eecdcec47bcd1edb9c Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 10:44:57 -0800 Subject: [PATCH 07/54] Complete rewrite of local dir collision detection using regular files. --- streaming/base/dataset.py | 114 ++++--- streaming/base/interproc/__init__.py | 4 + streaming/base/interproc/registry.py | 439 +++++++++++++++++++++++++++ 3 files changed, 516 insertions(+), 41 deletions(-) create mode 100644 streaming/base/interproc/__init__.py create mode 100644 streaming/base/interproc/registry.py diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index 294e14c64..a9fe50b37 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -12,6 +12,7 @@ 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 @@ -29,9 +30,9 @@ SHARD_ACCESS_TIMES, SHARD_STATES, TICK) from streaming.base.distributed import maybe_init_dist from streaming.base.format import get_index_basename +from streaming.base.interproc.registry import JobDir, JobRegistry from streaming.base.sampling import get_sampling -from streaming.base.shared import (SharedArray, SharedBarrier, SharedMemory, SharedScalar, - _get_path, get_shm_prefix) +from streaming.base.shared import SharedArray, SharedBarrier, SharedMemory, SharedScalar, _get_path from streaming.base.spanner import Spanner from streaming.base.stream import Stream from streaming.base.util import bytes_to_int, number_abbrev_to_int @@ -167,6 +168,34 @@ def on_exit(self) -> None: self._num_exited += 1 +def _test_config_root(config_root: str) -> None: + """Validate that the provided config root is usable. + + 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. + + Args: + config_root (str): Streaming configuration root directory. + """ + 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.') + + +def _get_default_config_root() -> str: + """Get the default Streaming configuration root directory. + + Returns: + str: Default Streaming configuration root directory. + """ + return os.path.join(gettempdir(), 'streaming') + + class StreamingDataset(Array, IterableDataset): """A mid-epoch-resumable streaming/caching pytorch IterableDataset. @@ -235,6 +264,10 @@ class StreamingDataset(Array, IterableDataset): * ``batching_method`` + * Configuration: + + * ``config_root`` + 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 @@ -305,32 +338,38 @@ class StreamingDataset(Array, IterableDataset): ``None``. batching_method (str): Which batching method to use, either ``random``, ``stratified``, or ``per_stream``. Defaults to ``random``. + 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. """ - 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, - 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: + 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, + 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', + config_root: str = _get_default_config_root(), + ) -> None: # Global arguments (which do not live in Streams). self.predownload = predownload self.cache_limit = cache_limit @@ -345,6 +384,9 @@ def __init__(self, self.shuffle_block_size = shuffle_block_size self.batching_method = batching_method + _test_config_root(config_root) + self.config_root = config_root + # 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. self.initial_physical_nodes = None @@ -501,13 +543,11 @@ def __init__(self, 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') + self.registry = JobRegistry(config_root) + self.job_dir = JobDir(self.registry, streams, world) + self._shm_prefix_int = int(self.job_dir.job_hash, 16) + + self._filelock_root = os.path.join(self.registry.config_root, self.job_dir.job_hash) os.makedirs(self._filelock_root, exist_ok=True) # Create the shared memory-backed barrier, without its lock, which is unpickleable. @@ -583,14 +623,6 @@ def __init__(self, del self._shared_barrier.lock # Remote the lock that makes it unpickleable. - 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 - @property def size(self) -> int: """Get the size of the dataset in samples. diff --git a/streaming/base/interproc/__init__.py b/streaming/base/interproc/__init__.py new file mode 100644 index 000000000..40b3649aa --- /dev/null +++ b/streaming/base/interproc/__init__.py @@ -0,0 +1,4 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Inter-process utilities.""" diff --git a/streaming/base/interproc/registry.py b/streaming/base/interproc/registry.py new file mode 100644 index 000000000..ee0102f81 --- /dev/null +++ b/streaming/base/interproc/registry.py @@ -0,0 +1,439 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Streaming job registry: local dir reuse detection.""" + +import json +import os +from hashlib import sha3_224 +from shutil import rmtree +from typing import Any, Dict, List, Optional, Sequence, Tuple + +from filelock import FileLock +from psutil import process_iter +from typing_extensions import Self + +from streaming.base.stream import Stream +from streaming.base.world import World + +__all__ = ['JobRegistry', 'JobDir'] + + +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. + create_time (int): Process creation time. + """ + + def __init__( + self, + *, + index: Optional[int] = None, + job_hash: str, + stream_hashes: List[str], + stream_locals: Optional[List[str]] = None, + process_id: int, + create_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.create_time = create_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'], + create_time=obj['create_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, + 'create_time': self.create_time, + } + + +class JobRegistryFile: + """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 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. + """ + job_hashes = [] + for job in self.jobs: + if job.create_time != pid2create_time.get(job.process_id): + self.remove(job.job_hash) + job_hashes.append(job.job_hash) + return job_hashes + + +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. + """ + + def __init__(self, config_root: str) -> None: + self.config_root = config_root + self._filelock_filename = os.path.join(config_root, 'filelock.bin') + 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()) + 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. + """ + 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, + create_time=create_time) + + with FileLock(self._filelock_filename): + reg = JobRegistryFile.read(self._registry_filename) + reg.add(entry) + del_job_hashes = reg.filter(pid2create_time) + reg.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: + return self._register(streams) + else: + return self._lookup(streams) + + 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 FileLock(self._filelock_filename): + reg = JobRegistryFile.read(self._registry_filename) + reg.remove(job_hash) + del_job_hashes = reg.filter(pid2create_time) + reg.write(self._registry_filename) + map(self._remove_dir, [job_hash] + del_job_hashes) + + 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: + pass + + +class JobDir: + """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) From eff80e6d4d4f0827bc2b3158ed026a013d6f9d1e Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 10:53:44 -0800 Subject: [PATCH 08/54] Add psutil. --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index dbe68f892..32946be36 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 = {} From 8fb9dca1237cd0a3b6501dd3bd9546e26a4e4883 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 11:00:46 -0800 Subject: [PATCH 09/54] Fix. --- streaming/base/dataset.py | 1 + 1 file changed, 1 insertion(+) diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index a9fe50b37..1df673909 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -179,6 +179,7 @@ def _test_config_root(config_root: str) -> None: 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: From 223f3ca81e6b9d8c2b5ebb2bd267c44242ab1c3b Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 11:11:13 -0800 Subject: [PATCH 10/54] Fix. --- streaming/base/interproc/registry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/streaming/base/interproc/registry.py b/streaming/base/interproc/registry.py index ee0102f81..492eb45bc 100644 --- a/streaming/base/interproc/registry.py +++ b/streaming/base/interproc/registry.py @@ -100,7 +100,7 @@ def read(cls, filename: str) -> Self: return cls(jobs) def write(self, filename: str) -> None: - jobs = [job.to_json() for job in self.jobs] + jobs = [job.to_json() for job in filter(bool, self.jobs)] obj = {'jobs': jobs} with open(filename, 'w') as out: json.dump(obj, out) From 664cbc1f959e39abdf3bdaf4c4b38bba1bf6103c Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 11:21:42 -0800 Subject: [PATCH 11/54] Fix. --- streaming/base/interproc/registry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/streaming/base/interproc/registry.py b/streaming/base/interproc/registry.py index 492eb45bc..ff543b1d3 100644 --- a/streaming/base/interproc/registry.py +++ b/streaming/base/interproc/registry.py @@ -215,7 +215,7 @@ def _get_live_procs(self) -> Dict[int, int]: """ ret = {} for obj in process_iter(['pid', 'create_time']): - ret[obj.pid] = int(obj.create_time()) + ret[obj.pid] = int(obj.create_time() * 1e9) return ret def _hash(self, data: bytes) -> str: From 02416d7d21c31ed65c00b89f03a4fb3dc6237a5e Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 11:46:05 -0800 Subject: [PATCH 12/54] Fix. --- streaming/base/interproc/registry.py | 38 +++++++++++++++++++++++++--- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/streaming/base/interproc/registry.py b/streaming/base/interproc/registry.py index ff543b1d3..d3a6c4906 100644 --- a/streaming/base/interproc/registry.py +++ b/streaming/base/interproc/registry.py @@ -7,6 +7,7 @@ import os from hashlib import sha3_224 from shutil import rmtree +from time import sleep from typing import Any, Dict, List, Optional, Sequence, Tuple from filelock import FileLock @@ -182,7 +183,7 @@ def filter(self, pid2create_time: Dict[int, int]) -> List[str]: List[str]: List of hashes of removed datasets. """ job_hashes = [] - for job in self.jobs: + for job in filter(bool, self.jobs): if job.create_time != pid2create_time.get(job.process_id): self.remove(job.job_hash) job_hashes.append(job.job_hash) @@ -200,8 +201,9 @@ class JobRegistry: your system. """ - def __init__(self, config_root: str) -> None: + def __init__(self, config_root: str, tick: float = 0.007) -> None: self.config_root = config_root + self._tick = tick self._filelock_filename = os.path.join(config_root, 'filelock.bin') self._registry_filename = os.path.join(config_root, 'registry.json') @@ -336,6 +338,32 @@ def _register(self, streams: Sequence[Stream]) -> str: return job_hash + def _wait_for_existence(self, job_hash: str) -> None: + """Wait for a directory to be created. + + Args: + job_hash (str): Job hash of directory. + """ + dirname = os.path.join(self.config_root, job_hash) + while True: + with FileLock(self._filelock_filename): + if os.path.exists(dirname): + break + sleep(self._tick) + + def _wait_for_removal(self, job_hash: str) -> None: + """Wait for a directory to be removed. + + Args: + job_hash (str): Job hash of directory. + """ + dirname = os.path.join(self.config_root, job_hash) + while True: + with FileLock(self._filelock_filename): + if not os.path.exists(dirname): + break + sleep(self._tick) + def _lookup(self, streams: Sequence[Stream]) -> str: """Look up this collection of StreamingDataset replicas. @@ -350,6 +378,7 @@ def _lookup(self, streams: Sequence[Stream]) -> str: str: Streaming config subdir for this job. """ _, _, job_hash = self._hash_streams(streams) + self._wait_for_existence(job_hash) return job_hash def register(self, streams: Sequence[Stream], world: World) -> str: @@ -386,7 +415,8 @@ def _unregister(self, job_hash: str) -> None: reg.remove(job_hash) del_job_hashes = reg.filter(pid2create_time) reg.write(self._registry_filename) - map(self._remove_dir, [job_hash] + del_job_hashes) + 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. @@ -400,7 +430,7 @@ def unregister(self, job_hash: str, world: World) -> None: if world.is_local_leader: self._unregister(job_hash) else: - pass + self._wait_for_removal(job_hash) class JobDir: From b0b3b5610a4247c36897ead3b6c67a5df28cb24e Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 13:26:26 -0800 Subject: [PATCH 13/54] Fix. --- tests/test_streaming.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 18dfef45e..3b8d6a664 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -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_dir # TODO: Gross hack. 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) From 23505c2939db7fc4aac4dcd57b13994d9361ef4d Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 13:47:57 -0800 Subject: [PATCH 14/54] Fix. --- streaming/base/interproc/registry.py | 62 ++++++++++++++-------------- 1 file changed, 32 insertions(+), 30 deletions(-) diff --git a/streaming/base/interproc/registry.py b/streaming/base/interproc/registry.py index d3a6c4906..331a96e7f 100644 --- a/streaming/base/interproc/registry.py +++ b/streaming/base/interproc/registry.py @@ -301,6 +301,32 @@ def _remove_dir(self, job_hash: str) -> None: dirname = os.path.join(self.config_root, job_hash) rmtree(dirname) + def _wait_for_existence(self, job_hash: str) -> None: + """Wait for a directory to be created. + + Args: + job_hash (str): Job hash of directory. + """ + dirname = os.path.join(self.config_root, job_hash) + while True: + with FileLock(self._filelock_filename): + if os.path.exists(dirname): + break + sleep(self._tick) + + def _wait_for_removal(self, job_hash: str) -> None: + """Wait for a directory to be removed. + + Args: + job_hash (str): Job hash of directory. + """ + dirname = os.path.join(self.config_root, job_hash) + while True: + with FileLock(self._filelock_filename): + if not os.path.exists(dirname): + break + sleep(self._tick) + def _register(self, streams: Sequence[Stream]) -> str: """Register this collection of StreamingDataset replicas. @@ -338,32 +364,6 @@ def _register(self, streams: Sequence[Stream]) -> str: return job_hash - def _wait_for_existence(self, job_hash: str) -> None: - """Wait for a directory to be created. - - Args: - job_hash (str): Job hash of directory. - """ - dirname = os.path.join(self.config_root, job_hash) - while True: - with FileLock(self._filelock_filename): - if os.path.exists(dirname): - break - sleep(self._tick) - - def _wait_for_removal(self, job_hash: str) -> None: - """Wait for a directory to be removed. - - Args: - job_hash (str): Job hash of directory. - """ - dirname = os.path.join(self.config_root, job_hash) - while True: - with FileLock(self._filelock_filename): - if not os.path.exists(dirname): - break - sleep(self._tick) - def _lookup(self, streams: Sequence[Stream]) -> str: """Look up this collection of StreamingDataset replicas. @@ -378,7 +378,6 @@ def _lookup(self, streams: Sequence[Stream]) -> str: str: Streaming config subdir for this job. """ _, _, job_hash = self._hash_streams(streams) - self._wait_for_existence(job_hash) return job_hash def register(self, streams: Sequence[Stream], world: World) -> str: @@ -396,9 +395,11 @@ def register(self, streams: Sequence[Stream], world: World) -> str: str: Subdir for this collection of StreamingDataset replicas. """ if world.is_local_leader: - return self._register(streams) + job_hash = self._register(streams) else: - return self._lookup(streams) + job_hash = self._lookup(streams) + self._wait_for_existence(job_hash) + return job_hash def _unregister(self, job_hash: str) -> None: """Unregister this collection of StreamingDataset replicas. @@ -430,7 +431,8 @@ def unregister(self, job_hash: str, world: World) -> None: if world.is_local_leader: self._unregister(job_hash) else: - self._wait_for_removal(job_hash) + pass + self._wait_for_removal(job_hash) class JobDir: From 67b936f306b61900ed4d89de00da339fca590e4a Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 14:13:02 -0800 Subject: [PATCH 15/54] Fix. --- streaming/base/interproc/registry.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/streaming/base/interproc/registry.py b/streaming/base/interproc/registry.py index 331a96e7f..9940fe8d8 100644 --- a/streaming/base/interproc/registry.py +++ b/streaming/base/interproc/registry.py @@ -7,7 +7,7 @@ import os from hashlib import sha3_224 from shutil import rmtree -from time import sleep +from time import sleep, time_ns from typing import Any, Dict, List, Optional, Sequence, Tuple from filelock import FileLock @@ -29,7 +29,7 @@ class JobEntry: 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. - create_time (int): Process creation time. + register_time (int): Process registration time. """ def __init__( @@ -40,14 +40,14 @@ def __init__( stream_hashes: List[str], stream_locals: Optional[List[str]] = None, process_id: int, - create_time: 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.create_time = create_time + self.register_time = register_time @classmethod def from_json(cls, obj: Dict[str, Any]) -> Self: @@ -63,7 +63,7 @@ def from_json(cls, obj: Dict[str, Any]) -> Self: stream_hashes=obj['stream_hashes'], stream_locals=obj.get('stream_locals'), process_id=obj['process_id'], - create_time=obj['create_time']) + register_time=obj['register_time']) def to_json(self) -> Dict[str, Any]: return { @@ -71,7 +71,7 @@ def to_json(self) -> Dict[str, Any]: 'stream_hashes': self.stream_hashes, # stream_locals is not saved, only their hashes. 'process_id': self.process_id, - 'create_time': self.create_time, + 'register_time': self.register_time, } @@ -182,12 +182,13 @@ def filter(self, pid2create_time: Dict[int, int]) -> List[str]: Returns: List[str]: List of hashes of removed datasets. """ - job_hashes = [] + del_job_hashes = [] for job in filter(bool, self.jobs): - if job.create_time != pid2create_time.get(job.process_id): + create_time = pid2create_time.get(job.process_id) + if not create_time or job.register_time < create_time: self.remove(job.job_hash) - job_hashes.append(job.job_hash) - return job_hashes + del_job_hashes.append(job.job_hash) + return del_job_hashes class JobRegistry: @@ -340,6 +341,7 @@ def _register(self, streams: Sequence[Stream]) -> str: 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) @@ -352,7 +354,7 @@ def _register(self, streams: Sequence[Stream]) -> str: stream_hashes=stream_hashes, stream_locals=stream_locals, process_id=pid, - create_time=create_time) + register_time=register_time) with FileLock(self._filelock_filename): reg = JobRegistryFile.read(self._registry_filename) From fa14130a27448b647b947374063a50f350ca47fe Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 14:35:41 -0800 Subject: [PATCH 16/54] Fix. --- streaming/base/interproc/registry.py | 5 ++--- tests/test_streaming.py | 1 - 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/streaming/base/interproc/registry.py b/streaming/base/interproc/registry.py index 9940fe8d8..810cc16f1 100644 --- a/streaming/base/interproc/registry.py +++ b/streaming/base/interproc/registry.py @@ -380,6 +380,7 @@ def _lookup(self, streams: Sequence[Stream]) -> str: str: Streaming config subdir for this job. """ _, _, job_hash = self._hash_streams(streams) + self._wait_for_existence(job_hash) return job_hash def register(self, streams: Sequence[Stream], world: World) -> str: @@ -400,7 +401,6 @@ def register(self, streams: Sequence[Stream], world: World) -> str: job_hash = self._register(streams) else: job_hash = self._lookup(streams) - self._wait_for_existence(job_hash) return job_hash def _unregister(self, job_hash: str) -> None: @@ -433,8 +433,7 @@ def unregister(self, job_hash: str, world: World) -> None: if world.is_local_leader: self._unregister(job_hash) else: - pass - self._wait_for_removal(job_hash) + self._wait_for_removal(job_hash) class JobDir: diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 3b8d6a664..42975766f 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -782,7 +782,6 @@ def test_streamingdataloader_mid_epoch_resumption(local_remote_dir: Any, batch_s sample_order.extend(batch['id'][:]) del dataloader - del dataset.job_dir # TODO: Gross hack. del dataset clean_stale_shared_memory() From c54887be4a846fa887ffa236f1790d5a842a4da4 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 14:47:29 -0800 Subject: [PATCH 17/54] Fix. --- streaming/base/interproc/registry.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/streaming/base/interproc/registry.py b/streaming/base/interproc/registry.py index 810cc16f1..9940fe8d8 100644 --- a/streaming/base/interproc/registry.py +++ b/streaming/base/interproc/registry.py @@ -380,7 +380,6 @@ def _lookup(self, streams: Sequence[Stream]) -> str: str: Streaming config subdir for this job. """ _, _, job_hash = self._hash_streams(streams) - self._wait_for_existence(job_hash) return job_hash def register(self, streams: Sequence[Stream], world: World) -> str: @@ -401,6 +400,7 @@ def register(self, streams: Sequence[Stream], world: World) -> str: job_hash = self._register(streams) else: job_hash = self._lookup(streams) + self._wait_for_existence(job_hash) return job_hash def _unregister(self, job_hash: str) -> None: @@ -433,7 +433,8 @@ def unregister(self, job_hash: str, world: World) -> None: if world.is_local_leader: self._unregister(job_hash) else: - self._wait_for_removal(job_hash) + pass + self._wait_for_removal(job_hash) class JobDir: From c21458900b9f5a3b0be5628101cf5074874fe92f Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 14:52:57 -0800 Subject: [PATCH 18/54] Fix. --- tests/test_streaming.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 42975766f..582bfe43a 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -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_dir del dataset clean_stale_shared_memory() From c0c82bd68d1df86985f0d62d242577f6e57f1b7d Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 15:10:39 -0800 Subject: [PATCH 19/54] Remove dist from StreamingDataset init. --- streaming/base/dataset.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index 1df673909..fefd98134 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -20,7 +20,6 @@ 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 @@ -28,7 +27,6 @@ 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.format import get_index_basename from streaming.base.interproc.registry import JobDir, JobRegistry from streaming.base.sampling import get_sampling @@ -448,9 +446,6 @@ def __init__( 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: for stream in streams: @@ -548,13 +543,17 @@ def __init__( self.job_dir = JobDir(self.registry, streams, world) self._shm_prefix_int = int(self.job_dir.job_hash, 16) + init_done_filename = self.job_dir.get_filename('init_done.txt') + if world.is_local_leader: + if os.path.exists(init_done_filename): + os.remove(init_done_filename) + self._filelock_root = os.path.join(self.registry.config_root, self.job_dir.job_hash) 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)) + self._shared_barrier = SharedBarrier(self.job_dir.get_filename('barrier_filelock.bin'), + _get_path(self._shm_prefix_int, BARRIER)) # Epoch counter. # @@ -564,8 +563,7 @@ def __init__( self._next_epoch = SharedScalar(np.int64, _get_path(self._shm_prefix_int, NEXT_EPOCH)) # 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_path = self.job_dir.get_filename('cache_filelock.bin') self._cache_filelock: FileLock # Cache usage in bytes. @@ -605,11 +603,13 @@ def __init__( 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() - - if destroy_dist: - dist.destroy_process_group() + dirname = os.path.dirname(init_done_filename) + os.makedirs(dirname, exist_ok=True) + with open(init_done_filename, 'wb') as out: + out.write(b'') + else: + wait_for_file_to_exist(init_done_filename, TICK, 300, + 'Waited too long for initialization') # Placeholder for a shared memory object where load_state_dict() saves its data to be # picked up by __iter__(). From 105ee1669bb09850d69f39438cabdbeb1402307f Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 15:10:50 -0800 Subject: [PATCH 20/54] Sleep first out of race paranoia. --- streaming/base/interproc/registry.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/streaming/base/interproc/registry.py b/streaming/base/interproc/registry.py index 9940fe8d8..31ed68029 100644 --- a/streaming/base/interproc/registry.py +++ b/streaming/base/interproc/registry.py @@ -310,10 +310,10 @@ def _wait_for_existence(self, job_hash: str) -> None: """ dirname = os.path.join(self.config_root, job_hash) while True: + sleep(self._tick) with FileLock(self._filelock_filename): if os.path.exists(dirname): break - sleep(self._tick) def _wait_for_removal(self, job_hash: str) -> None: """Wait for a directory to be removed. @@ -323,10 +323,10 @@ def _wait_for_removal(self, job_hash: str) -> None: """ dirname = os.path.join(self.config_root, job_hash) while True: + sleep(self._tick) with FileLock(self._filelock_filename): if not os.path.exists(dirname): break - sleep(self._tick) def _register(self, streams: Sequence[Stream]) -> str: """Register this collection of StreamingDataset replicas. From afafcd8007dc3e356e934e5322d0297c7de55cfb Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 26 Dec 2023 15:20:03 -0800 Subject: [PATCH 21/54] Fix. --- streaming/base/dataset.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index fefd98134..b2653aa4c 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -24,16 +24,15 @@ 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.constant import (BARRIER, CACHE_FILELOCK, CACHE_USAGE, EPOCH_DATA, EPOCH_SHAPE, + NEXT_EPOCH, RESUME, SHARD_ACCESS_TIMES, SHARD_STATES, TICK) from streaming.base.format import get_index_basename from streaming.base.interproc.registry import JobDir, JobRegistry from streaming.base.sampling import get_sampling from streaming.base.shared import SharedArray, SharedBarrier, SharedMemory, SharedScalar, _get_path 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.util import bytes_to_int, number_abbrev_to_int, wait_for_file_to_exist from streaming.base.world import World # An arbitrary time in the future, used for cold shard eviction. From 1012b21973d25caf77c79fb2a7909e7c37a81cc7 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Fri, 29 Dec 2023 22:30:47 -0800 Subject: [PATCH 22/54] Organize world, job/, and shmem/ into streaming/base/coord/. --- docs/source/conf.py | 3 +- simulation/core/sim_dataset.py | 1 + simulation/core/sim_world.py | 2 +- simulation/core/yaml_processing.py | 30 ++- streaming/base/batching/__init__.py | 2 +- streaming/base/batching/per_stream.py | 2 +- streaming/base/batching/random.py | 2 +- streaming/base/batching/stratified.py | 2 +- streaming/base/coord/__init__.py | 15 ++ streaming/base/coord/job/__init__.py | 9 + streaming/base/coord/job/directory.py | 51 ++++ streaming/base/coord/job/entry.py | 65 +++++ streaming/base/coord/job/file.py | 130 ++++++++++ .../base/{interproc => coord/job}/registry.py | 227 +----------------- streaming/base/coord/shmem/__init__.py | 17 ++ .../base/{shared => coord/shmem}/array.py | 2 +- .../base/{shared => coord/shmem}/barrier.py | 2 +- .../base/{shared => coord/shmem}/memory.py | 0 .../base/{shared => coord/shmem}/prefix.py | 4 +- .../base/{shared => coord/shmem}/scalar.py | 2 +- streaming/base/{ => coord}/world.py | 0 streaming/base/dataloader.py | 2 +- streaming/base/dataset.py | 9 +- streaming/base/interproc/__init__.py | 4 - streaming/base/shared/__init__.py | 17 -- streaming/base/stream.py | 2 +- streaming/base/util.py | 2 +- tests/test_barrier.py | 2 +- tests/test_reader.py | 3 +- tests/test_shared.py | 4 +- tests/test_streaming.py | 6 +- tests/test_util.py | 2 +- 32 files changed, 351 insertions(+), 270 deletions(-) create mode 100644 streaming/base/coord/__init__.py create mode 100644 streaming/base/coord/job/__init__.py create mode 100644 streaming/base/coord/job/directory.py create mode 100644 streaming/base/coord/job/entry.py create mode 100644 streaming/base/coord/job/file.py rename streaming/base/{interproc => coord/job}/registry.py (54%) create mode 100644 streaming/base/coord/shmem/__init__.py rename streaming/base/{shared => coord/shmem}/array.py (97%) rename streaming/base/{shared => coord/shmem}/barrier.py (98%) rename streaming/base/{shared => coord/shmem}/memory.py (100%) rename streaming/base/{shared => coord/shmem}/prefix.py (98%) rename streaming/base/{shared => coord/shmem}/scalar.py (93%) rename streaming/base/{ => coord}/world.py (100%) delete mode 100644 streaming/base/interproc/__init__.py delete mode 100644 streaming/base/shared/__init__.py diff --git a/docs/source/conf.py b/docs/source/conf.py index e25dc24ba..fd78ff9d5 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/simulation/core/sim_dataset.py b/simulation/core/sim_dataset.py index bc3646b2c..1b78f03c5 100644 --- a/simulation/core/sim_dataset.py +++ b/simulation/core/sim_dataset.py @@ -105,6 +105,7 @@ class SimulationDataset(StreamingDataset): """ def __init__(self, + *, nodes: int, devices: int, workers: int, diff --git a/simulation/core/sim_world.py b/simulation/core/sim_world.py index 6c607b8ad..f7a08743e 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 86e74dc3b..ae5aa6bfc 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 f4fd7f788..fdb81d273 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 d12b61a2c..99944aa7c 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 diff --git a/streaming/base/batching/random.py b/streaming/base/batching/random.py index 48e803acb..76050848a 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 diff --git a/streaming/base/batching/stratified.py b/streaming/base/batching/stratified.py index 2eef06fd5..4dfed207a 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 diff --git a/streaming/base/coord/__init__.py b/streaming/base/coord/__init__.py new file mode 100644 index 000000000..af0a17173 --- /dev/null +++ b/streaming/base/coord/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2023 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.mmap import MMapArray, MMapBarrier, MMapBuffer, MMapNumber +from streaming.base.coord.shmem import (SharedArray, SharedBarrier, SharedMemory, SharedScalar, + get_shm_prefix) +from streaming.base.coord.world import World + +__all__ = [ + 'JobDirectory', 'JobRegistry', 'MMapArray', 'MMapBarrier', 'MMapBuffer', 'MMapNumber', + 'SharedArray', 'SharedBarrier', 'SharedMemory', 'get_shm_prefix', 'SharedScalar', 'World' +] diff --git a/streaming/base/coord/job/__init__.py b/streaming/base/coord/job/__init__.py new file mode 100644 index 000000000..cd5f75465 --- /dev/null +++ b/streaming/base/coord/job/__init__.py @@ -0,0 +1,9 @@ +# Copyright 2023 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..6d077f4cd --- /dev/null +++ b/streaming/base/coord/job/directory.py @@ -0,0 +1,51 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""A directory containing all dataset-wide filesystem state for a Streaming job.""" + +import os +from pathlib import Path +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) + self.dirname = Path(os.path.join(registry.config_root, self.job_hash)) + + 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..c39305e6c --- /dev/null +++ b/streaming/base/coord/job/entry.py @@ -0,0 +1,65 @@ +# Copyright 2023 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..3383cd468 --- /dev/null +++ b/streaming/base/coord/job/file.py @@ -0,0 +1,130 @@ +# Copyright 2023 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__ = ['JobFile'] + + +class JobFile: + """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/interproc/registry.py b/streaming/base/coord/job/registry.py similarity index 54% rename from streaming/base/interproc/registry.py rename to streaming/base/coord/job/registry.py index 31ed68029..607425313 100644 --- a/streaming/base/interproc/registry.py +++ b/streaming/base/coord/job/registry.py @@ -1,194 +1,26 @@ # Copyright 2023 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 -"""Streaming job registry: local dir reuse detection.""" +"""A directory containing all Streaming-wide filesystem state. + +Useful for detecting collisions between different jobs' local dirs. +""" -import json import os from hashlib import sha3_224 from shutil import rmtree from time import sleep, time_ns -from typing import Any, Dict, List, Optional, Sequence, Tuple +from typing import Dict, List, Sequence, Tuple from filelock import FileLock from psutil import process_iter -from typing_extensions import Self +from streaming.base.coord.job.entry import JobEntry +from streaming.base.coord.job.file import JobFile +from streaming.base.coord.world import World from streaming.base.stream import Stream -from streaming.base.world import World - -__all__ = ['JobRegistry', 'JobDir'] - - -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, - } - - -class JobRegistryFile: - """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 +__all__ = ['JobRegistry'] class JobRegistry: @@ -203,6 +35,7 @@ class JobRegistry: """ def __init__(self, config_root: str, tick: float = 0.007) -> None: + os.makedirs(config_root, exist_ok=True) self.config_root = config_root self._tick = tick self._filelock_filename = os.path.join(config_root, 'filelock.bin') @@ -357,7 +190,7 @@ def _register(self, streams: Sequence[Stream]) -> str: register_time=register_time) with FileLock(self._filelock_filename): - reg = JobRegistryFile.read(self._registry_filename) + reg = JobFile.read(self._registry_filename) reg.add(entry) del_job_hashes = reg.filter(pid2create_time) reg.write(self._registry_filename) @@ -414,7 +247,7 @@ def _unregister(self, job_hash: str) -> None: pid2create_time = self._get_live_procs() with FileLock(self._filelock_filename): - reg = JobRegistryFile.read(self._registry_filename) + reg = JobFile.read(self._registry_filename) reg.remove(job_hash) del_job_hashes = reg.filter(pid2create_time) reg.write(self._registry_filename) @@ -435,39 +268,3 @@ def unregister(self, job_hash: str, world: World) -> None: else: pass self._wait_for_removal(job_hash) - - -class JobDir: - """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/shmem/__init__.py b/streaming/base/coord/shmem/__init__.py new file mode 100644 index 000000000..991be052c --- /dev/null +++ b/streaming/base/coord/shmem/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2023 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 20689d125..543dc7163 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 98% rename from streaming/base/shared/barrier.py rename to streaming/base/coord/shmem/barrier.py index ceeb3ec43..6cc9988af 100644 --- a/streaming/base/shared/barrier.py +++ b/streaming/base/coord/shmem/barrier.py @@ -12,7 +12,7 @@ from filelock import FileLock from streaming.base.constant import TICK -from streaming.base.shared.array import SharedArray +from streaming.base.coord.shmem.array import SharedArray # Time out to wait before raising exception TIMEOUT = 60 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 48d2aaa6c..69ab2031a 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 14cd5e7fa..03c142074 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/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 89cdb0026..266762fba 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 b2653aa4c..8dbee1835 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -26,14 +26,15 @@ from streaming.base.batching import generate_work from streaming.base.constant import (BARRIER, CACHE_FILELOCK, CACHE_USAGE, EPOCH_DATA, EPOCH_SHAPE, NEXT_EPOCH, RESUME, SHARD_ACCESS_TIMES, SHARD_STATES, TICK) +from streaming.base.coord.job import JobDirectory, JobRegistry +from streaming.base.coord.shmem import (SharedArray, SharedBarrier, SharedMemory, SharedScalar, + _get_path) +from streaming.base.coord.world import World from streaming.base.format import get_index_basename -from streaming.base.interproc.registry import JobDir, JobRegistry from streaming.base.sampling import get_sampling -from streaming.base.shared import SharedArray, SharedBarrier, SharedMemory, SharedScalar, _get_path from streaming.base.spanner import Spanner from streaming.base.stream import Stream from streaming.base.util import bytes_to_int, number_abbrev_to_int, wait_for_file_to_exist -from streaming.base.world import World # An arbitrary time in the future, used for cold shard eviction. NEVER = np.iinfo(np.uint64).max @@ -539,7 +540,7 @@ def __init__( # Register/lookup our shared memory prefix and filelock root directory. self.registry = JobRegistry(config_root) - self.job_dir = JobDir(self.registry, streams, world) + self.job_dir = JobDirectory(self.registry, streams, world) self._shm_prefix_int = int(self.job_dir.job_hash, 16) init_done_filename = self.job_dir.get_filename('init_done.txt') diff --git a/streaming/base/interproc/__init__.py b/streaming/base/interproc/__init__.py deleted file mode 100644 index 40b3649aa..000000000 --- a/streaming/base/interproc/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# Copyright 2023 MosaicML Streaming authors -# SPDX-License-Identifier: Apache-2.0 - -"""Inter-process utilities.""" diff --git a/streaming/base/shared/__init__.py b/streaming/base/shared/__init__.py deleted file mode 100644 index cf507c4fe..000000000 --- a/streaming/base/shared/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright 2023 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 bbd7f1fbb..e50e9b2fd 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: diff --git a/streaming/base/util.py b/streaming/base/util.py index e86876ee1..2e6a3be73 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 fdc5eb87d..0d8f206be 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: diff --git a/tests/test_reader.py b/tests/test_reader.py index fbe7ff723..24066a8d0 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 c28229472..ea711a76c 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 582bfe43a..f06496dbe 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, diff --git a/tests/test_util.py b/tests/test_util.py index e59f75911..98df2d719 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, From 0eb33274d49ab3ae2aff3abd2267d4129f507c26 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Fri, 29 Dec 2023 22:35:47 -0800 Subject: [PATCH 23/54] Fix. --- streaming/base/coord/__init__.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/streaming/base/coord/__init__.py b/streaming/base/coord/__init__.py index af0a17173..75b47257e 100644 --- a/streaming/base/coord/__init__.py +++ b/streaming/base/coord/__init__.py @@ -4,12 +4,11 @@ """Coordination among ranks and workers.""" from streaming.base.coord.job import JobDirectory, JobRegistry -from streaming.base.coord.mmap import MMapArray, MMapBarrier, MMapBuffer, MMapNumber from streaming.base.coord.shmem import (SharedArray, SharedBarrier, SharedMemory, SharedScalar, get_shm_prefix) from streaming.base.coord.world import World __all__ = [ - 'JobDirectory', 'JobRegistry', 'MMapArray', 'MMapBarrier', 'MMapBuffer', 'MMapNumber', - 'SharedArray', 'SharedBarrier', 'SharedMemory', 'get_shm_prefix', 'SharedScalar', 'World' + 'JobDirectory', 'JobRegistry', 'SharedArray', 'SharedBarrier', 'SharedMemory', + 'get_shm_prefix', 'SharedScalar', 'World' ] From ea7c7ada9f84ab27e6a5e8545f5383f6112e0fa2 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Fri, 29 Dec 2023 22:49:07 -0800 Subject: [PATCH 24/54] Fix. --- streaming/base/dataset.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index 8dbee1835..66ab9680c 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -428,10 +428,10 @@ def __init__( # 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: + 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.') + f'This may result in slower batch time. The 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 ' + From 5267bbb3ea80960df2c71643aeaf787da5cfd021 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Fri, 29 Dec 2023 23:13:09 -0800 Subject: [PATCH 25/54] Fix. --- streaming/base/coord/job/directory.py | 2 -- streaming/base/coord/job/registry.py | 1 - 2 files changed, 3 deletions(-) diff --git a/streaming/base/coord/job/directory.py b/streaming/base/coord/job/directory.py index 6d077f4cd..e41b14276 100644 --- a/streaming/base/coord/job/directory.py +++ b/streaming/base/coord/job/directory.py @@ -4,7 +4,6 @@ """A directory containing all dataset-wide filesystem state for a Streaming job.""" import os -from pathlib import Path from typing import Sequence from streaming.base.coord.job.registry import JobRegistry @@ -33,7 +32,6 @@ def __init__(self, registry: JobRegistry, streams: Sequence[Stream], world: Worl self.streams = streams self.world = world self.job_hash = registry.register(streams, world) - self.dirname = Path(os.path.join(registry.config_root, self.job_hash)) def get_filename(self, path: str) -> str: """Get a filename by relative path under its job dir. diff --git a/streaming/base/coord/job/registry.py b/streaming/base/coord/job/registry.py index 607425313..30e739eda 100644 --- a/streaming/base/coord/job/registry.py +++ b/streaming/base/coord/job/registry.py @@ -35,7 +35,6 @@ class JobRegistry: """ def __init__(self, config_root: str, tick: float = 0.007) -> None: - os.makedirs(config_root, exist_ok=True) self.config_root = config_root self._tick = tick self._filelock_filename = os.path.join(config_root, 'filelock.bin') From 74646a2d58a9969ef2e99e76da9567a3c349fa07 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Fri, 29 Dec 2023 23:51:12 -0800 Subject: [PATCH 26/54] Keep around the lazily initialized FileLock. --- streaming/base/coord/job/registry.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/streaming/base/coord/job/registry.py b/streaming/base/coord/job/registry.py index 30e739eda..f9ca325f8 100644 --- a/streaming/base/coord/job/registry.py +++ b/streaming/base/coord/job/registry.py @@ -38,6 +38,7 @@ def __init__(self, config_root: str, tick: float = 0.007) -> None: self.config_root = config_root self._tick = tick self._filelock_filename = os.path.join(config_root, 'filelock.bin') + self._filelock: FileLock self._registry_filename = os.path.join(config_root, 'registry.json') def _get_live_procs(self) -> Dict[int, int]: @@ -116,6 +117,11 @@ def _hash_streams(self, streams: Sequence[Stream]) -> Tuple[List[str], List[str] return stream_locals, stream_hashes, job_hash + def _ensure_filelock(self) -> None: + """Ensure the lazily-created file lock exists.""" + if not hasattr(self, '_filelock'): + self._filelock = FileLock(self._filelock_filename) + def _make_dir(self, job_hash: str) -> None: """Create a Streaming job config dir. @@ -143,7 +149,7 @@ def _wait_for_existence(self, job_hash: str) -> None: dirname = os.path.join(self.config_root, job_hash) while True: sleep(self._tick) - with FileLock(self._filelock_filename): + with self._filelock: if os.path.exists(dirname): break @@ -156,7 +162,7 @@ def _wait_for_removal(self, job_hash: str) -> None: dirname = os.path.join(self.config_root, job_hash) while True: sleep(self._tick) - with FileLock(self._filelock_filename): + with self._filelock: if not os.path.exists(dirname): break @@ -173,6 +179,8 @@ def _register(self, streams: Sequence[Stream]) -> str: Returns: str: Streaming config subdir for this job. """ + self._ensure_filelock() + register_time = time_ns() pid2create_time = self._get_live_procs() pid = os.getpid() @@ -188,7 +196,7 @@ def _register(self, streams: Sequence[Stream]) -> str: process_id=pid, register_time=register_time) - with FileLock(self._filelock_filename): + with self._filelock: reg = JobFile.read(self._registry_filename) reg.add(entry) del_job_hashes = reg.filter(pid2create_time) @@ -243,9 +251,11 @@ def _unregister(self, job_hash: str) -> None: Args: job_hash (str): Subdir identifying this Streaming job. """ + self._ensure_filelock() + pid2create_time = self._get_live_procs() - with FileLock(self._filelock_filename): + with self._filelock: reg = JobFile.read(self._registry_filename) reg.remove(job_hash) del_job_hashes = reg.filter(pid2create_time) From 6b3783f2c4c220ec577e12d7028d71febcbe2682 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 30 Dec 2023 00:44:25 -0800 Subject: [PATCH 27/54] What if it's the filelock? --- streaming/base/coord/file.py | 60 ++++++++++++++++++++++++++++ streaming/base/coord/job/registry.py | 12 +++--- 2 files changed, 66 insertions(+), 6 deletions(-) create mode 100644 streaming/base/coord/file.py diff --git a/streaming/base/coord/file.py b/streaming/base/coord/file.py new file mode 100644 index 000000000..253dfbbac --- /dev/null +++ b/streaming/base/coord/file.py @@ -0,0 +1,60 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""File locking via file open mode 'x'.""" + +import os +from contextlib import contextmanager +from time import sleep, time +from typing import Iterator, Optional + +__all__ = ['acquire_file_lock', 'release_file_lock', 'file_lock'] + + +def acquire_file_lock(filename: str, timeout: Optional[float] = 60, tick: float = 0.007) -> float: + """Acquire a file lock.""" + start = time() + + 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.') + + while True: + try: + file = open(filename, 'xb', 0) + file.close() + break + except: + pass + + if timeout is not None and start + timeout < time(): + raise ValueError(f'Timed out while attempting to acquire file lock: file ' + + f'{filename}, timeout {timeout} sec.') + + sleep(tick) + + return time() - start + + +def release_file_lock(filename: str) -> None: + """Release a file lock.""" + if os.path.exists(filename): + if os.path.isfile(filename): + os.remove(filename) + else: + raise ValueError(f'Path exists, but is not a file: {filename}.') + else: + raise ValueError(f'Path does not exist: {filename}.') + + +@contextmanager +def file_lock(filename: str, timeout: Optional[float] = 60, tick: float = 0.007) -> Iterator[None]: + """A with statement protected by a file lock.""" + acquire_file_lock(filename, timeout, tick) + + yield + + release_file_lock(filename) diff --git a/streaming/base/coord/job/registry.py b/streaming/base/coord/job/registry.py index 30e739eda..260ab3256 100644 --- a/streaming/base/coord/job/registry.py +++ b/streaming/base/coord/job/registry.py @@ -12,9 +12,9 @@ from time import sleep, time_ns from typing import Dict, List, Sequence, Tuple -from filelock import FileLock from psutil import process_iter +from streaming.base.coord.file import file_lock from streaming.base.coord.job.entry import JobEntry from streaming.base.coord.job.file import JobFile from streaming.base.coord.world import World @@ -37,7 +37,7 @@ class JobRegistry: def __init__(self, config_root: str, tick: float = 0.007) -> None: self.config_root = config_root self._tick = tick - self._filelock_filename = os.path.join(config_root, 'filelock.bin') + self._lock_filename = os.path.join(config_root, 'registry.lock') self._registry_filename = os.path.join(config_root, 'registry.json') def _get_live_procs(self) -> Dict[int, int]: @@ -143,7 +143,7 @@ def _wait_for_existence(self, job_hash: str) -> None: dirname = os.path.join(self.config_root, job_hash) while True: sleep(self._tick) - with FileLock(self._filelock_filename): + with file_lock(self._lock_filename): if os.path.exists(dirname): break @@ -156,7 +156,7 @@ def _wait_for_removal(self, job_hash: str) -> None: dirname = os.path.join(self.config_root, job_hash) while True: sleep(self._tick) - with FileLock(self._filelock_filename): + with file_lock(self._lock_filename): if not os.path.exists(dirname): break @@ -188,7 +188,7 @@ def _register(self, streams: Sequence[Stream]) -> str: process_id=pid, register_time=register_time) - with FileLock(self._filelock_filename): + with file_lock(self._lock_filename): reg = JobFile.read(self._registry_filename) reg.add(entry) del_job_hashes = reg.filter(pid2create_time) @@ -245,7 +245,7 @@ def _unregister(self, job_hash: str) -> None: """ pid2create_time = self._get_live_procs() - with FileLock(self._filelock_filename): + with file_lock(self._lock_filename): reg = JobFile.read(self._registry_filename) reg.remove(job_hash) del_job_hashes = reg.filter(pid2create_time) From 94d413608ad58cd847844e7c1c03281a294389e2 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 30 Dec 2023 01:01:43 -0800 Subject: [PATCH 28/54] Handle case where a process dies while holding a soft file lock. --- streaming/base/coord/file.py | 16 ++++++++++++++-- streaming/base/coord/job/registry.py | 5 ++--- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/streaming/base/coord/file.py b/streaming/base/coord/file.py index 253dfbbac..e84c9789c 100644 --- a/streaming/base/coord/file.py +++ b/streaming/base/coord/file.py @@ -24,12 +24,24 @@ def acquire_file_lock(filename: str, timeout: Optional[float] = 60, tick: float while True: try: - file = open(filename, 'xb', 0) - file.close() + text = str(os.getpid()) + with open(filename, 'x') as file: + file.write(text) break except: pass + try: + with open(filename, 'r') as file: + text = file.read() + pid = int(text) + + if pid != os.getpid(): + os.remove(filename) + continue + except: + pass + if timeout is not None and start + timeout < time(): raise ValueError(f'Timed out while attempting to acquire file lock: file ' + f'{filename}, timeout {timeout} sec.') diff --git a/streaming/base/coord/job/registry.py b/streaming/base/coord/job/registry.py index 260ab3256..e2ed2d2aa 100644 --- a/streaming/base/coord/job/registry.py +++ b/streaming/base/coord/job/registry.py @@ -232,7 +232,7 @@ def register(self, streams: Sequence[Stream], world: World) -> str: job_hash = self._register(streams) else: job_hash = self._lookup(streams) - self._wait_for_existence(job_hash) + self._wait_for_existence(job_hash) return job_hash def _unregister(self, job_hash: str) -> None: @@ -265,5 +265,4 @@ def unregister(self, job_hash: str, world: World) -> None: if world.is_local_leader: self._unregister(job_hash) else: - pass - self._wait_for_removal(job_hash) + self._wait_for_removal(job_hash) From f38e99f0f3681c428b5279c84968608c629b9607 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 30 Dec 2023 01:06:10 -0800 Subject: [PATCH 29/54] Fix. --- streaming/base/coord/job/registry.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/streaming/base/coord/job/registry.py b/streaming/base/coord/job/registry.py index f8b1e0835..e2ed2d2aa 100644 --- a/streaming/base/coord/job/registry.py +++ b/streaming/base/coord/job/registry.py @@ -116,11 +116,6 @@ def _hash_streams(self, streams: Sequence[Stream]) -> Tuple[List[str], List[str] return stream_locals, stream_hashes, job_hash - def _ensure_filelock(self) -> None: - """Ensure the lazily-created file lock exists.""" - if not hasattr(self, '_filelock'): - self._filelock = FileLock(self._filelock_filename) - def _make_dir(self, job_hash: str) -> None: """Create a Streaming job config dir. @@ -178,8 +173,6 @@ def _register(self, streams: Sequence[Stream]) -> str: Returns: str: Streaming config subdir for this job. """ - self._ensure_filelock() - register_time = time_ns() pid2create_time = self._get_live_procs() pid = os.getpid() @@ -250,8 +243,6 @@ def _unregister(self, job_hash: str) -> None: Args: job_hash (str): Subdir identifying this Streaming job. """ - self._ensure_filelock() - pid2create_time = self._get_live_procs() with file_lock(self._lock_filename): From 9743cd6d07e533ab45ffb06149173ddb7af17cfa Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 30 Dec 2023 03:03:34 -0800 Subject: [PATCH 30/54] Rewrite the homebrew soft file lock. --- streaming/base/coord/file.py | 170 +++++++++++++++++++-------- streaming/base/coord/job/registry.py | 11 +- streaming/base/coord/process.py | 20 ++++ 3 files changed, 147 insertions(+), 54 deletions(-) create mode 100644 streaming/base/coord/process.py diff --git a/streaming/base/coord/file.py b/streaming/base/coord/file.py index e84c9789c..c6d0ef70c 100644 --- a/streaming/base/coord/file.py +++ b/streaming/base/coord/file.py @@ -1,72 +1,144 @@ # Copyright 2023 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 -"""File locking via file open mode 'x'.""" +"""Soft file locking via file open mode 'x'.""" import os -from contextlib import contextmanager from time import sleep, time -from typing import Iterator, Optional +from types import TracebackType +from typing import Optional, Type -__all__ = ['acquire_file_lock', 'release_file_lock', 'file_lock'] +from typing_extensions import Self +from streaming.base.coord.process import get_live_processes -def acquire_file_lock(filename: str, timeout: Optional[float] = 60, tick: float = 0.007) -> float: - """Acquire a file lock.""" - start = time() +__all__ = ['SoftFileLock'] - 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.') +class SoftFileLock: + """Soft file locking via file open mode 'x'. - while True: - try: - text = str(os.getpid()) - with open(filename, 'x') as file: - file.write(text) - break - except: - pass + Args: + filename (str): Path to lock. + timeout (float, optional): How long to wait in seconds before raising an exception. + Defaults to ``10``. + tick (float): Polling interval in seconds. Defaults to ``0.007``. + """ - try: - with open(filename, 'r') as file: - text = file.read() - pid = int(text) + def __init__(self, filename: str, timeout: Optional[float] = 10, tick: float = 0.007) -> None: + if not filename: + raise ValueError('Path to file lock is empty.') - if pid != os.getpid(): - os.remove(filename) - continue - except: - pass + if timeout is not None: + if timeout <= 0: + raise ValueError( + f'Timeout must be positive float seconds, but got: {timeout} sec.') - if timeout is not None and start + timeout < time(): - raise ValueError(f'Timed out while attempting to acquire file lock: file ' + - f'{filename}, timeout {timeout} sec.') + if tick <= 0: + raise ValueError(f'Tick must be positive float seconds, but got: {tick} sec.') - sleep(tick) + self.filename = filename + self.timeout = timeout + self.tick = tick - return time() - start + self._garbage_collect(filename) + dirname = os.path.dirname(filename) + os.makedirs(dirname, exist_ok=True) -def release_file_lock(filename: str) -> None: - """Release a file lock.""" - if os.path.exists(filename): - if os.path.isfile(filename): - os.remove(filename) - else: - raise ValueError(f'Path exists, but is not a file: {filename}.') - else: - raise ValueError(f'Path does not exist: {filename}.') + @classmethod + def _set_pid(cls, filename: str, pid: int) -> None: + """Set the locking process's pid. + + Args: + filename (str): Path to lock. + """ + with open(filename, 'x') as file: + file.write(str(pid)) + @classmethod + def _get_pid(cls, filename: str) -> int: + """Get the locking process's pid. -@contextmanager -def file_lock(filename: str, timeout: Optional[float] = 60, tick: float = 0.007) -> Iterator[None]: - """A with statement protected by a file lock.""" - acquire_file_lock(filename, timeout, tick) + Args: + filename (str): Path to lock. + """ + with open(filename, 'r') as file: + text = file.read() + return int(text) - yield + @classmethod + def _try_remove(cls, filename: str) -> None: + """Try to remove this lock. - release_file_lock(filename) + Args: + filename (str): Path to lock. + """ + try: + os.remove(filename) + except: + pass + + @classmethod + def _garbage_collect(cls, filename: str) -> None: + """Release this lock if held by a dead process. + + Args: + filename (str): Path to lock. + """ + try: + pid = cls._get_pid(filename) + except: + cls._try_remove(filename) + return + + if pid not in get_live_processes(): + cls._try_remove(filename) + + def acquire(self) -> None: + """Acquire this lock.""" + start = time() + + while True: + try: + self._set_pid(self.filename, os.getpid()) + break + except: + pass + + if self.timeout is not None and start + self.timeout < time(): + raise ValueError(f'Timed out while attempting to acquire file lock: file ' + + f'{self.filename}, timeout {self.timeout} sec.') + + sleep(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 object. + """ + 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/job/registry.py b/streaming/base/coord/job/registry.py index e2ed2d2aa..048ddb84b 100644 --- a/streaming/base/coord/job/registry.py +++ b/streaming/base/coord/job/registry.py @@ -14,7 +14,7 @@ from psutil import process_iter -from streaming.base.coord.file import file_lock +from streaming.base.coord.file import SoftFileLock from streaming.base.coord.job.entry import JobEntry from streaming.base.coord.job.file import JobFile from streaming.base.coord.world import World @@ -38,6 +38,7 @@ def __init__(self, config_root: str, tick: float = 0.007) -> None: self.config_root = config_root 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]: @@ -143,7 +144,7 @@ def _wait_for_existence(self, job_hash: str) -> None: dirname = os.path.join(self.config_root, job_hash) while True: sleep(self._tick) - with file_lock(self._lock_filename): + with self._lock: if os.path.exists(dirname): break @@ -156,7 +157,7 @@ def _wait_for_removal(self, job_hash: str) -> None: dirname = os.path.join(self.config_root, job_hash) while True: sleep(self._tick) - with file_lock(self._lock_filename): + with self._lock: if not os.path.exists(dirname): break @@ -188,7 +189,7 @@ def _register(self, streams: Sequence[Stream]) -> str: process_id=pid, register_time=register_time) - with file_lock(self._lock_filename): + with self._lock: reg = JobFile.read(self._registry_filename) reg.add(entry) del_job_hashes = reg.filter(pid2create_time) @@ -245,7 +246,7 @@ def _unregister(self, job_hash: str) -> None: """ pid2create_time = self._get_live_procs() - with file_lock(self._lock_filename): + with self._lock: reg = JobFile.read(self._registry_filename) reg.remove(job_hash) del_job_hashes = reg.filter(pid2create_time) diff --git a/streaming/base/coord/process.py b/streaming/base/coord/process.py new file mode 100644 index 000000000..f78f6386c --- /dev/null +++ b/streaming/base/coord/process.py @@ -0,0 +1,20 @@ +# Copyright 2023 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 From b32785871e8fc84cd6a9c8fab8470b7a48928705 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 30 Dec 2023 04:30:14 -0800 Subject: [PATCH 31/54] Docstrings. --- streaming/base/coord/file.py | 95 +++++++++++++++++++++++------------- 1 file changed, 61 insertions(+), 34 deletions(-) diff --git a/streaming/base/coord/file.py b/streaming/base/coord/file.py index c6d0ef70c..ae339e95d 100644 --- a/streaming/base/coord/file.py +++ b/streaming/base/coord/file.py @@ -6,7 +6,7 @@ import os from time import sleep, time from types import TracebackType -from typing import Optional, Type +from typing import Optional, Type, Union from typing_extensions import Self @@ -21,11 +21,11 @@ class SoftFileLock: Args: filename (str): Path to lock. timeout (float, optional): How long to wait in seconds before raising an exception. - Defaults to ``10``. + Set to ``None`` to never time out. Defaults to ``30``. tick (float): Polling interval in seconds. Defaults to ``0.007``. """ - def __init__(self, filename: str, timeout: Optional[float] = 10, tick: float = 0.007) -> None: + 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.') @@ -41,14 +41,11 @@ def __init__(self, filename: str, timeout: Optional[float] = 10, tick: float = 0 self.timeout = timeout self.tick = tick - self._garbage_collect(filename) - - dirname = os.path.dirname(filename) - os.makedirs(dirname, exist_ok=True) + self._normalize(filename) @classmethod - def _set_pid(cls, filename: str, pid: int) -> None: - """Set the locking process's pid. + def _write(cls, filename: str, pid: int) -> None: + """Write the locking process's pid. Args: filename (str): Path to lock. @@ -57,59 +54,89 @@ def _set_pid(cls, filename: str, pid: int) -> None: file.write(str(pid)) @classmethod - def _get_pid(cls, filename: str) -> int: - """Get the locking process's pid. + def _read(cls, filename: str) -> int: + """Read the locking process's pid. Args: filename (str): Path to lock. """ with open(filename, 'r') as file: - text = file.read() - return int(text) + return int(file.read()) @classmethod - def _try_remove(cls, filename: str) -> None: - """Try to remove this lock. + 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) + 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: - os.remove(filename) + pid = cls._read(filename) except: - pass + 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 _garbage_collect(cls, filename: str) -> None: - """Release this lock if held by a dead process. + def _get_timeout(cls, + init_timeout: Optional[float], + timeout: Optional[Union[str, float]] = 'auto') -> Optional[float]: + """Determine the timeout for a given acquire(). Args: - filename (str): Path to lock. + 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. """ - try: - pid = cls._get_pid(filename) - except: - cls._try_remove(filename) - return + 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 - if pid not in get_live_processes(): - cls._try_remove(filename) + def acquire(self, timeout: Optional[Union[str, float]] = 'auto') -> None: + """Acquire this lock. - def acquire(self) -> None: - """Acquire this lock.""" + Args: + timeout (str | float, optional): Override timeout for just this method call. + """ start = time() - + timeout = self._get_timeout(self.timeout, timeout) while True: try: - self._set_pid(self.filename, os.getpid()) + self._write(self.filename, os.getpid()) break except: pass - if self.timeout is not None and start + self.timeout < time(): raise ValueError(f'Timed out while attempting to acquire file lock: file ' + f'{self.filename}, timeout {self.timeout} sec.') - sleep(self.tick) def release(self) -> None: @@ -125,7 +152,7 @@ def __enter__(self) -> Self: """Enter context manager. Returns: - Self: This object. + Self: This lock. """ self.acquire() return self From 0e475454692528e2e646dcd24da907a7e6fb0a44 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 30 Dec 2023 04:54:30 -0800 Subject: [PATCH 32/54] Switch all filelock.FileLock to streaming.base.coord.file.SoftFileLock. --- streaming/base/coord/shmem/barrier.py | 87 +++++++++++++-------------- streaming/base/dataset.py | 69 +++++++-------------- tests/test_barrier.py | 18 +++--- 3 files changed, 73 insertions(+), 101 deletions(-) diff --git a/streaming/base/coord/shmem/barrier.py b/streaming/base/coord/shmem/barrier.py index 6cc9988af..20a3085e8 100644 --- a/streaming/base/coord/shmem/barrier.py +++ b/streaming/base/coord/shmem/barrier.py @@ -9,9 +9,9 @@ from time import sleep import numpy as np -from filelock import FileLock from streaming.base.constant import TICK +from streaming.base.coord.file import SoftFileLock from streaming.base.coord.shmem.array import SharedArray # Time out to wait before raising exception @@ -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/dataset.py b/streaming/base/dataset.py index 21a276ac5..4bf088286 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -18,14 +18,15 @@ from typing import Any, Dict, Iterator, Optional, Sequence, Tuple, Union import numpy as np -from filelock import FileLock from numpy.typing import NDArray 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, CACHE_FILELOCK, CACHE_USAGE, EPOCH_DATA, EPOCH_SHAPE, - NEXT_EPOCH, RESUME, SHARD_ACCESS_TIMES, SHARD_STATES, TICK) +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.coord.file import SoftFileLock from streaming.base.coord.job import JobDirectory, JobRegistry from streaming.base.coord.shmem import (SharedArray, SharedBarrier, SharedMemory, SharedScalar, _get_path) @@ -552,7 +553,7 @@ def __init__( os.makedirs(self._filelock_root, exist_ok=True) # Create the shared memory-backed barrier, without its lock, which is unpickleable. - self._shared_barrier = SharedBarrier(self.job_dir.get_filename('barrier_filelock.bin'), + self._shared_barrier = SharedBarrier(self.job_dir.get_filename(BARRIER_FILELOCK), _get_path(self._shm_prefix_int, BARRIER)) # Epoch counter. @@ -563,8 +564,8 @@ def __init__( self._next_epoch = SharedScalar(np.int64, _get_path(self._shm_prefix_int, NEXT_EPOCH)) # Cache filelock. Protects downloading and evicting shards. - self._cache_filelock_path = self.job_dir.get_filename('cache_filelock.bin') - self._cache_filelock: FileLock + filename = self.job_dir.get_filename(CACHE_FILELOCK) + self._cache_lock = SoftFileLock(filename) # Cache usage in bytes. self._cache_usage = SharedScalar(np.int64, _get_path(self._shm_prefix_int, CACHE_USAGE)) @@ -622,8 +623,6 @@ def __init__( self._executor: ThreadPoolExecutor self._event: Event - del self._shared_barrier.lock # Remote the lock that makes it unpickleable. - @property def size(self) -> int: """Get the size of the dataset in samples. @@ -763,11 +762,6 @@ 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) @@ -980,11 +974,6 @@ 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) - # 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, @@ -1010,7 +999,7 @@ def _get_work(self, world: World, epoch: int, sample_in_epoch: int) -> NDArray[n 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. @@ -1075,12 +1064,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: @@ -1088,12 +1072,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: @@ -1109,12 +1088,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] @@ -1138,21 +1112,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: @@ -1174,16 +1147,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/tests/test_barrier.py b/tests/test_barrier.py index 0d8f206be..72e303e97 100644 --- a/tests/test_barrier.py +++ b/tests/test_barrier.py @@ -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()) From 0658ab9cb1cdbd596083889424b72970af778fb4 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 30 Dec 2023 04:59:34 -0800 Subject: [PATCH 33/54] Fix. --- streaming/base/coord/file.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/streaming/base/coord/file.py b/streaming/base/coord/file.py index ae339e95d..54af41055 100644 --- a/streaming/base/coord/file.py +++ b/streaming/base/coord/file.py @@ -72,7 +72,8 @@ def _normalize(cls, filename: str) -> None: """ # Ensure the file's parent directory exists so we can write it in one shot. dirname = os.path.dirname(filename) - os.makedirs(dirname, exist_ok=True) + if dirname: + os.makedirs(dirname, exist_ok=True) # If no file, we don't need to do anything. if not os.path.exists(filename): From 1d90d42256a1e29eeeba65ab467afe160582c702 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 30 Dec 2023 05:56:46 -0800 Subject: [PATCH 34/54] Rewrite StreamingDataset/SimulationDataset args handling to be rigorous. --- simulation/core/sim_dataset.py | 198 +++++----- streaming/base/batching/per_stream.py | 3 - streaming/base/batching/random.py | 3 - streaming/base/batching/stratified.py | 3 - streaming/base/dataset.py | 516 ++++++++++++++++---------- tests/test_streaming.py | 2 +- 6 files changed, 409 insertions(+), 316 deletions(-) diff --git a/simulation/core/sim_dataset.py b/simulation/core/sim_dataset.py index 1b78f03c5..a4ca7d16e 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,52 +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() @@ -139,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. @@ -211,15 +189,14 @@ def __init__(self, keep_zip=keep_zip, allow_unsafe_types=allow_unsafe_types) else: - stream = Stream(remote=remote, - local=local, - split=split, - download_retry=download_retry, - download_timeout=download_timeout, - validate_hash=validate_hash, - keep_zip=keep_zip, - allow_unsafe_types=allow_unsafe_types) - streams = [stream] + streams = Stream(remote=remote, + local=local, + split=split, + download_retry=download_retry, + download_timeout=download_timeout, + validate_hash=validate_hash, + 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. @@ -419,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: @@ -457,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]: @@ -529,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/streaming/base/batching/per_stream.py b/streaming/base/batching/per_stream.py index 99944aa7c..c313c5dc3 100644 --- a/streaming/base/batching/per_stream.py +++ b/streaming/base/batching/per_stream.py @@ -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 76050848a..113d5c360 100644 --- a/streaming/base/batching/random.py +++ b/streaming/base/batching/random.py @@ -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 4dfed207a..cecfb787b 100644 --- a/streaming/base/batching/stratified.py +++ b/streaming/base/batching/stratified.py @@ -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/dataset.py b/streaming/base/dataset.py index 4bf088286..93eefb3de 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -7,7 +7,6 @@ import logging import os import sys -import warnings from concurrent.futures import ThreadPoolExecutor, wait from concurrent.futures._base import Future from enum import IntEnum @@ -16,6 +15,7 @@ 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 numpy.typing import NDArray @@ -167,35 +167,6 @@ def on_exit(self) -> None: self._num_exited += 1 -def _test_config_root(config_root: str) -> None: - """Validate that the provided config root is usable. - - 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. - - 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.') - - -def _get_default_config_root() -> str: - """Get the default Streaming configuration root directory. - - Returns: - str: Default Streaming configuration root directory. - """ - return os.path.join(gettempdir(), 'streaming') - - class StreamingDataset(Array, IterableDataset): """A mid-epoch-resumable streaming/caching pytorch IterableDataset. @@ -237,6 +208,10 @@ class StreamingDataset(Array, IterableDataset): * How to iterate (the StreamingDataset arguments): + * Configuration: + + * ``config_root`` + * Shard lifecycle: * ``predownload`` @@ -264,10 +239,6 @@ class StreamingDataset(Array, IterableDataset): * ``batching_method`` - * Configuration: - - * ``config_root`` - 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 @@ -296,6 +267,9 @@ class StreamingDataset(Array, IterableDataset): 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 @@ -338,54 +312,64 @@ class StreamingDataset(Array, IterableDataset): ``None``. batching_method (str): Which batching method to use, either ``random``, ``stratified``, or ``per_stream``. Defaults to ``random``. - 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. """ 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, - 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', - config_root: str = _get_default_config_root(), + 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: - # 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 + # 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 - - _test_config_root(config_root) - self.config_root = config_root + 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. @@ -396,50 +380,6 @@ def __init__( 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. The 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: @@ -457,32 +397,22 @@ def __init__( keep_zip=keep_zip, allow_unsafe_types=allow_unsafe_types) else: - stream = Stream(remote=remote, - local=local, - split=split, - download_retry=download_retry, - download_timeout=download_timeout, - validate_hash=validate_hash, - keep_zip=keep_zip, - allow_unsafe_types=allow_unsafe_types) - streams = [stream] + streams = Stream(remote=remote, + local=local, + split=split, + download_retry=download_retry, + download_timeout=download_timeout, + validate_hash=validate_hash, + 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 = [] @@ -508,8 +438,6 @@ def __init__( # 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 ' + @@ -540,20 +468,20 @@ def __init__( self.length = ceil(self.epoch_size / world.num_ranks) # Register/lookup our shared memory prefix and filelock root directory. - self.registry = JobRegistry(config_root) - self.job_dir = JobDirectory(self.registry, streams, world) - self._shm_prefix_int = int(self.job_dir.job_hash, 16) + self.registry = JobRegistry(self.config_root) + self.job = JobDirectory(self.registry, streams, world) + self._shm_prefix_int = int(self.job.job_hash, 16) - init_done_filename = self.job_dir.get_filename('init_done.txt') + init_done_filename = self.job.get_filename('init_done.txt') if world.is_local_leader: if os.path.exists(init_done_filename): os.remove(init_done_filename) - self._filelock_root = os.path.join(self.registry.config_root, self.job_dir.job_hash) + self._filelock_root = os.path.join(self.config_root, self.job.job_hash) os.makedirs(self._filelock_root, exist_ok=True) # Create the shared memory-backed barrier, without its lock, which is unpickleable. - self._shared_barrier = SharedBarrier(self.job_dir.get_filename(BARRIER_FILELOCK), + self._shared_barrier = SharedBarrier(self.job.get_filename(BARRIER_FILELOCK), _get_path(self._shm_prefix_int, BARRIER)) # Epoch counter. @@ -564,7 +492,7 @@ def __init__( self._next_epoch = SharedScalar(np.int64, _get_path(self._shm_prefix_int, NEXT_EPOCH)) # Cache filelock. Protects downloading and evicting shards. - filename = self.job_dir.get_filename(CACHE_FILELOCK) + filename = self.job.get_filename(CACHE_FILELOCK) self._cache_lock = SoftFileLock(filename) # Cache usage in bytes. @@ -623,6 +551,250 @@ def __init__( self._executor: ThreadPoolExecutor self._event: Event + @classmethod + def _test_config_root(cls, config_root: str) -> None: + """Validate that the provided config root is usable. + + 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. + + 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.') + + @classmethod + def _get_config_root(cls, config_root: Optional[str]) -> str: + """Get the default Streaming configuration root directory. + + 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: """Get the size of the dataset in samples. @@ -676,17 +848,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. @@ -703,19 +864,6 @@ def _resume(self, world: World, epoch: int) -> Tuple[int, int]: 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) return epoch, 0 # SharedMemory buffers may contain additional null bytes at the end. @@ -726,31 +874,17 @@ def _resume(self, world: World, epoch: int) -> Tuple[int, int]: # Check if the resume state is stale. 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. 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]: diff --git a/tests/test_streaming.py b/tests/test_streaming.py index f06496dbe..409e4ffe9 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -782,7 +782,7 @@ def test_streamingdataloader_mid_epoch_resumption(local_remote_dir: Any, batch_s sample_order.extend(batch['id'][:]) del dataloader - del dataset.job_dir + del dataset.job del dataset clean_stale_shared_memory() From da9693f8ae2c60cebc9b1ef5bd3ed49e3e43cc28 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 30 Dec 2023 06:09:00 -0800 Subject: [PATCH 35/54] MMap-based cross-process Array, Barrier, Buffer, Number. --- streaming/base/coord/mmap/__init__.py | 11 +++ streaming/base/coord/mmap/array.py | 84 +++++++++++++++++ streaming/base/coord/mmap/barrier.py | 130 ++++++++++++++++++++++++++ streaming/base/coord/mmap/base.py | 116 +++++++++++++++++++++++ streaming/base/coord/mmap/buffer.py | 42 +++++++++ streaming/base/coord/mmap/number.py | 54 +++++++++++ 6 files changed, 437 insertions(+) create mode 100644 streaming/base/coord/mmap/__init__.py create mode 100644 streaming/base/coord/mmap/array.py create mode 100644 streaming/base/coord/mmap/barrier.py create mode 100644 streaming/base/coord/mmap/base.py create mode 100644 streaming/base/coord/mmap/buffer.py create mode 100644 streaming/base/coord/mmap/number.py diff --git a/streaming/base/coord/mmap/__init__.py b/streaming/base/coord/mmap/__init__.py new file mode 100644 index 000000000..7608cfed0 --- /dev/null +++ b/streaming/base/coord/mmap/__init__.py @@ -0,0 +1,11 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Share data across processes with mmap().""" + +from streaming.base.coord.mmap.array import MMapArray +from streaming.base.coord.mmap.barrier import MMapBarrier +from streaming.base.coord.mmap.buffer import MMapBuffer +from streaming.base.coord.mmap.number import MMapNumber + +__all__ = ['MMapArray', 'MMapBarrier', 'MMapBuffer', 'MMapNumber'] diff --git a/streaming/base/coord/mmap/array.py b/streaming/base/coord/mmap/array.py new file mode 100644 index 000000000..403c96b2a --- /dev/null +++ b/streaming/base/coord/mmap/array.py @@ -0,0 +1,84 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Share an array across processes using mmap().""" + +from mmap import mmap +from typing import Generic, Optional, Tuple, TypeVar, Union + +import numpy as np +from numpy.typing import NDArray + +from streaming.base.coord.mmap.base import ensure_file + +__all__ = ['MMapArray'] + +DType = TypeVar('DType', bound=np.number) + +IndexType = Union[int, slice, NDArray[np.integer]] +DataType = Union[DType, NDArray[DType]] + + +class MMapArray(Generic[DType]): + """Share an array across processes using mmap(). + + Args: + mode (str): Whether to ``create``, ``replace``, or ``attach``. Defaults to ``attach``. + filename (str): Path to memory-mapped file. + shape (int | Tuple[int], optional): Exact required shape, if known in advance. At most one + wildcard ``-1`` is acceptable. + dtype (DType): Data type of the number. + """ + + def __init__( + self, + *, + mode: str = 'attach', + filename: str, + shape: Optional[Union[int, Tuple[int]]] = None, + dtype: DType, + ) -> None: + self.mode = mode + self.filename = filename + self.shape = ensure_file(mode, filename, shape, 1) + self.dtype = dtype + self.file = open(filename, 'r+b', 0) + self.data = mmap(self.file.fileno(), 0) + + def __len__(self) -> int: + """Get the number of elements in the first axis of the array. + + Returns: + int: Length of the first axis of the array. + """ + return int(self.shape[0]) + + def as_array(self) -> NDArray[DType]: + """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[DType]: Our internal buffer as an ndarray. + """ + return np.ndarray(self.shape, buffer=self.data, dtype=self.dtype) + + def __getitem__(self, index: IndexType) -> DataType: + """Get the item at the index. + + Args: + index (IndexType): The index(es). + + Returns: + DataType; The item(s). + """ + return self.as_array()[index] + + def __setitem__(self, index: IndexType, item: DataType) -> None: + """Set the item at the index. + + Args: + index (IndexType): The index(es). + item (DataType): The item(s). + """ + self.as_array()[index] = item diff --git a/streaming/base/coord/mmap/barrier.py b/streaming/base/coord/mmap/barrier.py new file mode 100644 index 000000000..f311649e2 --- /dev/null +++ b/streaming/base/coord/mmap/barrier.py @@ -0,0 +1,130 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Share a barrier across processes using mmap().""" + +from time import sleep + +import numpy as np + +from streaming.base.coord.file import SoftFileLock +from streaming.base.coord.mmap.array import MMapArray + +__all__ = ['MMapBarrier'] + + +class MMapBarrier: + """Share a barrier across processes using mmap(). + + Args: + mode (str): Whether to ``create``, ``replace``, or ``attach``. Defaults to ``attach``. + mmap_filename (str): Path to memory-mapped file. + lock_filename (str): Path to SoftFileLock file. + tick (float): Polling interval in seconds. Defaults to ``0.007``. + """ + + def __init__( + self, + *, + mode: str = 'attach', + mmap_filename: str, + lock_filename: str, + tick: float = 0.007, + ) -> None: + self._lock = SoftFileLock(lock_filename) + self._arr = MMapArray(mode=mode, filename=mmap_filename, shape=3, dtype=np.int32()) + self._tick = tick + + self._num_enter = 0 + self._num_exit = -1 + self._flag = True + + @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] = np.int32(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] = np.int32(num_exit) + + @property + def _flag(self) -> bool: + """Getter for _flag. + + Returns: + bool: Flag value. + """ + return bool(self._arr[2]) + + @_flag.setter + def _flag(self, flag: bool) -> None: + """Setter for _flag. + + Args: + flag (bool): Flag value. + """ + self._arr[2] = np.int32(flag) + + def __call__(self, total: int) -> None: + + # 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 "don't go". + self._lock.acquire() + if not self._num_enter: + self._lock.release() + while self._num_exit != total: + sleep(self._tick) + self._lock.acquire() + self._flag = False + + # Note that we entered. + 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 = True + self._lock.release() + + # Everybody waits until the flag is set to "go". + while not self._flag: + sleep(self._tick) + + # Note that we exited. + with self._lock: + self._num_exit += 1 + if self._num_exit == total: + self._num_exit = -1 diff --git a/streaming/base/coord/mmap/base.py b/streaming/base/coord/mmap/base.py new file mode 100644 index 000000000..199bd2bbe --- /dev/null +++ b/streaming/base/coord/mmap/base.py @@ -0,0 +1,116 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Base functionality for sharing data across processes using mmap().""" + +import os +from typing import Optional, Tuple, Union + +import numpy as np + +__all__ = ['ensure_file'] + + +def _normalize_shape(shape: Optional[Union[int, Tuple[int]]]) -> \ + Tuple[Tuple[int], int, Optional[int]]: + """Normalize and validate a shape argument. + + Args: + shape (int | Tuple[int], optional): Input shape. + + Returns: + Tuple[Tuple[int], int, Optional[int]]: Normalized shape, number of elements without the + wildcard if present, and bytes per element. + """ + if shape is None: + shape = -1, + elif isinstance(shape, int): + shape = shape, + + num_wild = 0 + for dim in 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}.') + + numel = int(np.prod(shape)) + if numel < 0: + numel = -numel + wild_index = shape.index(-1) + else: + wild_index = None + + return shape, numel, wild_index + + +def ensure_file(mode: str, filename: str, shape: Optional[Union[int, Tuple[int]]], + unit: int) -> Tuple[int]: + """Ensure file existence and size according to mode. + + Args: + mode (str): Whether to ``create``, ``replace``, or ``attach``. Defaults to ``attach``. + filename (str): Path to memory-mapped file. + shape (int | Tuple[int], optional): Exact required number of units, along each axis, if + known in advance. At most one wildcard ``-1`` is acceptable. + unit (int): Stride of a single value in bytes. + + Returns: + int: Resulting exact shape. + """ + want_shape, want_numel, want_wild_index = _normalize_shape(shape) + + if unit < 1: + raise ValueError(f'{unit} must be a positive integer, but got: {unit}.') + + # Normalize file existence by mode. + if mode == 'create': + if os.path.exists(filename): + raise ValueError(f'File alreadfy exists: {filename}.') + elif mode == 'replace': + if os.path.exists(filename): + os.remove(filename) + elif mode == 'attach': + if not os.path.exists(filename): + raise ValueError(f'File does not exist: {filename}.') + else: + modes = {'create', 'replace', 'attach'} + raise ValueError(f'`mode` must be either replace,one of {sorted(modes)}, but got: {mode}.') + + # Perform the work. + if os.path.exists(filename): + # Use size info to validate the pre-existing file. + got_size = os.stat(filename).st_size + if want_wild_index is None: + want_size = want_numel * unit + if got_size != want_size: + raise ValueError(f'File is the wrong size: file {filename}, expected shape ' + + f'{want_shape}, expected unit {unit}, expected size ' + + f'{want_size}, actual size {got_size}.') + got_shape = want_numel, + else: + want_size = want_numel * unit + if got_size % want_size: + raise ValueError(f'File size is not evenly divisible: file {filename}, expected ' + + f'shape {want_shape}, expected unit {unit}, expected size to ' + + f'be divisible by {want_size}.') + wild_value = got_size // want_size + got_shape = list(want_shape) + got_shape[want_wild_index] = wild_value + got_shape = tuple(got_shape) + else: + # Use size info to create the (initially sparse) file. + if want_wild_index is not None: + raise ValueError(f'You must provide `shape`, without wildcards, in order to size ' + + f'the file: {filename}.') + with open(filename, 'wb') as out: + out.write(b'') + os.truncate(filename, want_numel * unit) + got_shape = want_shape + + # Return resulting exact shape. + return got_shape diff --git a/streaming/base/coord/mmap/buffer.py b/streaming/base/coord/mmap/buffer.py new file mode 100644 index 000000000..789f2383c --- /dev/null +++ b/streaming/base/coord/mmap/buffer.py @@ -0,0 +1,42 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Share a buffer across processes using mmap().""" + +from mmap import mmap +from typing import Optional + +from streaming.base.coord.mmap.base import ensure_file + +__all__ = ['MMapBuffer'] + + +class MMapBuffer: + """Share a buffer across processes using mmap(). + + Args: + mode (str): Whether to ``create``, ``replace``, or ``attach``. Defaults to ``attach``. + filename (str): Path to memory-mapped file. + size (int, optional): Exact required size, if known in advance. Defaults to ``None``. + """ + + def __init__( + self, + *, + mode: str = 'attach', + filename: str, + size: Optional[int] = None, + ) -> None: + self.mode = mode + self.filename = filename + self.size, = ensure_file(mode, filename, size, 1) + self.file = open(filename, 'r+b', 0) + self.data = mmap(self.file.fileno(), 0) + + def __len__(self) -> int: + """Get the number of bytes in the buffer. + + Returns: + int: Number of bytes in the buffer. + """ + return self.size diff --git a/streaming/base/coord/mmap/number.py b/streaming/base/coord/mmap/number.py new file mode 100644 index 000000000..f5d0fe03c --- /dev/null +++ b/streaming/base/coord/mmap/number.py @@ -0,0 +1,54 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Share a single number across processes using mmap().""" + +from mmap import mmap +from typing import Generic + +import numpy as np + +from streaming.base.coord.mmap.array import DType +from streaming.base.coord.mmap.base import ensure_file + +__init__ = ['MMapNumber'] + + +class MMapNumber(Generic[DType]): + """Share a single number across processes using mmap(). + + Args: + mode (str): Whether to ``create``, ``replace``, or ``attach``. Defaults to ``attach``. + filename (str): Path to memory-mapped file. + dtype (DType): Data type of the number. + """ + + def __init__( + self, + *, + mode: str = 'attach', + filename: str, + dtype: DType, + ) -> None: + self.mode = mode + self.filename = filename + ensure_file(mode, filename, 1, dtype.nbytes) + self.dtype = dtype + self.file = open(filename, 'r+b', 0) + self.data = mmap(self.file.fileno(), 0) + + def get(self) -> DType: + """Get our value. + + Returns: + DType: Our value. + """ + return np.frombuffer(self.data, self.dtype)[0] + + def set(self, value: DType) -> None: + """Set our value. + + Args: + value (DType): Our new value. + """ + self.data[:] = value.tobytes() From 3cd5fc78eb8321fdea21e3a9d5b779eee7a18570 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 30 Dec 2023 07:25:47 -0800 Subject: [PATCH 36/54] First attempt at replacing all SD shmem -> mmap. --- streaming/base/coord/mmap/array.py | 21 +++- streaming/base/dataset.py | 178 +++++++++++++++-------------- 2 files changed, 111 insertions(+), 88 deletions(-) diff --git a/streaming/base/coord/mmap/array.py b/streaming/base/coord/mmap/array.py index 403c96b2a..796f68f61 100644 --- a/streaming/base/coord/mmap/array.py +++ b/streaming/base/coord/mmap/array.py @@ -3,6 +3,7 @@ """Share an array across processes using mmap().""" +import os from mmap import mmap from typing import Generic, Optional, Tuple, TypeVar, Union @@ -40,8 +41,9 @@ def __init__( ) -> None: self.mode = mode self.filename = filename - self.shape = ensure_file(mode, filename, shape, 1) + self.shape = ensure_file(mode, filename, shape, dtype.nbytes) self.dtype = dtype + self.file = open(filename, 'r+b', 0) self.data = mmap(self.file.fileno(), 0) @@ -82,3 +84,20 @@ def __setitem__(self, index: IndexType, item: DataType) -> None: item (DataType): The item(s). """ self.as_array()[index] = item + + def flush(self) -> None: + """Flush the mmap.""" + self.data.flush() + + def close(self) -> None: + """Close the mmap and file handle.""" + if not self.data.closed: + self.data.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) diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index 93eefb3de..e41ff3dd0 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -23,13 +23,10 @@ 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.constant import TICK from streaming.base.coord.file import SoftFileLock from streaming.base.coord.job import JobDirectory, JobRegistry -from streaming.base.coord.shmem import (SharedArray, SharedBarrier, SharedMemory, SharedScalar, - _get_path) +from streaming.base.coord.mmap import MMapArray, MMapBarrier, MMapNumber from streaming.base.coord.world import World from streaming.base.format import get_index_basename from streaming.base.sampling import get_sampling @@ -469,44 +466,52 @@ def __init__( # Register/lookup our shared memory prefix and filelock root directory. self.registry = JobRegistry(self.config_root) - self.job = JobDirectory(self.registry, streams, world) - self._shm_prefix_int = int(self.job.job_hash, 16) + self.job = job = JobDirectory(self.registry, streams, world) init_done_filename = self.job.get_filename('init_done.txt') if world.is_local_leader: if os.path.exists(init_done_filename): os.remove(init_done_filename) - self._filelock_root = os.path.join(self.config_root, self.job.job_hash) - os.makedirs(self._filelock_root, exist_ok=True) + mode = 'create' if world.is_local_leader else 'attach' - # Create the shared memory-backed barrier, without its lock, which is unpickleable. - self._shared_barrier = SharedBarrier(self.job.get_filename(BARRIER_FILELOCK), - _get_path(self._shm_prefix_int, BARRIER)) + # Worker barrier. + self._worker_barrier = MMapBarrier(mode=mode, + mmap_filename=job.get_filename('worker_barrier.npy'), + lock_filename=job.get_filename('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)) + self._next_epoch = MMapNumber(mode=mode, + filename=job.get_filename('epoch.npy'), + dtype=np.int64()) - # Cache filelock. Protects downloading and evicting shards. - filename = self.job.get_filename(CACHE_FILELOCK) - self._cache_lock = SoftFileLock(filename) + # Cache filelock. + # + # Protects downloading and evicting shards. + self._cache_lock = SoftFileLock(job.get_filename('cache.lock')) # Cache usage in bytes. - self._cache_usage = SharedScalar(np.int64, _get_path(self._shm_prefix_int, CACHE_USAGE)) + self._cache_usage = MMapNumber(mode=mode, + filename=job.get_filename('cache_usage.npy'), + dtype=np.int64()) # 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 = MMapArray(mode=mode, + filename=job.get_filename('shard_states.npy'), + shape=self.num_shards, + dtype=np.uint8()) # 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 = MMapArray(mode=mode, + filename=job.get_filename('shard_access_times.npy'), + shape=self.num_shards, + dtype=np.uint64()) # Initialize shared memory objects. if world.is_local_leader: @@ -532,22 +537,16 @@ def __init__( self._shard_states[shard_id] = _ShardState.LOCAL if size else _ShardState.REMOTE self._shard_access_times[shard_id] = time_ns() - dirname = os.path.dirname(init_done_filename) - os.makedirs(dirname, exist_ok=True) - with open(init_done_filename, 'wb') as out: - out.write(b'') + with open(init_done_filename, 'x'): + pass else: wait_for_file_to_exist(init_done_filename, TICK, 300, 'Waited too long for initialization') - # Placeholder for a shared memory object where load_state_dict() saves its data to be - # picked up by __iter__(). - self._resume_shm: SharedMemory - # Placeholder for an _Iterator which tracks state during __iter__(). self._iterator: _Iterator - # For exception handling in __iter__ threads. + # Placeholder for exception handling in __iter__ threads. self._executor: ThreadPoolExecutor self._event: Event @@ -820,7 +819,7 @@ def next_epoch(self, next_epoch: int) -> None: Args: next_epoch (int): Next epoch. """ - self._next_epoch.set(next_epoch) + self._next_epoch.set(np.int64(next_epoch)) @property def cache_usage(self) -> int: @@ -838,7 +837,7 @@ def cache_usage(self, cache_usage: int) -> None: Args: cache_usage (int): Cache usage in bytes. """ - self._cache_usage.set(cache_usage) + self._cache_usage.set(np.int64(cache_usage)) def __len__(self) -> int: """Get the length as a PyTorch IterableDataset. @@ -858,25 +857,17 @@ 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 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: return epoch, 0 - # Load the correct resumption meta data. + # Load from checkpoint. epoch = obj['epoch'] sample_in_epoch = obj['sample_in_epoch'] self.shuffle_seed = obj['shuffle_seed'] @@ -901,7 +892,7 @@ def _resume_incr_epoch(self, world: World) -> Tuple[int, int]: 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: @@ -959,13 +950,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, @@ -1044,14 +1031,14 @@ 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]: + def _share_work(self, sample_ids: NDArray[np.int64]) -> Tuple[MMapArray, MMapArray]: """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. + Tuple[MMapArray, MMapArray]: Memory-mapped arrays containing shape and data. """ ndim = 5 @@ -1062,40 +1049,45 @@ def _share_work(self, sample_ids: NDArray[np.int64]) -> Tuple[SharedMemory, Shar 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() + work_shape = np.asarray(sample_ids.shape, np.int64) + work_shape_mmap = MMapArray(mode='create', + filename=self.job.get_filename('work_shape.npy'), + shape=work_shape.shape, + dtype=np.int64()) + work_shape_mmap[:] = work_shape + work_shape_mmap.flush() # 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() + work_mmap = MMapArray(mode='create', + filename=self.job.get_filename('work.npy'), + shape=sample_ids.shape, + dtype=np.int64()) + work_mmap[:] = sample_ids + work_mmap.flush() - return shape_shm, data_shm + return work_shape_mmap, work_mmap - def _attach_work(self) -> Tuple[NDArray[np.int64], SharedMemory, SharedMemory]: + def _attach_work(self) -> Tuple[MMapArray, MMapArray]: """Get an epoch's sample ordering from shared memory. Returns: - NDArray[np.int64]: Sample IDs. + Tuple[MMapArray, MMapArray]: Memory-mapped arrays containing shape and data. """ 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)) + work_shape_mmap = MMapArray(mode='attach', + filename=self.job.get_filename('work_shape.npy'), + shape=ndim, + dtype=np.int64()) + work_shape = tuple(work_shape_mmap.as_array()) - # 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) + work_mmap = MMapArray(mode='attach', + filename=self.job.get_filename('work.npy'), + shape=work_shape, + dtype=np.int64()) - return sample_ids, shape_shm, data_shm + return work_shape_mmap, work_mmap 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. @@ -1112,22 +1104,34 @@ def _get_work(self, world: World, epoch: int, sample_in_epoch: int) -> NDArray[n 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) + work_shape_mmap, work_mmap = self._share_work(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) + work_shape_mmap, work_mmap = self._attach_work() + epoch_sample_ids = work_mmap.as_array() # 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) + # Wait for each worker to extract their slice of the work. + self._worker_barrier(world.workers_per_node) + + # Close. + if not world.is_local_leader: + work_shape_mmap.close() + work_mmap.close() - # Now clean up after ourselves. - shape_shm.cleanup() - data_shm.cleanup() + # Wait for close. + self._worker_barrier(world.workers_per_node) + + # Delete. + if world.is_local_leader: + work_shape_mmap.delete() + work_mmap.delete() + # Proceed, not waiting for delete. return worker_sample_ids def _evict_shard(self, shard_id: int) -> None: @@ -1164,12 +1168,12 @@ def _evict_coldest_shard(self) -> None: """ while True: # Find the shard with the oldest last access time. - shard_id = int(self._shard_access_times.numpy().argmin()) + shard_id = int(self._shard_access_times.as_array().argmin()) # Check the shard's last access time. If it is NEVER, there are no downloaded shards to # evict. If any shards are currently being downloaded, wait, else raise an error. if self._shard_access_times[shard_id] == NEVER: - if (self._shard_states.numpy() == _ShardState.PREPARING).any(): + if (self._shard_states.as_array() == _ShardState.PREPARING).any(): sleep(TICK) continue else: From 8b994258522876db8b386bf97ef65df729471212 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 2 Jan 2024 08:33:22 -0800 Subject: [PATCH 37/54] Complete rewrite of all the mmap stuff --- streaming/base/coord/mmap/__init__.py | 25 +- streaming/base/coord/mmap/array.py | 103 ----- streaming/base/coord/mmap/barrier.py | 82 ++-- streaming/base/coord/mmap/base.py | 539 ++++++++++++++++++++++---- streaming/base/coord/mmap/buffer.py | 72 +++- streaming/base/coord/mmap/ndarray.py | 89 +++++ streaming/base/coord/mmap/number.py | 408 +++++++++++++++++-- streaming/base/dataset.py | 118 ++---- 8 files changed, 1088 insertions(+), 348 deletions(-) delete mode 100644 streaming/base/coord/mmap/array.py create mode 100644 streaming/base/coord/mmap/ndarray.py diff --git a/streaming/base/coord/mmap/__init__.py b/streaming/base/coord/mmap/__init__.py index 7608cfed0..34e34d921 100644 --- a/streaming/base/coord/mmap/__init__.py +++ b/streaming/base/coord/mmap/__init__.py @@ -3,9 +3,24 @@ """Share data across processes with mmap().""" -from streaming.base.coord.mmap.array import MMapArray -from streaming.base.coord.mmap.barrier import MMapBarrier -from streaming.base.coord.mmap.buffer import MMapBuffer -from streaming.base.coord.mmap.number import MMapNumber +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 +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) -__all__ = ['MMapArray', 'MMapBarrier', 'MMapBuffer', 'MMapNumber'] +__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/array.py b/streaming/base/coord/mmap/array.py deleted file mode 100644 index 796f68f61..000000000 --- a/streaming/base/coord/mmap/array.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright 2023 MosaicML Streaming authors -# SPDX-License-Identifier: Apache-2.0 - -"""Share an array 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 NDArray - -from streaming.base.coord.mmap.base import ensure_file - -__all__ = ['MMapArray'] - -DType = TypeVar('DType', bound=np.number) - -IndexType = Union[int, slice, NDArray[np.integer]] -DataType = Union[DType, NDArray[DType]] - - -class MMapArray(Generic[DType]): - """Share an array across processes using mmap(). - - Args: - mode (str): Whether to ``create``, ``replace``, or ``attach``. Defaults to ``attach``. - filename (str): Path to memory-mapped file. - shape (int | Tuple[int], optional): Exact required shape, if known in advance. At most one - wildcard ``-1`` is acceptable. - dtype (DType): Data type of the number. - """ - - def __init__( - self, - *, - mode: str = 'attach', - filename: str, - shape: Optional[Union[int, Tuple[int]]] = None, - dtype: DType, - ) -> None: - self.mode = mode - self.filename = filename - self.shape = ensure_file(mode, filename, shape, dtype.nbytes) - self.dtype = dtype - - self.file = open(filename, 'r+b', 0) - self.data = mmap(self.file.fileno(), 0) - - def __len__(self) -> int: - """Get the number of elements in the first axis of the array. - - Returns: - int: Length of the first axis of the array. - """ - return int(self.shape[0]) - - def as_array(self) -> NDArray[DType]: - """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[DType]: Our internal buffer as an ndarray. - """ - return np.ndarray(self.shape, buffer=self.data, dtype=self.dtype) - - def __getitem__(self, index: IndexType) -> DataType: - """Get the item at the index. - - Args: - index (IndexType): The index(es). - - Returns: - DataType; The item(s). - """ - return self.as_array()[index] - - def __setitem__(self, index: IndexType, item: DataType) -> None: - """Set the item at the index. - - Args: - index (IndexType): The index(es). - item (DataType): The item(s). - """ - self.as_array()[index] = item - - def flush(self) -> None: - """Flush the mmap.""" - self.data.flush() - - def close(self) -> None: - """Close the mmap and file handle.""" - if not self.data.closed: - self.data.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) diff --git a/streaming/base/coord/mmap/barrier.py b/streaming/base/coord/mmap/barrier.py index f311649e2..ada93a1c4 100644 --- a/streaming/base/coord/mmap/barrier.py +++ b/streaming/base/coord/mmap/barrier.py @@ -3,21 +3,29 @@ """Share a barrier across processes using mmap().""" +from enum import IntEnum from time import sleep import numpy as np from streaming.base.coord.file import SoftFileLock -from streaming.base.coord.mmap.array import MMapArray +from streaming.base.coord.mmap.ndarray import MemMapNDArray -__all__ = ['MMapBarrier'] +__all__ = ['MemMapBarrier', 'barrier'] -class MMapBarrier: - """Share a barrier across processes using mmap(). +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: - mode (str): Whether to ``create``, ``replace``, or ``attach``. Defaults to ``attach``. + create (bool): If ``True``, create. If ``False``, attach. mmap_filename (str): Path to memory-mapped file. lock_filename (str): Path to SoftFileLock file. tick (float): Polling interval in seconds. Defaults to ``0.007``. @@ -25,19 +33,18 @@ class MMapBarrier: def __init__( self, - *, - mode: str = 'attach', + create: bool, mmap_filename: str, lock_filename: str, tick: float = 0.007, ) -> None: - self._lock = SoftFileLock(lock_filename) - self._arr = MMapArray(mode=mode, filename=mmap_filename, shape=3, dtype=np.int32()) - self._tick = tick - + 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 = True + self._flag = BarrierFlag.GO + self._lock = SoftFileLock(lock_filename) + self._tick = tick @property def _num_enter(self) -> int: @@ -55,7 +62,7 @@ def _num_enter(self, num_enter: int) -> None: Args: num_enter (int): Entered process count. """ - self._arr[0] = np.int32(num_enter) + self._arr[0] = num_enter @property def _num_exit(self) -> int: @@ -73,53 +80,57 @@ def _num_exit(self, num_exit: int) -> None: Args: num_exit (int): Exited process count. """ - self._arr[1] = np.int32(num_exit) + self._arr[1] = num_exit @property - def _flag(self) -> bool: + def _flag(self) -> BarrierFlag: """Getter for _flag. Returns: - bool: Flag value. + BarrierFlat: Flag value. """ - return bool(self._arr[2]) + return BarrierFlag(self._arr[2]) @_flag.setter - def _flag(self, flag: bool) -> None: + def _flag(self, flag: BarrierFlag) -> None: """Setter for _flag. Args: - flag (bool): Flag value. + flag (BarrierFlag): Flag value. """ - self._arr[2] = np.int32(flag) + self._arr[2] = flag def __call__(self, total: int) -> None: + """A set number of processes enter, wait, and exit the barrier. - # Initialize num_exit to the number of processes. + 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 "don't go". + # 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() while self._num_exit != total: sleep(self._tick) self._lock.acquire() - self._flag = False + self._flag = BarrierFlag.STOP - # Note that we entered. + # 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 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 = True + self._flag = BarrierFlag.GO self._lock.release() - # Everybody waits until the flag is set to "go". + # Everybody waits until `_flag` is set to `GO`. while not self._flag: sleep(self._tick) @@ -128,3 +139,20 @@ def __call__(self, total: int) -> None: self._num_exit += 1 if self._num_exit == total: self._num_exit = -1 + + +def barrier( + create: bool, + mmap_filename: str, + lock_filename: str, + 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. + tick (float): Polling interval in seconds. Defaults to ``0.007``. + """ + return MemMapBarrier(create, mmap_filename, lock_filename, tick) diff --git a/streaming/base/coord/mmap/base.py b/streaming/base/coord/mmap/base.py index 199bd2bbe..034aefb38 100644 --- a/streaming/base/coord/mmap/base.py +++ b/streaming/base/coord/mmap/base.py @@ -4,31 +4,35 @@ """Base functionality for sharing data across processes using mmap().""" import os -from typing import Optional, Tuple, Union +from mmap import mmap +from typing import IO, Generic, Optional, Tuple, TypeVar, Union import numpy as np +from numpy.typing import DTypeLike, NDArray -__all__ = ['ensure_file'] +__all__ = ['T', 'MemMap', 'Number'] +Number = Union[int, float, complex, np.number] -def _normalize_shape(shape: Optional[Union[int, Tuple[int]]]) -> \ - Tuple[Tuple[int], int, Optional[int]]: + +def _get_wildcarded_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[Tuple[int], int, Optional[int]]: Normalized shape, number of elements without the - wildcard if present, and bytes per element. + Tuple[int]: Normalized shape containing at least one dimension, and at most one wildcard. """ if shape is None: - shape = -1, + wild_shape = -1, elif isinstance(shape, int): - shape = shape, + wild_shape = shape, + else: + wild_shape = shape num_wild = 0 - for dim in shape: + for dim in wild_shape: if dim == -1: num_wild += 1 elif dim < 1: @@ -38,79 +42,468 @@ def _normalize_shape(shape: Optional[Union[int, Tuple[int]]]) -> \ if 1 < num_wild: raise ValueError(f'Shape contains multiple ({num_wild}) wildcards: {shape}.') - numel = int(np.prod(shape)) - if numel < 0: - numel = -numel - wild_index = shape.index(-1) + 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: - wild_index = None + 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 shape, numel, wild_index + return exact_shape -def ensure_file(mode: str, filename: str, shape: Optional[Union[int, Tuple[int]]], - unit: int) -> Tuple[int]: - """Ensure file existence and size according to mode. +def _normalize_ndarray(value: NDArray[np.number]) -> NDArray[np.number]: + """Normalize an input ndarray to rule out the zero ndim case. Args: - mode (str): Whether to ``create``, ``replace``, or ``attach``. Defaults to ``attach``. - filename (str): Path to memory-mapped file. - shape (int | Tuple[int], optional): Exact required number of units, along each axis, if - known in advance. At most one wildcard ``-1`` is acceptable. - unit (int): Stride of a single value in bytes. + value (NDArray[np.number]): Input ndarray. Returns: - int: Resulting exact shape. - """ - want_shape, want_numel, want_wild_index = _normalize_shape(shape) - - if unit < 1: - raise ValueError(f'{unit} must be a positive integer, but got: {unit}.') - - # Normalize file existence by mode. - if mode == 'create': - if os.path.exists(filename): - raise ValueError(f'File alreadfy exists: {filename}.') - elif mode == 'replace': - if os.path.exists(filename): - os.remove(filename) - elif mode == 'attach': - if not os.path.exists(filename): - raise ValueError(f'File does not exist: {filename}.') + 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: - modes = {'create', 'replace', 'attach'} - raise ValueError(f'`mode` must be either replace,one of {sorted(modes)}, but got: {mode}.') + 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: Tuple[int]) -> bool: + """Tell whether the shape restriction accepts the observed shape. + + Args: + shape (Tuple[int]): Observed shape. + wildcarded_shape (Tuple[int]): Shape restriction. + + Returns: + bool: Whether acceptable. + """ + 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 + + +# 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 info. + + Args: + shape (Tuple[int]): Normalized ndarray shape. + dtype (DTypeLike): 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 _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:] + ndim, = np.frombuffer(part, np.uint64) + 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 + - # Perform the work. +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. + """ + # Write header and body to a temp file. + tmp_filename = filename + '.tmp' + with open(tmp_filename, 'wb') as out: + header = _encode_file_header(arr.shape, arr.dtype) + out.write(header) + out.write(arr.tobytes()) + + # 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: + shape (Tuple[int]): Normalized ndarray shape. + dtype (np.dtype): Normalized ndarray dtype. + filename (str): Path to file. + """ + if (arr == 0).all(): + _write_sparse_file(shape, arr.dtype, filename) + else: + arr = arr.repeat(np.prod(shape)) + arr = arr.reshape(shape) + _write_dense_file(arr, filename) + + +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 _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 _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, int]: + """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, dtype, and offset. + """ + # The file must not already exist. if os.path.exists(filename): - # Use size info to validate the pre-existing file. - got_size = os.stat(filename).st_size - if want_wild_index is None: - want_size = want_numel * unit - if got_size != want_size: - raise ValueError(f'File is the wrong size: file {filename}, expected shape ' + - f'{want_shape}, expected unit {unit}, expected size ' + - f'{want_size}, actual size {got_size}.') - got_shape = want_numel, - else: - want_size = want_numel * unit - if got_size % want_size: - raise ValueError(f'File size is not evenly divisible: file {filename}, expected ' + - f'shape {want_shape}, expected unit {unit}, expected size to ' + - f'be divisible by {want_size}.') - wild_value = got_size // want_size - got_shape = list(want_shape) - got_shape[want_wild_index] = wild_value - got_shape = tuple(got_shape) + 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) + + # Get the offset of the array part. + offset = _get_file_header_size(exact_shape) + + return exact_shape, arr.dtype, offset + + +def _check_file( + filename: str, + shape: Optional[Union[int, Tuple[int]]] = None, + dtype: Optional[np.dtype] = None, +) -> Tuple[Tuple[int], np.dtype, int]: + """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``. + + Returns: + Tuple[Tuple[int], np.dtype, int]: The file's exact shape, dtype, and offset. + """ + # The file must already exist. + if not os.path.exists(filename): + raise ValueError(f'`value` was not provided, so attaching the file, but it does not ' + + f'exist: {filename}.') + + # 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}.') + + # Get the offset of the array part. + offset = _get_file_header_size(got_shape) + + return got_shape, got_dtype, offset + + +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, +) -> Tuple[Tuple[int], np.dtype, int]: + """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``. + + Returns: + Tuple[Tuple[int], np.dtype, int]: The file's exact shape, dtype, and offset. + """ + if value is not None: + return _create_file(filename, value, shape, dtype) else: - # Use size info to create the (initially sparse) file. - if want_wild_index is not None: - raise ValueError(f'You must provide `shape`, without wildcards, in order to size ' + - f'the file: {filename}.') - with open(filename, 'wb') as out: - out.write(b'') - os.truncate(filename, want_numel * unit) - got_shape = want_shape - - # Return resulting exact shape. - return got_shape + return _check_file(filename, shape, dtype) + + +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``. + """ + + 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: + norm_dtype = np.dtype(dtype) if dtype is not None else None + self.shape, self.dtype, self.offset = _ensure_file(filename, shape, norm_dtype, value) + self.filename = filename + self.file = open(filename, 'w+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 index 789f2383c..e52d166e5 100644 --- a/streaming/base/coord/mmap/buffer.py +++ b/streaming/base/coord/mmap/buffer.py @@ -3,35 +3,35 @@ """Share a buffer across processes using mmap().""" -from mmap import mmap from typing import Optional -from streaming.base.coord.mmap.base import ensure_file +import numpy as np +from numpy.dtypes import UInt8DType -__all__ = ['MMapBuffer'] +from streaming.base.coord.mmap.base import MemMap +__all__ = ['MemMapBuffer', 'buffer'] -class MMapBuffer: - """Share a buffer across processes using mmap(). + +class MemMapBuffer(MemMap[UInt8DType]): + """A buffer backed by a memory-mapped file. Args: - mode (str): Whether to ``create``, ``replace``, or ``attach``. Defaults to ``attach``. - filename (str): Path to memory-mapped file. - size (int, optional): Exact required size, if known in advance. Defaults to ``None``. + 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, - *, - mode: str = 'attach', filename: str, size: Optional[int] = None, + value: Optional[bytes] = None, ) -> None: - self.mode = mode - self.filename = filename - self.size, = ensure_file(mode, filename, size, 1) - self.file = open(filename, 'r+b', 0) - self.data = mmap(self.file.fileno(), 0) + 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. @@ -39,4 +39,44 @@ def __len__(self) -> int: Returns: int: Number of bytes in the buffer. """ - return self.size + 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/ndarray.py b/streaming/base/coord/mmap/ndarray.py new file mode 100644 index 000000000..fece0d4c0 --- /dev/null +++ b/streaming/base/coord/mmap/ndarray.py @@ -0,0 +1,89 @@ +# Copyright 2023 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, buffer=self.mmap, offset=self.offset, dtype=self.dtype) + + 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 index f5d0fe03c..7573e24c6 100644 --- a/streaming/base/coord/mmap/number.py +++ b/streaming/base/coord/mmap/number.py @@ -1,54 +1,406 @@ # Copyright 2023 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 -"""Share a single number across processes using mmap().""" +"""Share a single number across processes using mmap(). -from mmap import mmap -from typing import Generic +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.array import DType -from streaming.base.coord.mmap.base import ensure_file +from streaming.base.coord.mmap.base import MemMap, Number, T -__init__ = ['MMapNumber'] +__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 MMapNumber(Generic[DType]): - """Share a single number across processes using mmap(). +class MemMapNumber(MemMap[T]): + """A number backed by a memory-mapped file. Args: - mode (str): Whether to ``create``, ``replace``, or ``attach``. Defaults to ``attach``. - filename (str): Path to memory-mapped file. - dtype (DType): Data type of the number. + 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, - *, - mode: str = 'attach', filename: str, - dtype: DType, + dtype: Optional[DTypeLike] = None, + value: Optional[Union[Number, NDArray[np.number]]] = None, ) -> None: - self.mode = mode - self.filename = filename - ensure_file(mode, filename, 1, dtype.nbytes) - self.dtype = dtype - self.file = open(filename, 'r+b', 0) - self.data = mmap(self.file.fileno(), 0) + super().__init__(filename, 1, dtype, value) - def get(self) -> DType: - """Get our value. + def get(self) -> T: + """Get value. Returns: - DType: Our value. + np.number: The value. """ - return np.frombuffer(self.data, self.dtype)[0] + return np.frombuffer(self.mmap, self.dtype)[0] - def set(self, value: DType) -> None: - """Set our value. + def set(self, value: Number) -> None: + """Set value. Args: - value (DType): Our new value. + value (Number): The value. """ - self.data[:] = value.tobytes() + dtype_class = getattr(np, self.dtype.name) + self.mmap[:] = dtype_class.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/dataset.py b/streaming/base/dataset.py index e41ff3dd0..830ee29af 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -24,9 +24,9 @@ from streaming.base.array import Array from streaming.base.batching import generate_work 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.mmap import MMapArray, MMapBarrier, MMapNumber from streaming.base.coord.world import World from streaming.base.format import get_index_basename from streaming.base.sampling import get_sampling @@ -466,52 +466,43 @@ def __init__( # Register/lookup our shared memory prefix and filelock root directory. self.registry = JobRegistry(self.config_root) - self.job = job = JobDirectory(self.registry, streams, world) + self.job = JobDirectory(self.registry, streams, world) + job_file = self.job.get_filename init_done_filename = self.job.get_filename('init_done.txt') if world.is_local_leader: if os.path.exists(init_done_filename): os.remove(init_done_filename) - mode = 'create' if world.is_local_leader else 'attach' - # Worker barrier. - self._worker_barrier = MMapBarrier(mode=mode, - mmap_filename=job.get_filename('worker_barrier.npy'), - lock_filename=job.get_filename('worker_barrier.lock')) + 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 = MMapNumber(mode=mode, - filename=job.get_filename('epoch.npy'), - dtype=np.int64()) + 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_lock = SoftFileLock(job.get_filename('cache.lock')) + self._cache_lock = SoftFileLock(job_file('cache.lock')) # Cache usage in bytes. - self._cache_usage = MMapNumber(mode=mode, - filename=job.get_filename('cache_usage.npy'), - dtype=np.int64()) + 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 = MMapArray(mode=mode, - filename=job.get_filename('shard_states.npy'), - shape=self.num_shards, - dtype=np.uint8()) + 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 = MMapArray(mode=mode, - filename=job.get_filename('shard_access_times.npy'), - shape=self.num_shards, - dtype=np.uint64()) + self._shard_access_times = mm.ndarray(job_file('shard_access_times.npy'), self.num_shards, + np.uint64, value) # Initialize shared memory objects. if world.is_local_leader: @@ -1031,64 +1022,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[MMapArray, MMapArray]: - """Put an epoch's sample ordering into shared memory. - - Args: - sample_ids (NDArray[np.int64]): Sample IDs. - - Returns: - Tuple[MMapArray, MMapArray]: Memory-mapped 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. - work_shape = np.asarray(sample_ids.shape, np.int64) - work_shape_mmap = MMapArray(mode='create', - filename=self.job.get_filename('work_shape.npy'), - shape=work_shape.shape, - dtype=np.int64()) - work_shape_mmap[:] = work_shape - work_shape_mmap.flush() - - # Save the generated epoch data to shared memory. - work_mmap = MMapArray(mode='create', - filename=self.job.get_filename('work.npy'), - shape=sample_ids.shape, - dtype=np.int64()) - work_mmap[:] = sample_ids - work_mmap.flush() - - return work_shape_mmap, work_mmap - - def _attach_work(self) -> Tuple[MMapArray, MMapArray]: - """Get an epoch's sample ordering from shared memory. - - Returns: - Tuple[MMapArray, MMapArray]: Memory-mapped arrays containing shape and data. - """ - ndim = 5 - - # Load the generated epoch shape from shared memory. - work_shape_mmap = MMapArray(mode='attach', - filename=self.job.get_filename('work_shape.npy'), - shape=ndim, - dtype=np.int64()) - work_shape = tuple(work_shape_mmap.as_array()) - - work_mmap = MMapArray(mode='attach', - filename=self.job.get_filename('work.npy'), - shape=work_shape, - dtype=np.int64()) - - return work_shape_mmap, work_mmap - 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. @@ -1100,38 +1033,31 @@ def _get_work(self, world: World, epoch: int, sample_in_epoch: int) -> NDArray[n Returns: Optional[NDArray[np.int64]]: Our partition of the epoch. """ - # Do expensive work that may use a lot of cores/memory just once, in the local leader. + filename = self.job.get_filename('epoch_sample_ids.npy') + if world.is_local_leader: epoch_sample_ids = generate_work(self.batching_method, self, world, epoch, sample_in_epoch) - work_shape_mmap, work_mmap = self._share_work(epoch_sample_ids) + io = mm.ndarray(filename, value=epoch_sample_ids) self._worker_barrier(world.workers_per_node) else: self._worker_barrier(world.workers_per_node) - work_shape_mmap, work_mmap = self._attach_work() - epoch_sample_ids = work_mmap.as_array() + 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() - # Wait for each worker to extract their slice of the work. self._worker_barrier(world.workers_per_node) - # Close. if not world.is_local_leader: - work_shape_mmap.close() - work_mmap.close() + io.close() - # Wait for close. self._worker_barrier(world.workers_per_node) - # Delete. if world.is_local_leader: - work_shape_mmap.delete() - work_mmap.delete() + io.delete() - # Proceed, not waiting for delete. return worker_sample_ids def _evict_shard(self, shard_id: int) -> None: @@ -1168,12 +1094,12 @@ def _evict_coldest_shard(self) -> None: """ while True: # Find the shard with the oldest last access time. - shard_id = int(self._shard_access_times.as_array().argmin()) + shard_id = int(self._shard_access_times.numpy().argmin()) # Check the shard's last access time. If it is NEVER, there are no downloaded shards to # evict. If any shards are currently being downloaded, wait, else raise an error. if self._shard_access_times[shard_id] == NEVER: - if (self._shard_states.as_array() == _ShardState.PREPARING).any(): + if (self._shard_states.numpy() == _ShardState.PREPARING).any(): sleep(TICK) continue else: From f60d2b87dfb130aa9d2096865e7277fcbb0f4923 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 2 Jan 2024 08:45:58 -0800 Subject: [PATCH 38/54] Tweak. --- streaming/base/dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index 830ee29af..e14b97809 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -810,7 +810,7 @@ def next_epoch(self, next_epoch: int) -> None: Args: next_epoch (int): Next epoch. """ - self._next_epoch.set(np.int64(next_epoch)) + self._next_epoch.set(next_epoch) @property def cache_usage(self) -> int: @@ -828,7 +828,7 @@ def cache_usage(self, cache_usage: int) -> None: Args: cache_usage (int): Cache usage in bytes. """ - self._cache_usage.set(np.int64(cache_usage)) + self._cache_usage.set(cache_usage) def __len__(self) -> int: """Get the length as a PyTorch IterableDataset. From 85d48109cbd4c9ef1ceefe64218e37f40608ad15 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 2 Jan 2024 12:56:57 -0800 Subject: [PATCH 39/54] Fix. --- streaming/base/coord/file/__init__.py | 9 ++ streaming/base/coord/file/barrier.py | 90 +++++++++++++++++++ .../base/coord/{file.py => file/lock.py} | 0 streaming/base/coord/mmap/base.py | 30 +++++-- streaming/base/coord/mmap/number.py | 4 +- 5 files changed, 124 insertions(+), 9 deletions(-) create mode 100644 streaming/base/coord/file/__init__.py create mode 100644 streaming/base/coord/file/barrier.py rename streaming/base/coord/{file.py => file/lock.py} (100%) diff --git a/streaming/base/coord/file/__init__.py b/streaming/base/coord/file/__init__.py new file mode 100644 index 000000000..3553cc011 --- /dev/null +++ b/streaming/base/coord/file/__init__.py @@ -0,0 +1,9 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Coordinating processes using files.""" + +from streaming.base.coord.file.barrier import create_file, wait_for_creation, wait_for_deletion +from streaming.base.coord.file.lock import SoftFileLock + +__all__ = ['wait_for_creation', 'wait_for_deletion', 'create_file', 'SoftFileLock'] diff --git a/streaming/base/coord/file/barrier.py b/streaming/base/coord/file/barrier.py new file mode 100644 index 000000000..9cce1e66a --- /dev/null +++ b/streaming/base/coord/file/barrier.py @@ -0,0 +1,90 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +"""Utilities for doing barrier-like work with files.""" + +import os +from time import sleep, time +from typing import Optional + +__all__ = ['wait_for_creation', 'wait_for_deletion', 'create_file'] + + +def _wait_for_path(path: str, + exist: bool, + timeout: Optional[float] = 60, + poll_interval: float = 0.007) -> None: + """Wait for the creation or deletion of a path on the local filesystem. + + Args: + path (str): Local path to wait on. + exist (bool): Wait for existence if ``True`` or non-existence if ``False``. + timeout (float, optional): How long to wait before raising an error, in seconds. Defaults + to ``60``. + poll_interval (float): Check interval, in seconds. Defaults to ``0.007``. + """ + if timeout is not None and timeout <= 0: + timeout_str = f'{timeout:.3f}'.rstrip('0').rstrip('.') + raise ValueError(f'Timeout must be positive if provided, but got: {timeout_str}.') + + if poll_interval <= 0: + poll_interval_str = f'{poll_interval:.3f}'.rstrip('0').rstrip('.') + raise ValueError(f'Poll interval must be positive if provided, but got: ' + + f'{poll_interval_str}.') + + start = time() + while True: + if os.path.exists(path) == exist: + break + + if timeout is not None: + elapsed = time() - start + if timeout <= elapsed: + timeout_str = f'{timeout:.3f}'.rstrip('0').rstrip('.') + elapsed_str = f'{elapsed:.3f}'.rstrip('0').rstrip('.') + raise RuntimeError(f'Timed out while waiting for path to exist: path {path}, ' + + f'timeout {timeout_str} sec, elapsed {elapsed_str} sec.') + + sleep(poll_interval) + + +def wait_for_creation(path: str, + timeout: Optional[float] = 60, + poll_interval: float = 0.007) -> None: + """Wait for the creation of a path on the local filesystem. + + Args: + path (str): Local path to wait on. + timeout (float, optional): How long to wait before raising an error, in seconds. Defaults + to ``60``. + poll_interval (float): Check interval, in seconds. Defaults to ``0.007``. + """ + _wait_for_path(path, True, timeout, poll_interval) + + +def wait_for_deletion(path: str, + timeout: Optional[float] = 60, + poll_interval: float = 0.007) -> None: + """Wait for the deletion of a path on the local filesystem. + + Args: + path (str): Local path to wait on. + timeout (float, optional): How long to wait before raising an error, in seconds. Defaults + to ``60``. + poll_interval (float): Check interval, in seconds. Defaults to ``0.007``. + """ + _wait_for_path(path, False, timeout, poll_interval) + + +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) + file = open(filename, 'x') + file.close() diff --git a/streaming/base/coord/file.py b/streaming/base/coord/file/lock.py similarity index 100% rename from streaming/base/coord/file.py rename to streaming/base/coord/file/lock.py diff --git a/streaming/base/coord/mmap/base.py b/streaming/base/coord/mmap/base.py index 034aefb38..319254adb 100644 --- a/streaming/base/coord/mmap/base.py +++ b/streaming/base/coord/mmap/base.py @@ -10,22 +10,28 @@ import numpy as np from numpy.typing import DTypeLike, NDArray +from streaming.base.coord.file.barrier import wait_for_creation + __all__ = ['T', 'MemMap', 'Number'] Number = Union[int, float, complex, np.number] -def _get_wildcarded_shape(shape: Optional[Union[int, Tuple[int]]]) -> Tuple[int]: +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]: Normalized shape containing at least one dimension, and at most one wildcard. + 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: - wild_shape = -1, + return None + + if shape == (): + wild_shape = 1, elif isinstance(shape, int): wild_shape = shape, else: @@ -123,11 +129,14 @@ def _accepts_shape(shape: Tuple[int], wildcarded_shape: Tuple[int]) -> bool: Args: shape (Tuple[int]): Observed shape. - wildcarded_shape (Tuple[int]): Shape restriction. + wildcarded_shape (Optional[Tuple[int]]): Shape restriction, if provided. Returns: bool: Whether acceptable. """ + if wildcarded_shape is None: + return True + if len(shape) != len(wildcarded_shape): return False @@ -281,8 +290,9 @@ def _write_file(arr: NDArray[np.number], shape: Tuple[int], filename: str) -> No if (arr == 0).all(): _write_sparse_file(shape, arr.dtype, filename) else: - arr = arr.repeat(np.prod(shape)) - arr = arr.reshape(shape) + if arr.size == 1: + arr = arr.repeat(np.prod(shape)) + arr = arr.reshape(shape) _write_dense_file(arr, filename) @@ -481,11 +491,17 @@ def __init__( shape: Optional[Union[int, Tuple[int]]] = None, dtype: Optional[DTypeLike] = None, value: Optional[Union[Number, NDArray[np.number]]] = None, + timeout: Optional[float] = 60, + poll_interval: float = 0.007, ) -> None: norm_dtype = np.dtype(dtype) if dtype is not None else None self.shape, self.dtype, self.offset = _ensure_file(filename, shape, norm_dtype, value) + + if value is None: + wait_for_creation(filename, timeout, poll_interval) + self.filename = filename - self.file = open(filename, 'w+b', 0) + self.file = open(filename, 'r+b', 0) self.mmap = mmap(self.file.fileno(), 0) def flush(self) -> None: diff --git a/streaming/base/coord/mmap/number.py b/streaming/base/coord/mmap/number.py index 7573e24c6..1fe40a4c3 100644 --- a/streaming/base/coord/mmap/number.py +++ b/streaming/base/coord/mmap/number.py @@ -75,8 +75,8 @@ def set(self, value: Number) -> None: Args: value (Number): The value. """ - dtype_class = getattr(np, self.dtype.name) - self.mmap[:] = dtype_class.tobytes() + cls = getattr(np, self.dtype.name) + self.mmap[self.offset:] = cls(value).tobytes() class MemMapInteger(MemMapNumber): From 7d69d662113674a1e0cd34e80f67451af232be8b Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 2 Jan 2024 13:17:37 -0800 Subject: [PATCH 40/54] Fix. --- streaming/base/coord/mmap/base.py | 5 +++-- streaming/base/coord/mmap/buffer.py | 3 +-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/streaming/base/coord/mmap/base.py b/streaming/base/coord/mmap/base.py index 319254adb..ac9307fe9 100644 --- a/streaming/base/coord/mmap/base.py +++ b/streaming/base/coord/mmap/base.py @@ -124,12 +124,12 @@ def _is_shape_nonempty(shape: Tuple[int]) -> bool: return True -def _accepts_shape(shape: Tuple[int], wildcarded_shape: Tuple[int]) -> bool: +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 (Optional[Tuple[int]]): Shape restriction, if provided. + wildcarded_shape (Tuple[int], optional): Shape restriction, if provided. Returns: bool: Whether acceptable. @@ -291,6 +291,7 @@ def _write_file(arr: NDArray[np.number], shape: Tuple[int], filename: str) -> No _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/buffer.py b/streaming/base/coord/mmap/buffer.py index e52d166e5..ca346db56 100644 --- a/streaming/base/coord/mmap/buffer.py +++ b/streaming/base/coord/mmap/buffer.py @@ -6,14 +6,13 @@ from typing import Optional import numpy as np -from numpy.dtypes import UInt8DType from streaming.base.coord.mmap.base import MemMap __all__ = ['MemMapBuffer', 'buffer'] -class MemMapBuffer(MemMap[UInt8DType]): +class MemMapBuffer(MemMap[np.dtype[np.uint8]]): """A buffer backed by a memory-mapped file. Args: From 43e5a7202ea45018612fa500e60b08d240fa7803 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 2 Jan 2024 13:51:12 -0800 Subject: [PATCH 41/54] Fix. --- streaming/base/coord/mmap/buffer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/streaming/base/coord/mmap/buffer.py b/streaming/base/coord/mmap/buffer.py index ca346db56..a7a3672a4 100644 --- a/streaming/base/coord/mmap/buffer.py +++ b/streaming/base/coord/mmap/buffer.py @@ -12,7 +12,7 @@ __all__ = ['MemMapBuffer', 'buffer'] -class MemMapBuffer(MemMap[np.dtype[np.uint8]]): +class MemMapBuffer(MemMap[np.dtype]): """A buffer backed by a memory-mapped file. Args: From 4fa51148127efb7453c4788a82869f0f3495d873 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 2 Jan 2024 13:58:31 -0800 Subject: [PATCH 42/54] Fix. --- streaming/base/coord/mmap/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/streaming/base/coord/mmap/base.py b/streaming/base/coord/mmap/base.py index ac9307fe9..a89c27ab0 100644 --- a/streaming/base/coord/mmap/base.py +++ b/streaming/base/coord/mmap/base.py @@ -181,7 +181,7 @@ def _encode_file_header(shape: Tuple[int], dtype: np.dtype) -> bytes: Args: shape (Tuple[int]): Normalized ndarray shape. - dtype (DTypeLike): Normalized ndarray dtype. + dtype (np.dtype): Normalized ndarray dtype. Returns: bytes: Header data. From 372e5eb0ea1177c12ca6dd37d4c8378b2e0b102a Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 2 Jan 2024 14:13:56 -0800 Subject: [PATCH 43/54] Fix. --- streaming/base/coord/mmap/__init__.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/streaming/base/coord/mmap/__init__.py b/streaming/base/coord/mmap/__init__.py index 34e34d921..6ae1ff58a 100644 --- a/streaming/base/coord/mmap/__init__.py +++ b/streaming/base/coord/mmap/__init__.py @@ -7,14 +7,11 @@ 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 -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) +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) __all__ = [ 'MemMapBarrier', 'barrier', 'MemMap', 'MemMapBuffer', 'buffer', 'MemMapNumber', From fd37096830cbe6615410bf012f34d6cd5a74cd1b Mon Sep 17 00:00:00 2001 From: James Knighton Date: Tue, 2 Jan 2024 14:25:46 -0800 Subject: [PATCH 44/54] Fix. --- streaming/base/coord/mmap/__init__.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/streaming/base/coord/mmap/__init__.py b/streaming/base/coord/mmap/__init__.py index 6ae1ff58a..3e0e84a6a 100644 --- a/streaming/base/coord/mmap/__init__.py +++ b/streaming/base/coord/mmap/__init__.py @@ -7,11 +7,14 @@ 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', From f85bb6b40da3dac880c8bb2f0cf242a2a45a83fe Mon Sep 17 00:00:00 2001 From: James Knighton Date: Wed, 3 Jan 2024 00:12:35 -0800 Subject: [PATCH 45/54] Add lock usage to wait_for_existence, wait_for_removal. --- streaming/base/coord/file/barrier.py | 81 +++++++++++++++++++--------- streaming/base/coord/job/file.py | 4 +- streaming/base/coord/job/registry.py | 79 +++++++++++---------------- 3 files changed, 90 insertions(+), 74 deletions(-) diff --git a/streaming/base/coord/file/barrier.py b/streaming/base/coord/file/barrier.py index 9cce1e66a..2dfc3aecf 100644 --- a/streaming/base/coord/file/barrier.py +++ b/streaming/base/coord/file/barrier.py @@ -4,53 +4,81 @@ """Utilities for doing barrier-like work with files.""" import os +from contextlib import nullcontext from time import sleep, time -from typing import Optional +from typing import Any, Optional __all__ = ['wait_for_creation', 'wait_for_deletion', 'create_file'] -def _wait_for_path(path: str, - exist: bool, - timeout: Optional[float] = 60, - poll_interval: float = 0.007) -> None: +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_for_path( + path: str, + exists: bool, + timeout: Optional[float] = 30, + poll_interval: float = 0.007, + lock: Optional[Any] = None, +) -> None: """Wait for the creation or deletion of a path on the local filesystem. Args: path (str): Local path to wait on. - exist (bool): Wait for existence if ``True`` or non-existence if ``False``. + exists (bool): Wait for existence if ``True`` or non-existence if ``False``. timeout (float, optional): How long to wait before raising an error, in seconds. Defaults to ``60``. poll_interval (float): Check interval, in seconds. Defaults to ``0.007``. + lock (Any, optional): Lock that must be held when probing for the path. Defaults to + ``None``. """ + start = time() + if timeout is not None and timeout <= 0: - timeout_str = f'{timeout:.3f}'.rstrip('0').rstrip('.') - raise ValueError(f'Timeout must be positive if provided, but got: {timeout_str}.') + raise ValueError(f'Timeout must be positive if provided, but got: ' + + f'{_say_duration(timeout)} sec.') if poll_interval <= 0: - poll_interval_str = f'{poll_interval:.3f}'.rstrip('0').rstrip('.') raise ValueError(f'Poll interval must be positive if provided, but got: ' + - f'{poll_interval_str}.') + f'{_say_duration(poll_interval)} 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}.') - start = time() while True: - if os.path.exists(path) == exist: - break + with lock or nullcontext(): + if os.path.exists(path) == exists: + break if timeout is not None: - elapsed = time() - start - if timeout <= elapsed: - timeout_str = f'{timeout:.3f}'.rstrip('0').rstrip('.') - elapsed_str = f'{elapsed:.3f}'.rstrip('0').rstrip('.') + now = time() + if timeout <= now - start: raise RuntimeError(f'Timed out while waiting for path to exist: path {path}, ' + - f'timeout {timeout_str} sec, elapsed {elapsed_str} sec.') + f'timeout {_say_duration(timeout)} sec, elapsed ' + + f'{_say_duration(now - start)} sec.') sleep(poll_interval) -def wait_for_creation(path: str, - timeout: Optional[float] = 60, - poll_interval: float = 0.007) -> None: +def wait_for_creation( + path: str, + timeout: Optional[float] = 30, + poll_interval: float = 0.007, + lock: Optional[Any] = None, +) -> None: """Wait for the creation of a path on the local filesystem. Args: @@ -58,13 +86,15 @@ def wait_for_creation(path: str, timeout (float, optional): How long to wait before raising an error, in seconds. Defaults to ``60``. poll_interval (float): Check interval, in seconds. Defaults to ``0.007``. + lock (Any): Lock that must be held when probing for the path. """ - _wait_for_path(path, True, timeout, poll_interval) + _wait_for_path(path, True, timeout, poll_interval, lock) def wait_for_deletion(path: str, - timeout: Optional[float] = 60, - poll_interval: float = 0.007) -> None: + timeout: Optional[float] = 30, + poll_interval: float = 0.007, + lock: Optional[Any] = None) -> None: """Wait for the deletion of a path on the local filesystem. Args: @@ -72,8 +102,9 @@ def wait_for_deletion(path: str, timeout (float, optional): How long to wait before raising an error, in seconds. Defaults to ``60``. poll_interval (float): Check interval, in seconds. Defaults to ``0.007``. + lock (Any): Lock that must be held when probing for the path. """ - _wait_for_path(path, False, timeout, poll_interval) + _wait_for_path(path, False, timeout, poll_interval, lock) def create_file(filename: str) -> None: diff --git a/streaming/base/coord/job/file.py b/streaming/base/coord/job/file.py index 3383cd468..5fb669bb6 100644 --- a/streaming/base/coord/job/file.py +++ b/streaming/base/coord/job/file.py @@ -11,10 +11,10 @@ from streaming.base.coord.job.entry import JobEntry -__all__ = ['JobFile'] +__all__ = ['RegistryFile'] -class JobFile: +class RegistryFile: """StreamingDataset job registry, which is backed by a JSON file. Args: diff --git a/streaming/base/coord/job/registry.py b/streaming/base/coord/job/registry.py index 048ddb84b..e2458609c 100644 --- a/streaming/base/coord/job/registry.py +++ b/streaming/base/coord/job/registry.py @@ -9,14 +9,15 @@ import os from hashlib import sha3_224 from shutil import rmtree -from time import sleep, time_ns -from typing import Dict, List, Sequence, Tuple +from time import time_ns +from typing import Dict, List, Optional, Sequence, Tuple from psutil import process_iter -from streaming.base.coord.file import SoftFileLock +from streaming.base.coord.file.barrier import wait_for_creation, wait_for_deletion +from streaming.base.coord.file.lock import SoftFileLock from streaming.base.coord.job.entry import JobEntry -from streaming.base.coord.job.file import JobFile +from streaming.base.coord.job.file import RegistryFile from streaming.base.coord.world import World from streaming.base.stream import Stream @@ -34,12 +35,20 @@ class JobRegistry: your system. """ - def __init__(self, config_root: str, tick: float = 0.007) -> None: + def __init__( + self, + config_root: str, + timeout: Optional[float] = 30, + poll_interval: float = 0.007, + ) -> None: self.config_root = config_root - 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') + self.timeout = timeout + self.poll_interval = poll_interval + + 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. @@ -135,32 +144,6 @@ def _remove_dir(self, job_hash: str) -> None: dirname = os.path.join(self.config_root, job_hash) rmtree(dirname) - def _wait_for_existence(self, job_hash: str) -> None: - """Wait for a directory to be created. - - Args: - job_hash (str): Job hash of directory. - """ - dirname = os.path.join(self.config_root, job_hash) - while True: - sleep(self._tick) - with self._lock: - if os.path.exists(dirname): - break - - def _wait_for_removal(self, job_hash: str) -> None: - """Wait for a directory to be removed. - - Args: - job_hash (str): Job hash of directory. - """ - dirname = os.path.join(self.config_root, job_hash) - while True: - sleep(self._tick) - with self._lock: - if not os.path.exists(dirname): - break - def _register(self, streams: Sequence[Stream]) -> str: """Register this collection of StreamingDataset replicas. @@ -189,11 +172,11 @@ def _register(self, streams: Sequence[Stream]) -> str: process_id=pid, register_time=register_time) - with self._lock: - reg = JobFile.read(self._registry_filename) - reg.add(entry) - del_job_hashes = reg.filter(pid2create_time) - reg.write(self._registry_filename) + 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) @@ -233,7 +216,8 @@ def register(self, streams: Sequence[Stream], world: World) -> str: job_hash = self._register(streams) else: job_hash = self._lookup(streams) - self._wait_for_existence(job_hash) + dirname = os.path.join(self.config_root, job_hash) + wait_for_creation(dirname, self.timeout, self.poll_interval, self.lock) return job_hash def _unregister(self, job_hash: str) -> None: @@ -246,11 +230,11 @@ def _unregister(self, job_hash: str) -> None: """ pid2create_time = self._get_live_procs() - with self._lock: - reg = JobFile.read(self._registry_filename) - reg.remove(job_hash) - del_job_hashes = reg.filter(pid2create_time) - reg.write(self._registry_filename) + 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) @@ -266,4 +250,5 @@ def unregister(self, job_hash: str, world: World) -> None: if world.is_local_leader: self._unregister(job_hash) else: - self._wait_for_removal(job_hash) + dirname = os.path.join(self.config_root, job_hash) + wait_for_deletion(dirname, self.timeout, self.poll_interval, self.lock) From e660c962033b4120689e4f5cf715b910d457de56 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Wed, 3 Jan 2024 03:12:18 -0800 Subject: [PATCH 46/54] Generalize waiting, etc. --- streaming/base/coord/file/__init__.py | 4 +- streaming/base/coord/file/barrier.py | 121 -------------------------- streaming/base/coord/file/lock.py | 52 ++++++----- streaming/base/coord/file/waiting.py | 71 +++++++++++++++ streaming/base/coord/job/registry.py | 13 +-- streaming/base/coord/mmap/barrier.py | 22 +++-- streaming/base/coord/mmap/base.py | 2 +- streaming/base/coord/waiting.py | 73 ++++++++++++++++ 8 files changed, 201 insertions(+), 157 deletions(-) delete mode 100644 streaming/base/coord/file/barrier.py create mode 100644 streaming/base/coord/file/waiting.py create mode 100644 streaming/base/coord/waiting.py diff --git a/streaming/base/coord/file/__init__.py b/streaming/base/coord/file/__init__.py index 3553cc011..ec6aeb7d2 100644 --- a/streaming/base/coord/file/__init__.py +++ b/streaming/base/coord/file/__init__.py @@ -3,7 +3,7 @@ """Coordinating processes using files.""" -from streaming.base.coord.file.barrier import create_file, wait_for_creation, wait_for_deletion from streaming.base.coord.file.lock import SoftFileLock +from streaming.base.coord.file.waiting import create_file, wait_for_creation, wait_for_deletion -__all__ = ['wait_for_creation', 'wait_for_deletion', 'create_file', 'SoftFileLock'] +__all__ = ['create_file', 'wait_for_creation', 'wait_for_deletion', 'SoftFileLock'] diff --git a/streaming/base/coord/file/barrier.py b/streaming/base/coord/file/barrier.py deleted file mode 100644 index 2dfc3aecf..000000000 --- a/streaming/base/coord/file/barrier.py +++ /dev/null @@ -1,121 +0,0 @@ -# Copyright 2023 MosaicML Streaming authors -# SPDX-License-Identifier: Apache-2.0 - -"""Utilities for doing barrier-like work with files.""" - -import os -from contextlib import nullcontext -from time import sleep, time -from typing import Any, Optional - -__all__ = ['wait_for_creation', 'wait_for_deletion', 'create_file'] - - -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_for_path( - path: str, - exists: bool, - timeout: Optional[float] = 30, - poll_interval: float = 0.007, - lock: Optional[Any] = None, -) -> None: - """Wait for the creation or deletion of a path on the local filesystem. - - Args: - path (str): Local path to wait on. - exists (bool): Wait for existence if ``True`` or non-existence if ``False``. - timeout (float, optional): How long to wait before raising an error, in seconds. Defaults - to ``60``. - poll_interval (float): Check interval, in seconds. Defaults to ``0.007``. - lock (Any, optional): Lock that must be held when probing for the path. 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 poll_interval <= 0: - raise ValueError(f'Poll interval must be positive if provided, but got: ' + - f'{_say_duration(poll_interval)} 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}.') - - while True: - with lock or nullcontext(): - if os.path.exists(path) == exists: - break - - if timeout is not None: - now = time() - if timeout <= now - start: - raise RuntimeError(f'Timed out while waiting for path to exist: path {path}, ' + - f'timeout {_say_duration(timeout)} sec, elapsed ' + - f'{_say_duration(now - start)} sec.') - - sleep(poll_interval) - - -def wait_for_creation( - path: str, - timeout: Optional[float] = 30, - poll_interval: 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. - timeout (float, optional): How long to wait before raising an error, in seconds. Defaults - to ``60``. - poll_interval (float): Check interval, in seconds. Defaults to ``0.007``. - lock (Any): Lock that must be held when probing for the path. - """ - _wait_for_path(path, True, timeout, poll_interval, lock) - - -def wait_for_deletion(path: str, - timeout: Optional[float] = 30, - poll_interval: 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. - timeout (float, optional): How long to wait before raising an error, in seconds. Defaults - to ``60``. - poll_interval (float): Check interval, in seconds. Defaults to ``0.007``. - lock (Any): Lock that must be held when probing for the path. - """ - _wait_for_path(path, False, timeout, poll_interval, 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) - file = open(filename, 'x') - file.close() diff --git a/streaming/base/coord/file/lock.py b/streaming/base/coord/file/lock.py index 54af41055..37912b549 100644 --- a/streaming/base/coord/file/lock.py +++ b/streaming/base/coord/file/lock.py @@ -4,13 +4,13 @@ """Soft file locking via file open mode 'x'.""" import os -from time import sleep, time 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'] @@ -25,7 +25,12 @@ class SoftFileLock: tick (float): Polling interval in seconds. Defaults to ``0.007``. """ - def __init__(self, filename: str, timeout: Optional[float] = 30, tick: float = 0.007) -> None: + 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.') @@ -91,9 +96,11 @@ def _normalize(cls, filename: str) -> None: os.remove(filename) @classmethod - def _get_timeout(cls, - init_timeout: Optional[float], - timeout: Optional[Union[str, float]] = 'auto') -> Optional[float]: + def _get_timeout( + cls, + init_timeout: Optional[float], + timeout: Optional[Union[str, float]] = 'auto', + ) -> Optional[float]: """Determine the timeout for a given acquire(). Args: @@ -121,24 +128,27 @@ def _get_timeout(cls, f'init, but got: {timeout}.') return ret - def acquire(self, timeout: Optional[Union[str, float]] = 'auto') -> None: + 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. """ - start = time() - timeout = self._get_timeout(self.timeout, timeout) - while True: + + def stop() -> bool: try: - self._write(self.filename, os.getpid()) - break + with open(self.filename, 'x') as out: + text = str(os.getpid()) + out.write(text) + return True except: - pass - if self.timeout is not None and start + self.timeout < time(): - raise ValueError(f'Timed out while attempting to acquire file lock: file ' + - f'{self.filename}, timeout {self.timeout} sec.') - sleep(self.tick) + return False + + norm_timeout = self._get_timeout(self.timeout, timeout) + wait(stop, norm_timeout, self.tick) def release(self) -> None: """Release this lock.""" @@ -158,10 +168,12 @@ def __enter__(self) -> Self: self.acquire() return self - def __exit__(self, - err_type: Optional[Type[BaseException]] = None, - err: Optional[BaseException] = None, - trace: Optional[TracebackType] = None) -> None: + def __exit__( + self, + err_type: Optional[Type[BaseException]] = None, + err: Optional[BaseException] = None, + trace: Optional[TracebackType] = None, + ) -> None: """Exit context manager. Args: diff --git a/streaming/base/coord/file/waiting.py b/streaming/base/coord/file/waiting.py new file mode 100644 index 000000000..2bef4a847 --- /dev/null +++ b/streaming/base/coord/file/waiting.py @@ -0,0 +1,71 @@ +# Copyright 2023 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/registry.py b/streaming/base/coord/job/registry.py index e2458609c..b69b70540 100644 --- a/streaming/base/coord/job/registry.py +++ b/streaming/base/coord/job/registry.py @@ -14,8 +14,8 @@ from psutil import process_iter -from streaming.base.coord.file.barrier import wait_for_creation, wait_for_deletion 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 @@ -33,17 +33,20 @@ class JobRegistry: 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, - poll_interval: float = 0.007, + tick: float = 0.007, ) -> None: self.config_root = config_root self.timeout = timeout - self.poll_interval = poll_interval + self.tick = tick self.lock_filename = os.path.join(config_root, 'registry.lock') self.lock = SoftFileLock(self.lock_filename) @@ -217,7 +220,7 @@ def register(self, streams: Sequence[Stream], world: World) -> str: else: job_hash = self._lookup(streams) dirname = os.path.join(self.config_root, job_hash) - wait_for_creation(dirname, self.timeout, self.poll_interval, self.lock) + wait_for_creation(dirname, self.timeout, self.tick, self.lock) return job_hash def _unregister(self, job_hash: str) -> None: @@ -251,4 +254,4 @@ def unregister(self, job_hash: str, world: World) -> None: self._unregister(job_hash) else: dirname = os.path.join(self.config_root, job_hash) - wait_for_deletion(dirname, self.timeout, self.poll_interval, self.lock) + wait_for_deletion(dirname, self.timeout, self.tick, self.lock) diff --git a/streaming/base/coord/mmap/barrier.py b/streaming/base/coord/mmap/barrier.py index ada93a1c4..064db664a 100644 --- a/streaming/base/coord/mmap/barrier.py +++ b/streaming/base/coord/mmap/barrier.py @@ -4,12 +4,13 @@ """Share a barrier across processes using mmap().""" from enum import IntEnum -from time import sleep +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'] @@ -28,7 +29,9 @@ class MemMapBarrier: create (bool): If ``True``, create. If ``False``, attach. mmap_filename (str): Path to memory-mapped file. lock_filename (str): Path to SoftFileLock file. - tick (float): Polling interval in seconds. Defaults to ``0.007``. + 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__( @@ -36,6 +39,7 @@ def __init__( create: bool, mmap_filename: str, lock_filename: str, + timeout: Optional[float] = 30, tick: float = 0.007, ) -> None: value = 0 if create else None @@ -44,6 +48,7 @@ def __init__( self._num_exit = -1 self._flag = BarrierFlag.GO self._lock = SoftFileLock(lock_filename) + self._timeout = timeout self._tick = tick @property @@ -115,8 +120,7 @@ def __call__(self, total: int) -> None: self._lock.acquire() if not self._num_enter: self._lock.release() - while self._num_exit != total: - sleep(self._tick) + wait(lambda: self._num_exit == total, self._timeout, self._tick, self._lock) self._lock.acquire() self._flag = BarrierFlag.STOP @@ -131,8 +135,7 @@ def __call__(self, total: int) -> None: self._lock.release() # Everybody waits until `_flag` is set to `GO`. - while not self._flag: - sleep(self._tick) + wait(lambda: self._flag == BarrierFlag.GO, self._timeout, self._tick, self._lock) # Note that we exited. with self._lock: @@ -145,6 +148,7 @@ 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. @@ -153,6 +157,8 @@ def barrier( create (bool): If ``True``, create. If ``False``, attach. mmap_filename (str): Path to memory-mapped file. lock_filename (str): Path to SoftFileLock file. - tick (float): Polling interval in seconds. Defaults to ``0.007``. + 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, tick) + 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 index a89c27ab0..905ecb766 100644 --- a/streaming/base/coord/mmap/base.py +++ b/streaming/base/coord/mmap/base.py @@ -10,7 +10,7 @@ import numpy as np from numpy.typing import DTypeLike, NDArray -from streaming.base.coord.file.barrier import wait_for_creation +from streaming.base.coord.file.waiting import wait_for_creation __all__ = ['T', 'MemMap', 'Number'] diff --git a/streaming/base/coord/waiting.py b/streaming/base/coord/waiting.py new file mode 100644 index 000000000..da6237446 --- /dev/null +++ b/streaming/base/coord/waiting.py @@ -0,0 +1,73 @@ +# Copyright 2023 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, + interval: 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``. + interval (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 interval <= 0: + raise ValueError(f'Poll interval must be positive if provided, but got: ' + + f'{_say_duration(interval)} 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(interval) From 61a5d42348c27eb39592db08a8b327ce0a5e8f76 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Wed, 3 Jan 2024 04:45:23 -0800 Subject: [PATCH 47/54] Fix. --- streaming/base/coord/file/lock.py | 2 +- streaming/base/coord/mmap/base.py | 34 ++++++++++++++++++---------- streaming/base/coord/mmap/ndarray.py | 2 +- streaming/base/coord/mmap/number.py | 2 +- streaming/base/coord/waiting.py | 11 ++++----- 5 files changed, 30 insertions(+), 21 deletions(-) diff --git a/streaming/base/coord/file/lock.py b/streaming/base/coord/file/lock.py index 37912b549..95988a821 100644 --- a/streaming/base/coord/file/lock.py +++ b/streaming/base/coord/file/lock.py @@ -22,7 +22,7 @@ class SoftFileLock: 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): Polling interval in seconds. Defaults to ``0.007``. + tick (float): Check interval in seconds. Defaults to ``0.007``. """ def __init__( diff --git a/streaming/base/coord/mmap/base.py b/streaming/base/coord/mmap/base.py index 905ecb766..7a3786002 100644 --- a/streaming/base/coord/mmap/base.py +++ b/streaming/base/coord/mmap/base.py @@ -220,7 +220,8 @@ def _read_file_header(file: IO[bytes]) -> Tuple[Tuple[int], np.dtype]: # Get ndim. part = data[max_itemsize:] - ndim, = np.frombuffer(part, np.uint64) + arr = np.frombuffer(part, np.uint64) + ndim = int(arr[0]) if ndim < 0: raise ValueError(f'Header ndim is negative: {ndim}.') @@ -388,6 +389,8 @@ 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, int]: """Validate the file backing the memory mapping given optional explicit shape and dtype. @@ -397,14 +400,15 @@ def _check_file( 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, int]: The file's exact shape, dtype, and offset. """ - # The file must already exist. - if not os.path.exists(filename): - raise ValueError(f'`value` was not provided, so attaching the file, but it does not ' + - f'exist: {filename}.') + # 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: @@ -442,6 +446,8 @@ def _ensure_file( 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, int]: """Create the file backing the memory mapping given optional explicit shape and dtype. @@ -453,6 +459,9 @@ def _ensure_file( ``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, dtype, and offset. @@ -460,7 +469,7 @@ def _ensure_file( if value is not None: return _create_file(filename, value, shape, dtype) else: - return _check_file(filename, shape, dtype) + return _check_file(filename, shape, dtype, timeout, tick) T = TypeVar('T', bound=np.dtype) @@ -484,6 +493,9 @@ class MemMap(Generic[T]): ``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__( @@ -492,14 +504,12 @@ def __init__( shape: Optional[Union[int, Tuple[int]]] = None, dtype: Optional[DTypeLike] = None, value: Optional[Union[Number, NDArray[np.number]]] = None, - timeout: Optional[float] = 60, - poll_interval: float = 0.007, + 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, self.offset = _ensure_file(filename, shape, norm_dtype, value) - - if value is None: - wait_for_creation(filename, timeout, poll_interval) + self.shape, self.dtype, self.offset = _ensure_file(filename, shape, norm_dtype, value, + timeout, tick) self.filename = filename self.file = open(filename, 'r+b', 0) diff --git a/streaming/base/coord/mmap/ndarray.py b/streaming/base/coord/mmap/ndarray.py index fece0d4c0..eb9c92dce 100644 --- a/streaming/base/coord/mmap/ndarray.py +++ b/streaming/base/coord/mmap/ndarray.py @@ -44,7 +44,7 @@ def numpy(self) -> NDArray: Returns: NDArray[T]: Our internal buffer as an ndarray. """ - return np.ndarray(self.shape, buffer=self.mmap, offset=self.offset, dtype=self.dtype) + 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. diff --git a/streaming/base/coord/mmap/number.py b/streaming/base/coord/mmap/number.py index 1fe40a4c3..544927ab9 100644 --- a/streaming/base/coord/mmap/number.py +++ b/streaming/base/coord/mmap/number.py @@ -67,7 +67,7 @@ def get(self) -> T: Returns: np.number: The value. """ - return np.frombuffer(self.mmap, self.dtype)[0] + return np.frombuffer(self.mmap, self.dtype, 1, self.offset)[0] def set(self, value: Number) -> None: """Set value. diff --git a/streaming/base/coord/waiting.py b/streaming/base/coord/waiting.py index da6237446..6dd8bf4bb 100644 --- a/streaming/base/coord/waiting.py +++ b/streaming/base/coord/waiting.py @@ -25,7 +25,7 @@ def _say_duration(duration: float) -> str: def wait( stop: Callable[[], bool], timeout: Optional[float] = 30, - interval: float = 0.007, + tick: float = 0.007, lock: Optional[Any] = None, ) -> None: """Wait for the predicate to succeed. @@ -34,7 +34,7 @@ def wait( 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``. - interval (float): Check interval, in seconds. Defaults to ``0.007``. + 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``. """ @@ -44,9 +44,8 @@ def wait( raise ValueError(f'Timeout must be positive if provided, but got: ' + f'{_say_duration(timeout)} sec.') - if interval <= 0: - raise ValueError(f'Poll interval must be positive if provided, but got: ' + - f'{_say_duration(interval)} 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__'): @@ -70,4 +69,4 @@ def wait( raise RuntimeError(f'Wait timed out: timeout {_say_duration(timeout)} sec vs ' + f'elapsed {_say_duration(now - start)} sec.') - sleep(interval) + sleep(tick) From 507744e536418e9d8e8d05b89c03378d9b684422 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Wed, 3 Jan 2024 04:57:07 -0800 Subject: [PATCH 48/54] Stop doing pytest in ten parts, as the file handle issue is now fixed too. --- .github/workflows/pytest.yaml | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) 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 From 0e6d717684925d3563164442559de14e308e80f8 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Wed, 3 Jan 2024 05:13:50 -0800 Subject: [PATCH 49/54] Fix (docstring). --- streaming/base/coord/mmap/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/streaming/base/coord/mmap/base.py b/streaming/base/coord/mmap/base.py index 7a3786002..c73e8e656 100644 --- a/streaming/base/coord/mmap/base.py +++ b/streaming/base/coord/mmap/base.py @@ -284,8 +284,8 @@ def _write_file(arr: NDArray[np.number], shape: Tuple[int], filename: str) -> No """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. - dtype (np.dtype): Normalized ndarray dtype. filename (str): Path to file. """ if (arr == 0).all(): From 986bef83cefcb6b7b3e013054054e873190d6f32 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Wed, 3 Jan 2024 05:37:25 -0800 Subject: [PATCH 50/54] Split file format functionality out of base.py into file.py --- streaming/base/coord/mmap/base.py | 196 +++-------------------------- streaming/base/coord/mmap/file.py | 201 ++++++++++++++++++++++++++++++ 2 files changed, 217 insertions(+), 180 deletions(-) create mode 100644 streaming/base/coord/mmap/file.py diff --git a/streaming/base/coord/mmap/base.py b/streaming/base/coord/mmap/base.py index c73e8e656..af9994687 100644 --- a/streaming/base/coord/mmap/base.py +++ b/streaming/base/coord/mmap/base.py @@ -5,12 +5,14 @@ import os from mmap import mmap -from typing import IO, Generic, Optional, Tuple, TypeVar, Union +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'] @@ -163,171 +165,12 @@ def _accepts_dtype(have: np.dtype, want: Optional[np.dtype]) -> bool: return have == want -# 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 info. - - 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 _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. - """ - # Write header and body to a temp file. - tmp_filename = filename + '.tmp' - with open(tmp_filename, 'wb') as out: - header = _encode_file_header(arr.shape, arr.dtype) - out.write(header) - out.write(arr.tobytes()) - - # 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) - - -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 _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 _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, int]: +) -> Tuple[Tuple[int], np.dtype]: """Create the file backing the memory mapping given optional explicit shape and dtype. Args: @@ -339,7 +182,7 @@ def _create_file( ``None``. Returns: - Tuple[Tuple[int], np.dtype, int]: The file's exact shape, dtype, and offset. + Tuple[Tuple[int], np.dtype, int]: The file's exact shape and dtype. """ # The file must not already exist. if os.path.exists(filename): @@ -377,12 +220,9 @@ def _create_file( f'`value` dtype {arr.dtype} vs explicit `dtype` {dtype}.') # Create the file (dense or sparse). - _write_file(arr, exact_shape, filename) + write_file(arr, exact_shape, filename) - # Get the offset of the array part. - offset = _get_file_header_size(exact_shape) - - return exact_shape, arr.dtype, offset + return exact_shape, arr.dtype def _check_file( @@ -391,7 +231,7 @@ def _check_file( dtype: Optional[np.dtype] = None, timeout: Optional[float] = 30, tick: float = 0.007, -) -> Tuple[Tuple[int], np.dtype, int]: +) -> Tuple[Tuple[int], np.dtype]: """Validate the file backing the memory mapping given optional explicit shape and dtype. Args: @@ -405,17 +245,17 @@ def _check_file( tick (float): Check interval, in seconds. Defaults to ``0.007``. Returns: - Tuple[Tuple[int], np.dtype, int]: The file's exact shape, dtype, and offset. + 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) + 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) + 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 ' + @@ -435,10 +275,7 @@ def _check_file( f'file: `filename` {filename}, `dtype` {dtype_name}, actual dtype ' + f'{got_dtype.name}.') - # Get the offset of the array part. - offset = _get_file_header_size(got_shape) - - return got_shape, got_dtype, offset + return got_shape, got_dtype def _ensure_file( @@ -448,7 +285,7 @@ def _ensure_file( value: Optional[Union[Number, NDArray[np.number]]] = None, timeout: Optional[float] = 30, tick: float = 0.007, -) -> Tuple[Tuple[int], np.dtype, int]: +) -> Tuple[Tuple[int], np.dtype]: """Create the file backing the memory mapping given optional explicit shape and dtype. Args: @@ -464,7 +301,7 @@ def _ensure_file( tick (float): Check interval, in seconds. Defaults to ``0.007``. Returns: - Tuple[Tuple[int], np.dtype, int]: The file's exact shape, dtype, and offset. + 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) @@ -508,9 +345,8 @@ def __init__( tick: float = 0.007, ) -> None: norm_dtype = np.dtype(dtype) if dtype is not None else None - self.shape, self.dtype, self.offset = _ensure_file(filename, shape, norm_dtype, value, - timeout, tick) - + 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) diff --git a/streaming/base/coord/mmap/file.py b/streaming/base/coord/mmap/file.py new file mode 100644 index 000000000..7fd453368 --- /dev/null +++ b/streaming/base/coord/mmap/file.py @@ -0,0 +1,201 @@ +# Copyright 2023 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) From 82993ec3b366ae34c254f8e9714c0934231824d2 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Fri, 5 Jan 2024 02:31:58 -0800 Subject: [PATCH 51/54] SD init barrier: cruft flag file -> proper MemMapBarrier. --- streaming/base/dataset.py | 30 ++++++++++++------------------ 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index e14b97809..478322c1e 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -32,7 +32,7 @@ from streaming.base.sampling import get_sampling from streaming.base.spanner import Spanner from streaming.base.stream import Stream -from streaming.base.util import bytes_to_int, number_abbrev_to_int, wait_for_file_to_exist +from streaming.base.util import bytes_to_int, number_abbrev_to_int # An arbitrary time in the future, used for cold shard eviction. NEVER = np.iinfo(np.uint64).max @@ -464,15 +464,14 @@ def __init__( # 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. + # 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 - init_done_filename = self.job.get_filename('init_done.txt') - if world.is_local_leader: - if os.path.exists(init_done_filename): - os.remove(init_done_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'), @@ -504,7 +503,7 @@ def __init__( 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 @@ -528,18 +527,13 @@ def __init__( self._shard_states[shard_id] = _ShardState.LOCAL if size else _ShardState.REMOTE self._shard_access_times[shard_id] = time_ns() - with open(init_done_filename, 'x'): - pass - else: - wait_for_file_to_exist(init_done_filename, TICK, 300, - 'Waited too long for initialization') - - # Placeholder for an _Iterator which tracks state during __iter__(). - self._iterator: _Iterator + # These fields are set each __iter__(). + self._iterator: _Iterator # Tracks thread positions. + self._executor: ThreadPoolExecutor # Multi-threading. + self._event: Event # Exception handling. - # Placeholder for exception handling in __iter__ threads. - self._executor: ThreadPoolExecutor - self._event: Event + # Init is not done for anyone until all interprocess state is populated by local leader. + self._rank_barrier(world.ranks_per_node) @classmethod def _test_config_root(cls, config_root: str) -> None: From 34efcf7ae2397261c3d68516556715c6a9437dcf Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 20 Jan 2024 01:35:22 -0800 Subject: [PATCH 52/54] Misc. --- benchmarks/locking/bench.py | 209 ++++++++++++++++++++++++++++++++++++ 1 file changed, 209 insertions(+) create mode 100644 benchmarks/locking/bench.py diff --git a/benchmarks/locking/bench.py b/benchmarks/locking/bench.py new file mode 100644 index 000000000..dc69611b2 --- /dev/null +++ b/benchmarks/locking/bench.py @@ -0,0 +1,209 @@ +# Copyright 2023 MosaicML Streaming authors +# SPDX-License-Identifier: Apache-2.0 + +import os +from argparse import ArgumentParser, Namespace +from functools import partial +from multiprocessing import Process +from tempfile import TemporaryDirectory +from time import time +from typing import List, Type, TypeVar, Union + +import numpy as np +from numpy.typing import NDArray +from tqdm import tqdm + +from streaming.base.coord.file.lock import SoftFileLock + + +def parse_args() -> Namespace: + """Parse command-line arguments. + + Returns: + Namespace: Command-line arguments. + """ + args = ArgumentParser() + args.add_argument('--tick_ms', type=str, default='1,2,4,8,16,32,64') + args.add_argument('--num_procs', type=str, default='1,2,4,8,16,32,64') + args.add_argument('--turns_per_proc', type=float, default=1024) + args.add_argument('--turn_mean_ms', type=float, default=42) + args.add_argument('--turn_std_ms', type=float, default=7) + args.add_argument('--progress_bar', type=int, default=1) + return args.parse_args() + + +domain2triple = { + 'neg': [True, False, False], + 'nonpos': [True, True, False], + 'negneg': [False, True, True], + 'pos': [False, False, True], +} + +Number = Union[int, float] + + +def get_domain_index(value: Number) -> int: + """Get the domain index of a number: 0 (neg), 1 (zero), or 2 (pos). + + Args: + value (Number): Value. + + Returns: + int: Domain index. + """ + if value < 0: + idx = 0 + elif value == 0: + idx = 1 + else: + idx = 2 + return idx + + +def check_number( + name: str, + domain: str, + value: Number, +) -> None: + """Check a number. + + Args: + name (str): Argument command-line name. + domain (str): Output argument value valid range: one of {neg, nonpos, nonneg, pos}.' + value (Number): Parsed argument value. + """ + is_ok = domain2triple[domain] + idx = get_domain_index(value) + if not is_ok[idx]: + raise ValueError(f'Argument --{name} must be {domain}, but got: {value}.') + + +T = TypeVar('T', bound=Number) + + +def parse_number( + name: str, + kind: Type[T], + domain: str, + text: str, +) -> T: + """Parse a number passed in a str command-line argument. + + Args: + name (str): Argument command-line name. + kind (Type[T]): Type of output argument values. + domain (str): Output argument value valid range: one of {neg, nonpos, nonneg, pos}.' + text (str): Argument value as a str. + + Returns: + T: Output argument values. + """ + value = kind(text) + check_number(name, domain, value) + return value + + +def parse_numbers( + name: str, + kind: Type[T], + domain: str, + text: str, +) -> List[T]: + """Parse a list of numbers passed in a str command-line argument. + + Args: + name (str): Argument command-line name. + kind (Type[T]): Type of output argument values. + domain (str): Output argument value valid range: one of {neg, nonpos, nonneg, pos}.' + text (str): Argument value as a str. + + Returns: + List[T]: Output argument values. + """ + parse = partial(parse_number, name, kind, domain) + texts = text.split(',') if text else [] + return list(map(parse, texts)) + + +def get_lock_filename(dirname: str) -> str: + return os.path.join(dirname, f'lock') + + +def get_times_filename(dirname: str, proc_id: int) -> str: + return os.path.join(dirname, f'{proc_id}_times.npy') + + +def bench_process( + dirname: str, + tick: float, + turns_per_proc: int, + proc_id: int, +) -> None: + lock_filename = get_lock_filename(dirname) + lock = SoftFileLock(lock_filename, 60, tick) + times = np.zeros(turns_per_proc + 1, np.float64) + for i in range(turns_per_proc): + times[i] = time() + with lock: + pass + times[-1] = time() + times_filename = get_times_filename(dirname, proc_id) + times.tofile(times_filename) + + +def bench( + dirname: str, + tick: float, + num_procs: int, + turns_per_proc: int, +) -> NDArray[np.float64]: + procs = [] + for proc_id in range(num_procs): + proc = Process(target=bench_process, args=(dirname, tick, turns_per_proc, proc_id)) + procs.append(proc) + + for proc in procs: + proc.start() + + for proc in procs: + proc.join() + + times = np.zeros((num_procs, turns_per_proc + 1), np.float64) + for proc_id in range(num_procs): + times_filename = get_times_filename(dirname, proc_id) + times[proc_id] = np.fromfile(times_filename, np.float64) + return times + + +def main(args: Namespace) -> None: + """Main. + + Args: + args (Namespace): Command-line arguments. + """ + tick_mss = parse_numbers('tick_ms', float, 'pos', args.tick_ms) + ticks = list(map(lambda tick_ms: tick_ms / 1000, tick_mss)) + proc_counts = parse_numbers('num_procs', int, 'pos', args.num_procs) + turns_per_proc = parse_number('turns_per_proc', int, 'pos', args.turns_per_proc) + + durs = [] + with TemporaryDirectory() as dirname: + assert os.path.isdir(dirname) + total = len(proc_counts) * len(ticks) + progress_bar = tqdm(total=total, leave=False) if args.progress_bar else None + for tick in ticks: + for num_procs in proc_counts: + arr = bench(dirname, tick, num_procs, turns_per_proc) + dur = arr.max() - arr.min() + durs.append(dur) + if progress_bar is not None: + progress_bar.update(1) + assert os.path.isdir(dirname) + durs = np.array(durs, np.float64) + durs = durs.reshape(len(ticks), len(proc_counts)) + durs *= 1000 + print(durs) + + +if __name__ == '__main__': + main(parse_args()) From 15808f3826a8e94bfb41055031dfbdc43b686104 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 20 Jan 2024 01:46:24 -0800 Subject: [PATCH 53/54] Update the new files' copyright years. --- streaming/base/coord/__init__.py | 2 +- streaming/base/coord/file/__init__.py | 2 +- streaming/base/coord/file/lock.py | 2 +- streaming/base/coord/file/waiting.py | 2 +- streaming/base/coord/job/__init__.py | 2 +- streaming/base/coord/job/directory.py | 2 +- streaming/base/coord/job/entry.py | 2 +- streaming/base/coord/job/file.py | 2 +- streaming/base/coord/job/registry.py | 2 +- streaming/base/coord/mmap/__init__.py | 2 +- streaming/base/coord/mmap/barrier.py | 2 +- streaming/base/coord/mmap/base.py | 2 +- streaming/base/coord/mmap/buffer.py | 2 +- streaming/base/coord/mmap/file.py | 2 +- streaming/base/coord/mmap/ndarray.py | 2 +- streaming/base/coord/mmap/number.py | 2 +- streaming/base/coord/process.py | 2 +- streaming/base/coord/shmem/__init__.py | 2 +- streaming/base/coord/waiting.py | 2 +- 19 files changed, 19 insertions(+), 19 deletions(-) diff --git a/streaming/base/coord/__init__.py b/streaming/base/coord/__init__.py index 75b47257e..430dfc136 100644 --- a/streaming/base/coord/__init__.py +++ b/streaming/base/coord/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Coordination among ranks and workers.""" diff --git a/streaming/base/coord/file/__init__.py b/streaming/base/coord/file/__init__.py index ec6aeb7d2..90f2459d8 100644 --- a/streaming/base/coord/file/__init__.py +++ b/streaming/base/coord/file/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Coordinating processes using files.""" diff --git a/streaming/base/coord/file/lock.py b/streaming/base/coord/file/lock.py index 95988a821..0c8c59689 100644 --- a/streaming/base/coord/file/lock.py +++ b/streaming/base/coord/file/lock.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Soft file locking via file open mode 'x'.""" diff --git a/streaming/base/coord/file/waiting.py b/streaming/base/coord/file/waiting.py index 2bef4a847..daf4a8996 100644 --- a/streaming/base/coord/file/waiting.py +++ b/streaming/base/coord/file/waiting.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Waiting on files.""" diff --git a/streaming/base/coord/job/__init__.py b/streaming/base/coord/job/__init__.py index cd5f75465..82e765f4b 100644 --- a/streaming/base/coord/job/__init__.py +++ b/streaming/base/coord/job/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# 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.""" diff --git a/streaming/base/coord/job/directory.py b/streaming/base/coord/job/directory.py index e41b14276..03adb2a6a 100644 --- a/streaming/base/coord/job/directory.py +++ b/streaming/base/coord/job/directory.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """A directory containing all dataset-wide filesystem state for a Streaming job.""" diff --git a/streaming/base/coord/job/entry.py b/streaming/base/coord/job/entry.py index c39305e6c..6ebf88a6f 100644 --- a/streaming/base/coord/job/entry.py +++ b/streaming/base/coord/job/entry.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """An entry in a Streaming job registry file.""" diff --git a/streaming/base/coord/job/file.py b/streaming/base/coord/job/file.py index 5fb669bb6..213394eac 100644 --- a/streaming/base/coord/job/file.py +++ b/streaming/base/coord/job/file.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """A Streaming job registry file.""" diff --git a/streaming/base/coord/job/registry.py b/streaming/base/coord/job/registry.py index b69b70540..56b46b0a2 100644 --- a/streaming/base/coord/job/registry.py +++ b/streaming/base/coord/job/registry.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """A directory containing all Streaming-wide filesystem state. diff --git a/streaming/base/coord/mmap/__init__.py b/streaming/base/coord/mmap/__init__.py index 3e0e84a6a..c208a57a0 100644 --- a/streaming/base/coord/mmap/__init__.py +++ b/streaming/base/coord/mmap/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Share data across processes with mmap().""" diff --git a/streaming/base/coord/mmap/barrier.py b/streaming/base/coord/mmap/barrier.py index 064db664a..f7eae5f61 100644 --- a/streaming/base/coord/mmap/barrier.py +++ b/streaming/base/coord/mmap/barrier.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Share a barrier across processes using mmap().""" diff --git a/streaming/base/coord/mmap/base.py b/streaming/base/coord/mmap/base.py index af9994687..5f21b8822 100644 --- a/streaming/base/coord/mmap/base.py +++ b/streaming/base/coord/mmap/base.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Base functionality for sharing data across processes using mmap().""" diff --git a/streaming/base/coord/mmap/buffer.py b/streaming/base/coord/mmap/buffer.py index a7a3672a4..d0fe113d8 100644 --- a/streaming/base/coord/mmap/buffer.py +++ b/streaming/base/coord/mmap/buffer.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Share a buffer across processes using mmap().""" diff --git a/streaming/base/coord/mmap/file.py b/streaming/base/coord/mmap/file.py index 7fd453368..07591a9fe 100644 --- a/streaming/base/coord/mmap/file.py +++ b/streaming/base/coord/mmap/file.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Base functionality for sharing data across processes using mmap().""" diff --git a/streaming/base/coord/mmap/ndarray.py b/streaming/base/coord/mmap/ndarray.py index eb9c92dce..0994ee4ce 100644 --- a/streaming/base/coord/mmap/ndarray.py +++ b/streaming/base/coord/mmap/ndarray.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Share an ndarray across processes using mmap().""" diff --git a/streaming/base/coord/mmap/number.py b/streaming/base/coord/mmap/number.py index 544927ab9..dd35acb80 100644 --- a/streaming/base/coord/mmap/number.py +++ b/streaming/base/coord/mmap/number.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Share a single number across processes using mmap(). diff --git a/streaming/base/coord/process.py b/streaming/base/coord/process.py index f78f6386c..facedb420 100644 --- a/streaming/base/coord/process.py +++ b/streaming/base/coord/process.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Utility methods for coordination related to processes.""" diff --git a/streaming/base/coord/shmem/__init__.py b/streaming/base/coord/shmem/__init__.py index 991be052c..66a487542 100644 --- a/streaming/base/coord/shmem/__init__.py +++ b/streaming/base/coord/shmem/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Objects that live in shared memory. diff --git a/streaming/base/coord/waiting.py b/streaming/base/coord/waiting.py index 6dd8bf4bb..92a640630 100644 --- a/streaming/base/coord/waiting.py +++ b/streaming/base/coord/waiting.py @@ -1,4 +1,4 @@ -# Copyright 2023 MosaicML Streaming authors +# Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 """Waiting on predicates.""" From 32c3e5e6bb5efa83fe25998282c8df5fb2cea945 Mon Sep 17 00:00:00 2001 From: James Knighton Date: Sat, 20 Jan 2024 02:01:24 -0800 Subject: [PATCH 54/54] Benchmark separately. --- benchmarks/locking/bench.py | 209 ------------------------------------ 1 file changed, 209 deletions(-) delete mode 100644 benchmarks/locking/bench.py diff --git a/benchmarks/locking/bench.py b/benchmarks/locking/bench.py deleted file mode 100644 index dc69611b2..000000000 --- a/benchmarks/locking/bench.py +++ /dev/null @@ -1,209 +0,0 @@ -# Copyright 2023 MosaicML Streaming authors -# SPDX-License-Identifier: Apache-2.0 - -import os -from argparse import ArgumentParser, Namespace -from functools import partial -from multiprocessing import Process -from tempfile import TemporaryDirectory -from time import time -from typing import List, Type, TypeVar, Union - -import numpy as np -from numpy.typing import NDArray -from tqdm import tqdm - -from streaming.base.coord.file.lock import SoftFileLock - - -def parse_args() -> Namespace: - """Parse command-line arguments. - - Returns: - Namespace: Command-line arguments. - """ - args = ArgumentParser() - args.add_argument('--tick_ms', type=str, default='1,2,4,8,16,32,64') - args.add_argument('--num_procs', type=str, default='1,2,4,8,16,32,64') - args.add_argument('--turns_per_proc', type=float, default=1024) - args.add_argument('--turn_mean_ms', type=float, default=42) - args.add_argument('--turn_std_ms', type=float, default=7) - args.add_argument('--progress_bar', type=int, default=1) - return args.parse_args() - - -domain2triple = { - 'neg': [True, False, False], - 'nonpos': [True, True, False], - 'negneg': [False, True, True], - 'pos': [False, False, True], -} - -Number = Union[int, float] - - -def get_domain_index(value: Number) -> int: - """Get the domain index of a number: 0 (neg), 1 (zero), or 2 (pos). - - Args: - value (Number): Value. - - Returns: - int: Domain index. - """ - if value < 0: - idx = 0 - elif value == 0: - idx = 1 - else: - idx = 2 - return idx - - -def check_number( - name: str, - domain: str, - value: Number, -) -> None: - """Check a number. - - Args: - name (str): Argument command-line name. - domain (str): Output argument value valid range: one of {neg, nonpos, nonneg, pos}.' - value (Number): Parsed argument value. - """ - is_ok = domain2triple[domain] - idx = get_domain_index(value) - if not is_ok[idx]: - raise ValueError(f'Argument --{name} must be {domain}, but got: {value}.') - - -T = TypeVar('T', bound=Number) - - -def parse_number( - name: str, - kind: Type[T], - domain: str, - text: str, -) -> T: - """Parse a number passed in a str command-line argument. - - Args: - name (str): Argument command-line name. - kind (Type[T]): Type of output argument values. - domain (str): Output argument value valid range: one of {neg, nonpos, nonneg, pos}.' - text (str): Argument value as a str. - - Returns: - T: Output argument values. - """ - value = kind(text) - check_number(name, domain, value) - return value - - -def parse_numbers( - name: str, - kind: Type[T], - domain: str, - text: str, -) -> List[T]: - """Parse a list of numbers passed in a str command-line argument. - - Args: - name (str): Argument command-line name. - kind (Type[T]): Type of output argument values. - domain (str): Output argument value valid range: one of {neg, nonpos, nonneg, pos}.' - text (str): Argument value as a str. - - Returns: - List[T]: Output argument values. - """ - parse = partial(parse_number, name, kind, domain) - texts = text.split(',') if text else [] - return list(map(parse, texts)) - - -def get_lock_filename(dirname: str) -> str: - return os.path.join(dirname, f'lock') - - -def get_times_filename(dirname: str, proc_id: int) -> str: - return os.path.join(dirname, f'{proc_id}_times.npy') - - -def bench_process( - dirname: str, - tick: float, - turns_per_proc: int, - proc_id: int, -) -> None: - lock_filename = get_lock_filename(dirname) - lock = SoftFileLock(lock_filename, 60, tick) - times = np.zeros(turns_per_proc + 1, np.float64) - for i in range(turns_per_proc): - times[i] = time() - with lock: - pass - times[-1] = time() - times_filename = get_times_filename(dirname, proc_id) - times.tofile(times_filename) - - -def bench( - dirname: str, - tick: float, - num_procs: int, - turns_per_proc: int, -) -> NDArray[np.float64]: - procs = [] - for proc_id in range(num_procs): - proc = Process(target=bench_process, args=(dirname, tick, turns_per_proc, proc_id)) - procs.append(proc) - - for proc in procs: - proc.start() - - for proc in procs: - proc.join() - - times = np.zeros((num_procs, turns_per_proc + 1), np.float64) - for proc_id in range(num_procs): - times_filename = get_times_filename(dirname, proc_id) - times[proc_id] = np.fromfile(times_filename, np.float64) - return times - - -def main(args: Namespace) -> None: - """Main. - - Args: - args (Namespace): Command-line arguments. - """ - tick_mss = parse_numbers('tick_ms', float, 'pos', args.tick_ms) - ticks = list(map(lambda tick_ms: tick_ms / 1000, tick_mss)) - proc_counts = parse_numbers('num_procs', int, 'pos', args.num_procs) - turns_per_proc = parse_number('turns_per_proc', int, 'pos', args.turns_per_proc) - - durs = [] - with TemporaryDirectory() as dirname: - assert os.path.isdir(dirname) - total = len(proc_counts) * len(ticks) - progress_bar = tqdm(total=total, leave=False) if args.progress_bar else None - for tick in ticks: - for num_procs in proc_counts: - arr = bench(dirname, tick, num_procs, turns_per_proc) - dur = arr.max() - arr.min() - durs.append(dur) - if progress_bar is not None: - progress_bar.update(1) - assert os.path.isdir(dirname) - durs = np.array(durs, np.float64) - durs = durs.reshape(len(ticks), len(proc_counts)) - durs *= 1000 - print(durs) - - -if __name__ == '__main__': - main(parse_args())