diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 45551178e..a7d5f1280 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -109,10 +109,24 @@ class _StreamingDatasetAdapter(IterableDataset): class _ConsumerCommitIterator: - def __init__(self, iterator, dataset: "StreamingDataset", *, require_uniform: bool): + def __init__( + self, + iterator, + dataset: "StreamingDataset", + *, + initial_state: dict, + require_uniform: bool, + restore_on_exhaustion: bool, + ): self._iterator = iterator self._dataset = dataset self._require_uniform = require_uniform + self._restore_on_exhaustion = restore_on_exhaustion + self._released = False + self._last_delivered_state = initial_state + self._last_delivered_requires_uniform = ( + dataset._consumer_checkpoint_requires_uniform + ) def __iter__(self): return self @@ -121,6 +135,22 @@ class _ConsumerCommitIterator: try: batch = next(self._iterator) except StopIteration: + try: + if self._restore_on_exhaustion: + # A single-process DataLoader advances the parent dataset + # while assembling a batch. Restore the last delivered + # state when drop_last discards an incomplete tail. + self._dataset.load_state_dict(self._last_delivered_state) + self._dataset._consumer_checkpoint_requires_uniform = ( + self._last_delivered_requires_uniform + ) + except BaseException as exc: + self._dataset._invalidate_checkpoint( + f"discarded DataLoader tail state could not be restored: {exc}" + ) + raise + finally: + self._release() raise except BaseException as exc: self._dataset._invalidate_checkpoint( @@ -135,6 +165,10 @@ class _ConsumerCommitIterator: self._dataset._commit_worker_state( batch.state, require_uniform=self._require_uniform ) + self._last_delivered_state = self._dataset._checkpoint_snapshot() + self._last_delivered_requires_uniform = ( + self._dataset._consumer_checkpoint_requires_uniform + ) return batch.data except BaseException as exc: self._dataset._invalidate_checkpoint( @@ -142,6 +176,29 @@ class _ConsumerCommitIterator: ) raise + def _release(self) -> None: + if self.__dict__.get("_released", True): + return + self._released = True + dataset = self.__dict__.get("_dataset") + if dataset is not None: + dataset._release_consumer_iterator() + + def _shutdown_workers(self): + iterator = self.__dict__.get("_iterator") + shutdown = getattr(iterator, "_shutdown_workers", None) + try: + if shutdown is not None: + return shutdown() + finally: + self._release() + + def __del__(self): + try: + self._shutdown_workers() + except BaseException: + self._release() + def __getattr__(self, name): return getattr(self._iterator, name) @@ -487,6 +544,8 @@ class StreamingDataset(IterableDataset): # reach the same logical step boundary. self._checkpoint_invalid_reason: Optional[str] = None self._consumer_checkpoint_requires_uniform = False + self._consumer_iterator_lock = threading.Lock() + self._consumer_iterator_active = False # Cumulative bytes of Arrow buffer data fetched across all iterations. self._bytes_loaded: int = 0 @@ -634,7 +693,9 @@ class StreamingDataset(IterableDataset): def _iter( self, *, consumer_checkpoint_transport: bool = False ) -> Iterator[dict[str, Any]]: - if self._raw_batches_ref is not None: + if self._raw_batches_ref is not None or ( + self._consumer_iterator_active and not consumer_checkpoint_transport + ): raise RuntimeError( "StreamingDataset does not support concurrent iteration. " "Only one active iterator per dataset instance is allowed." @@ -974,6 +1035,9 @@ class StreamingDataset(IterableDataset): def _yield_row(i: int): pos, row = cooked[i].popleft() + # Surface any completed prefetched failure before the + # current row becomes durable checkpoint progress. + _advance(i) local_consumed[i] += 1 pos_consumed[i] = pos + 1 split_idx = my_splits[i] @@ -981,7 +1045,6 @@ class StreamingDataset(IterableDataset): initial_samples[i] + local_consumed[i] ) self._resume_positions[split_idx] = pos_consumed[i] - _advance(i) return row def _update_progress_stats() -> None: @@ -1222,6 +1285,7 @@ class StreamingDataset(IterableDataset): "_local_consumed_ref", ): state[key] = None + state["_consumer_iterator_lock"] = None return state def __setstate__(self, state): @@ -1232,6 +1296,7 @@ class StreamingDataset(IterableDataset): table_state = state.pop("_table") perm_name, perm_data = state.pop("_perm_table") self.__dict__.update(state) + self._consumer_iterator_lock = threading.Lock() if self._connection_factory is not None: self._table = self._connection_factory(table_name) else: @@ -1331,6 +1396,20 @@ class StreamingDataset(IterableDataset): if self._checkpoint_invalid_reason is None: self._checkpoint_invalid_reason = reason + def _acquire_consumer_iterator(self) -> None: + """Reserve this parent dataset for one checkpoint-aware iterator.""" + with self._consumer_iterator_lock: + if self._consumer_iterator_active or self._raw_batches_ref is not None: + raise RuntimeError( + "StreamingDataset does not support concurrent iteration. " + "Only one active iterator per dataset instance is allowed." + ) + self._consumer_iterator_active = True + + def _release_consumer_iterator(self) -> None: + with self._consumer_iterator_lock: + self._consumer_iterator_active = False + def _commit_worker_state(self, state: dict, *, require_uniform: bool) -> None: """Merge one trainer-consumed worker batch into parent state.""" for key, expected in ( @@ -1619,6 +1698,10 @@ class StreamingDataLoader(DataLoader): parent-committed state. If batch collation raises, checkpointing remains invalid for that dataset instance; restore the last valid checkpoint into a fresh dataset before continuing. + Only one active iterator may own a dataset at a time, including when worker + processes are used. Exhausting or explicitly shutting down the iterator + releases that ownership. With ``drop_last=True``, an incomplete discarded + tail is rolled back to the last batch returned to the trainer. Parameters are the same as ``torch.utils.data.DataLoader`` except that ``dataset`` must be a @@ -1683,8 +1766,15 @@ class StreamingDataLoader(DataLoader): "partial logical step; resume from a checkpoint whose splits " "assigned to this rank have equal consumed-sample counts" ) - return _ConsumerCommitIterator( - super().__iter__(), - self._streaming_dataset, - require_uniform=self.num_workers > 1 or packed, - ) + self._streaming_dataset._acquire_consumer_iterator() + try: + return _ConsumerCommitIterator( + super().__iter__(), + self._streaming_dataset, + initial_state=state, + require_uniform=self.num_workers > 1 or packed, + restore_on_exhaustion=self.drop_last, + ) + except BaseException: + self._streaming_dataset._release_consumer_iterator() + raise diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index a7e7f46c2..dddcb0b1c 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -32,6 +32,7 @@ Parameters used throughout: import dataclasses import logging +import threading from unittest.mock import patch import lancedb @@ -1244,6 +1245,108 @@ def test_parent_commit_base_exception_invalidates_consumer_checkpoint(tmp_path): dataset.state_dict() +def test_direct_iteration_surfaces_prefetch_failure_before_committing_row( + tmp_path, monkeypatch +): + db = lancedb.connect(tmp_path) + table = db.create_table("prefetch_failure", pa.table({"id": list(range(4))})) + release = threading.Event() + failed = threading.Event() + real_getitems = streaming.Permutation.__getitems__ + + def controlled_getitems(permutation, indices): + if indices and indices[0] >= 2: + assert release.wait(timeout=5) + failed.set() + raise RuntimeError("later prefetched I/O failed") + return real_getitems(permutation, indices) + + class SignalDict(dict): + def __setitem__(self, key, value): + super().__setitem__(key, value) + release.set() + assert failed.wait(timeout=5) + + monkeypatch.setattr(streaming.Permutation, "__getitems__", controlled_getitems) + dataset = StreamingDataset( + table, + num_splits=1, + shuffle=False, + read_batch_size=2, + io_queue_depth=2, + ) + dataset._resume_positions = SignalDict() + iterator = iter(dataset) + + assert next(iterator)["id"] == 0 + with pytest.raises(RuntimeError, match="later prefetched I/O failed"): + next(iterator) + + checkpoint = dataset.state_dict() + assert checkpoint["samples_consumed_per_split"] == [1] + assert checkpoint["positions_consumed_per_split"] == [1] + + +def test_streaming_dataloader_drop_last_restores_discarded_tail(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table("drop_last", pa.table({"id": [0, 1, 2]})) + results = [] + + for workers in (0, 1): + dataset = StreamingDataset(table, num_splits=1, shuffle=False) + worker_options = {"multiprocessing_context": "spawn"} if workers else {} + loader = StreamingDataLoader( + dataset, + batch_size=2, + num_workers=workers, + drop_last=True, + **worker_options, + ) + batches = [batch["id"].tolist() for batch in loader] + results.append( + (workers, batches, dataset.state_dict()["samples_consumed_per_split"]) + ) + + assert results == [ + (0, [[0, 1]], [2]), + (1, [[0, 1]], [2]), + ] + + +def test_streaming_dataloader_owns_one_iterator_until_teardown(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table("iterator_owner", pa.table({"id": list(range(4))})) + dataset = StreamingDataset(table, num_splits=1, shuffle=False) + loader = StreamingDataLoader( + dataset, + batch_size=2, + num_workers=1, + multiprocessing_context="spawn", + ) + + first = iter(loader) + try: + assert next(first)["id"].tolist() == [0, 1] + with pytest.raises(RuntimeError, match="concurrent iteration"): + iter(loader) + finally: + first._shutdown_workers() + + second = iter(loader) + try: + assert [batch["id"].tolist() for batch in second] == [[2, 3]] + except BaseException: + second._shutdown_workers() + raise + + # Natural exhaustion releases ownership too. + third = iter(loader) + try: + assert list(third) == [] + finally: + third._shutdown_workers() + + def test_streaming_dataloader_rejects_dataset_iter_override(tmp_path): class CustomizedDataset(StreamingDataset): def __iter__(self):