mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 21:19:40 +00:00
fix(datasets): stop frame errors being treated as shard exhaustion in StreamingLeRobotDataset (#4237)
* fix(datasets): stop frame errors being treated as shard exhaustion in StreamingLeRobotDataset StreamingLeRobotDataset.__iter__ caught every RuntimeError and treated all as exhausted shard. Real errors like video decode failure made each shard get dropped on the first frame, so iteration ended while yeilding zero frames with no errors. Shard exahustion is StopIteration raised from make_frame generator, which python converts to RuntimeError with StopIteration as __cause__. Added check to tell StopIteration from everything else, consuming real shard exhaustion while re-raising everything else. Added a test that injects a decode failure and asserts iteration raises instead of returning on empty stream. Fixes #4066 * refactor(datasets): exception streaming --------- Co-authored-by: Mohit Yadav <mohitydv09@gmail.com>
This commit is contained in:
@@ -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",
|
||||
[
|
||||
|
||||
Reference in New Issue
Block a user