streaming-community 0.14.0__py3-none-any.whl

simulation/__init__.py

@@ -0,0 +1,4 @@
+# Copyright 2022-2024 MosaicML Streaming authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Streaming simulation for throughput, network downloads, and shuffle quality."""

simulation/core/__init__.py

@@ -0,0 +1,4 @@
+# Copyright 2022-2024 MosaicML Streaming authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Core classes and functions for streaming simulation."""

simulation/core/create_index.py

@@ -0,0 +1,86 @@
+# Copyright 2022-2024 MosaicML Streaming authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Create a dataset index file from input parameters."""
+
+import json
+import logging
+import os
+import random
+import string
+from typing import Optional
+
+from streaming.base.format import get_index_basename
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
+def get_random_foldername() -> str:
+    """Generate random folder name to store the index file in.
+
+    Returns:
+        str: random alphanumeric folder name.
+    """
+    return ''.join(
+        random.choice(string.ascii_uppercase + string.ascii_lowercase + string.digits)
+        for _ in range(16))
+
+
+def create_stream_index(shards: int, samples_per_shard: int, avg_raw_shard_size: int,
+                        avg_zip_shard_size: Optional[int]) -> str:
+    """Create dataset index file from input parameters.
+
+    Args:
+        shards (int): Number of shards.
+        samples_per_shard (int): Number of samples per shard.
+        avg_raw_shard_size (int): Average raw shard size.
+        avg_zip_shard_size (int): Average compressed shard size.
+
+    Returns:
+        local path to created index file for stream.
+    """
+    index_data = {'version': 2, 'shards': []}
+
+    shards_list = []
+    for _ in range(shards):
+        shard_data = {
+            'column_encodings': [],
+            'column_names': [],
+            'column_sizes': [],
+            'format': 'mds',
+            'raw_data': {
+                'basename': '',
+                'bytes': avg_raw_shard_size,
+                'hashes': {}
+            },
+            'hashes': [],
+            'samples': samples_per_shard,
+            'size_limit': avg_raw_shard_size,
+            'version': 2,
+            'zip_data': None,
+            'compression': None
+        }
+        if avg_zip_shard_size is not None:
+            shard_data['zip_data'] = {'basename': '', 'bytes': avg_zip_shard_size, 'hashes': {}}
+            shard_data['compression'] = ''
+        shards_list.append(shard_data)
+
+    index_data['shards'] = shards_list
+
+    # Try making the directory for the stream's index.json file
+    foldername = get_random_foldername() + '_indexcreated'
+    try:
+        os.mkdir(foldername)
+    except FileExistsError:
+        logger.warning(' Folder already exists, trying again...')
+        foldername = get_random_foldername()
+        os.mkdir(foldername)
+
+    index_basename = get_index_basename()
+    index_path = os.path.join(foldername, index_basename)
+
+    with open(index_path, 'w') as f:
+        json.dump(index_data, f)
+
+    return index_path

simulation/core/last_used_ordered_set.py

@@ -0,0 +1,37 @@
+# Copyright 2022-2024 MosaicML Streaming authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""An ordered set that can be used as an LRU cache."""
+
+from collections import OrderedDict
+from typing import Any
+
+
+class LastUsedOrderedSet(OrderedDict):
+    """An ordered dict that can be used as an LRU cache.
+
+    This is a subclass of OrderedDict, with some LRU-specific functions and all values as ``None``.
+    """
+
+    def setitem(self, key: Any, move_to_end: bool = True):
+        """Set/add an item.
+
+        Args:
+            key (Any): key to be added.
+            move_to_end (bool, optional): whether to move the item to the end, signifying most
+                recent access. Defaults to ``True``.
+        """
+        super().__setitem__(key, None)
+        self.move_to_end(key, last=move_to_end)
+
+    def popLRU(self):
+        """Pop the least recently used item (located at the front)."""
+        return self.popitem(last=False)[0]
+
+    def setuse(self, key: Any):
+        """Mark an item as used, moving it to the end.
+
+        Args:
+            key (Any): key of element to move to the end, signifying most recent access.
+        """
+        self.setitem(key)

simulation/core/main.py

