Raise a failed download thread's error from iteration and reject a cache_limit that local-only shards would fill

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. 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
 

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.
@@ -1433,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.
@@ -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()

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,36 @@ 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)
+
+    # 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))

tests/test_streaming.py

@@ -7,13 +7,15 @@
 import pickle
 import shutil
 from multiprocessing import Process
+from time import sleep
 from typing import Any, Optional
 
 import pytest
 from torch.utils.data import DataLoader
 
 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
@@ -733,6 +735,63 @@ 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.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])