Close prefix-lookup shared memory probes and reset a shard's state when its download fails

CHANGELOG.md

@@ -8,6 +8,8 @@ streaming-community is a fork of MosaicML Streaming 0.13.0. Each entry names the
 
 - Start DataLoader workers with `fork` after the parent iterated the dataset: a forked child inherited the parent's thread pool without its threads and waited forever. A dataset that is in, or has finished, an epoch also pickles again, so `spawn` and `forkserver` workers no longer fail on it either (#48).
 - Raise the error of a failed download thread from iteration, with the cause attached, instead of ending the epoch early or waiting forever for the thread. In 0.14.0, a `cache_limit` too tight for a mix of local-only and remote streams hung or returned an empty epoch this way; that case now fails at construction with the limit to set. The ready thread also downloads a shard again if it was evicted while being waited for, and stops at interpreter exit, so a stalled epoch no longer blocks shutdown (#49).
+- Close each shared memory block that the prefix lookup opens to read another dataset's local directories. Every probe stayed open until the process exited, two file descriptors each, so they piled up in a process that creates many datasets (#54).
+- Return a shard to its previous state when its download or decompression raises. It stayed `PREPARING`, so every worker waiting for it waited forever and no later call downloaded it again; now the next call retries (#54).
 
 ### Docs
 

streaming/base/dataset.py

@@ -1221,7 +1221,15 @@ def prepare_shard(self, shard_id: int, blocking: bool = True) -> None:
             lock.release()
 
             # Perform the download (shard will not be modified by others in PREPARING state).
-            delta = stream.prepare_shard(shard)
+            # If it fails, put the shard back to REMOTE so that waiters stop waiting and the next
+            # call downloads it again, instead of leaving it in PREPARING forever.
+            try:
+                delta = stream.prepare_shard(shard)
+            except BaseException:
+                lock.acquire()
+                self._shard_states[shard_id] = _ShardState.REMOTE
+                lock.release()
+                raise
 
             # Download completed, so note the time and transition shard state to LOCAL.
             lock.acquire()
@@ -1258,8 +1266,14 @@ def prepare_shard(self, shard_id: int, blocking: bool = True) -> None:
                 # Lock the shard.
                 self._shard_states[shard_id] = _ShardState.PREPARING
                 lock.release()  # Unblock other workers.
-                # Decompress and remove zip.
-                delta = stream.prepare_shard(shard)
+                # Decompress and remove zip. If that fails, restore LOCAL so the next call retries.
+                try:
+                    delta = stream.prepare_shard(shard)
+                except BaseException:
+                    lock.acquire()
+                    self._shard_states[shard_id] = _ShardState.LOCAL
+                    lock.release()
+                    raise
                 lock.acquire()  # Briefly take the lock back.
                 # Restore shard state.
                 self._shard_states[shard_id] = _ShardState.LOCAL

streaming/base/shared/prefix.py

@@ -144,16 +144,21 @@ def _check_and_find(streams_local: list[str], streams_remote: list[Union[str, No
 
         # Attempt to access shared memory by name. Use prefix_int if files do not exist
         try:
-            shm = SharedMemory(name, False)
+            shm = SharedMemory(name, False, auto_cleanup=False)
         except PermissionError:
             continue
         except FileNotFoundError:
             break
 
+        # The probe only reads the block: copy what is needed and close it now, instead of
+        # leaving it to an atexit handler that keeps it open until the process exits.
+        data = bytes(shm.buf) if shm_name == LOCALS else b''
+        shm.cleanup()
+
         if shm_name != LOCALS:
             continue
 
-        their_locals, _ = _unpack_locals(bytes(shm.buf))
+        their_locals, _ = _unpack_locals(data)
 
         # Do not check for a conflicting local directories across existing shared memory if
         # remote directories are None. Get the next prefix.

tests/test_eviction.py

@@ -8,11 +8,13 @@
 import pickle
 from shutil import copytree, rmtree
 from typing import Any
+from unittest.mock import patch
 
 import pytest
 from torch.utils.data import DataLoader
 
 from streaming import MDSWriter, Stream, StreamingDataset
+from streaming.base.dataset import _ShardState
 from tests.common.utils import convert_to_mds
 
 
@@ -388,3 +390,54 @@ def streams() -> list[Stream]:
                                cache_limit=floor + 4 * largest)
     assert sum(1 for _ in dataset) == 400
     assert set(os.listdir(local_only)) == set(os.listdir(remote))
+
+
+def _write_small_shards(remote: str, compression: Any = None) -> None:
+    with MDSWriter(out=remote, columns={'data': 'bytes'}, compression=compression,
+                   size_limit=500) as out:
+        for _ in range(200):
+            out.write({'data': b'\0'})
+
+
+def test_failed_download_returns_shard_to_remote(local_remote_dir: tuple[str, str]):
+    """A failed download puts the shard back to REMOTE, so the next call downloads it again."""
+    local, remote = local_remote_dir
+    _write_small_shards(remote)
+    dataset = StreamingDataset(remote=remote, local=local, batch_size=1)
+    download = Stream.prepare_shard
+    calls = []
+
+    def fail_first(self: Stream, shard: Any) -> int:
+        calls.append(shard)
+        if len(calls) == 1:
+            raise OSError('simulated download failure')
+        return download(self, shard)
+
+    with patch.object(Stream, 'prepare_shard', fail_first):
+        with pytest.raises(OSError, match='simulated download failure'):
+            dataset.prepare_shard(0)
+        assert dataset._shard_states[0] == _ShardState.REMOTE
+        dataset.prepare_shard(0, blocking=False)
+    assert dataset._shard_states[0] == _ShardState.LOCAL
+    assert len(calls) == 2
+
+
+def test_failed_decompression_keeps_shard_local(local_remote_dir: tuple[str, str]):
+    """A failed decompression restores LOCAL, so the next call decompresses again."""
+    local, remote = local_remote_dir
+    _write_small_shards(remote, compression='zstd')
+    dataset = StreamingDataset(remote=remote, local=local, batch_size=1, keep_zip=True)
+    dataset.prepare_shard(0)
+    stream = dataset.streams[dataset.stream_per_shard[0]]
+    raw_info, _ = dataset.shards[0].file_pairs[0]
+    raw_filename = os.path.join(stream.local, stream.split, raw_info.basename)
+    os.remove(raw_filename)
+
+    with patch.object(Stream,
+                      'prepare_shard',
+                      side_effect=OSError('simulated decompression failure')):
+        with pytest.raises(OSError, match='simulated decompression failure'):
+            dataset.prepare_shard(0)
+    assert dataset._shard_states[0] == _ShardState.LOCAL
+    dataset.prepare_shard(0, blocking=False)
+    assert os.path.isfile(raw_filename)

tests/test_shared.py

@@ -404,3 +404,22 @@ def test_forced_race():
             clean_stale_shared_memory()
         except:
             pass
+
+
+@pytest.mark.skipif(not os.path.isdir('/dev/fd'), reason='counts open descriptors in /dev/fd')
+def test_check_and_find_closes_probes(tmp_path: Any):
+    """Probing an existing prefix leaves no shared memory descriptor open."""
+    clean_stale_shared_memory()
+    data = _pack_locals([str(tmp_path / 'theirs')], 0)
+    owner = SharedMemory(_get_path(0, LOCALS), True, len(data), auto_cleanup=False)
+    owner.buf[:len(data)] = data
+    try:
+        # No filelocks in this directory, so every call probes prefix 0.
+        with patch('streaming.base.shared.prefix.gettempdir', return_value=str(tmp_path)):
+            before = len(os.listdir('/dev/fd'))
+            for _ in range(20):
+                assert _check_and_find([str(tmp_path / 'ours')], [None], LOCALS) == 1
+            after = len(os.listdir('/dev/fd'))
+    finally:
+        owner.cleanup()
+    assert after == before