@@ -0,0 +1,243 @@
+# Copyright 2022-2024 MosaicML Streaming authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Main simulation function, simulating bytes downloaded and time taken each training step."""
+
+import logging
+import time
+from typing import Generator, Union
+
+import numpy as np
+from numpy.typing import NDArray
+
+from simulation.core.node_tracker import NodeTracker
+from simulation.core.shard_downloads import run_cache_limit, simulate_shard_downloads
+from simulation.core.sim_dataset import SimulationDataset
+from simulation.core.sim_time import Time
+from simulation.core.utils import bytes_to_time, get_batches_epochs, time_to_bytes
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
+def simulate(
+    dataset: SimulationDataset,
+    time_per_sample: float,
+    node_network_bandwidth: int,
+    max_duration: Time,
+    generator: bool = False
+) -> Generator[Union[tuple[int, float, int], tuple[float, int], tuple[NDArray, NDArray, float,
+                                                                      int]], None, None]:
+    """Simulates step time and downloads using streaming for the specified input parameters.
+
+    At each training step, the simulation does the following:
+       * gets the shards containing the current batch's samples.
+       * for each node, downloads shards if they are not present (round-robin through workers).
+       * predownloads more shards during model process time plus extra time from previous step.
+       * tracks the time the the step took as well as downloaded bytes.
+
+    Key Notes and Assumptions:
+       * assume that batch time is solely made up of two things: batch shard download wait time
+         and batch processing time
+       * loop through workers in a node round-robin style for batches and for downloads
+       * assume each node has a separate, uniform network bandwidth
+       * the batch has to wait until all nodes have downloaded the shards containing batch samples.
+       * for shard eviction itself, use LRU shard eviction to take out the least recently used
+         shard, per node.
+       * shards are shared across devices on a single node, but nodes do not share shards between
+         each other.
+       * if a shard is unavailable, we download shards round-robin until we have it.
+       * if a shard is available in a node, we just use it.
+
+    Args:
+        dataset (SimulationDataset): SimulationDataset object created based on input params/yaml
+        time_per_sample (float): time to process one sample on one device (seconds)
+        node_network_bandwidth (int): network bandwidth per node (bytes/s)
+        max_duration (Time): max duration of simulation. Defaults to ``None``.
+        generator (bool): True if we yield throughput and shard_download one step at a time.
+
+    Returns:
+        Generator[Union[Tuple[int, float, int],
+            Tuple[NDArray, NDArray, float, int],
+            Tuple[float, int]], None, None]: either a tuple of step number, step time, and
+            downloaded bytes, a tuple of startup time and min needed cache limit,
+            (both when generator=True), or a tuple of all step times, downloaded bytes,
+            startup_time, and min needed cache limit.
+    """
+    # tracking startup time, which includes SimulationDataset instantiation time.
+    start_time = time.time()
+    startup_time = dataset.get_instantiation_time()
+
+    # Get batches, epochs, total batches from dataset and provided time info.
+    batches_per_epoch, epochs, total_batches = get_batches_epochs(dataset, max_duration)
+
+    # Retrieve streaming and dataset information from SimulationDataset.
+    physical_nodes = dataset.get_nodes()
+    devices = dataset.get_devices()
+    workers = dataset.get_workers()
+    device_batch_size = dataset.get_batch_size()
+    predownload = dataset.get_predownload()
+    cache_limit = dataset.get_cache_limit()
+    total_shards = dataset.get_num_shards()
+    raw_shard_sizes = dataset.get_raw_shard_sizes()
+    zip_shard_sizes = dataset.get_zip_shard_sizes()
+    # dataset's spanner object maps global sample id to shard id.
+    sample_to_shard = dataset.get_spanner()
+
+    # Initialize NodeTracker objects for each node. These keep track of shards, worker downloads,
+    # cache usage, shard usage ranges, etc. for each node.
+    nodes = []
+    for _ in range(physical_nodes):
+        nodes.append(
+            NodeTracker(workers, devices, predownload, device_batch_size, total_shards,
+                        cache_limit))
+
+    # Time for the global batch is just device batch size * time per sample.
+    # We assume all devices process their microbatch perfectly in parallel.
+    avg_batch_process_time = device_batch_size * time_per_sample
+
+    notification_batches = int(batches_per_epoch) / 20
+
+    # Simulate training by looping over epochs and batches.
+
+    # Track time and downloads at each step.
+    step_times = []
+    step_downloads = []
+
+    for epoch in range(epochs):
+
+        # Get the samples, divided up per node, for this epoch.
+        samples_per_node = dataset.get_samples_per_node(epoch, 0)
+
+        # Set the samples for each node for this epoch.
+        for node_id, node in enumerate(nodes):
+            node.samples = samples_per_node[node_id]
+            node.initialize_worker_downloads(sample_to_shard)
+
+        # Track which sample we are currently on, as a worker id. We round-robin over workers.
+        worker_sample_index = 0
+
+        # if first epoch, add time so far to startup time
+        if epoch == 0:
+            startup_time += time.time() - start_time
+
+        # Iterate over batches
+        for batch in range(batches_per_epoch):
+
+            step_num = batch + (batches_per_epoch * epoch)
+
+            # If we are at the last batch, exit.
+            if step_num >= total_batches:
+                break
+
+            # Log progress every notification_batches interval.
+            if (batch + 1) % notification_batches == 0:
+                logger.info(
+                    f' Epoch: {str(epoch + 1)} | Batch {str(batch + 1)}/{str(batches_per_epoch)}')
+
+            # We round-robin over workers per device. The current batch's worker is the same
+            # across every device.
+            curr_worker = batch % workers
+            # Track how long each node takes to download the shards that the current batch needs.
+            node_batch_download_times = np.array([0] * physical_nodes)
+            # Track how many total bytes are downloaded in this batch by all nodes.
+            downloaded_bytes = 0
+
+            # Get current samples and download samples for each node, for this batch.
+            for node_id, node in enumerate(nodes):
+
+                shards_needed, shards_present = node.get_current_batch_shards(
+                    curr_worker, worker_sample_index, sample_to_shard)
+                # Mark all shards present as accessed most recently in this node.
+                node.set_shards_used(shards_present, step_num)
+                # Push the predownload for the current batch workers ahead by device_batch_size.
+                node.update_worker_predownloads(curr_worker, worker_sample_index, sample_to_shard)
+                # Track bytes downloaded by this node.
+                node_downloaded_bytes = 0
+
+                # Because we assume downloads also round-robin over workers, we can download shards
+                # other than the current batch's shards while looking for current batch's shards.
+                while len(shards_needed) > 0:
+                    download_outcome, download_size = \
+                        simulate_shard_downloads(node,
+                                                 raw_shard_sizes,
+                                                 zip_shard_sizes,
+                                                 current_batch_downloads=True,
+                                                 step_num=step_num,
+                                                 cache_limit=cache_limit,
+                                                 shards_needed=shards_needed)
+                    if download_outcome == 'downloaded':
+                        node_downloaded_bytes += download_size
+                        downloaded_bytes += download_size
+                    elif download_outcome == 'present':
+                        # If the shard was already present in the node, continue downloading.
+                        pass
+                    else:
+                        # If no shard downloads are left in the node, stop downloading.
+                        break
+
+                # Calculate how much time this node spent downloading shards
+                node_batch_download_times[node_id] = bytes_to_time(node_downloaded_bytes,
+                                                                   node_network_bandwidth)
+
+            # The node that took the longest to download shards is the bottleneck. All other nodes
+            # use the extra time to continue downloading.
+            slowest_download_time = np.max(node_batch_download_times)
+
+            # If we are on the first step, add slowest_download_time to startup time
+            if epoch == 0 and batch == 0:
+                startup_time += slowest_download_time
+
+            # Iterate over nodes again to continue downloading shards.
+            for node_id, node in enumerate(nodes):
+                # The download time each node has is the avg_batch_process_time plus the extra time
+                # the node has from finishing earlier than the slowest node.
+                download_time_left = avg_batch_process_time + (slowest_download_time -
+                                                               node_batch_download_times[node_id])
+                # Get number of bytes we can download in download_time_left.
+                # We also include any partially downloaded bytes from previous steps.
+                download_bytes_left = time_to_bytes(download_time_left, node_network_bandwidth) + \
+                    node.partial_shard_bytes
+
+                while True:
+                    download_outcome, download_size = \
+                        simulate_shard_downloads(node,
+                                                 raw_shard_sizes,
+                                                 zip_shard_sizes,
+                                                 current_batch_downloads=False,
+                                                 step_num=step_num,
+                                                 cache_limit=cache_limit,
+                                                 download_bytes_left=download_bytes_left)
+                    if download_outcome == 'downloaded':
+                        downloaded_bytes += download_size
+                        download_bytes_left -= download_size
+                    elif download_outcome == 'present':
+                        pass
+                    else:
+                        # If no shard downloads are left in the node, or we could only partially
+                        # download a shard, stop downloading.
+                        break
+
+            # Yield or store step number, time and download for this step
+            if generator:
+                yield step_num, slowest_download_time + avg_batch_process_time, downloaded_bytes
+            else:
+                step_times.append(slowest_download_time + avg_batch_process_time)
+                step_downloads.append(downloaded_bytes)
+
+            # Increment worker_sample_index by device_batch_size if we are at the last worker.
+            # As we round-robin over workers, the sample index per worker is increased as we loop
+            # through all workers.
+            if curr_worker == workers - 1:
+                worker_sample_index += device_batch_size
+
+    # Simulation is finished. Calculate needed cache limit.
+    min_cache_limit = run_cache_limit(nodes, raw_shard_sizes)
+
+    # Yield results.
+    if not generator:
+        step_times = np.array(step_times)
+        step_downloads = np.array(step_downloads)
+        yield step_times, step_downloads, startup_time, min_cache_limit
+    else:
+        yield startup_time, min_cache_limit

