Remove streaming benchmark prototypes

2026-07-28 20:26:05 +00:00 · 2026-07-28 11:14:10 +02:00
parent 4db03ed535
commit fdf6014214
7 changed files with 0 additions and 1742 deletions
@@ -172,22 +172,6 @@ lerobot-train \
 Use `--dataset.streaming_data_root=hf://buckets/OWNER/BUCKET/PREFIX` when metadata lives in a
 dataset repository but Parquet and MP4 data are mirrored in an HF Bucket.
 To capture comparable data-pipeline results on a training host:
 ```bash
 uv run python scripts/bench_streaming_dataset.py \
  --repo-id=yaak-ai/L2D-v3 \
  --batch-size=16 \
  --num-workers=1 \
  --fetch-workers=4 \
  --decode-threads=2 \
  --summary-json=streaming-benchmark.json
 ```
 The JSON records the code and dataset revisions, initialization/first-batch time, steady-state
 samples per second, batch-wait p50/p95, duplicate indices, and the exact cache/worker settings.
 Run end-to-end `lerobot-train` separately to measure GPU utilization and full training step time.
 The Python API uses the same implementation:
 ```python
@@ -105,32 +105,6 @@ Streaming checkpoints created by the earlier multi-worker sampler record a diffe
 topology and are rejected for sample-exact resume. Start a new run from the saved policy weights
 rather than claiming that its dataset stream resumes exactly.
 ## Benchmarking
 Measure the integrated path on the same hosts used for training:
 ```bash
 uv run python scripts/bench_streaming_dataset.py \
  --repo-id=OWNER/DATASET \
  --revision=COMMIT_SHA \
  --batch-size=16 \
  --num-workers=1 \
  --fetch-workers=4 \
  --decode-threads=2 \
  --decoded-queue-size=8 \
  --summary-json=streaming-benchmark.json
 ```
 The output records source revisions, startup and first-batch latency, steady-state throughput,
 batch-wait percentiles, duplicate indices, memory high-water marks, and the exact settings. Run a
 separate end-to-end training A/B to include policy compute, device transfer, and optimizer time.
 Do not compare results unless the code revision, dataset revision, hardware, and settings match.
 For lower-level byte-fetch, retry, decoder-open, and refill diagnostics, build the sidecar explicitly
 and run `scripts/bench_episode_byte_cache.py` with the same `--repo-id`, `--revision`,
 `--data-root`, and `--sidecar-path`. Treat this as a stage profiler; throughput claims should come
 from the integrated dataset benchmark and end-to-end training A/B.
 ## Troubleshooting
 - **The first batch takes a long time:** check logs for a sidecar build. Reuse a shared cache or
@@ -1,214 +0,0 @@
 #!/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 collections.abc import Sequence
 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,
        choices=(0, 1),
        default=1,
        help="Rank-level DataLoader process count. Use 1 for the production pipeline.",
    )
    parser.add_argument("--fetch-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("--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=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)
    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: Sequence[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()
    if args.fetch_workers <= 0:
        raise ValueError("--fetch-workers must be positive")
    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,
        decode_threads=args.decode_threads,
        decoded_queue_size=args.decoded_queue_size,
        max_open_decoders=args.max_open_decoders,
        native_http_connections=args.native_http_connections,
        native_http_subranges=args.native_http_subranges,
        max_num_shards=max(1, args.fetch_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] = []
    unique_episodes_per_batch: 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()
                episode_indices = batch["episode_index"].reshape(-1).tolist()
                measured_indices.extend(int(index) for index in indices)
                unique_episodes_per_batch.append(len({int(index) for index in episode_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)),
        "unique_episodes_per_batch_mean": (
            statistics.fmean(unique_episodes_per_batch) if unique_episodes_per_batch else 0.0
        ),
        "unique_episodes_per_batch_p50": percentile(unique_episodes_per_batch, 0.50),
        "unique_episodes_per_batch_p95": percentile(unique_episodes_per_batch, 0.95),
        "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,
            "fetch_workers": args.fetch_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,
            "decode_threads": args.decode_threads,
            "decoded_queue_size": args.decoded_queue_size,
            "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,
            "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()
@@ -1,65 +0,0 @@
 #!/usr/bin/env python
 from __future__ import annotations
 import argparse
 import json
 from pathlib import Path
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Summarize distributed episode pool benchmark JSON files.")
    parser.add_argument("summaries", nargs="+", help="Rank summary JSON files.")
    return parser.parse_args()
 def _load(path: str) -> dict:
    return json.loads(Path(path).read_text())
 def _fmt(value: float) -> str:
    return f"{value:.1f}"
 def main() -> None:
    args = parse_args()
    rows = [_load(path) for path in args.summaries]
    rows.sort(key=lambda row: int(row.get("distributed_shard_index", 0)))
    total_bytes = sum(float(row.get("fetch_bytes", 0.0)) for row in rows)
    max_fetch_s = max(float(row.get("fetch_s", 0.0)) for row in rows)
    aggregate_mib_s = total_bytes / max_fetch_s / 1024**2 if max_fetch_s > 0 else float("inf")
    summed_rank_mib_s = sum(float(row.get("fetch_mib_s", 0.0)) for row in rows)
    total_decode_samples_s = sum(float(row.get("pool_decode_training_samples_s", 0.0)) for row in rows)
    total_stream_samples_s = sum(float(row.get("pool_stream_actual_samples_s", 0.0)) for row in rows)
    kept_up = all(bool(row.get("pool_stream_kept_up", 0.0)) for row in rows)
    print("| Aggregate | value |")
    print("|---|---:|")
    print(f"| ranks | {len(rows)} |")
    print(f"| total fetched GiB | {total_bytes / 1024**3:.2f} |")
    print(f"| aggregate fetch MiB/s | {_fmt(aggregate_mib_s)} |")
    print(f"| summed rank fetch MiB/s | {_fmt(summed_rank_mib_s)} |")
    if total_decode_samples_s:
        print(f"| aggregate resident decode samples/s | {_fmt(total_decode_samples_s)} |")
    if total_stream_samples_s:
        print(f"| aggregate stream samples/s | {_fmt(total_stream_samples_s)} |")
        print(f"| all ranks kept up | {'yes' if kept_up else 'no'} |")
    print()
    print("| Rank | host | fetch MiB/s | fetch s | GiB | decode samples/s | stream samples/s | kept up |")
    print("|---:|---|---:|---:|---:|---:|---:|---|")
    for row in rows:
        rank = int(row.get("distributed_shard_index", 0))
        print(
            f"| {rank} | {row.get('hostname', '')} | "
            f"{_fmt(float(row.get('fetch_mib_s', 0.0)))} | "
            f"{_fmt(float(row.get('fetch_s', 0.0)))} | "
            f"{float(row.get('fetch_gib', 0.0)):.2f} | "
            f"{_fmt(float(row.get('pool_decode_training_samples_s', 0.0)))} | "
            f"{_fmt(float(row.get('pool_stream_actual_samples_s', 0.0)))} | "
            f"{'yes' if row.get('pool_stream_kept_up', 0.0) else 'no'} |"
        )
 if __name__ == "__main__":
    main()
@@ -1,36 +0,0 @@
 # 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
 """Compatibility imports for the split training-time episode streaming stack.
 New internal code should import from `episode_cache`, `episode_pool`, `manifest`, or
 `range_fetch` directly. This module keeps the original prototype import path working.
 """
 from lerobot.streaming.episode_cache import EpisodeByteCache, open_video_decoder
 from lerobot.streaming.episode_pool import ExactCoveragePool
 from lerobot.streaming.manifest import EpisodeVideoManifest, EpisodeVideoSpan, VideoFileRecord
 from lerobot.streaming.range_fetch import (
    NativeHTTPRangeFetcher,
    ThreadLocalRangeFetcher,
    _log_http_failure as _log_http_failure,
    make_range_fetcher,
 )
 __all__ = [
    "EpisodeByteCache",
    "EpisodeVideoManifest",
    "EpisodeVideoSpan",
    "ExactCoveragePool",
    "NativeHTTPRangeFetcher",
    "ThreadLocalRangeFetcher",
    "VideoFileRecord",
    "make_range_fetcher",
    "open_video_decoder",
 ]
@@ -39,15 +39,3 @@ from lerobot.streaming.mp4 import Mp4Index
    )
    assert result.returncode == 0, result.stderr
 def test_episode_video_compatibility_imports() -> None:
    from lerobot.streaming.episode_cache import EpisodeByteCache
    from lerobot.streaming.episode_pool import ExactCoveragePool
    from lerobot.streaming.episode_video import (
        EpisodeByteCache as CompatibilityEpisodeByteCache,
        ExactCoveragePool as CompatibilityExactCoveragePool,
    )
    assert CompatibilityEpisodeByteCache is EpisodeByteCache
    assert CompatibilityExactCoveragePool is ExactCoveragePool