feat(dataset): integrate episode streaming into training

scripts/bench_episode_byte_cache.py

@@ -16,7 +16,6 @@ import os
 import random
 import resource
 import socket
-import tempfile
 import threading
 import time
 from collections.abc import Sequence
@@ -30,27 +29,24 @@ import pyarrow.compute as pc
 import pyarrow.parquet as pq
 
 from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
-from lerobot.datasets.episode_video_streaming import (
-    EpisodeByteCache,
-    ExactCoveragePool,
-    EpisodeVideoManifest,
-    NativeHTTPRangeFetcher,
-    assert_hf_hub_range_cache_branch,
-)
 from lerobot.datasets.video_utils import VideoDecoderCache, decode_video_frames_torchcodec
-
-DEFAULT_REPO = "allenai/MolmoAct2-BimanualYAM-Dataset"
-DEFAULT_REVISION = "e9f21ae15074330839f2ac25ed4b49d76dfa1f9c"
-DEFAULT_DATA_ROOT = "hf://buckets/pepijn223/MolmoAct2-BimanualYAM-Dataset-bucket"
-SIDECAR_CACHE_DIR = Path(tempfile.gettempdir()) / "lerobot-sidecars"
-FULL_SIDECAR_NAME = "molmoact2-full.npz"
+from lerobot.streaming.episode_video import (
+    EpisodeByteCache,
+    EpisodeVideoManifest,
+    ExactCoveragePool,
+)
 
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Benchmark episode-level streaming mini-MP4 cache.")
-    parser.add_argument("--repo-id", default=DEFAULT_REPO)
-    parser.add_argument("--revision", default=DEFAULT_REVISION)
-    parser.add_argument("--data-root", default=DEFAULT_DATA_ROOT)
+    parser = argparse.ArgumentParser(description="Profile the episode-level streaming mini-MP4 cache.")
+    parser.add_argument("--repo-id", required=True)
+    parser.add_argument("--revision", default=None)
+    parser.add_argument("--data-root", required=True)
+    parser.add_argument(
+        "--sidecar-path",
+        default=None,
+        help="Optional validated sidecar built by scripts/build_mp4_sidecar.py.",
+    )
     parser.add_argument(
         "--strategy",
         choices=("both", "full", "indexed", "remote-decoder", "native-http"),
@@ -144,7 +140,6 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--in-memory", action="store_true", help="Accepted for compatibility; manifest is always in memory."
     )
-    parser.add_argument("--no-hub-branch-assert", action="store_true")
     return parser.parse_args()
 
 
