From 8928c96b10df00dbe416f52a8b78e2611c670096 Mon Sep 17 00:00:00 2001 From: August Vi Date: Sun, 11 Oct 2026 19:23:03 -0700 Subject: [PATCH 1/2] Raise a failed download thread's error from iteration instead of a short epoch or a hang, and reject a cache_limit that local-only shards would fill at construction. --- CHANGELOG.md | 1 + streaming/base/dataset.py | 71 +++++++++++++++++++++++++++------------ tests/test_eviction.py | 29 ++++++++++++++++ tests/test_streaming.py | 27 +++++++++++++++ 4 files changed, 106 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b2831b..23419d7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ streaming-community is a fork of MosaicML Streaming 0.13.0. Each entry names the ### Fixes - 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 (#49). ### Docs diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index 29c400b..5572f29 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -9,13 +9,12 @@ import os import warnings from concurrent.futures import ThreadPoolExecutor, wait -from concurrent.futures._base import Future from enum import IntEnum from math import ceil from tempfile import gettempdir from threading import Event, Lock from time import sleep, time_ns -from typing import Any, Iterator, Optional, Sequence, Union +from typing import Any, Callable, Iterator, Optional, Sequence, Union import numpy as np from filelock import FileLock @@ -518,6 +517,21 @@ def __init__(self, f'bytes) which includes raw (decompressed) and zip ' + f'(compressed) file size. Recommendation is to provide a ' + f'`cache_limit` as high as possible to avoid thrashing.') + # Shards of a stream without a remote are never evicted, because ``local`` is their + # only copy. Like the index files, they occupy the cache for good, so the limit must + # leave room above them for the shards that are downloaded. + local_only_bytes = sum(shard.get_max_size() + for shard, stream_id in zip(self.shards, stream_per_shard) + if streams[stream_id].remote is None) + if local_only_bytes and self.cache_limit < (min_cache_usage + local_only_bytes + 4 * + self.max_shard_size_across_all_streams): + raise ValueError(f'Cache limit ({self.cache_limit} bytes) is too low. The index ' + + f'files ({min_cache_usage} bytes) and the shards of streams ' + + f'with no remote ({local_only_bytes} bytes) are never evicted, ' + + f'so the limit must be at least their size plus four times ' + + f'the largest shard size ' + + f'({self.max_shard_size_across_all_streams} bytes). Raise ' + + f'`cache_limit`, or leave it unset if every stream is local.') # Build the shard index (for partitioning and mapping samples to shards). self.samples_per_shard = np.array([shard.samples for shard in self.shards], np.int64) @@ -1269,9 +1283,8 @@ def get_item(self, sample_id: int, retry: int = 7) -> Any: Returns: Dict[str, Any]: Mapping of column name to column data. """ - # Background thread crashed, terminate the main process - if hasattr(self, '_event') and self._event.is_set(): - raise RuntimeError('Background thread failed. Check other traceback.') + # A background thread crashed, so stop the main thread too. + self._raise_if_background_failed() # Locate the shard and sample offset within that shard where the sample lives. shard_id, shard_sample_id = self.spanner[sample_id] shard = self.shards[shard_id] @@ -1315,23 +1328,38 @@ def get_item(self, sample_id: int, retry: int = 7) -> Any: return sample - def on_exception(self, future: Future) -> None: - """Raise an exception to the caller if an exception was generated by a thread. + def _run_thread(self, fn: Callable[[_Iterator], None], it: _Iterator) -> None: + """Run a thread of __iter__, recording a failure so the yield loop stops and reports it. - Also, set the thread event to let the other threads know about the exception. + A thread that raises skips its own ``it.on_exit()``, so it is called here. Otherwise + ``_Iterator.exit()`` would wait for that thread forever. Args: - future (Future): The status of the task. + fn (Callable[[_Iterator], None]): Thread body. + it (_Iterator): State of __iter__. + """ + try: + fn(it) + except BaseException as exc: + self._background_error = exc + self._event.set() + it.on_exit() + raise + + def _raise_if_background_failed(self) -> None: + """Raise if a thread of __iter__ has failed in this process, chaining its exception. Raises: - Exception: re-raises the exception. + RuntimeError: If a thread has failed. """ - exception = future.exception() - if exception: - # Set the event to let the other threadpool threads know about the exception. - self._event.set() - # Re-raise the exception. - raise exception + event = getattr(self, '_event', None) + if event is None or not event.is_set(): + return + error = getattr(self, '_background_error', None) + if error is None: + raise RuntimeError('Background thread failed. Check other traceback.') + raise RuntimeError(f'StreamingDataset background thread failed: {type(error).__name__}: ' + + f'{error}') from error def _prepare_thread(self, it: _Iterator) -> None: """Download the relevant shards in the background while we are being iterated. @@ -1516,8 +1544,7 @@ def __iter__(self) -> Iterator[dict[str, Any]]: # Exit the threads that are pre-downloading and iterating the shards for the previous # epoch. self._iterator.exit() - if self._event.is_set(): - raise RuntimeError('Background thread failed. Check other traceback.') + self._raise_if_background_failed() # Discover where we left off, if there is a checkpoint, or start at the next epoch. # Also pre-increment the epoch counter. @@ -1532,10 +1559,10 @@ def __iter__(self) -> Iterator[dict[str, Any]]: # Iterate over the samples while downloading ahead. self._iterator = it = _Iterator(sample_ids) - prepare_future = self._executor.submit(self._prepare_thread, it) - prepare_future.add_done_callback(self.on_exception) - ready_future = self._executor.submit(self._ready_thread, it) - ready_future.add_done_callback(self.on_exception) + prepare_future = self._executor.submit(self._run_thread, self._prepare_thread, it) + ready_future = self._executor.submit(self._run_thread, self._ready_thread, it) yield from map(self.__getitem__, self._each_sample_id(it)) wait([prepare_future, ready_future], return_when='FIRST_EXCEPTION') it.exit() + # A thread failure ends the yield loop early: report it instead of a short epoch. + self._raise_if_background_failed() diff --git a/tests/test_eviction.py b/tests/test_eviction.py index 1e1add5..f718a25 100644 --- a/tests/test_eviction.py +++ b/tests/test_eviction.py @@ -1,6 +1,7 @@ # Copyright 2022-2024 MosaicML Streaming authors # SPDX-License-Identifier: Apache-2.0 +import json import multiprocessing as mp import operator import os @@ -354,3 +355,31 @@ def test_evict_coldest_skips_local_only_stream(tmp_path: Any): dataset.evict_coldest_shard() assert set(os.listdir(local_only)) == before + + +def test_cache_limit_leaves_room_above_local_only_streams(tmp_path: Any): + """A cache limit that local-only shards would fill fails at init, not in a download thread.""" + remote = tmp_path / 'remote' + local_only = tmp_path / 'local_only' + with MDSWriter(out=str(remote), columns={'data': 'bytes'}, size_limit=500) as out: + for _ in range(200): + out.write({'data': b'\0' * 100}) + copytree(remote, local_only) + shards = json.load(open(remote / 'index.json'))['shards'] + largest = max(shard['raw_data']['bytes'] for shard in shards) + # Both index files and the local-only shards stay in the cache for good. + floor = 2 * os.path.getsize(remote / 'index.json') + sum( + shard['raw_data']['bytes'] for shard in shards) + + def streams() -> list[Stream]: + return [ + Stream(local=str(local_only), remote=None), + Stream(local=str(tmp_path / 'cache'), remote=str(remote)), + ] + + with pytest.raises(ValueError, match='never evicted'): + StreamingDataset(streams=streams(), batch_size=1, cache_limit=floor + 4 * largest - 1) + + dataset = StreamingDataset(streams=streams(), batch_size=1, cache_limit=floor + 4 * largest) + assert sum(1 for _ in dataset) == 400 + assert set(os.listdir(local_only)) == set(os.listdir(remote)) diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 9742887..19681c3 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -733,6 +733,33 @@ def test_dataloader_after_parent_use(local_remote_dir: tuple[str, str], start_me assert samples_seen == 117 +@pytest.mark.timeout(60) +@pytest.mark.usefixtures('local_remote_dir') +def test_background_thread_failure_raises(local_remote_dir: tuple[str, str], + monkeypatch: pytest.MonkeyPatch): + """A failing download thread raises from iteration, instead of a short epoch or a hang.""" + local, remote = local_remote_dir + convert_to_mds(out_root=remote, + dataset_name='sequencedataset', + num_samples=117, + size_limit=1 << 8) + dataset = StreamingDataset(local=local, remote=remote, batch_size=2) + + def fail(shard_id: int, blocking: bool = True) -> None: + raise OSError(f'download of shard {shard_id} failed') + + # Both the download thread and the ready thread hit this, so both die without `on_exit()`. + monkeypatch.setattr(dataset, 'prepare_shard', fail) + with pytest.raises(RuntimeError, match='download of shard') as info: + for _ in dataset: + pass + assert isinstance(info.value.__cause__, OSError) + + # The failure is permanent for this dataset, and the next epoch reports it right away. + with pytest.raises(RuntimeError, match='download of shard'): + next(iter(dataset)) + + @pytest.mark.parametrize('batch_size', [4]) @pytest.mark.parametrize('size_limit', [256, 512]) @pytest.mark.parametrize('seed', [1111]) From fe3ea85f16f3542fab29c13c55f742164d7f6edd Mon Sep 17 00:00:00 2001 From: August Vi Date: Sun, 11 Oct 2026 22:47:34 -0700 Subject: [PATCH 2/2] Download a shard again from the ready thread if it was evicted while being waited for, and let that wait stop at exit. --- CHANGELOG.md | 2 +- streaming/base/dataset.py | 14 +++++++------- tests/test_eviction.py | 7 ++++++- tests/test_streaming.py | 32 ++++++++++++++++++++++++++++++++ 4 files changed, 46 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 23419d7..28e2be4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,7 +7,7 @@ streaming-community is a fork of MosaicML Streaming 0.13.0. Each entry names the ### Fixes - 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 (#49). +- 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). ### Docs diff --git a/streaming/base/dataset.py b/streaming/base/dataset.py index 5572f29..9403614 100644 --- a/streaming/base/dataset.py +++ b/streaming/base/dataset.py @@ -1461,15 +1461,15 @@ def _ready_thread(self, it: _Iterator) -> None: # Wait for the shard for this sample to be downloaded and decompressed, if not already. shard_id, _ = self.spanner[sample_id] - # During cold shard eviction, shard state might go in the reverse direction. If a shard - # is missing while fetching a sample, download it. - if self._shard_states[shard_id] == _ShardState.REMOTE: - self.prepare_shard(shard_id, False) - # Wait for a shard file to download completely. while self._shard_states[shard_id] != _ShardState.LOCAL: - # Background thread or a main process crashed, terminate this thread. - if self._event.is_set(): + # A new epoch started, a background thread or the main process crashed: exit. + if it.should_exit() or self._event.is_set(): break + # During cold shard eviction, shard state might go in the reverse direction. If + # the shard is missing, or was evicted while we waited for it, download it: nobody + # else will, since the download thread is past it and the yield loop is behind us. + if self._shard_states[shard_id] == _ShardState.REMOTE: + self.prepare_shard(shard_id, False) sleep(TICK) # Step forward one sample. diff --git a/tests/test_eviction.py b/tests/test_eviction.py index f718a25..23d76df 100644 --- a/tests/test_eviction.py +++ b/tests/test_eviction.py @@ -380,6 +380,11 @@ def streams() -> list[Stream]: with pytest.raises(ValueError, match='never evicted'): StreamingDataset(streams=streams(), batch_size=1, cache_limit=floor + 4 * largest - 1) - dataset = StreamingDataset(streams=streams(), batch_size=1, cache_limit=floor + 4 * largest) + # At the bound there is room for four shards: the one being read, the one the ready thread + # waits for, and what the download thread fetches ahead, so keep the look-ahead to one shard. + dataset = StreamingDataset(streams=streams(), + batch_size=1, + predownload=4, + cache_limit=floor + 4 * largest) assert sum(1 for _ in dataset) == 400 assert set(os.listdir(local_only)) == set(os.listdir(remote)) diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 19681c3..74f7782 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -7,6 +7,7 @@ import pickle import shutil from multiprocessing import Process +from time import sleep from typing import Any, Optional import pytest @@ -14,6 +15,7 @@ from streaming.base import Stream, StreamingDataLoader, StreamingDataset from streaming.base.batching import generate_work +from streaming.base.dataset import _ShardState from streaming.base.stream import streams_registry from streaming.base.util import clean_stale_shared_memory from streaming.base.world import World @@ -760,6 +762,36 @@ def fail(shard_id: int, blocking: bool = True) -> None: next(iter(dataset)) +@pytest.mark.timeout(60) +@pytest.mark.usefixtures('local_remote_dir') +def test_ready_thread_requests_evicted_shard(local_remote_dir: tuple[str, str]): + """The ready thread must download a shard again if it is evicted while being waited for.""" + local, remote = local_remote_dir + convert_to_mds(out_root=remote, + dataset_name='sequencedataset', + num_samples=117, + size_limit=1 << 8) + dataset = StreamingDataset(local=local, remote=remote, batch_size=2, predownload=32) + assert dataset.num_shards > 2 + first, second = dataset.shards[0].samples, dataset.shards[1].samples + + # Pretend another worker is downloading the second shard. The download thread skips it and + # moves on to the shards after it, while the ready thread waits for it after the first shard. + dataset._shard_states[1] = _ShardState.PREPARING + iterator = iter(dataset) + assert len([next(iterator) for _ in range(first)]) == first + for _ in range(200): + if dataset._iterator.prepare_index > first + second: + break + sleep(0.05) + assert dataset._iterator.prepare_index > first + second + + # That worker evicted the shard instead. Only the ready thread can request it again. + dataset._shard_states[1] = _ShardState.REMOTE + assert next(iterator) is not None + assert dataset._shard_states[1] == _ShardState.LOCAL + + @pytest.mark.parametrize('batch_size', [4]) @pytest.mark.parametrize('size_limit', [256, 512]) @pytest.mark.parametrize('seed', [1111])