Harden production dataset streaming pipeline

tests/configs/test_default.py

@@ -44,6 +44,11 @@ def test_dataset_config_empty_episodes_ok():
         ("streaming_episode_pool_size", 0, "episode_pool_size"),
         ("streaming_prefetch_episodes", -1, "prefetch_episodes"),
         ("streaming_byte_budget_gb", 0, "byte_budget_gb"),
+        ("streaming_decode_threads", 0, "decode_threads"),
+        ("streaming_decoded_queue_size", 0, "decoded_queue_size"),
+        ("streaming_max_open_decoders", 0, "max_open_decoders"),
+        ("streaming_native_http_connections", 0, "native_http_connections"),
+        ("streaming_native_http_subranges", 0, "native_http_subranges"),
     ],
 )
 def test_dataset_config_rejects_invalid_streaming_resource_limits(field, value, message):

tests/datasets/test_episode_parquet_reader.py

@@ -122,3 +122,33 @@ def test_reader_forwards_explicit_token_to_hf_filesystem(monkeypatch) -> None:
     )
 
     url_to_fs.assert_called_once_with("hf://datasets/private@revision", token="hf_test_token")
+
+
+def test_reader_retries_transient_remote_file_not_found(tmp_path: Path, monkeypatch) -> None:
+    path = tmp_path / "episode.parquet"
+    pq.write_table(_table([0, 0]), path)
+    filesystem = Mock()
+    attempts = 0
+
+    def open_remote(*_args, **_kwargs):
+        nonlocal attempts
+        attempts += 1
+        if attempts == 1:
+            raise FileNotFoundError("partial HfFileSystem dircache")
+        return path.open("rb")
+
+    filesystem.open.side_effect = open_remote
+    filesystem.invalidate_cache = Mock()
+    monkeypatch.setattr(fsspec.core, "url_to_fs", Mock(return_value=(filesystem, "root")))
+    reader = EpisodeParquetReader(
+        "hf://datasets/private@revision",
+        columns=("episode_index", "frame_index"),
+        max_retries=1,
+        retry_backoff_s=0,
+    )
+
+    table = reader.read_episode("data/file.parquet", episode_index=0, expected_rows=2)
+
+    assert len(table) == 2
+    assert attempts == 2
+    filesystem.invalidate_cache.assert_called_once()

tests/datasets/test_episode_video_streaming.py

@@ -11,6 +11,7 @@
 import json
 import struct
 import threading
+import time
 from concurrent.futures import ThreadPoolExecutor
 
 import numpy as np
