rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-11-14 06:36:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(python): commit streaming worker checkpoints on consumption

Browse Source
This commit is contained in:
Gatefixer
2026-08-21 22:23:15 +00:00
parent 29822306d2
commit 865b167de2
3 changed files with 341 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/streaming.py
+253 -53
View File
@@ -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)
python/python/tests/test_elastic_dataloader.py
+86
View File
@@ -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)."""