From 865b167de2192dcdbc7c868d8fff1d6234eac951 Mon Sep 17 00:00:00 2001
From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Date: Fri, 21 Aug 2026 22:23:15 +0000
Subject: [PATCH] fix(python): commit streaming worker checkpoints on
 consumption

---
 docs/src/python/python.md                     |   2 +
 python/python/lancedb/streaming.py            | 306 +++++++++++++++---
 .../python/tests/test_elastic_dataloader.py   |  86 +++++
 3 files changed, 341 insertions(+), 53 deletions(-)

diff --git a/docs/src/python/python.md b/docs/src/python/python.md
index 70b2a7207..fc5cc3b17 100644
--- a/docs/src/python/python.md
+++ b/docs/src/python/python.md
@@ -255,6 +255,8 @@ instead of being materialized with the rest of the row.
 
 ::: lancedb.streaming.StreamingDataset
 
+::: lancedb.streaming.StreamingDataLoader
+
 ::: lancedb.permutation.permutation_builder
 
 ::: lancedb.permutation.PermutationBuilder
diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py
index b27e606a4..159dea5d6 100644
--- a/python/python/lancedb/streaming.py
+++ b/python/python/lancedb/streaming.py
@@ -19,6 +19,7 @@ above.
 """
 
 import ctypes
+import heapq
 import logging
 import os
 import random
@@ -27,9 +28,9 @@ import time
 from collections import deque
 from concurrent.futures import ThreadPoolExecutor
 from multiprocessing import RawArray
-from typing import Any, Callable, Iterator, Optional, Union
+from typing import Any, Callable, Iterator, NamedTuple, Optional, Union
 
-from torch.utils.data import IterableDataset, get_worker_info
+from torch.utils.data import DataLoader, IterableDataset, get_worker_info
 
 from .permutation import (
     Permutation,
@@ -50,6 +51,75 @@ DEFAULT_READ_BATCH_SIZE = 64
 DEFAULT_PREFETCH_BATCHES = 4
 
 
+class _WorkerSample(NamedTuple):
+    data: Any
+    dataset: "StreamingDataset"
+
+
+class _WorkerBatch(NamedTuple):
+    data: Any
+    state: dict
+
+
+class _CheckpointCollate:
+    """Attach the worker's post-fetch state to a collated batch."""
+
+    def __init__(self, collate_fn: Callable):
+        self._collate_fn = collate_fn
+
+    def __call__(self, samples):
+        if isinstance(samples, list):
+            if not samples:
+                return _WorkerBatch(self._collate_fn(samples), {})
+            worker_samples = samples
+            data = self._collate_fn([sample.data for sample in worker_samples])
+            dataset = worker_samples[-1].dataset
+        else:
+            data = self._collate_fn(samples.data)
+            dataset = samples.dataset
+        return _WorkerBatch(data, dataset._checkpoint_snapshot())
+
+
+class _StreamingDatasetAdapter(IterableDataset):
+    """Yield private sample wrappers for :class:`StreamingDataLoader`."""
+
+    def __init__(self, dataset: "StreamingDataset"):
+        super().__init__()
+        self.dataset = dataset
+
+    def __iter__(self):
+        self.dataset._consumer_checkpoint_transport = True
+        for sample in self.dataset:
+            yield _WorkerSample(sample, self.dataset)
+
+    def __getattr__(self, name):
+        dataset = self.__dict__.get("dataset")
+        if dataset is None:
+            raise AttributeError(name)
+        return getattr(dataset, name)
+
+
+class _ConsumerCommitIterator:
+    def __init__(self, iterator, dataset: "StreamingDataset"):
+        self._iterator = iterator
+        self._dataset = dataset
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        batch = next(self._iterator)
+        if not isinstance(batch, _WorkerBatch):
+            raise RuntimeError(
+                "StreamingDataLoader did not receive worker checkpoint metadata"
+            )
+        self._dataset._commit_worker_state(batch.state)
+        return batch.data
+
+    def __getattr__(self, name):
+        return getattr(self._iterator, name)
+
+
 class StreamingDataset(IterableDataset):
     """An elastic, resumable PyTorch IterableDataset backed by a LanceDB table.
 
@@ -260,6 +330,13 @@ class StreamingDataset(IterableDataset):
         #          rows_skipped]
         self._worker_stats: RawArray = RawArray(ctypes.c_int64, 8)
 
+        # A standard multi-process DataLoader cannot report which prefetched
+        # batches were actually returned to its consumer.  Workers set this
+        # shared flag so state_dict() can reject a stale parent checkpoint
+        # unless StreamingDataLoader installed the consumer-commit transport.
+        self._untracked_worker_iteration: RawArray = RawArray(ctypes.c_int64, 1)
+        self._consumer_checkpoint_transport = False
+
         # Cumulative bytes of Arrow buffer data fetched across all iterations.
         self._bytes_loaded: int = 0
         # Cumulative seconds spent in LanceDB I/O and in transform functions.
@@ -272,6 +349,10 @@ class StreamingDataset(IterableDataset):
         # step boundaries all splits have consumed this many samples, so a
         # single scalar captures the topology-independent checkpoint state.
         self._resume_offset: int = 0
+        # Exact yielded-sample counts for splits this process has advanced.
+        # Missing entries use _resume_offset, which remains the lower-bound
+        # checkpoint inherited from an earlier uniform/global state.
+        self._resume_samples: dict[int, int] = {}
         # Permutation position each split has consumed through, keyed by
         # global split index.  Equal to _resume_offset for every split unless
         # on_transform_error skipped rows, in which case skipped positions
@@ -334,6 +415,10 @@ class StreamingDataset(IterableDataset):
                 "StreamingDataset does not support concurrent iteration. "
                 "Only one active iterator per dataset instance is allowed."
             )
+        real_worker = get_worker_info() is not None
+        if real_worker and not self._consumer_checkpoint_transport:
+            self._untracked_worker_iteration[0] = 1
+
         my_splits = self._resolve_my_splits()
         if not my_splits:
             return
@@ -341,6 +426,7 @@ class StreamingDataset(IterableDataset):
         # Set identity transform on each Permutation so __getitems__ returns
         # the raw RecordBatch.  Stage 2 applies the real transform.
         permutations: list[Permutation] = []
+        initial_samples: list[int] = []
         initial_positions: list[int] = []
         for split_idx in my_splits:
             perm = Permutation.from_tables(
@@ -349,15 +435,16 @@ class StreamingDataset(IterableDataset):
             if self._columns is not None:
                 perm = perm.select_columns(self._columns)
             perm = perm.with_transform(lambda batch: batch)
-            start_pos = self._resume_positions.get(split_idx, self._resume_offset)
+            sample_count = self._resume_samples.get(split_idx, self._resume_offset)
+            start_pos = self._resume_positions.get(split_idx, sample_count)
             if start_pos > 0:
                 perm = perm.with_skip(start_pos)
+            initial_samples.append(sample_count)
             initial_positions.append(start_pos)
             permutations.append(perm)
 
         n = len(permutations)
         split_sizes = [perm.num_rows for perm in permutations]
-        initial_offset = self._resume_offset
         local_consumed = [0] * n
         # Permutation position each split has consumed through (absolute,
         # i.e. counted from the start of the unskipped split).  Runs ahead of
@@ -542,6 +629,55 @@ class StreamingDataset(IterableDataset):
                     for i in range(n):
                         _fill_io(i)
 
+                    def _yield_row(i: int):
+                        pos, row = cooked[i].popleft()
+                        local_consumed[i] += 1
+                        pos_consumed[i] = pos + 1
+                        split_idx = my_splits[i]
+                        self._resume_samples[split_idx] = (
+                            initial_samples[i] + local_consumed[i]
+                        )
+                        self._resume_positions[split_idx] = pos_consumed[i]
+                        _advance(i)
+                        return row
+
+                    def _update_progress_stats() -> None:
+                        if not real_worker:
+                            self._resume_offset = min(
+                                initial_samples[j] + local_consumed[j] for j in range(n)
+                            )
+                        ws = self._worker_stats
+                        ws[0] = sum(split_sizes[j] - fetch_head[j] for j in range(n))
+                        ws[1] = sum(
+                            batch.num_rows for q in raw_batches for _, batch in q
+                        )
+                        ws[2] = sum(len(q) for q in cooked)
+                        ws[3] = sum(local_consumed)
+                        ws[4] = self._bytes_loaded
+                        ws[5] = int(self._fetch_time * 1_000_000)
+                        ws[6] = int(self._transform_time * 1_000_000)
+                        ws[7] = self._rows_skipped
+
+                    # A checkpoint taken between round-robin split turns has
+                    # non-uniform counts.  Resume lagging splits first so the
+                    # exact canonical sequence continues without replaying
+                    # already-consumed rows.
+                    if len(set(initial_samples)) > 1:
+                        pending = [
+                            (initial_samples[i], my_splits[i], i) for i in range(n)
+                        ]
+                        heapq.heapify(pending)
+                        while pending:
+                            consumed, _, i = heapq.heappop(pending)
+                            _ensure_cooked(i)
+                            if not cooked[i]:
+                                continue
+                            row = _yield_row(i)
+                            heapq.heappush(pending, (consumed + 1, my_splits[i], i))
+                            _update_progress_stats()
+                            yield row
+                        return
+
                     while True:
                         # A cycle only runs if every split can still produce a
                         # row.  Without skips all splits exhaust simultaneously
@@ -562,34 +698,14 @@ class StreamingDataset(IterableDataset):
                             break
 
                         for i in range(n):
-                            pos, row = cooked[i].popleft()
-                            local_consumed[i] += 1
-                            pos_consumed[i] = pos + 1
-                            _advance(i)
+                            row = _yield_row(i)
 
                             # After the last split in each cycle: update the
                             # global offset and refresh the shared-memory stats
                             # so the main process can observe pipeline depth
                             # even when __iter__ runs in a worker process.
                             if i == n - 1:
-                                self._resume_offset = initial_offset + local_consumed[i]
-                                for j, split_idx in enumerate(my_splits):
-                                    self._resume_positions[split_idx] = pos_consumed[j]
-                                ws = self._worker_stats
-                                ws[0] = sum(
-                                    split_sizes[j] - fetch_head[j] for j in range(n)
-                                )
-                                ws[1] = sum(
-                                    batch.num_rows
-                                    for q in raw_batches
-                                    for _, batch in q
-                                )
-                                ws[2] = sum(len(q) for q in cooked)
-                                ws[3] = sum(local_consumed)
-                                ws[4] = self._bytes_loaded
-                                ws[5] = int(self._fetch_time * 1_000_000)
-                                ws[6] = int(self._transform_time * 1_000_000)
-                                ws[7] = self._rows_skipped
+                                _update_progress_stats()
 
                             yield row
                 finally:
@@ -763,33 +879,80 @@ class StreamingDataset(IterableDataset):
     def state_dict(self) -> dict:
         """Snapshot the dataset's consumption state.
 
