diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index e02550e2f..b63e34dc8 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -58,6 +58,10 @@ class LookAheadError(Exception): pass +class _ShardExhaustedError(Exception): + """Raised when a streaming dataset shard has no more items.""" + + class Backtrackable[T]: """ Wrap any iterator/iterable so you can step back up to `history` items @@ -422,10 +426,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): else: frames_buffer.append(frame) break # random shard sampled, switch shard - except ( - RuntimeError, - StopIteration, - ): # NOTE: StopIteration inside a generator throws a RuntimeError since python 3.7 + except _ShardExhaustedError: del idx_to_backtrack_dataset[shard_key] # Remove exhausted shard, onto another shard # Once shards are all exhausted, shuffle the buffer and yield the remaining frames @@ -503,7 +504,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): def make_frame(self, dataset_iterator: Backtrackable) -> Generator: """Makes a frame starting from a dataset iterator""" - item = next(dataset_iterator) + try: + item = next(dataset_iterator) + except StopIteration as e: + # Translate exhaustion here, before PEP 479 turns it into an indistinguishable RuntimeError. + raise _ShardExhaustedError from e item = item_to_torch(item) updates = [] # list of "updates" to apply to the item retrieved from hf_dataset (w/o camera features) diff --git a/tests/datasets/test_streaming.py b/tests/datasets/test_streaming.py index cae4be5b6..08544326f 100644 --- a/tests/datasets/test_streaming.py +++ b/tests/datasets/test_streaming.py @@ -252,6 +252,54 @@ def test_frames_order_with_shards(tmp_path, lerobot_dataset_factory, shuffle): assert frames_match +def test_iter_raises_on_frame_error(tmp_path, lerobot_dataset_factory, monkeypatch): + """Video decode failures must propagate instead of being silently ignored.""" + ds_num_frames = 20 + ds_num_episodes = 2 + buffer_size = 10 + + local_path = tmp_path / "test" + repo_id = f"{DUMMY_REPO_ID}" + + lerobot_dataset_factory( + root=local_path, + repo_id=repo_id, + total_episodes=ds_num_episodes, + total_frames=ds_num_frames, + ) + + streaming_ds = StreamingLeRobotDataset(repo_id=repo_id, root=local_path, buffer_size=buffer_size) + + def broken_video_decode(*args, **kwargs): + raise RuntimeError("Could not load libtorchcodec") + + monkeypatch.setattr(streaming_dataset_module, "decode_video_frames_torchcodec", broken_video_decode) + + with pytest.raises(RuntimeError, match="libtorchcodec"): + next(iter(streaming_ds)) + + +def test_iter_raises_on_nested_generator_error(tmp_path, lerobot_dataset_factory, monkeypatch): + """PEP 479 errors below frame construction must not be mistaken for shard exhaustion.""" + local_path = tmp_path / "test" + repo_id = DUMMY_REPO_ID + + lerobot_dataset_factory(root=local_path, repo_id=repo_id, total_episodes=2, total_frames=20) + streaming_ds = StreamingLeRobotDataset(repo_id=repo_id, root=local_path, buffer_size=10) + + def broken_frame_generator(): + raise StopIteration("decoder internal failure") + yield + + def broken_video_decode(*args, **kwargs): + return next(broken_frame_generator()) + + monkeypatch.setattr(streaming_dataset_module, "decode_video_frames_torchcodec", broken_video_decode) + + with pytest.raises(RuntimeError, match="generator raised StopIteration"): + next(iter(streaming_ds)) + + @pytest.mark.parametrize( "state_deltas, action_deltas", [