From d808cd04100c1632ca91a0154c0dc30ca248c133 Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Mon, 24 Aug 2026 09:44:12 +0000 Subject: [PATCH] fix(python): make streaming iterator ownership atomic --- python/python/lancedb/streaming.py | 189 +++++++++++------- .../python/tests/test_elastic_dataloader.py | 170 ++++++++++++++-- 2 files changed, 270 insertions(+), 89 deletions(-) diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index a7d5f1280..f3d9ead95 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -114,43 +114,27 @@ class _ConsumerCommitIterator: iterator, dataset: "StreamingDataset", *, - initial_state: dict, + owner_token: int, require_uniform: bool, - restore_on_exhaustion: bool, ): self._iterator = iterator self._dataset = dataset + self._owner_token = owner_token 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 - ) + self._terminal = False def __iter__(self): return self def __next__(self): + if self._terminal: + raise StopIteration 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() + self._terminal = True + self._release() raise except BaseException as exc: self._dataset._invalidate_checkpoint( @@ -165,10 +149,6 @@ 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( @@ -182,22 +162,39 @@ class _ConsumerCommitIterator: self._released = True dataset = self.__dict__.get("_dataset") if dataset is not None: - dataset._release_consumer_iterator() + dataset._release_consumer_iterator(self._owner_token) def _shutdown_workers(self): + if self.__dict__.get("_released", True): + return None + self._terminal = True iterator = self.__dict__.get("_iterator") shutdown = getattr(iterator, "_shutdown_workers", None) try: if shutdown is not None: - return shutdown() - finally: + shutdown() + else: + fetcher = getattr(iterator, "_dataset_fetcher", None) + dataset_iterator = getattr(fetcher, "dataset_iter", None) + close = getattr(dataset_iterator, "close", None) + if close is None: + raise RuntimeError( + "StreamingDataLoader could not close its inner iterator" + ) + close() + except BaseException as exc: + self._dataset._invalidate_checkpoint( + f"a DataLoader iterator could not be shut down safely: {exc}" + ) + raise + else: self._release() def __del__(self): try: self._shutdown_workers() except BaseException: - self._release() + pass def __getattr__(self, name): return getattr(self._iterator, name) @@ -546,6 +543,9 @@ class StreamingDataset(IterableDataset): self._consumer_checkpoint_requires_uniform = False self._consumer_iterator_lock = threading.Lock() self._consumer_iterator_active = False + self._consumer_iterator_generation = 0 + self._consumer_iterator_owner: Optional[int] = None + self._consumer_iterator_owner_thread: Optional[int] = None # Cumulative bytes of Arrow buffer data fetched across all iterations. self._bytes_loaded: int = 0 @@ -693,9 +693,32 @@ class StreamingDataset(IterableDataset): def _iter( self, *, consumer_checkpoint_transport: bool = False ) -> Iterator[dict[str, Any]]: - if self._raw_batches_ref is not None or ( - self._consumer_iterator_active and not consumer_checkpoint_transport - ): + owner_token = None + previous_owner = self._consumer_iterator_owner + if consumer_checkpoint_transport: + if not self._consumer_iterator_active: + raise RuntimeError( + "StreamingDataLoader worker transport requires an active " + "parent iterator reservation" + ) + else: + try: + owner_token = self._acquire_consumer_iterator() + except BaseException: + self._release_consumer_iterator_after_failed_acquire(previous_owner) + raise + try: + yield from self._iter_owned( + consumer_checkpoint_transport=consumer_checkpoint_transport + ) + finally: + if owner_token is not None: + self._release_consumer_iterator(owner_token) + + def _iter_owned( + self, *, consumer_checkpoint_transport: bool + ) -> Iterator[dict[str, Any]]: + if self._raw_batches_ref is not None: raise RuntimeError( "StreamingDataset does not support concurrent iteration. " "Only one active iterator per dataset instance is allowed." @@ -1396,7 +1419,7 @@ class StreamingDataset(IterableDataset): if self._checkpoint_invalid_reason is None: self._checkpoint_invalid_reason = reason - def _acquire_consumer_iterator(self) -> None: + def _acquire_consumer_iterator(self) -> int: """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: @@ -1404,11 +1427,32 @@ class StreamingDataset(IterableDataset): "StreamingDataset does not support concurrent iteration. " "Only one active iterator per dataset instance is allowed." ) + self._consumer_iterator_generation += 1 + owner_token = self._consumer_iterator_generation self._consumer_iterator_active = True + self._consumer_iterator_owner = owner_token + self._consumer_iterator_owner_thread = threading.get_ident() + return owner_token - def _release_consumer_iterator(self) -> None: + def _release_consumer_iterator(self, owner_token: int) -> None: with self._consumer_iterator_lock: - self._consumer_iterator_active = False + if self._consumer_iterator_owner == owner_token: + self._consumer_iterator_active = False + self._consumer_iterator_owner = None + self._consumer_iterator_owner_thread = None + + def _release_consumer_iterator_after_failed_acquire( + self, previous_owner: Optional[int] + ) -> None: + """Clean up when an interrupted acquire set a lease but did not return it.""" + with self._consumer_iterator_lock: + if ( + self._consumer_iterator_owner != previous_owner + and self._consumer_iterator_owner_thread == threading.get_ident() + ): + self._consumer_iterator_active = False + self._consumer_iterator_owner = None + self._consumer_iterator_owner_thread = None def _commit_worker_state(self, state: dict, *, require_uniform: bool) -> None: """Merge one trainer-consumed worker batch into parent state.""" @@ -1700,8 +1744,9 @@ class StreamingDataLoader(DataLoader): 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. + releases that ownership. ``drop_last=True`` is not supported because worker + replicas discard incomplete tails independently, which cannot produce a + topology-independent checkpoint. Parameters are the same as ``torch.utils.data.DataLoader`` except that ``dataset`` must be a @@ -1739,42 +1784,50 @@ class StreamingDataLoader(DataLoader): ) self._streaming_dataset = dataset super().__init__(_StreamingDatasetAdapter(dataset), *args, **kwargs) + if self.drop_last: + raise ValueError( + "StreamingDataLoader does not support drop_last=True because " + "discarded worker tails cannot be checkpointed " + "topology-independently" + ) self.collate_fn = _CheckpointCollate(self.collate_fn) def __iter__(self): - state = self._streaming_dataset._checkpoint_snapshot() - packed = self._streaming_dataset._pack_sequences is not None - if packed: - blocks = state["blocks_emitted_per_split"] - rank_blocks = [ - blocks[split] for split in self._streaming_dataset._rank_splits - ] - if len(set(rank_blocks)) > 1: - raise RuntimeError( - "StreamingDataLoader cannot start from a partial packed " - "logical step; resume from a checkpoint whose splits assigned " - "to this rank have equal emitted-block counts" - ) - elif self.num_workers > 1: - samples = state["samples_consumed_per_split"] - rank_samples = [ - samples[split] for split in self._streaming_dataset._rank_splits - ] - if len(set(rank_samples)) > 1: - raise RuntimeError( - "StreamingDataLoader cannot start multiple workers from a " - "partial logical step; resume from a checkpoint whose splits " - "assigned to this rank have equal consumed-sample counts" - ) - self._streaming_dataset._acquire_consumer_iterator() + dataset = self._streaming_dataset + previous_owner = dataset._consumer_iterator_owner + owner_token = None try: + owner_token = dataset._acquire_consumer_iterator() + state = dataset._checkpoint_snapshot() + packed = dataset._pack_sequences is not None + if packed: + blocks = state["blocks_emitted_per_split"] + rank_blocks = [blocks[split] for split in dataset._rank_splits] + if len(set(rank_blocks)) > 1: + raise RuntimeError( + "StreamingDataLoader cannot start from a partial packed " + "logical step; resume from a checkpoint whose splits " + "assigned to this rank have equal emitted-block counts" + ) + elif self.num_workers > 1: + samples = state["samples_consumed_per_split"] + rank_samples = [samples[split] for split in dataset._rank_splits] + if len(set(rank_samples)) > 1: + raise RuntimeError( + "StreamingDataLoader cannot start multiple workers from a " + "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, - initial_state=state, + dataset, + owner_token=owner_token, require_uniform=self.num_workers > 1 or packed, - restore_on_exhaustion=self.drop_last, ) except BaseException: - self._streaming_dataset._release_consumer_iterator() + if owner_token is not None: + dataset._release_consumer_iterator(owner_token) + else: + dataset._release_consumer_iterator_after_failed_acquire(previous_owner) raise diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index dddcb0b1c..853a37496 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -1287,30 +1287,21 @@ def test_direct_iteration_surfaces_prefetch_failure_before_committing_row( assert checkpoint["positions_consumed_per_split"] == [1] -def test_streaming_dataloader_drop_last_restores_discarded_tail(tmp_path): +@pytest.mark.parametrize("workers", [0, 1, 2]) +def test_streaming_dataloader_rejects_drop_last(tmp_path, workers): db = lancedb.connect(tmp_path) table = db.create_table("drop_last", pa.table({"id": [0, 1, 2]})) - results = [] + dataset = StreamingDataset(table, num_splits=1, shuffle=False) + worker_options = {"multiprocessing_context": "spawn"} if workers else {} - for workers in (0, 1): - dataset = StreamingDataset(table, num_splits=1, shuffle=False) - worker_options = {"multiprocessing_context": "spawn"} if workers else {} - loader = StreamingDataLoader( + with pytest.raises(ValueError, match="drop_last=True"): + 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): @@ -1347,6 +1338,137 @@ def test_streaming_dataloader_owns_one_iterator_until_teardown(tmp_path): third._shutdown_workers() +def test_zero_worker_shutdown_closes_inner_iterator_before_release(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table("zero_worker_shutdown", pa.table({"id": list(range(6))})) + dataset = StreamingDataset(table, num_splits=1, shuffle=False) + loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0) + + first = iter(loader) + assert next(first)["id"].tolist() == [0, 1] + first._shutdown_workers() + + assert dataset._consumer_iterator_active is False + assert dataset._raw_batches_ref is None + second = iter(loader) + try: + with pytest.raises(StopIteration): + next(first) + assert next(second)["id"].tolist() == [2, 3] + finally: + second._shutdown_workers() + + +def test_direct_and_loader_admission_share_one_atomic_lease(tmp_path, monkeypatch): + db = lancedb.connect(tmp_path) + table = db.create_table("direct_loader_lease", pa.table({"id": list(range(4))})) + dataset = StreamingDataset(table, num_splits=1, shuffle=False) + loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0) + entered = threading.Event() + release = threading.Event() + direct_result = [] + direct_error = [] + contender = [] + real_resolve = dataset._resolve_my_splits + + def controlled_resolve(): + if threading.current_thread().name == "direct-start": + entered.set() + assert release.wait(timeout=5) + return real_resolve() + + def advance_direct(iterator): + try: + direct_result.append(next(iterator)["id"]) + except BaseException as exc: + direct_error.append(exc) + + monkeypatch.setattr(dataset, "_resolve_my_splits", controlled_resolve) + direct = iter(dataset) + thread = threading.Thread( + target=advance_direct, args=(direct,), name="direct-start" + ) + thread.start() + assert entered.wait(timeout=5) + try: + with pytest.raises(RuntimeError, match="concurrent iteration"): + contender.append(iter(loader)) + finally: + release.set() + thread.join(timeout=5) + if contender: + contender[0]._shutdown_workers() + direct.close() + + assert not thread.is_alive() + assert direct_error == [] + assert direct_result == [0] + + +def test_loader_acquires_before_snapshot_and_cleans_interrupted_acquire( + tmp_path, monkeypatch +): + db = lancedb.connect(tmp_path) + table = db.create_table("lease_snapshot", pa.table({"id": list(range(4))})) + dataset = StreamingDataset(table, num_splits=1, shuffle=False) + loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0) + first = iter(loader) + assert next(first)["id"].tolist() == [0, 1] + + entered = threading.Event() + release = threading.Event() + pending = [] + pending_errors = [] + observed_snapshots = [] + real_acquire = dataset._acquire_consumer_iterator + real_snapshot = dataset._checkpoint_snapshot + + def controlled_acquire(): + if threading.current_thread().name == "stale-start": + entered.set() + assert release.wait(timeout=5) + return real_acquire() + + def recording_snapshot(): + state = real_snapshot() + if threading.current_thread().name == "stale-start": + observed_snapshots.append(state["samples_consumed_per_split"]) + return state + + def create_pending_iterator(): + try: + pending.append(iter(loader)) + except BaseException as exc: + pending_errors.append(exc) + + monkeypatch.setattr(dataset, "_acquire_consumer_iterator", controlled_acquire) + monkeypatch.setattr(dataset, "_checkpoint_snapshot", recording_snapshot) + thread = threading.Thread(target=create_pending_iterator, name="stale-start") + thread.start() + assert entered.wait(timeout=5) + assert next(first)["id"].tolist() == [2, 3] + with pytest.raises(StopIteration): + next(first) + release.set() + thread.join(timeout=5) + + assert not thread.is_alive() + assert pending_errors == [] + assert observed_snapshots == [[4]] + assert len(pending) == 1 + assert list(pending[0]) == [] + assert dataset.state_dict()["samples_consumed_per_split"] == [4] + + def interrupted_acquire(): + real_acquire() + raise KeyboardInterrupt("after acquire") + + monkeypatch.setattr(dataset, "_acquire_consumer_iterator", interrupted_acquire) + with pytest.raises(KeyboardInterrupt, match="after acquire"): + iter(loader) + assert dataset._consumer_iterator_active is False + + def test_streaming_dataloader_rejects_dataset_iter_override(tmp_path): class CustomizedDataset(StreamingDataset): def __iter__(self): @@ -1374,12 +1496,18 @@ def test_interleaved_adapters_do_not_authorize_plain_iteration(tmp_path): 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] + owner_a = dataset_a._acquire_consumer_iterator() + owner_b = dataset_b._acquire_consumer_iterator() + try: + 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] + finally: + dataset_a._release_consumer_iterator(owner_a) + dataset_b._release_consumer_iterator(owner_b) dataset_a.load_state_dict(initial_state) with patch(