Cold shard eviction

streaming/base/format/base/reader.py

@@ -3,9 +3,10 @@
 
 """Read and decode sample from shards."""
 
+import os
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, Dict, Iterator, List, Optional
+from typing import Any, Dict, Iterator, List, Optional, Set
 
 from streaming.base.array import Array
 
@@ -74,6 +75,152 @@ def __len__(self) -> int:
         """
         return self.samples
 
+    def _evict_raw(self) -> None:
+        """Remove all raw files belonging to this shard."""
+        for raw_info, _ in self.file_pairs:
+            filename = os.path.join(self.dirname, self.split, raw_info.basename)
+            if os.path.exists(filename):
+                os.remove(filename)
+
+    def _evict_zip(self) -> None:
+        """Remove all zip files belonging to this shard."""
+        for _, zip_info in self.file_pairs:
+            if zip_info:
+                filename = os.path.join(self.dirname, self.split, zip_info.basename)
+                if os.path.exists(filename):
+                    os.remove(filename)
+
+    def evict(self) -> None:
+        """Remove all files belonging to this shard."""
+        self._evict_raw()
+        self._evict_zip()
+
+    def init_local_dir(self, filenames_present: Set[str], keep_zip: bool) -> bool:
+        """Bring what shard files are present to a consistent state, returning whether present.
+
+        Args:
+            filenames_present (Set[str]): The listing of all files under dirname/[split/]. This is
+                listed once and then saved because there could potentially be very many shard
+                files.
+            keep_zip (bool): Whether to keep zip files when decompressing. Possible when
+                compression was used. Necessary when local is the remote or there is no remote.
+
+        Returns:
+            bool: Whether the shard is present.
+        """
+        # For raw/zip to be considered present, each raw/zip file must be present.
+        raw_files_present = 0
+        zip_files_present = 0
+        for raw_info, zip_info in self.file_pairs:
+            if raw_info:
+                filename = os.path.join(self.dirname, self.split, raw_info.basename)
+                if filename in filenames_present:
+                    raw_files_present += 1
+            if zip_info:
+                filename = os.path.join(self.dirname, self.split, zip_info.basename)
+                if filename in filenames_present:
+                    zip_files_present += 1
+
+        # If the shard raw files are partially present, garbage collect the present ones and mark
+        # the shard raw as not present, in order to achieve consistency.
+        if not raw_files_present:
+            is_raw_present = False
+        elif raw_files_present < len(self.file_pairs):
+            is_raw_present = False
+            self._evict_raw()
+        else:
+            is_raw_present = True
+
+        # Same as the above, but for shard zip files.
+        if not zip_files_present:
+            is_zip_present = False
+        elif zip_files_present < len(self.file_pairs):
+            is_zip_present = False
+            self._evict_zip()
+        else:
+            is_zip_present = True
+
+        # Do we keep_zip?
+        if keep_zip:
+            # If we can keep_zip, and we do, and have either raw or zip, we must have the other one
+            # too, because they are downloaded and decompressed together.
+            if self.compression and (is_zip_present != is_raw_present):
+                if is_raw_present:
+                    is_raw_present = False
+                    self._evict_raw()
+                elif is_zip_present:
+                    is_zip_present = False
+                    self._evict_zip()
+        else:
+            # If we don't keep_zip, drop any zip files.
+            if is_zip_present:
+                is_zip_present = False
+                self._evict_zip()
+
+        # Now, the shard is either entirely or not at all present given keep_zip.
+        return is_raw_present
+
+    def get_raw_size(self) -> int:
+        """Get the raw (uncompressed) size of this shard.
+
+        Returns:
+            int: Size in bytes.
+        """
+        size = 0
+        for info, _ in self.file_pairs:
+            size += info.bytes
+        return size
+
+    def get_zip_size(self) -> Optional[int]:
+        """Get the zip (compressed) size of this shard, if compression was used.
+
+        Returns:
+            Optional[int]: Size in bytes, or ``None`` if does not exist.
+        """
+        size = 0
+        for _, info in self.file_pairs:
+            if info is None:
+                return None
+            size += info.bytes
+        return size
+
+    def get_full_size(self) -> int:
+        """Get the full size of this shard.
+
+        "Full" in this case means both the raw (decompressed) and zip (compressed) versions are
+        resident (assuming it has a zip form). This is the maximum disk usage the shard can reach.
+        When compressed was used, even if keep_zip is ``False``, the zip form must still be
+        resident at the same time as the raw form during shard decompression.
+
+        Returns:
+            int: Size in bytes.
+        """
+        raw_size = self.get_raw_size()
+        zip_size = self.get_zip_size() or 0
+        return raw_size + zip_size
+
+    def get_persistent_size(self, keep_zip: bool) -> int:
+        """Get the persistent size of this shard.
+
+        "Persistent" in this case means whether both raw and zip are present is subject to
+        keep_zip. If we are not keeping zip files after decompression, they don't count to the
+        shard's persistent size on disk.
+
+        Args:
+            keep_zip (bool): Whether to keep zip files after decompressing.
+
+        Returns:
+            int: Size in bytes.
+        """
+        if self.compression:
+            if keep_zip:
+                size = self.get_full_size()
+            else:
+                size = self.get_raw_size()
+        else:
+            size = self.get_raw_size()
+        return size
+
     @abstractmethod
     def decode_sample(self, data: bytes) -> Dict[str, Any]:
         """Decode a sample dict from bytes.

