Let's blow away dist, and also shared memory

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__:

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 = {}

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

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

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

streaming/base/batching/per_stream.py

@@ -10,9 +10,9 @@
 import numpy as np
 from numpy.typing import NDArray
 
+from streaming.base.coord.world import World
 from streaming.base.partition import get_partitions
 from streaming.base.shuffle import get_shuffle
-from streaming.base.world import World
 
 if TYPE_CHECKING:
     from streaming.base.dataset import StreamingDataset
@@ -63,9 +63,6 @@ def generate_work_per_stream_batching(dataset: StreamingDataset, world: World, e
             # same as the ratio of the stream's samples to overall samples.
             # This ensures that the overall training shuffle block size is still approximately
             # equal to what is set by the user, and allows for reasoning about cache_limit as well.
-            if not isinstance(dataset.shuffle_block_size, int):
-                raise TypeError(f'Dataset `shuffle_block_size` must be an integer. ' +
-                                f'Got {type(dataset.shuffle_block_size)} instead.')
             shuffle_block_portion = int(dataset.shuffle_block_size * stream.proportion)
             stream_shuffle = get_shuffle(dataset.shuffle_algo, shuffle_units,
                                          dataset.num_canonical_nodes, dataset.shuffle_seed, epoch,

streaming/base/batching/random.py

@@ -10,9 +10,9 @@
 import numpy as np
 from numpy.typing import NDArray
 
+from streaming.base.coord.world import World
 from streaming.base.partition import get_partitions
 from streaming.base.shuffle import get_shuffle
-from streaming.base.world import World
 
 if TYPE_CHECKING:
     from streaming.base.dataset import StreamingDataset
@@ -58,9 +58,6 @@ def generate_work_random_batching(dataset: StreamingDataset, world: World, epoch
 
     # If we need to shuffle, shuffle in a node-aware and *underlying* shard-aware way.
     if dataset.shuffle:
-        if not isinstance(dataset.shuffle_block_size, int):
-            raise TypeError(f'Dataset `shuffle_block_size` must be an integer. ' +
-                            f'Got {type(dataset.shuffle_block_size)} instead.')
         shuffle = get_shuffle(dataset.shuffle_algo, shuffle_units, dataset.num_canonical_nodes,
                               dataset.shuffle_seed, epoch, dataset.shuffle_block_size)
         big_ids = np.where(big_ids != -1, shuffle[big_ids], -1)

streaming/base/batching/stratified.py

@@ -11,9 +11,9 @@
 import numpy as np
 from numpy.typing import NDArray
 
+from streaming.base.coord.world import World
 from streaming.base.partition import get_partitions
 from streaming.base.shuffle import get_shuffle
-from streaming.base.world import World
 
 if TYPE_CHECKING:
     from streaming.base.dataset import StreamingDataset
@@ -75,9 +75,6 @@ def generate_work_stratified_batching(dataset: StreamingDataset, world: World, e
             # same as the ratio of the stream's samples to overall samples.
             # This ensures that the overall training shuffle block size is still approximately
             # equal to what is set by the user, and allows for reasoning about cache_limit as well.
-            if not isinstance(dataset.shuffle_block_size, int):
-                raise TypeError(f'Dataset `shuffle_block_size` must be an integer. ' +
-                                f'Got {type(dataset.shuffle_block_size)} instead.')
             shuffle_block_portion = int(dataset.shuffle_block_size * stream.proportion)
             stream_shuffle = get_shuffle(dataset.shuffle_algo, shuffle_units,
                                          dataset.num_canonical_nodes, dataset.shuffle_seed, epoch,

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'
+]

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']

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)

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,
+        }