diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index dd0632f82..8f72090f5 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -27,7 +27,6 @@ import threading import time from collections import deque from concurrent.futures import ThreadPoolExecutor -from contextvars import ContextVar from multiprocessing import RawArray from typing import Any, Callable, Iterator, NamedTuple, Optional, Union @@ -51,10 +50,6 @@ _EPOCH_PRIME = 100003 DEFAULT_READ_BATCH_SIZE = 64 DEFAULT_PREFETCH_BATCHES = 4 -_CONSUMER_CHECKPOINT_DATASET: ContextVar[Optional[object]] = ContextVar( - "lancedb_streaming_consumer_checkpoint_dataset", default=None -) - class _WorkerSample(NamedTuple): data: Any @@ -98,12 +93,8 @@ class _StreamingDatasetAdapter(IterableDataset): self.dataset = dataset def __iter__(self): - token = _CONSUMER_CHECKPOINT_DATASET.set(self.dataset) - try: - for sample in self.dataset: - yield _WorkerSample(sample, self.dataset) - finally: - _CONSUMER_CHECKPOINT_DATASET.reset(token) + for sample in self.dataset._iter(consumer_checkpoint_transport=True): + yield _WorkerSample(sample, self.dataset) def __getattr__(self, name): dataset = self.__dict__.get("dataset") @@ -126,7 +117,7 @@ class _ConsumerCommitIterator: batch = next(self._iterator) except StopIteration: raise - except Exception as exc: + except BaseException as exc: self._dataset._invalidate_checkpoint( f"a DataLoader batch failed before it was returned: {exc}" ) @@ -443,9 +434,7 @@ class StreamingDataset(IterableDataset): return self._rank_splits[start : start + splits_per_worker] def __iter__(self) -> Iterator[dict[str, Any]]: - return self._iter( - consumer_checkpoint_transport=_CONSUMER_CHECKPOINT_DATASET.get() is self - ) + return self._iter() def _iter( self, *, consumer_checkpoint_transport: bool = False @@ -1151,6 +1140,9 @@ class StreamingDataLoader(DataLoader): Parameters are the same as ``torch.utils.data.DataLoader`` except that ``dataset`` must be a [StreamingDataset][lancedb.streaming.StreamingDataset]. + Subclasses that override ``StreamingDataset.__iter__`` are not supported + because the custom iterator cannot provide the exact per-yield checkpoint + snapshots required by this loader. Examples -------- @@ -1163,6 +1155,12 @@ class StreamingDataLoader(DataLoader): def __init__(self, dataset: StreamingDataset, *args, **kwargs): if not isinstance(dataset, StreamingDataset): raise TypeError("StreamingDataLoader requires a StreamingDataset") + if type(dataset).__iter__ is not StreamingDataset.__iter__: + raise TypeError( + "StreamingDataLoader does not support StreamingDataset subclasses " + "that override __iter__ because they cannot provide exact " + "per-yield checkpoint state" + ) if kwargs.get("in_order", True) is False: raise ValueError( "StreamingDataLoader requires in_order=True for deterministic " diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index 34c4b95ea..06c24e13a 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -107,6 +107,13 @@ def _collate_with_first_batch_stop(samples): return ids +def _collate_with_first_batch_interrupt(samples): + ids = [sample["id"] for sample in samples] + if ids == [0, 1]: + raise KeyboardInterrupt("first batch interrupted") + return ids + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @@ -1194,7 +1201,27 @@ def test_collate_stop_iteration_invalidates_consumer_checkpoint(tmp_path): assert list(iterator) == [[2, 3], [4, 5]] -def test_streaming_dataloader_preserves_dataset_iter_override(tmp_path): +def test_batch_base_exception_invalidates_consumer_checkpoint(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table("collate_interrupt", pa.table({"id": list(range(6))})) + dataset = StreamingDataset(table, num_splits=1, shuffle=False) + loader = StreamingDataLoader( + dataset, + batch_size=2, + num_workers=0, + collate_fn=_collate_with_first_batch_interrupt, + ) + iterator = iter(loader) + + with pytest.raises(KeyboardInterrupt, match="first batch interrupted"): + next(iterator) + assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2] + with pytest.raises(RuntimeError, match="failed before it was returned"): + dataset.state_dict() + assert list(iterator) == [[2, 3], [4, 5]] + + +def test_streaming_dataloader_rejects_dataset_iter_override(tmp_path): class CustomizedDataset(StreamingDataset): def __iter__(self): return iter([1000, 1001]) @@ -1204,13 +1231,42 @@ def test_streaming_dataloader_preserves_dataset_iter_override(tmp_path): dataset = CustomizedDataset(table, num_splits=1, shuffle=False) assert list(dataset) == [1000, 1001] - loader = StreamingDataLoader( - dataset, - batch_size=2, - num_workers=0, - collate_fn=list, - ) - assert list(loader) == [[1000, 1001]] + with pytest.raises(TypeError, match="override __iter__"): + StreamingDataLoader( + dataset, + batch_size=2, + num_workers=0, + collate_fn=list, + ) + + +def test_interleaved_adapters_do_not_authorize_plain_iteration(tmp_path): + db = lancedb.connect(tmp_path) + table_a = db.create_table("adapter_a", pa.table({"id": [0, 1]})) + table_b = db.create_table("adapter_b", pa.table({"id": [10, 11]})) + dataset_a = StreamingDataset(table_a, num_splits=1, shuffle=False) + dataset_b = StreamingDataset(table_b, num_splits=1, shuffle=False) + initial_state = dataset_a.state_dict() + + iterator_a = iter(streaming._StreamingDatasetAdapter(dataset_a)) + iterator_b = iter(streaming._StreamingDatasetAdapter(dataset_b)) + assert next(iterator_a).data["id"] == 0 + assert next(iterator_b).data["id"] == 10 + assert [sample.data["id"] for sample in iterator_a] == [1] + assert [sample.data["id"] for sample in iterator_b] == [11] + + dataset_a.load_state_dict(initial_state) + with patch( + "lancedb.streaming.get_worker_info", + return_value=FakeWorkerInfo(id=0, num_workers=1), + ): + plain_iterator = iter(dataset_a) + assert next(plain_iterator)["id"] == 0 + plain_iterator.close() + + assert dataset_a._untracked_worker_iteration[0] == 1 + with pytest.raises(RuntimeError, match="Use StreamingDataLoader"): + dataset_a.state_dict() def test_resume_from_partial_split_cycle_preserves_remaining_order(tmp_path):