admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-28 12:15:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

Remove streaming benchmark prototypes

Browse Source
This commit is contained in:
Pepijn
2026-07-28 11:14:10 +02:00
parent 4db03ed535
commit fdf6014214
7 changed files with 0 additions and 1742 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/lerobot-dataset-v3.mdx
-16
View File
@@ -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
docs/source/training_dataset_streaming.mdx
-26
View File
@@ -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
scripts/bench_episode_byte_cache.py
-1373
View File
File diff suppressed because it is too large Load Diff
scripts/bench_streaming_dataset.py
-214
View File
@@ -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()
scripts/summarize_episode_pool_bench.py
-65
View File
@@ -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()
src/lerobot/streaming/episode_video.py
-36
View File
@@ -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",
]
tests/test_streaming_core_imports.py
-12
View File
@@ -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