simulation/core/node_tracker.py

@@ -0,0 +1,224 @@
+# Copyright 2022-2024 MosaicML Streaming authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Class for tracking node information during simulation."""
+
+from typing import Optional
+
+import numpy as np
+from numpy.typing import NDArray
+from sortedcollections import OrderedSet
+
+from simulation.core.last_used_ordered_set import LastUsedOrderedSet
+from simulation.core.utils import remove_padded_samples
+from streaming.base.spanner import Spanner
+
+
+class NodeTracker():
+    """Tracker for node information during simulation.
+
+    Args:
+        workers (int): The number of workers.
+        devices (int): The number of devices.
+        predownload (int): The number of samples to predownload.
+        device_batch_size (int): The device batch size.
+        total_shards (int): Total number of shards in the dataset.
+        cache_limit (Optional[int]): The cache limit for the node. Defaults to ``None``.
+    """
+
+    def __init__(self,
+                 workers: int,
+                 devices: int,
+                 predownload: int,
+                 device_batch_size: int,
+                 total_shards: int,
+                 cache_limit: Optional[int] = None):
+
+        self.shards = LastUsedOrderedSet()
+        self.all_shards = set()
+        self.cache_usage = 0
+        self.partial_shard_bytes = 0
+        self.worker_download_index = 0
+        self.devices = devices
+        self.workers = workers
+        self.total_workers = workers * devices
+        self.device_batch_size = device_batch_size
+        self.predownload = predownload
+        self.cache_limit = cache_limit
+        self.worker_downloads = []
+        self.shard_access_starts = np.full(total_shards, -1)
+        self.shard_access_ends = np.full(total_shards, -1)
+
+        # Use the set_epoch_samples method every epoch to set the node's samples.
+        self.samples = None
+
+    def initialize_worker_downloads(self, sample_to_shard: Spanner):
+        """Initialize worker downloads, making shards in the predownload sample range available.
+
+        Args:
+            sample_to_shard (Spanner): The mapping from samples to shards.
+        """
+        # For downloads, we round-robin over devices first, then workers.
+        if self.samples is not None:
+            for worker in range(self.workers):
+                for device in range(self.devices):
+                    download_samples = remove_padded_samples(
+                        self.samples[device, worker, :self.predownload])
+                    # Get the shards these samples correspond to, maintaining access order.
+                    download_shards = OrderedSet(
+                        [sample_to_shard[sample] for sample in download_samples])
+                    self.worker_downloads.append(download_shards)
+        else:
+            raise AttributeError('Must set node samples before accessing them.')
+
+    def set_shards_used(self, shards: set, step_num: int):
+        """Mark a set of shards as recently used.
+
+        Args:
+            shards (set): The shards to set as used.
+            step_num (int): The current step number.
+        """
+        for shard in shards:
+            self.shards.setuse(shard)
+            # For any shard access, we are accessing the shard so we need the shard until
+            # at least the next step begins. Adding 0.5 ensures that we evict shards
+            # after they are used for the last time, but before they are replaced by
+            # new downloads in the next step.
+            self.shard_access_ends[shard] = step_num + 0.5
+
+    def add_shard(self, shard: int, used: bool = True):
+        """Add a shard to the node.
+
+        Args:
+            shard (int): The shard to add.
+            used (bool): Whether the shard is used when added.
+        """
+        self.shards.setitem(shard, used)
+        self.all_shards.add(shard)
+
+    def get_all_shards(self):
+        """Get all the shards in the node."""
+        return self.all_shards
+
+    def evict_shard(self) -> int:
+        """Evict a shard.
+
+        Returns:
+            int: The evicted shard.
+        """
+        evicted_shard = self.shards.popLRU()
+        return evicted_shard
+
+    def evict_until_satisfied(self, incoming_shard_size: int, raw_shard_sizes: NDArray[np.int64]):
+        """Evict shards until the node has enough space to download the incoming shard.
+
+        Args:
+            incoming_shard_size (int): The size of the incoming shard.
+            raw_shard_sizes (NDArray[np.int64]): The raw shard sizes.
+        """
+        # We evict shards until the incoming shard fits into the node's cache.
+        if self.cache_limit is not None:
+            while self.cache_usage + incoming_shard_size > self.cache_limit:
+                evicted_shard = self.evict_shard()
+                self.cache_usage -= int(raw_shard_sizes[evicted_shard])
+
+    def increment_worker_download_index(self):
+        """Increment the worker download index."""
+        self.worker_download_index = (self.worker_download_index + 1) % \
+            (self.workers * self.devices)
+
+    def get_worker_download(self,
+                            worker: Optional[int] = None,
+                            device: Optional[int] = None,
+                            index: Optional[int] = None) -> OrderedSet:
+        """Get the shard downloads for a worker on a specific device.
+
+        Args:
+            worker (Optional[int]): The worker index.
+            device (Optional[int]): The device index the worker is on.
+            index (Optional[int]): The index of the worker download for direct access.
+
+        Returns:
+            OrderedSet: The shard downloads, in order, for this worker.
+        """
+        if index is not None:
+            # Directly access worker_downloads through an index.
+            return self.worker_downloads[index]
+        elif worker is not None and device is not None:
+            # Access worker_downloads through worker and device indices.
+            return self.worker_downloads[worker * self.devices + device]
+        else:
+            raise ValueError('Must specify either index, or worker and device.')
+
+    def get_current_batch_shards(self, worker: int, worker_sample_index: int,
+                                 sample_to_shard: Spanner) -> tuple[set, set]:
+        """Get this node's shards for the current batch.
+
+        Args:
+            worker (int): The worker.
+            worker_sample_index (int): The worker sample index.
+            sample_to_shard (Spanner): The mapping from samples to shards.
+
+        Returns:
+            Tuple[set, set]: shard ids needed by node, shard ids present in node.
+        """
+        if self.samples is not None:
+            batch_samples = remove_padded_samples(
+                self.samples[:, worker, worker_sample_index:worker_sample_index +
+                             self.device_batch_size].flatten())
+            batch_shards = {sample_to_shard[sample] for sample in batch_samples}
+            shards_needed = batch_shards.difference(self.shards.keys())
+            shards_present = batch_shards.difference(shards_needed)
+            return shards_needed, shards_present
+        else:
+            raise AttributeError('Must set node samples before accessing them.')
+
+    def get_next_worker_with_downloads(self) -> Optional[OrderedSet]:
+        """Get the next worker with samples to download.
+
+        Returns:
+            Optional[OrderedSet]: The next worker's shard downloads, or None if no workers have
+                samples to download.
+        """
+        empty_download_counter = 0
+        worker_download = self.get_worker_download(index=self.worker_download_index)
+        while len(worker_download) == 0:
+            empty_download_counter += 1
+            self.increment_worker_download_index()
+            worker_download = self.get_worker_download(index=self.worker_download_index)
+
+            # No workers in the node have samples to download.
+            if empty_download_counter >= self.total_workers:
+                return None
+
+        return worker_download
+
+    def update_worker_predownloads(self, worker: int, worker_sample_index: int,
+                                   sample_to_shard: Spanner):
+        """Update the worker predownload samples for a worker and device.
+
+        Args:
+            worker (int): The current batch worker index.
+            worker_sample_index (int): The worker sample index.
+            sample_to_shard (Spanner): The mapping from samples to shards.
+        """
+        if self.samples is not None:
+            for device in range(self.devices):
+                # Retrieve new samples that are now within predownload range of the worker.
+                new_download_samples = remove_padded_samples(
+                    self.samples[device, worker,
+                                 worker_sample_index + self.predownload:worker_sample_index +
+                                 self.device_batch_size + self.predownload])
+
+                # Want to maintain the shard access order.
+                new_download_shards = OrderedSet(
+                    [sample_to_shard[sample] for sample in new_download_samples])
+
+                worker_downloads = self.get_worker_download(worker=worker, device=device)
+
+                # Add in new shards to the worker's shard downloads only if node does not yet have it.
+                for shard in new_download_shards:
+                    if shard not in self.shards:
+                        worker_downloads.add(shard)
+        else:
+            raise AttributeError('Must set node samples before accessing them.')