From fe34fb847b25ca28c54ac27c32357726ce3ef51e Mon Sep 17 00:00:00 2001 From: August Vi Date: Tue, 20 Oct 2026 11:51:03 -0700 Subject: [PATCH] Close prefix-lookup shared memory probes, and return a shard to its previous state when its download or decompression fails. --- CHANGELOG.md | 2 ++ streaming/base/dataset.py | 20 +++++++++++-- streaming/base/shared/prefix.py | 9 ++++-- tests/test_eviction.py | 53 +++++++++++++++++++++++++++++++++ tests/test_shared.py | 19 ++++++++++++ 5 files changed, 98 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 28e2be4..528c72b 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index 9403614..6d6178d 100644 --- a/streaming/base/dataset.py +++ b/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 diff --git a/streaming/base/shared/prefix.py b/streaming/base/shared/prefix.py index 6f93e91..18925d0 100644 --- a/streaming/base/shared/prefix.py +++ b/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. diff --git a/tests/test_eviction.py b/tests/test_eviction.py index 23d76df..1c3d362 100644 --- a/tests/test_eviction.py +++ b/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) diff --git a/tests/test_shared.py b/tests/test_shared.py index 3db8687..14ff078 100644 --- a/tests/test_shared.py +++ b/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