@@ -461,7 +456,6 @@ def run_exact_coverage_stream(
 
     elapsed = time.perf_counter() - start
     result = {
-        "coverage_mode": "exact",
         "target_samples_s": target_samples_s,
         "actual_samples_s": samples_done / elapsed if elapsed > 0 else float("inf"),
         "stream_wall_s": elapsed,
@@ -584,74 +578,6 @@ def _root_join(data_root: str, relative_path: str) -> str:
     return str(Path(data_root) / relative_path)
 
 
-def _find_or_download_sidecar(data_root: str, manifest_episode_count: int) -> Path | None:
-    _ = manifest_episode_count
-    local = SIDECAR_CACHE_DIR / FULL_SIDECAR_NAME
-    if _valid_sidecar(local):
-        return local
-    if local.exists():
-        print(f"mp4_sidecar_invalid_local: {local}")
-        local.unlink()
-    remote_relative = f"meta/mp4-sidecars/{FULL_SIDECAR_NAME}"
-    remote = _root_join(data_root, remote_relative)
-    protocol = "hf" if data_root.startswith("hf://") else "file"
-    fs = fsspec.filesystem(protocol)
-    if not fs.exists(remote):
-        return None
-    local.parent.mkdir(parents=True, exist_ok=True)
-    print(f"downloading_mp4_sidecar: {remote} -> {local}")
-    if data_root.startswith("hf://"):
-        _download_sidecar_native_http(data_root, remote_relative, local)
-    else:
-        fs.get(remote, str(local))
-    return local
-
-
-def _valid_sidecar(path: Path) -> bool:
-    if not path.exists():
-        return False
-    try:
-        with np.load(path, allow_pickle=False) as data:
-            return "manifest_json" in data
-    except Exception:
-        return False
-
-
-def _download_sidecar_native_http(data_root: str, relative_path: str, local: Path) -> None:
-    fetcher = NativeHTTPRangeFetcher(data_root, max_connections=16)
-    tmp = local.with_suffix(local.suffix + ".tmp")
-    try:
-        size = fetcher.info_size(relative_path)
-        chunk_size = 16 * 1024 * 1024
-        ranges = [(offset, min(chunk_size, size - offset)) for offset in range(0, size, chunk_size)]
-        with tmp.open("wb") as out_file:
-            out_file.truncate(size)
-
-        def read_chunk(offset_length: tuple[int, int]) -> tuple[int, bytes]:
-            offset, length = offset_length
-            return offset, fetcher.read_range(relative_path, offset, length)
-
-        start = time.perf_counter()
-        done = 0
-        with ThreadPoolExecutor(max_workers=8) as pool:
-            futures = [pool.submit(read_chunk, item) for item in ranges]
-            with tmp.open("r+b") as rw_file:
-                for future in futures:
-                    offset, data = future.result()
-                    rw_file.seek(offset)
-                    rw_file.write(data)
-                    done += len(data)
-                    elapsed = max(time.perf_counter() - start, 1e-9)
-                    print(
-                        f"sidecar_download: {done / 1024**2:.1f}/{size / 1024**2:.1f} MiB "
-                        f"({done / elapsed / 1024**2:.1f} MiB/s)",
-                        flush=True,
-                    )
-        tmp.replace(local)
-    finally:
-        fetcher.close()
-
-
 class EpisodeParquetReader:
     def __init__(self, meta: LeRobotDatasetMetadata, data_root: str):
         self.meta = meta
@@ -1369,15 +1295,14 @@ def main() -> None:
         os.environ["LEROBOT_HTTP_FAILURE_LOG"] = args.http_failure_log
         print(f"http_failure_log: {args.http_failure_log}")
     data_root = args.data_root
-    if data_root.startswith("hf://") and not args.no_hub_branch_assert:
-        assert_hf_hub_range_cache_branch()
-
     meta = LeRobotDatasetMetadata(args.repo_id, revision=args.revision)
     meta.ensure_readable()
     parquet_reader = EpisodeParquetReader(meta, data_root)
     manifest_episode_count = args.manifest_episodes or int(meta.total_episodes)
     manifest_episode_count = min(manifest_episode_count, int(meta.total_episodes), args.num_episodes)
-    sidecar_path = _find_or_download_sidecar(data_root, manifest_episode_count)
+    sidecar_path = Path(args.sidecar_path).expanduser() if args.sidecar_path else None
+    if sidecar_path is not None and not sidecar_path.is_file():
+        raise FileNotFoundError(f"MP4 sidecar not found: {sidecar_path}")
 
     if sidecar_path is not None:
         print(f"using_mp4_sidecar: {sidecar_path}")
@@ -1419,15 +1344,7 @@ def main() -> None:
         )
         return
     if args.strategy == "both":
-        expected_sidecar = SIDECAR_CACHE_DIR / FULL_SIDECAR_NAME
-        expected_remote = _root_join(data_root, f"meta/mp4-sidecars/{FULL_SIDECAR_NAME}")
-        print(f"mp4_sidecar_missing_local: {expected_sidecar}")
-        print(f"mp4_sidecar_missing_remote: {expected_remote}")
-        print(
-            "build_mp4_sidecar: "
-            "uv run --no-sync python scripts/build_mp4_sidecar.py "
-            f"--workers {args.workers} --range-backend native-http --output {expected_sidecar}"
-        )
+        print("mp4_sidecar: none (pass --sidecar-path to profile a prebuilt sidecar)")
         print("running_without_mp4_sidecar: indexed variants will build MP4 indexes online")
         print()
 

scripts/bench_streaming_dataset.py

@@ -0,0 +1,179 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+"""Benchmark the production StreamingLeRobotDataset path used by lerobot-train."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import platform
+import resource
+import shutil
+import socket
+import statistics
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+from lerobot.datasets import StreamingLeRobotDataset
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--repo-id", required=True)
+    parser.add_argument("--revision", default=None)
+    parser.add_argument("--root", default=None)
+    parser.add_argument("--data-root", default=None)
+    parser.add_argument("--episodes", type=int, default=None, help="Use the first N episodes.")
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument("--prefetch-factor", type=int, default=2)
+    parser.add_argument("--episode-pool-size", type=int, default=32)
+    parser.add_argument("--prefetch-episodes", type=int, default=8)
+    parser.add_argument("--byte-budget-gb", type=float, default=8.0)
+    parser.add_argument("--warmup-batches", type=int, default=8)
+    parser.add_argument("--measure-batches", type=int, default=128)
+    parser.add_argument("--summary-json", type=Path, default=None)
+    return parser.parse_args()
+
+
+def percentile(values: list[float], quantile: float) -> float:
+    if not values:
+        return 0.0
+    ordered = sorted(values)
+    index = round((len(ordered) - 1) * quantile)
+    return ordered[index]
+
+
+def git_commit() -> str | None:
+    git = shutil.which("git")
+    if git is None:
+        return None
+    try:
+        return subprocess.run(
+            [git, "rev-parse", "HEAD"],
+            check=True,
+            capture_output=True,
+            text=True,
+        ).stdout.strip()
+    except (OSError, subprocess.CalledProcessError):
+        return None
+
+
+def main_process_max_rss_mb() -> float:
+    rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+    return rss / 1024**2 if sys.platform == "darwin" else rss / 1024
+
+
+def child_process_max_rss_mb() -> float:
+    rss = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss
+    return rss / 1024**2 if sys.platform == "darwin" else rss / 1024
+
+
+def main() -> None:
+    args = parse_args()
+    episodes = list(range(args.episodes)) if args.episodes is not None else None
+
+    init_start = time.perf_counter()
+    dataset = StreamingLeRobotDataset(
+        args.repo_id,
+        root=args.root,
+        episodes=episodes,
+        revision=args.revision,
+        data_root=args.data_root,
+        episode_pool_size=args.episode_pool_size,
+        prefetch_episodes=args.prefetch_episodes,
+        byte_budget_gb=args.byte_budget_gb,
+        max_num_shards=max(1, args.num_workers),
+        return_uint8=True,
+    )
+    dataset_init_s = time.perf_counter() - init_start
+
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=torch.cuda.is_available(),
+        prefetch_factor=args.prefetch_factor if args.num_workers else None,
+        persistent_workers=args.num_workers > 0,
+    )
+    iterator = iter(loader)
+    waits: list[float] = []
+    measured_samples = 0
+    measured_indices: list[int] = []
+    first_batch_s = 0.0
+    exhausted = False
+
+    try:
+        for batch_index in range(args.warmup_batches + args.measure_batches):
+            wait_start = time.perf_counter()
+            try:
+                batch = next(iterator)
+            except StopIteration:
+                exhausted = True
+                break
+            wait_s = time.perf_counter() - wait_start
+            if batch_index == 0:
+                first_batch_s = wait_s
+            if batch_index >= args.warmup_batches:
+                waits.append(wait_s)
+                indices = batch["index"].reshape(-1).tolist()
+                measured_indices.extend(int(index) for index in indices)
+                measured_samples += len(indices)
+    finally:
+        shutdown = getattr(iterator, "_shutdown_workers", None)
+        if shutdown is not None:
+            shutdown()
+
+    measured_wall_s = sum(waits)
+    summary = {
+        "repo_id": args.repo_id,
+        "revision": str(dataset.revision),
+        "git_commit": git_commit(),
+        "host": socket.gethostname(),
+        "platform": platform.platform(),
+        "torch_version": torch.__version__,
+        "dataset_init_s": dataset_init_s,
+        "first_batch_s": first_batch_s,
+        "measured_batches": len(waits),
+        "measured_samples": measured_samples,
+        "measured_wall_s": measured_wall_s,
+        "samples_s": measured_samples / measured_wall_s if measured_wall_s else 0.0,
+        "batch_wait_mean_ms": statistics.fmean(waits) * 1000 if waits else 0.0,
+        "batch_wait_p50_ms": percentile(waits, 0.50) * 1000,
+        "batch_wait_p95_ms": percentile(waits, 0.95) * 1000,
+        "batch_wait_p99_ms": percentile(waits, 0.99) * 1000,
+        "duplicate_indices": measured_samples - len(set(measured_indices)),
+        "epoch_exhausted": exhausted,
+        "main_process_max_rss_mb": main_process_max_rss_mb(),
+        "worker_process_max_rss_mb": child_process_max_rss_mb(),
+        "config": {
+            "batch_size": args.batch_size,
+            "num_workers": args.num_workers,
+            "prefetch_factor": args.prefetch_factor,
+            "episode_pool_size": args.episode_pool_size,
+            "prefetch_episodes": args.prefetch_episodes,
+            "byte_budget_gb": args.byte_budget_gb,
+            "warmup_batches": args.warmup_batches,
+            "measure_batches": args.measure_batches,
+        },
+    }
+    print(json.dumps(summary, indent=2, sort_keys=True))
+    if args.summary_json is not None:
+        args.summary_json.parent.mkdir(parents=True, exist_ok=True)
+        args.summary_json.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n")
+
+
+if __name__ == "__main__":
+    main()

scripts/build_mp4_sidecar.py

@@ -17,76 +17,78 @@ from pathlib import Path
 import fsspec
 
 from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
-from lerobot.datasets.episode_video_streaming import EpisodeVideoManifest, assert_hf_hub_range_cache_branch
-
-DEFAULT_REPO = "allenai/MolmoAct2-BimanualYAM-Dataset"
-DEFAULT_REVISION = "e9f21ae15074330839f2ac25ed4b49d76dfa1f9c"
-DEFAULT_DATA_ROOT = "hf://buckets/pepijn223/MolmoAct2-BimanualYAM-Dataset-bucket"
+from lerobot.datasets.streaming_sidecar import (
+    build_mp4_sidecar,
+    make_sidecar_spec,
+    published_sidecar_url,
+    range_backend_for_root,
+)
+from lerobot.streaming.sidecar import SidecarSpec
 
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Build a reusable MP4 byte-index sidecar for streaming.")
-    parser.add_argument("--repo-id", default=DEFAULT_REPO)
-    parser.add_argument("--revision", default=DEFAULT_REVISION)
-    parser.add_argument("--data-root", default=DEFAULT_DATA_ROOT)
+    parser.add_argument("--repo-id", required=True)
+    parser.add_argument("--revision", default=None)
+    parser.add_argument("--data-root", required=True)
     parser.add_argument("--output", required=True)
     parser.add_argument("--episodes", type=int, default=None)
     parser.add_argument("--workers", type=int, default=8)
-    parser.add_argument("--range-backend", choices=("fsspec", "native-http"), default="native-http")
+    parser.add_argument("--range-backend", choices=("fsspec", "native-http"), default=None)
     parser.add_argument("--max-probe-mb", type=int, default=64)
-    parser.add_argument(
-        "--no-push", action="store_true", help="Do not upload the sidecar to data_root/meta/mp4-sidecars."
-    )
-    parser.add_argument("--no-hub-branch-assert", action="store_true")
+    parser.add_argument("--push", action="store_true", help="Explicitly publish the sidecar to data_root.")
     return parser.parse_args()
 
 
-def push_sidecar(local_path: str, data_root: str) -> list[str]:
-    if not data_root.startswith("hf://"):
-        return []
+def push_sidecar(local_path: str, spec: SidecarSpec) -> list[str]:
+    if not spec.data_root.startswith("hf://"):
+        raise ValueError("--push currently supports only hf:// data roots")
 
-    local = Path(local_path)
     fs = fsspec.filesystem("hf")
-    remote_dir = f"{data_root.rstrip('/')}/meta/mp4-sidecars"
-    remote_paths = [f"{remote_dir}/{local.name}"]
-
-    for remote in remote_paths:
-        fs.put(str(local), remote)
-    return remote_paths
+    remote = published_sidecar_url(spec)
+    fs.put(str(Path(local_path)), remote)
+    return [remote]
 
 
 def main() -> None:
     args = parse_args()
-    if args.data_root.startswith("hf://") and not args.no_hub_branch_assert:
-        assert_hf_hub_range_cache_branch()
 
     meta = LeRobotDatasetMetadata(args.repo_id, revision=args.revision)
     meta.ensure_readable()
     total = (
         int(meta.total_episodes) if args.episodes is None else min(args.episodes, int(meta.total_episodes))
     )
-    rel_paths = sorted(
-        {str(meta.get_video_file_path(ep_idx, key)) for ep_idx in range(total) for key in meta.video_keys}
-    )
+    spec = make_sidecar_spec(meta, args.data_root)
+    if total != int(meta.total_episodes):
+        selected_paths = {
+            str(meta.get_video_file_path(ep_idx, key)) for ep_idx in range(total) for key in meta.video_keys
+        }
+        spec = SidecarSpec(
+            repo_id=spec.repo_id,
+            revision=spec.revision,
+            data_root=spec.data_root,
+            source_files=tuple(item for item in spec.source_files if item[0] in selected_paths),
+        )
 
     start = time.perf_counter()
-    EpisodeVideoManifest.write_file_sidecar(
+    build_mp4_sidecar(
         args.output,
-        rel_paths,
-        args.data_root,
-        range_backend=args.range_backend,
+        spec,
+        range_backend=args.range_backend or range_backend_for_root(args.data_root),
         workers=args.workers,
         max_probe_bytes=args.max_probe_mb * 1024 * 1024,
     )
     elapsed = time.perf_counter() - start
     print(f"wrote {args.output}")
-    print(f"episodes={total} files={len(rel_paths)} elapsed_s={elapsed:.2f}")
-    if args.no_push:
-        print("push_skipped: --no-push")
-    else:
-        pushed = push_sidecar(args.output, args.data_root)
+    print(f"episodes={total} files={len(spec.source_files)} elapsed_s={elapsed:.2f}")
+    if args.push:
+        if total != int(meta.total_episodes):
+            raise ValueError("Only a complete dataset sidecar can be published")
+        pushed = push_sidecar(args.output, spec)
         for remote in pushed:
             print(f"pushed {remote}")
+    else:
+        print("push_skipped: pass --push for explicit publication")
 
 
 if __name__ == "__main__":