diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 159dea5d6..f1d80c9c8 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -88,8 +88,7 @@ class _StreamingDatasetAdapter(IterableDataset): self.dataset = dataset def __iter__(self): - self.dataset._consumer_checkpoint_transport = True - for sample in self.dataset: + for sample in self.dataset._iter(consumer_checkpoint_transport=True): yield _WorkerSample(sample, self.dataset) def __getattr__(self, name): @@ -100,20 +99,34 @@ class _StreamingDatasetAdapter(IterableDataset): class _ConsumerCommitIterator: - def __init__(self, iterator, dataset: "StreamingDataset"): + def __init__(self, iterator, dataset: "StreamingDataset", *, require_uniform: bool): self._iterator = iterator self._dataset = dataset + self._require_uniform = require_uniform def __iter__(self): return self def __next__(self): - batch = next(self._iterator) + try: + batch = next(self._iterator) + except StopIteration: + raise + except Exception as exc: + self._dataset._invalidate_checkpoint( + f"a DataLoader batch failed before it was returned: {exc}" + ) + raise if not isinstance(batch, _WorkerBatch): + self._dataset._invalidate_checkpoint( + "worker checkpoint metadata was missing from a returned batch" + ) raise RuntimeError( "StreamingDataLoader did not receive worker checkpoint metadata" ) - self._dataset._commit_worker_state(batch.state) + self._dataset._commit_worker_state( + batch.state, require_uniform=self._require_uniform + ) return batch.data def __getattr__(self, name): @@ -335,7 +348,13 @@ class StreamingDataset(IterableDataset): # 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 + + # Parent-side checkpoint lifecycle. A failed DataLoader task creates + # a permanent hole in that iterator's delivery stream, while a + # multi-worker checkpoint is safe to restore only after all splits + # reach the same logical step boundary. + self._checkpoint_invalid_reason: Optional[str] = None + self._consumer_checkpoint_requires_uniform = False # Cumulative bytes of Arrow buffer data fetched across all iterations. self._bytes_loaded: int = 0 @@ -410,13 +429,18 @@ class StreamingDataset(IterableDataset): return self._rank_splits[start : start + splits_per_worker] def __iter__(self) -> Iterator[dict[str, Any]]: + return self._iter() + + def _iter( + self, *, consumer_checkpoint_transport: bool = False + ) -> 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." ) real_worker = get_worker_info() is not None - if real_worker and not self._consumer_checkpoint_transport: + if real_worker and not consumer_checkpoint_transport: self._untracked_worker_iteration[0] = 1 my_splits = self._resolve_my_splits() @@ -663,20 +687,23 @@ class StreamingDataset(IterableDataset): # exact canonical sequence continues without replaying # already-consumed rows. if len(set(initial_samples)) > 1: + catch_up_to = max(initial_samples) pending = [ - (initial_samples[i], my_splits[i], i) for i in range(n) + (initial_samples[i], my_splits[i], i) + for i in range(n) + if initial_samples[i] < catch_up_to ] heapq.heapify(pending) while pending: consumed, _, i = heapq.heappop(pending) _ensure_cooked(i) if not cooked[i]: - continue + return row = _yield_row(i) - heapq.heappush(pending, (consumed + 1, my_splits[i], i)) + if consumed + 1 < catch_up_to: + 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 @@ -897,17 +924,28 @@ class StreamingDataset(IterableDataset): 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 - ): + if self._untracked_worker_iteration[0] 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() + if self._checkpoint_invalid_reason is not None: + raise RuntimeError( + "StreamingDataset checkpointing is invalid because " + f"{self._checkpoint_invalid_reason}. Load the last valid " + "checkpoint into a fresh dataset before continuing." + ) + state = self._checkpoint_snapshot() + samples = state["samples_consumed_per_split"] + if self._consumer_checkpoint_requires_uniform and len(set(samples)) > 1: + raise RuntimeError( + "StreamingDataLoader checkpointing with multiple workers is only " + "safe at a complete logical step boundary, when every split has " + "the same consumed-sample count. Consume more batches before " + "calling state_dict()." + ) + return state def _checkpoint_snapshot(self) -> dict: samples = [ @@ -926,7 +964,11 @@ class StreamingDataset(IterableDataset): "positions_consumed_per_split": positions, } - def _commit_worker_state(self, state: dict) -> None: + def _invalidate_checkpoint(self, reason: str) -> None: + if self._checkpoint_invalid_reason is None: + self._checkpoint_invalid_reason = reason + + def _commit_worker_state(self, state: dict, *, require_uniform: bool) -> None: """Merge one trainer-consumed worker batch into parent state.""" for key, expected in ( ("shuffle_seed", self._shuffle_seed), @@ -952,6 +994,7 @@ class StreamingDataset(IterableDataset): self._resume_samples.get(split, self._resume_offset) for split in range(self._num_splits) ) + self._consumer_checkpoint_requires_uniform |= require_uniform def load_state_dict(self, state: dict) -> None: """Resume from a previously snapshotted state. @@ -970,6 +1013,7 @@ class StreamingDataset(IterableDataset): f"shuffle_seed mismatch: checkpoint has {state['shuffle_seed']}, " f"current dataset has {self._shuffle_seed}" ) + self._consumer_checkpoint_requires_uniform = False consumed = state["samples_consumed_per_split"] if isinstance(consumed, list): self._resume_offset = min(consumed) if consumed else 0 @@ -1079,6 +1123,13 @@ class StreamingDataLoader(DataLoader): The trainer receives the same collated batch it would receive from a standard ``torch.utils.data.DataLoader``. + With more than one worker, ``state_dict()`` is available only at complete + logical step boundaries, when every split has the same consumed-sample + count. ``persistent_workers=True`` is not supported because prefetched + worker copies cannot be restored from 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. + Parameters are the same as ``torch.utils.data.DataLoader`` except that ``dataset`` must be a [StreamingDataset][lancedb.streaming.StreamingDataset]. @@ -1099,10 +1150,28 @@ class StreamingDataLoader(DataLoader): "StreamingDataLoader requires in_order=True for deterministic " "consumer checkpoints" ) - dataset._consumer_checkpoint_transport = True + if kwargs.get("persistent_workers", False): + raise ValueError( + "StreamingDataLoader does not support persistent_workers=True " + "because worker prefetch state cannot be reset from a checkpoint" + ) 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) + if self.num_workers > 1: + samples = self._streaming_dataset._checkpoint_snapshot()[ + "samples_consumed_per_split" + ] + if len(set(samples)) > 1: + raise RuntimeError( + "StreamingDataLoader cannot start multiple workers from a " + "partial logical step; resume from a checkpoint whose splits " + "have equal consumed-sample counts" + ) + return _ConsumerCommitIterator( + super().__iter__(), + self._streaming_dataset, + require_uniform=self.num_workers > 1, + ) diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index 6742fe0d1..dbc2761d8 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -93,6 +93,13 @@ class FakeWorkerInfo: num_workers: int +def _collate_with_first_batch_error(samples): + ids = [sample["id"] for sample in samples] + if ids == [0, 1]: + raise ValueError("first batch fails") + return ids + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @@ -1012,7 +1019,9 @@ def test_multi_worker_elastic_det_across_worker_counts(lance_table): 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]})) + table = db.create_table( + "worker_commit", pa.table({"id": [1, 2, 3, 4, 10, 20, 30, 40]}) + ) dataset = StreamingDataset(table, num_splits=2, shuffle=False) loader = StreamingDataLoader( dataset, @@ -1024,21 +1033,35 @@ def test_streaming_dataloader_commits_only_consumed_worker_batches(tmp_path): 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] + assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2, 0] + with pytest.raises(RuntimeError, match="complete logical step boundary"): + dataset.state_dict() second = next(iterator)["id"].tolist() assert second == [10, 20] - assert dataset.state_dict()["samples_consumed_per_split"] == [2, 2] + checkpoint = dataset.state_dict() + assert checkpoint["samples_consumed_per_split"] == [2, 2] + uninterrupted = [batch["id"].tolist() for batch in iterator] 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] + resumed_loader = StreamingDataLoader( + resumed, + batch_size=2, + num_workers=2, + multiprocessing_context="spawn", + prefetch_factor=4, + ) + resumed_iterator = iter(resumed_loader) + try: + remaining = [batch["id"].tolist() for batch in resumed_iterator] + finally: + resumed_iterator._shutdown_workers() + assert remaining == uninterrupted == [[3, 4], [30, 40]] def test_standard_dataloader_rejects_stale_parent_checkpoint(tmp_path): @@ -1046,6 +1069,9 @@ def test_standard_dataloader_rejects_stale_parent_checkpoint(tmp_path): 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) + # Merely constructing the checkpoint-aware loader must not authorize a + # later plain DataLoader's worker progress. + StreamingDataLoader(dataset, batch_size=2, num_workers=0) loader = torch.utils.data.DataLoader( dataset, batch_size=2, @@ -1057,6 +1083,48 @@ def test_standard_dataloader_rejects_stale_parent_checkpoint(tmp_path): assert next(iterator)["id"].tolist() == [1, 2] with pytest.raises(RuntimeError, match="Use StreamingDataLoader"): dataset.state_dict() + list(iterator) + finally: + iterator._shutdown_workers() + + +def test_streaming_dataloader_rejects_persistent_workers(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table("persistent_workers", pa.table({"id": [1, 2]})) + dataset = StreamingDataset(table, num_splits=2, shuffle=False) + + with pytest.raises(ValueError, match="persistent_workers=True"): + StreamingDataLoader( + dataset, + batch_size=1, + num_workers=2, + persistent_workers=True, + ) + + +def test_collate_failure_invalidates_consumer_checkpoint(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table( + "collate_failure", pa.table({"id": [0, 1, 2, 3, 100, 101, 102, 103]}) + ) + dataset = StreamingDataset(table, num_splits=2, shuffle=False) + loader = StreamingDataLoader( + dataset, + batch_size=2, + num_workers=2, + multiprocessing_context="spawn", + collate_fn=_collate_with_first_batch_error, + prefetch_factor=2, + ) + iterator = iter(loader) + try: + with pytest.raises(ValueError, match="first batch fails"): + next(iterator) + assert next(iterator) == [100, 101] + assert next(iterator) == [2, 3] + with pytest.raises(RuntimeError, match="failed before it was returned"): + dataset.state_dict() + list(iterator) finally: iterator._shutdown_workers() @@ -1077,6 +1145,29 @@ def test_resume_from_partial_split_cycle_preserves_remaining_order(tmp_path): assert [row["id"] for row in resumed] == [10, 2, 20] +def test_partial_cycle_resume_preserves_skip_truncation(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table( + "partial_skip", pa.table({"id": [0, 1, 2, 3, 100, 101, 102, 103]}) + ) + kwargs = dict( + num_splits=2, + shuffle=False, + transform=_failing_transform({1, 2, 3}), + on_transform_error="skip", + ) + dataset = StreamingDataset(table, **kwargs) + iterator = iter(dataset) + + assert next(iterator)["id"] == 0 + checkpoint = dataset.state_dict() + uninterrupted = [row["id"] for row in iterator] + + resumed = StreamingDataset(table, **kwargs) + resumed.load_state_dict(checkpoint) + assert [row["id"] for row in resumed] == uninterrupted == [100] + + def test_multi_worker_resumability_same_topology(lance_table): """Checkpoint with num_workers=2, resume with num_workers=2: exact continuation.""" world_size = 1