-        The returned dict is topology-independent: at global step boundaries
-        every split has been consumed the same number of times (by the
-        round-robin design), so the per-split count is a single uniform value
-        that is identical across all ranks and DataLoader workers.
+        When using DataLoader workers, construct a
+        [StreamingDataLoader][lancedb.streaming.StreamingDataLoader].  It
+        commits worker state only when a prefetched batch is returned to the
+        trainer.  A standard multi-process ``DataLoader`` cannot expose that
+        boundary, so calling this method after one has started raises
+        ``RuntimeError`` instead of returning stale producer state.
 
         ``positions_consumed_per_split`` records how far into each split's
         permutation iteration has advanced.  It only differs from
         ``samples_consumed_per_split`` when ``on_transform_error`` skipped
         rows, in which case entries are exact for the splits this instance
         iterated and a lower bound (the sample count) for splits owned by
-        other ranks or workers.  Combine the state dicts from all ranks with
+        other ranks.  ``StreamingDataLoader`` combines worker state in its
+        parent process.  Combine the parent state dicts from all ranks with
         [merge_state_dicts][lancedb.streaming.StreamingDataset.merge_state_dicts]
         to recover the exact value for every split before resuming on a
         different topology.
         """
+        if (
+            self._untracked_worker_iteration[0]
+            and not self._consumer_checkpoint_transport
+            and get_worker_info() is None
+        ):
+            raise RuntimeError(
+                "StreamingDataset cannot checkpoint a standard DataLoader with "
+                "num_workers > 0 because prefetched worker progress is not "
+                "consumer-committed. Use StreamingDataLoader instead."
+            )
+        return self._checkpoint_snapshot()
+
+    def _checkpoint_snapshot(self) -> dict:
+        samples = [
+            self._resume_samples.get(split, self._resume_offset)
+            for split in range(self._num_splits)
+        ]
         positions = [
-            self._resume_positions.get(split, self._resume_offset)
+            self._resume_positions.get(split, samples[split])
             for split in range(self._num_splits)
         ]
         return {
             "shuffle_seed": self._shuffle_seed,
             "num_splits": self._num_splits,
             "epoch": self._epoch,
-            "samples_consumed_per_split": [self._resume_offset] * self._num_splits,
+            "samples_consumed_per_split": samples,
             "positions_consumed_per_split": positions,
         }
 
+    def _commit_worker_state(self, state: dict) -> None:
+        """Merge one trainer-consumed worker batch into parent state."""
+        for key, expected in (
+            ("shuffle_seed", self._shuffle_seed),
+            ("num_splits", self._num_splits),
+            ("epoch", self._epoch),
+        ):
+            if state.get(key) != expected:
+                raise ValueError(
+                    f"{key} mismatch in worker checkpoint: "
+                    f"{state.get(key)} != {expected}"
+                )
+        samples = state["samples_consumed_per_split"]
+        positions = state.get("positions_consumed_per_split", samples)
+        for split, count in enumerate(samples):
+            current = self._resume_samples.get(split, self._resume_offset)
+            self._resume_samples[split] = max(current, int(count))
+        for split, position in enumerate(positions):
+            current = self._resume_positions.get(
+                split, self._resume_samples.get(split, self._resume_offset)
+            )
+            self._resume_positions[split] = max(current, int(position))
+        self._resume_offset = min(
+            self._resume_samples.get(split, self._resume_offset)
+            for split in range(self._num_splits)
+        )
+
     def load_state_dict(self, state: dict) -> None:
         """Resume from a previously snapshotted state.
 