@@ -217,6 +218,66 @@ def test_decoder_falls_back_to_pyav_when_torchcodec_rejects_mini_mp4(monkeypatch
         assert cache.decoder_fallback_count == 1
 
 
+def test_torchcodec_frame_indices_are_clamped_to_decoder_bounds(monkeypatch, tmp_path):
+    requested_indices = []
+
+    class FakeDecoder:
+        metadata = type("Metadata", (), {"average_fps": 30.0, "num_frames": 10})()
+
+        def get_frames_at(self, *, indices):
+            requested_indices.extend(indices)
+            return type("Frames", (), {"data": indices})()
+
+    with _fake_cache(monkeypatch, tmp_path) as cache:
+        monkeypatch.setattr(
+            cache.manifest,
+            "lookup",
+            lambda *_args: type("Span", (), {"source_start_pts": 0.0})(),
+        )
+        monkeypatch.setattr(cache, "_decoder_for_frames", lambda *_args: (FakeDecoder(), None))
+
+        cache.get_frames(0, "camera", [-0.1, 1.0])
+
+    assert requested_indices == [0, 9]
+
+
+def test_frame_reads_serialize_access_to_each_decoder(monkeypatch, tmp_path):
+    state_lock = threading.Lock()
+    active = 0
+    max_active = 0
+
+    class FakeDecoder:
+        metadata = type("Metadata", (), {"average_fps": 30.0, "num_frames": 10})()
+
+        def get_frames_at(self, *, indices):
+            nonlocal active, max_active
+            with state_lock:
+                active += 1
+                max_active = max(max_active, active)
+            try:
+                time.sleep(0.01)
+                return type("Frames", (), {"data": indices})()
+            finally:
+                with state_lock:
+                    active -= 1
+
+    with _fake_cache(monkeypatch, tmp_path) as cache:
+        monkeypatch.setattr(
+            cache.manifest,
+            "lookup",
+            lambda *_args: type("Span", (), {"source_start_pts": 0.0})(),
+        )
+        decoder = FakeDecoder()
+        monkeypatch.setattr(cache, "_decoder_for_frames", lambda *_args: (decoder, None))
+
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            futures = [executor.submit(cache.get_frames, 0, "camera", [0.1]) for _ in range(2)]
+            for future in futures:
+                future.result()
+
+    assert max_active == 1
+
+
 def test_releasing_episode_allows_immediate_eviction(monkeypatch, tmp_path):
     with _fake_cache(monkeypatch, tmp_path, byte_budget=5) as cache:
         cache.retain_episode(0)

tests/datasets/test_streaming_factory.py

@@ -38,6 +38,11 @@ def test_factory_wires_production_streaming_settings(monkeypatch):
         streaming_episode_pool_size=7,
         streaming_prefetch_episodes=3,
         streaming_byte_budget_gb=2.5,
+        streaming_decode_threads=2,
+        streaming_decoded_queue_size=5,
+        streaming_max_open_decoders=17,
+        streaming_native_http_connections=9,
+        streaming_native_http_subranges=3,
     )
     cfg = SimpleNamespace(
         dataset=dataset_config,
@@ -54,6 +59,11 @@ def test_factory_wires_production_streaming_settings(monkeypatch):
     assert captured["kwargs"]["episode_pool_size"] == 7
     assert captured["kwargs"]["prefetch_episodes"] == 3
     assert captured["kwargs"]["byte_budget_gb"] == 2.5
+    assert captured["kwargs"]["decode_threads"] == 2
+    assert captured["kwargs"]["decoded_queue_size"] == 5
+    assert captured["kwargs"]["max_open_decoders"] == 17
+    assert captured["kwargs"]["native_http_connections"] == 9
+    assert captured["kwargs"]["native_http_subranges"] == 3
     assert captured["kwargs"]["max_num_shards"] == 1
     assert captured["kwargs"]["video_backend"] == "pyav"
     assert captured["kwargs"]["return_uint8"] is True

tests/datasets/test_streaming_production.py

@@ -10,6 +10,8 @@
 
 from __future__ import annotations
 
+import threading
+import time
 from itertools import islice
 from pathlib import Path
 
@@ -61,6 +63,61 @@ def test_streaming_matches_map_style_with_exact_coverage(tmp_path: Path, lerobot
         _assert_item_equal(sample, map_dataset[int(sample["index"])])
 
 
+def test_parallel_decode_queue_preserves_planner_order(
+    tmp_path: Path,
+    lerobot_dataset_factory,
+    monkeypatch,
+) -> None:
+    root = tmp_path / "dataset"
+    lerobot_dataset_factory(
+        root=root,
+        repo_id=DUMMY_REPO_ID,
+        total_episodes=4,
+        total_frames=40,
+        use_videos=False,
+    )
+    expected = _indices(
+        StreamingLeRobotDataset(
+            DUMMY_REPO_ID,
+            root=root,
+            seed=17,
+            buffer_size=3,
+            decode_threads=1,
+            decoded_queue_size=1,
+        )
+    )
+    parallel = StreamingLeRobotDataset(
+        DUMMY_REPO_ID,
+        root=root,
+        seed=17,
+        buffer_size=3,
+        decode_threads=3,
+        decoded_queue_size=5,
+    )
+    original_make_item = parallel._make_episode_item
+    state_lock = threading.Lock()
+    active = 0
+    max_active = 0
+
+    def delayed_make_item(*args, **kwargs):
+        nonlocal active, max_active
+        with state_lock:
+            active += 1
+            max_active = max(max_active, active)
+        try:
+            frame_index = int(args[2])
+            time.sleep(0.005 if frame_index % 3 == 0 else 0.001)
+            return original_make_item(*args, **kwargs)
+        finally:
+            with state_lock:
+                active -= 1
+
+    monkeypatch.setattr(parallel, "_make_episode_item", delayed_make_item)
+
+    assert _indices(parallel) == expected
+    assert 1 < max_active <= parallel.decode_threads
+
+
 @pytest.mark.parametrize("video_backend", ["torchcodec", "pyav"])
 def test_streaming_rgb_video_matches_map_style(
     tmp_path: Path,