diff --git a/deepspec/data/cuda_prefetcher.py b/deepspec/data/cuda_prefetcher.py index 69eac905..37be8ad8 100644 --- a/deepspec/data/cuda_prefetcher.py +++ b/deepspec/data/cuda_prefetcher.py @@ -29,6 +29,7 @@ def __iter__(self): self._done = False self._gpu_batch = None self._thread = None + self._background_error = None # First batch: fetch synchronously so __next__ has something to return. self._fetch_and_transfer() return self @@ -43,12 +44,24 @@ def _fetch_and_transfer(self): with torch.cuda.stream(self.stream): self._gpu_batch = move_batch_to_device(cpu_batch, self.device) + def _prefetch(self): + try: + self._fetch_and_transfer() + except BaseException as exc: + self._background_error = exc + def __next__(self): # Join the background thread kicked off in the previous __next__. if self._thread is not None: self._thread.join() self._thread = None + if self._background_error is not None: + error = self._background_error + self._background_error = None + self._gpu_batch = None + raise error + if self._done: raise StopIteration @@ -64,7 +77,7 @@ def __next__(self): # Kick off the next fetch and H2D in a background thread so it # overlaps with compute on the batch we are about to return. - self._thread = Thread(target=self._fetch_and_transfer, daemon=True) + self._thread = Thread(target=self._prefetch, daemon=True) self._thread.start() return batch diff --git a/tests/test_cuda_prefetcher.py b/tests/test_cuda_prefetcher.py new file mode 100644 index 00000000..da7388d4 --- /dev/null +++ b/tests/test_cuda_prefetcher.py @@ -0,0 +1,116 @@ +import contextlib +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + + +class FakeTensor: + def __init__(self, value): + self.value = value + self.dtype = "int64" + self.recorded_streams = [] + + def to(self, *args, **kwargs): + return self + + def record_stream(self, stream): + self.recorded_streams.append(stream) + + +class FakeStream: + def wait_stream(self, stream): + self.waited_for = stream + + +class ScriptedIterator: + def __init__(self, events): + self._events = iter(events) + + def __iter__(self): + return self + + def __next__(self): + event = next(self._events) + if isinstance(event, BaseException): + raise event + return event + + +class ResettableLoader: + def __init__(self, runs): + self._runs = iter(runs) + + def __iter__(self): + return ScriptedIterator(next(self._runs)) + + def __len__(self): + return 1 + + +@pytest.fixture +def prefetcher_module(monkeypatch): + current_stream = FakeStream() + fake_cuda = types.SimpleNamespace( + Stream=lambda device=None: FakeStream(), + stream=lambda stream: contextlib.nullcontext(), + current_stream=lambda device=None: current_stream, + ) + fake_torch = types.ModuleType("torch") + fake_torch.cuda = fake_cuda + fake_torch.long = "int64" + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + module_path = ( + Path(__file__).parents[1] / "deepspec" / "data" / "cuda_prefetcher.py" + ) + spec = importlib.util.spec_from_file_location("cuda_prefetcher_under_test", module_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def batch(value): + return {"input_ids": FakeTensor(value)} + + +def test_normal_batches_are_returned_in_order(prefetcher_module): + prefetcher = prefetcher_module.CUDAPrefetcher( + ScriptedIterator([batch(1), batch(2)]), + "cuda:0", + ) + + assert [item["input_ids"].value for item in prefetcher] == [1, 2] + + +def test_background_loader_error_is_reraised(prefetcher_module): + error = RuntimeError("bad batch") + prefetcher = prefetcher_module.CUDAPrefetcher( + ScriptedIterator([batch(1), error]), + "cuda:0", + ) + iterator = iter(prefetcher) + + assert next(iterator)["input_ids"].value == 1 + with pytest.raises(RuntimeError, match="bad batch") as raised: + next(iterator) + assert raised.value is error + + +def test_reiteration_resets_background_error(prefetcher_module): + loader = ResettableLoader( + [ + [batch(1), RuntimeError("bad batch")], + [batch(2)], + ] + ) + prefetcher = prefetcher_module.CUDAPrefetcher(loader, "cuda:0") + first = iter(prefetcher) + + assert next(first)["input_ids"].value == 1 + with pytest.raises(RuntimeError, match="bad batch"): + next(first) + + assert [item["input_ids"].value for item in prefetcher] == [2]