fix(python): make streaming iterator ownership atomic

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

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(