@@ -808,14 +971,17 @@ class StreamingDataset(IterableDataset):
                 f"current dataset has {self._shuffle_seed}"
             )
         consumed = state["samples_consumed_per_split"]
-        # All entries are equal at step boundaries; use the first.
         if isinstance(consumed, list):
-            self._resume_offset = consumed[0] if consumed else 0
+            self._resume_offset = min(consumed) if consumed else 0
+            self._resume_samples = {
+                split: int(count) for split, count in enumerate(consumed)
+            }
         else:
             self._resume_offset = int(consumed)
+            self._resume_samples = {}
         # Older checkpoints predate positions_consumed_per_split; without
         # skipped rows positions equal sample counts, so falling back to
-        # _resume_offset (the .get default in __iter__) is exact.
+        # the per-split sample count (the .get default in __iter__) is exact.
         positions = state.get("positions_consumed_per_split")
         if positions is None:
             self._resume_positions = {}
@@ -828,16 +994,13 @@ class StreamingDataset(IterableDataset):
     def merge_state_dicts(states: list[dict]) -> dict:
         """Merge state dicts saved by different ranks into one exact state.
 
-        Only needed when ``on_transform_error`` skips rows in multi-rank
-        training: each rank then knows the exact permutation position only for
-        its own splits, and records a lower bound for the rest.  Because
-        exactly one rank owns each split, the elementwise maximum across all
-        ranks' ``positions_consumed_per_split`` recovers the exact position of
-        every split.  Without skipped rows every rank's state is already
-        identical and merging is a no-op.
+        Each rank knows exact consumer-committed progress for its own splits
+        and records a lower bound for the rest.  Because exactly one rank owns
+        each split, the elementwise maximum across all ranks recovers both the
+        sample count and permutation position for every split.
 
         Raises ``ValueError`` if the states are empty or were not produced by
-        the same run (mismatched seed, split count, epoch, or sample counts).
+        the same run (mismatched seed, split count, or epoch).
 
         The merge is always all-to-all and topology-agnostic: collect the
         ``state_dict()`` from every rank of the *previous* run into one list,
@@ -886,16 +1049,13 @@ class StreamingDataset(IterableDataset):
                         f"{key} mismatch across state dicts: "
                         f"{state[key]} != {first[key]}"
                     )
-            if (
-                state["samples_consumed_per_split"]
-                != first["samples_consumed_per_split"]
-            ):
-                raise ValueError(
-                    "samples_consumed_per_split mismatch across state dicts; "
-                    "state_dict() must be called at the same global step "
-                    "boundary on every rank"
-                )
         merged = dict(first)
+        merged["samples_consumed_per_split"] = [
+            max(per_split)
+            for per_split in zip(
+                *(state["samples_consumed_per_split"] for state in states)
+            )
+        ]
         all_positions = [
             state.get(
                 "positions_consumed_per_split", state["samples_consumed_per_split"]
@@ -906,3 +1066,43 @@ class StreamingDataset(IterableDataset):
             max(per_split) for per_split in zip(*all_positions)
         ]
         return merged
+
+
+class StreamingDataLoader(DataLoader):
+    """A PyTorch DataLoader with consumer-committed dataset checkpoints.
+
+    PyTorch workers prefetch batches ahead of the trainer, so worker-local
+    producer progress is not a safe checkpoint.  This loader carries a state
+    snapshot alongside every internal batch and applies it to the parent
+    [StreamingDataset][lancedb.streaming.StreamingDataset] only when that batch
+    is returned by ``next()``.
+    The trainer receives the same collated batch it would receive from a
+    standard ``torch.utils.data.DataLoader``.
+
+    Parameters are the same as ``torch.utils.data.DataLoader`` except that
+    ``dataset`` must be a
+    [StreamingDataset][lancedb.streaming.StreamingDataset].
+
+    Examples
+    --------
+    >>> # dataset = StreamingDataset(table, num_splits=2)
+    >>> # loader = StreamingDataLoader(dataset, batch_size=8, num_workers=2)
+    >>> # batch = next(iter(loader))
+    >>> # checkpoint = dataset.state_dict()
+    """
+
+    def __init__(self, dataset: StreamingDataset, *args, **kwargs):
+        if not isinstance(dataset, StreamingDataset):
+            raise TypeError("StreamingDataLoader requires a StreamingDataset")
+        if kwargs.get("in_order", True) is False:
+            raise ValueError(
+                "StreamingDataLoader requires in_order=True for deterministic "
+                "consumer checkpoints"
+            )
+        dataset._consumer_checkpoint_transport = True
+        self._streaming_dataset = dataset
+        super().__init__(_StreamingDatasetAdapter(dataset), *args, **kwargs)
+        self.collate_fn = _CheckpointCollate(self.collate_fn)
+
+    def __iter__(self):
+        return _ConsumerCommitIterator(super().__iter__(), self._streaming_dataset)
diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py
index 0c1a70765..6742fe0d1 100644
--- a/python/python/tests/test_elastic_dataloader.py
+++ b/python/python/tests/test_elastic_dataloader.py
@@ -46,6 +46,7 @@ from utils import (
 torch = pytest.importorskip("torch")
 streaming = pytest.importorskip("lancedb.streaming")
 StreamingDataset = streaming.StreamingDataset
+StreamingDataLoader = streaming.StreamingDataLoader
 
 # ---------------------------------------------------------------------------
 # Dataset parameters
@@ -1008,6 +1009,74 @@ def test_multi_worker_elastic_det_across_worker_counts(lance_table):
 # ── Resumability with num_workers ─────────────────────────────────────────────
 
 
+def test_streaming_dataloader_commits_only_consumed_worker_batches(tmp_path):
+    """Prefetched worker state is committed only as the trainer receives it."""
+    db = lancedb.connect(tmp_path)
+    table = db.create_table("worker_commit", pa.table({"id": [1, 2, 10, 20]}))
+    dataset = StreamingDataset(table, num_splits=2, shuffle=False)
+    loader = StreamingDataLoader(
+        dataset,
+        batch_size=2,
+        num_workers=2,
+        multiprocessing_context="spawn",
+        prefetch_factor=4,
+    )
+    iterator = iter(loader)
+    try:
+        first = next(iterator)["id"].tolist()
+        checkpoint = dataset.state_dict()
+
+        assert first == [1, 2]
+        assert checkpoint["samples_consumed_per_split"] == [2, 0]
+        assert checkpoint["positions_consumed_per_split"] == [2, 0]
+
+        second = next(iterator)["id"].tolist()
+        assert second == [10, 20]
+        assert dataset.state_dict()["samples_consumed_per_split"] == [2, 2]
+    finally:
+        iterator._shutdown_workers()
+
+    resumed = StreamingDataset(table, num_splits=2, shuffle=False)
+    resumed.load_state_dict(checkpoint)
+    assert [row["id"] for row in resumed] == [10, 20]
+
+
+def test_standard_dataloader_rejects_stale_parent_checkpoint(tmp_path):
+    """A standard DataLoader must not expose prefetched producer progress."""
+    db = lancedb.connect(tmp_path)
+    table = db.create_table("untracked_workers", pa.table({"id": [1, 2, 10, 20]}))
+    dataset = StreamingDataset(table, num_splits=2, shuffle=False)
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=2,
+        num_workers=2,
+        multiprocessing_context="spawn",
+    )
+    iterator = iter(loader)
+    try:
+        assert next(iterator)["id"].tolist() == [1, 2]
+        with pytest.raises(RuntimeError, match="Use StreamingDataLoader"):
+            dataset.state_dict()
+    finally:
+        iterator._shutdown_workers()
+
+
+def test_resume_from_partial_split_cycle_preserves_remaining_order(tmp_path):
+    db = lancedb.connect(tmp_path)
+    table = db.create_table("partial_cycle", pa.table({"id": [1, 2, 10, 20]}))
+    dataset = StreamingDataset(table, num_splits=2, shuffle=False)
+    iterator = iter(dataset)
+
+    assert next(iterator)["id"] == 1
+    checkpoint = dataset.state_dict()
+    iterator.close()
+    assert checkpoint["samples_consumed_per_split"] == [1, 0]
+
+    resumed = StreamingDataset(table, num_splits=2, shuffle=False)
+    resumed.load_state_dict(checkpoint)
+    assert [row["id"] for row in resumed] == [10, 2, 20]
+
+
 def test_multi_worker_resumability_same_topology(lance_table):
     """Checkpoint with num_workers=2, resume with num_workers=2: exact continuation."""
     world_size = 1
@@ -1836,6 +1905,23 @@ def test_merge_state_dicts_validates_consistency(lance_table):
         StreamingDataset.merge_state_dicts([])
 
 
+def test_merge_state_dicts_combines_nonuniform_consumer_progress(lance_table):
+    dataset = StreamingDataset(
+        lance_table, num_splits=2, shuffle=False, shuffle_seed=SHUFFLE_SEED
+    )
+    rank0 = dataset.state_dict()
+    rank0["samples_consumed_per_split"] = [2, 0]
+    rank0["positions_consumed_per_split"] = [2, 0]
+    rank1 = dataset.state_dict()
+    rank1["samples_consumed_per_split"] = [0, 2]
+    rank1["positions_consumed_per_split"] = [0, 2]
+
+    merged = StreamingDataset.merge_state_dicts([rank0, rank1])
+
+    assert merged["samples_consumed_per_split"] == [2, 2]
+    assert merged["positions_consumed_per_split"] == [2, 2]
+
+
 def test_load_state_dict_without_positions_key(lance_table):
     """Checkpoints from before positions_consumed_per_split existed still
     resume exactly (positions equal sample counts when nothing is skipped)."""