diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index f3d9ead95..2c0e2d5c1 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -66,6 +66,11 @@ class _WorkerBatch(NamedTuple): state: dict +class _ConsumerIteratorLease(NamedTuple): + owner_token: int + owner_thread: int + + class _CheckpointCollate: """Attach the worker's post-fetch state to a collated batch.""" @@ -542,10 +547,8 @@ class StreamingDataset(IterableDataset): self._checkpoint_invalid_reason: Optional[str] = None 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 + self._consumer_iterator_lease: Optional[_ConsumerIteratorLease] = None # Cumulative bytes of Arrow buffer data fetched across all iterations. self._bytes_loaded: int = 0 @@ -694,7 +697,7 @@ class StreamingDataset(IterableDataset): self, *, consumer_checkpoint_transport: bool = False ) -> Iterator[dict[str, Any]]: owner_token = None - previous_owner = self._consumer_iterator_owner + previous_lease = self._consumer_iterator_lease if consumer_checkpoint_transport: if not self._consumer_iterator_active: raise RuntimeError( @@ -705,7 +708,7 @@ class StreamingDataset(IterableDataset): try: owner_token = self._acquire_consumer_iterator() except BaseException: - self._release_consumer_iterator_after_failed_acquire(previous_owner) + self._release_consumer_iterator_after_failed_acquire(previous_lease) raise try: yield from self._iter_owned( @@ -1419,6 +1422,20 @@ class StreamingDataset(IterableDataset): if self._checkpoint_invalid_reason is None: self._checkpoint_invalid_reason = reason + @property + def _consumer_iterator_active(self) -> bool: + return self._consumer_iterator_lease is not None + + @property + def _consumer_iterator_owner(self) -> Optional[int]: + lease = self._consumer_iterator_lease + return lease.owner_token if lease is not None else None + + @property + def _consumer_iterator_owner_thread(self) -> Optional[int]: + lease = self._consumer_iterator_lease + return lease.owner_thread if lease is not None else None + def _acquire_consumer_iterator(self) -> int: """Reserve this parent dataset for one checkpoint-aware iterator.""" with self._consumer_iterator_lock: @@ -1427,32 +1444,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() + owner_thread = threading.get_ident() + owner_token = self._consumer_iterator_generation + 1 + lease = _ConsumerIteratorLease(owner_token, owner_thread) + self._consumer_iterator_generation = owner_token + self._consumer_iterator_lease = lease return owner_token def _release_consumer_iterator(self, owner_token: int) -> None: with self._consumer_iterator_lock: - if self._consumer_iterator_owner == owner_token: - self._consumer_iterator_active = False - self._consumer_iterator_owner = None - self._consumer_iterator_owner_thread = None + lease = self._consumer_iterator_lease + if lease is not None and lease.owner_token == owner_token: + self._consumer_iterator_lease = None def _release_consumer_iterator_after_failed_acquire( - self, previous_owner: Optional[int] + self, previous_lease: Optional[_ConsumerIteratorLease] ) -> None: """Clean up when an interrupted acquire set a lease but did not return it.""" + owner_thread = threading.current_thread().ident with self._consumer_iterator_lock: + lease = self._consumer_iterator_lease if ( - self._consumer_iterator_owner != previous_owner - and self._consumer_iterator_owner_thread == threading.get_ident() + lease is not None + and lease is not previous_lease + and lease.owner_thread == owner_thread ): - self._consumer_iterator_active = False - self._consumer_iterator_owner = None - self._consumer_iterator_owner_thread = None + self._consumer_iterator_lease = None def _commit_worker_state(self, state: dict, *, require_uniform: bool) -> None: """Merge one trainer-consumed worker batch into parent state.""" @@ -1794,7 +1811,7 @@ class StreamingDataLoader(DataLoader): def __iter__(self): dataset = self._streaming_dataset - previous_owner = dataset._consumer_iterator_owner + previous_lease = dataset._consumer_iterator_lease owner_token = None try: owner_token = dataset._acquire_consumer_iterator() @@ -1829,5 +1846,5 @@ class StreamingDataLoader(DataLoader): if owner_token is not None: dataset._release_consumer_iterator(owner_token) else: - dataset._release_consumer_iterator_after_failed_acquire(previous_owner) + dataset._release_consumer_iterator_after_failed_acquire(previous_lease) raise diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index 853a37496..da27e5bfc 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -1469,6 +1469,34 @@ def test_loader_acquires_before_snapshot_and_cleans_interrupted_acquire( assert dataset._consumer_iterator_active is False +def test_consumer_iterator_lease_publication_is_atomic(tmp_path, monkeypatch): + db = lancedb.connect(tmp_path) + table = db.create_table("atomic_lease", pa.table({"id": [0, 1]})) + dataset = StreamingDataset(table, num_splits=1, shuffle=False) + loader = StreamingDataLoader(dataset, batch_size=1, num_workers=0) + real_get_ident = streaming.threading.get_ident + calls = 0 + + def interrupt_during_publication(): + nonlocal calls + calls += 1 + if calls == 1: + raise KeyboardInterrupt("during lease mutation") + return real_get_ident() + + monkeypatch.setattr(streaming.threading, "get_ident", interrupt_during_publication) + with pytest.raises(KeyboardInterrupt, match="during lease mutation"): + iter(loader) + monkeypatch.setattr(streaming.threading, "get_ident", real_get_ident) + + assert dataset._consumer_iterator_active is False + iterator = iter(loader) + try: + assert next(iterator)["id"].tolist() == [0] + finally: + iterator._shutdown_workers() + + def test_streaming_dataloader_rejects_dataset_iter_override(tmp_path): class CustomizedDataset(StreamingDataset): def __iter__(self):