fix(python): preserve streaming delivery boundaries

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

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):