streaming/base/shared/__init__.py

@@ -7,8 +7,10 @@
 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 CreateSharedMemory as CreateSharedMemory
+from streaming.base.shared.memory import SharedMemory as SharedMemory
 from streaming.base.shared.prefix import get_shm_prefix as get_shm_prefix
+from streaming.base.shared.scalar import SharedScalar as SharedScalar
 
-__all__ = ['SharedBarrier', 'CreateSharedMemory', 'get_shm_prefix']
+__all__ = ['SharedArray', 'SharedBarrier', 'SharedMemory', 'get_shm_prefix', 'SharedScalar']

streaming/base/shared/array.py

@@ -0,0 +1,58 @@
+# Copyright 2023 MosaicML Streaming authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""A numpy array of predetermined shape and dtype that lives in shared memory."""
+
+from typing import Any, Tuple, Union
+
+import numpy as np
+from numpy.typing import NDArray
+
+from streaming.base.shared.memory import SharedMemory
+
+
+class SharedArray:
+    """A numpy array of predetermined shape and dtype that lives in shared memory.
+
+    Args:
+        shape (Union[int, Tuple[int]]): Shape of the array.
+        dtype (type): Dtype of the array.
+        name (str): Its name in shared memory.
+    """
+
+    def __init__(self, shape: Union[int, Tuple[int]], dtype: type, name: str) -> None:
+        self.shape = np.empty(shape).shape
+        self.dtype = dtype
+        self.name = name
+        size = np.prod(shape) * dtype(0).nbytes
+        self.shm = SharedMemory(name=name, size=size)
+
+    def numpy(self) -> NDArray:
+        """Get as a numpy array.
+
+        We can't just internally store and use this numpy array shared memory wrapper because it's
+        not compatible with spawn.
+        """
+        return np.ndarray(self.shape, buffer=self.shm.buf, dtype=self.dtype)
+
+    def __getitem__(self, index: Any) -> Any:
+        """Get the scalar(s) at the given index, slice, or array of indices.
+
+        Args:
+            index (Any): The index, slice, or array of indices.
+
+        Returns:
+            The scalar(s) at the given location(s).
+        """
+        arr = self.numpy()
+        return arr[index]
+
+    def __setitem__(self, index: Any, value: Any) -> Any:
+        """Set the scalar(s) at the given index, slice, or array of indices.
+
+        Args:
+            index (Any): The index, slice, or array of indices.
+            value (Any): The scalar(s) at the given location(s).
+        """
+        arr = self.numpy()
+        arr[index] = value

streaming/base/shared/barrier.py

@@ -1,22 +1,20 @@
 # Copyright 2023 MosaicML Streaming authors
 # SPDX-License-Identifier: Apache-2.0
 
-"""Synchronization primitives that live in shared memory.
+"""Barrier that lives 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.
+Implemented with shared array and a filelock.
 """
 
 import atexit
 import os
 import shutil
-from multiprocessing import resource_tracker  # pyright: ignore
 from time import sleep
 
 import numpy as np
 from filelock import FileLock
 
-from streaming.base.shared.memory import CreateSharedMemory
+from streaming.base.shared.array import SharedArray
 
 # Time to wait, in seconds.
 TICK = 0.07
@@ -39,24 +37,19 @@ class SharedBarrier:
 
     def __init__(self, filelock_path: str, shm_path: str) -> None:
         self.filelock_path = filelock_path
-        self.created_shms = []
-        self.opened_shms = []
 
         # Create three int32 fields in shared memory: num_enter, num_exit, flag.
-        size = 3 * np.int32(0).nbytes
-        shared_barrier_shm = CreateSharedMemory(name=shm_path, size=size)
-        self._shm = shared_barrier_shm.shm
+        self._arr = SharedArray(3, np.int32, shm_path)
+        self.num_enter = 0
+        self.num_exit = -1
+        self.flag = True
 
         # Create filelock.
         dirname = os.path.dirname(filelock_path)
-        os.makedirs(dirname, exist_ok=True)
+        if dirname:
+            os.makedirs(dirname, exist_ok=True)
         self.lock = FileLock(filelock_path)
 
-        self._arr = np.ndarray(3, buffer=self._shm.buf, dtype=np.int32)
-        self._arr[0] = 0
-        self._arr[1] = -1
-        self._arr[2] = True
-
         def cleanup():
             """Directory clean up."""
             if os.path.islink(dirname):
@@ -125,10 +118,6 @@ def __call__(self, num_procs: int) -> None:
         Args:
             num_procs (int): How many processes are sharing this barrier.
         """
