diff --git a/docs/source/training_dataset_streaming.mdx b/docs/source/training_dataset_streaming.mdx index 7288637eb..7cd053ce6 100644 --- a/docs/source/training_dataset_streaming.mdx +++ b/docs/source/training_dataset_streaming.mdx @@ -58,17 +58,17 @@ Subset sidecars cannot be published. The production defaults are: -| Option | Default | Meaning | -| ----------------------------------- | ------: | --------------------------------------------------- | -| `streaming_episode_pool_size` | 32 | Maximum complete episodes mixed by each rank | -| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool | -| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per rank | -| `streaming_decode_threads` | 2 | Parallel sample assembly and video decode workers | -| `streaming_decoded_queue_size` | 8 | Decoded samples buffered ahead, in planner order | -| `streaming_max_open_decoders` | 64 | Independent open-decoder LRU cap per rank | -| `streaming_native_http_connections` | unset | Native HTTP connection cap per rank | -| `streaming_native_http_subranges` | 1 | Concurrent subranges per native HTTP range read | -| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root | +| Option | Default | Meaning | +| ----------------------------------- | -------------: | --------------------------------------------------- | +| `streaming_episode_pool_size` | 32 | Maximum complete episodes mixed by each rank | +| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool | +| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per rank | +| `streaming_decode_threads` | 2 | Parallel sample assembly and video decode workers | +| `streaming_decoded_queue_size` | 8 | Decoded samples buffered ahead, in planner order | +| `streaming_max_open_decoders` | pool × cameras | Independent open-decoder LRU cap per rank | +| `streaming_native_http_connections` | unset | Native HTTP connection cap per rank | +| `streaming_native_http_subranges` | 1 | Concurrent subranges per native HTTP range read | +| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root | The active episode set is capped by both episode count and the exact synthesized mini-MP4 sizes computed from the sidecar. An episode larger than the complete rank budget fails before training @@ -84,7 +84,6 @@ lerobot-train \ --dataset.streaming_byte_budget_gb=4 \ --dataset.streaming_decode_threads=2 \ --dataset.streaming_decoded_queue_size=8 \ - --dataset.streaming_max_open_decoders=64 \ --num_workers=4 \ --policy.type=act \ --output_dir=outputs/train/act_streaming diff --git a/scripts/bench_streaming_dataset.py b/scripts/bench_streaming_dataset.py index c19d99727..c2aa0c821 100644 --- a/scripts/bench_streaming_dataset.py +++ b/scripts/bench_streaming_dataset.py @@ -52,7 +52,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--byte-budget-gb", type=float, default=8.0) parser.add_argument("--decode-threads", type=int, default=2) parser.add_argument("--decoded-queue-size", type=int, default=8) - parser.add_argument("--max-open-decoders", type=int, default=64) + parser.add_argument("--max-open-decoders", type=int, default=None) parser.add_argument("--native-http-connections", type=int, default=None) parser.add_argument("--native-http-subranges", type=int, default=1) parser.add_argument("--warmup-batches", type=int, default=8) @@ -196,7 +196,8 @@ def main() -> None: "byte_budget_gb": args.byte_budget_gb, "decode_threads": args.decode_threads, "decoded_queue_size": args.decoded_queue_size, - "max_open_decoders": args.max_open_decoders, + "max_open_decoders": dataset.max_open_decoders, + "requested_max_open_decoders": args.max_open_decoders, "native_http_connections": args.native_http_connections, "native_http_subranges": args.native_http_subranges, "warmup_batches": args.warmup_batches, diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index 9ad8f047e..dbfc93bd0 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -55,8 +55,8 @@ class DatasetConfig: # Parallel sample assembly/decode workers and their bounded in-order result queue. streaming_decode_threads: int = 2 streaming_decoded_queue_size: int = 8 - # Independent decoder-state cap. - streaming_max_open_decoders: int = 64 + # Independent decoder-state cap. None covers every camera in the configured episode pool. + streaming_max_open_decoders: int | None = None # Per-rank native HTTP limits. None preserves the fetcher's worker-derived default. streaming_native_http_connections: int | None = None streaming_native_http_subranges: int = 1 @@ -80,7 +80,7 @@ class DatasetConfig: raise ValueError("streaming_decode_threads must be positive") if self.streaming_decoded_queue_size <= 0: raise ValueError("streaming_decoded_queue_size must be positive") - if self.streaming_max_open_decoders <= 0: + if self.streaming_max_open_decoders is not None and self.streaming_max_open_decoders <= 0: raise ValueError("streaming_max_open_decoders must be positive") if self.streaming_native_http_connections is not None and self.streaming_native_http_connections <= 0: raise ValueError("streaming_native_http_connections must be positive") diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index f9f1bee2f..e5c95a8db 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -125,7 +125,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): token: str | bool | None = None, decode_threads: int = 2, decoded_queue_size: int = 8, - max_open_decoders: int = 64, + max_open_decoders: int | None = None, native_http_connections: int | None = None, native_http_subranges: int = 1, ): @@ -161,7 +161,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): decode_threads (int, optional): Parallel sample assembly and video decode workers. decoded_queue_size (int, optional): Maximum number of decoded samples produced ahead of the consumer. Results are yielded in exact planner order. - max_open_decoders (int, optional): Maximum number of open video decoders per rank. + max_open_decoders (int | None, optional): Maximum number of open video decoders per + rank. By default every camera in the configured episode pool can remain open. native_http_connections (int | None, optional): Native HTTP connection limit per rank. ``None`` preserves the fetcher's worker-derived default. native_http_subranges (int, optional): Concurrent HTTP subranges per video range read. @@ -213,7 +214,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): raise ValueError("decode_threads must be positive") if decoded_queue_size <= 0: raise ValueError("decoded_queue_size must be positive") - if max_open_decoders <= 0: + if max_open_decoders is not None and max_open_decoders <= 0: raise ValueError("max_open_decoders must be positive") if native_http_connections is not None and native_http_connections <= 0: raise ValueError("native_http_connections must be positive") @@ -224,7 +225,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): self.byte_budget = int(byte_budget_gb * 1024**3) self.decode_threads = decode_threads self.decoded_queue_size = decoded_queue_size - self.max_open_decoders = max_open_decoders self.native_http_connections = native_http_connections self.native_http_subranges = native_http_subranges self.repeat = repeat @@ -250,6 +250,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): self.meta.rescale_depth_stats(self._depth_output_unit) # Check version check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION) + self.max_open_decoders = ( + max_open_decoders + if max_open_decoders is not None + else max(1, self.episode_pool_size * len(self.meta.video_keys)) + ) self._depth_encoder_configs: dict[str, DepthEncoderConfig] = { vid_key: DepthEncoderConfig.from_video_info(self.meta.features[vid_key].get("info")) diff --git a/tests/configs/test_default.py b/tests/configs/test_default.py index b5c4846c8..c4229303d 100644 --- a/tests/configs/test_default.py +++ b/tests/configs/test_default.py @@ -38,6 +38,10 @@ def test_dataset_config_empty_episodes_ok(): DatasetConfig(repo_id="user/repo", episodes=[]) +def test_dataset_config_derives_streaming_decoder_limit_by_default(): + assert DatasetConfig(repo_id="user/repo").streaming_max_open_decoders is None + + @pytest.mark.parametrize( ("field", "value", "message"), [ diff --git a/tests/datasets/test_streaming_production.py b/tests/datasets/test_streaming_production.py index e876dd9f1..ff244856e 100644 --- a/tests/datasets/test_streaming_production.py +++ b/tests/datasets/test_streaming_production.py @@ -118,6 +118,35 @@ def test_parallel_decode_queue_preserves_planner_order( assert 1 < max_active <= parallel.decode_threads +def test_default_decoder_limit_covers_the_configured_episode_pool( + tmp_path: Path, + lerobot_dataset_factory, +) -> None: + root = tmp_path / "dataset" + lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=2, + total_frames=20, + ) + + streaming = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + episode_pool_size=7, + ) + + assert streaming.max_open_decoders == 7 * len(streaming.meta.video_keys) + + overridden = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + episode_pool_size=7, + max_open_decoders=5, + ) + assert overridden.max_open_decoders == 5 + + @pytest.mark.parametrize("video_backend", ["torchcodec", "pyav"]) def test_streaming_rgb_video_matches_map_style( tmp_path: Path,