diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 8f72090f5..6367cc723 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -122,17 +122,20 @@ class _ConsumerCommitIterator: f"a DataLoader batch failed before it was returned: {exc}" ) raise - if not isinstance(batch, _WorkerBatch): + try: + if not isinstance(batch, _WorkerBatch): + raise RuntimeError( + "StreamingDataLoader did not receive worker checkpoint metadata" + ) + self._dataset._commit_worker_state( + batch.state, require_uniform=self._require_uniform + ) + return batch.data + except BaseException as exc: self._dataset._invalidate_checkpoint( - "worker checkpoint metadata was missing from a returned batch" + f"a DataLoader batch failed before it was returned: {exc}" ) - raise RuntimeError( - "StreamingDataLoader did not receive worker checkpoint metadata" - ) - self._dataset._commit_worker_state( - batch.state, require_uniform=self._require_uniform - ) - return batch.data + raise def __getattr__(self, name): return getattr(self._iterator, name) diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index 06c24e13a..3a29ef4ad 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -1221,6 +1221,29 @@ def test_batch_base_exception_invalidates_consumer_checkpoint(tmp_path): assert list(iterator) == [[2, 3], [4, 5]] +def test_parent_commit_base_exception_invalidates_consumer_checkpoint(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table("commit_interrupt", pa.table({"id": list(range(4))})) + dataset = StreamingDataset(table, num_splits=1, shuffle=False) + loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0) + iterator = iter(loader) + real_commit = dataset._commit_worker_state + + def interrupt_after_commit(state, *, require_uniform): + real_commit(state, require_uniform=require_uniform) + raise KeyboardInterrupt("after parent commit") + + with patch.object( + dataset, "_commit_worker_state", side_effect=interrupt_after_commit + ): + with pytest.raises(KeyboardInterrupt, match="after parent commit"): + next(iterator) + + assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2] + with pytest.raises(RuntimeError, match="failed before it was returned"): + dataset.state_dict() + + def test_streaming_dataloader_rejects_dataset_iter_override(tmp_path): class CustomizedDataset(StreamingDataset): def __iter__(self):