-        # Re-init the numpy array pointing to shared memory. Necessary when spawn is the
-        # multiprocessing method used.
-        self._arr = np.ndarray(3, buffer=self._shm.buf, dtype=np.int32)
-
         # Initialize num_exit to the number of processes.
         with self.lock:
             if self.num_exit == -1:
@@ -160,3 +149,5 @@ def __call__(self, num_procs: int) -> None:
         # Note that we exited.
         with self.lock:
             self.num_exit += 1
+            if self.num_exit == num_procs:
+                self.num_exit = -1

streaming/base/shared/memory.py

@@ -1,43 +1,35 @@
 # Copyright 2023 MosaicML Streaming authors
 # SPDX-License-Identifier: Apache-2.0
 
-"""Synchronization primitives 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.
-"""
+"""Improved quiet implementation of shared memory in pure python."""
 
 import atexit
 from multiprocessing import resource_tracker  # pyright: ignore
-from multiprocessing.shared_memory import SharedMemory
+from multiprocessing.shared_memory import SharedMemory as BuiltinSharedMemory
 from time import sleep
 from typing import Any, Optional
 
 # Time to wait, in seconds.
 TICK = 0.07
 
-# Time out to wait before raising exception
-TIMEOUT = 60
-
 
-class CreateSharedMemory:
-    """Create a new Shared Memory block or attach to an existing shared memory block.
+class SharedMemory:
+    """Improved quiet implementation of shared memory.
 
     Args:
-        name (Optional[str], optional): A unique shared memory block name. Defaults to ``None``.
-        create (Optional[bool], optional): Creates a new shared memory block or attaches to an
-            existing shared memory block. Defaults to ``None``.
+        name (str, optional): A unique shared memory block name. Defaults to ``None``.
+        create (bool, optional): Creates a new shared memory block or attaches to an existin
+            shared memory block. Defaults to ``None``.
         size (int, optional): A size of a shared memory block. Defaults to ``0``.
-        auto_cleanup (bool, optional): Register atexit handler for cleanup or not.
-            Defaults to ``True``.
+        auto_cleanup (bool, optional): Register atexit handler for cleanup or not. Defaults to
+            ``True``.
     """
 
     def __init__(self,
                  name: Optional[str] = None,
                  create: Optional[bool] = None,
                  size: int = 0,
                  auto_cleanup: bool = True):
-
         self.created_shms = []
         self.opened_shms = []
         shm = None
@@ -47,28 +39,29 @@ def __init__(self,
         try:
             if create is True:
                 # Creates a new shared memory block
-                shm = SharedMemory(name, create, size)
+                shm = BuiltinSharedMemory(name, create, size)
                 self.created_shms.append(shm)
             elif create is False:
                 # Avoid tracking shared memory resources in a process who attaches to an existing
                 # shared memory block because the process who created the shared memory is
                 # responsible for destroying the shared memory block.
                 resource_tracker.register = self.fix_register
                 # Attaches to an existing shared memory block
-                shm = SharedMemory(name, create, size)
+                shm = BuiltinSharedMemory(name, create, size)
                 self.opened_shms.append(shm)
             else:
                 try:
                     # Creates a new shared memory block
-                    shm = SharedMemory(name, True, size)
+                    shm = BuiltinSharedMemory(name, True, size)
                     self.created_shms.append(shm)
                 except FileExistsError:
                     sleep(TICK)
                     resource_tracker.register = self.fix_register
                     # Attaches to an existing shared memory block
-                    shm = SharedMemory(name, False, size)
+                    shm = BuiltinSharedMemory(name, False, size)
                     self.opened_shms.append(shm)
             self.shm = shm
+            self.buf = shm.buf
         finally:
             resource_tracker.register = original_rtracker_reg