From fdf6014214dc24e6a1e1bd8dc3a685e1b76f4f11 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Tue, 28 Jul 2026 11:14:10 +0200 Subject: [PATCH] Remove streaming benchmark prototypes --- docs/source/lerobot-dataset-v3.mdx | 16 - docs/source/training_dataset_streaming.mdx | 26 - scripts/bench_episode_byte_cache.py | 1373 -------------------- scripts/bench_streaming_dataset.py | 214 --- scripts/summarize_episode_pool_bench.py | 65 - src/lerobot/streaming/episode_video.py | 36 - tests/test_streaming_core_imports.py | 12 - 7 files changed, 1742 deletions(-) delete mode 100644 scripts/bench_episode_byte_cache.py delete mode 100644 scripts/bench_streaming_dataset.py delete mode 100644 scripts/summarize_episode_pool_bench.py delete mode 100644 src/lerobot/streaming/episode_video.py diff --git a/docs/source/lerobot-dataset-v3.mdx b/docs/source/lerobot-dataset-v3.mdx index b63b2a08d..91571e4cd 100644 --- a/docs/source/lerobot-dataset-v3.mdx +++ b/docs/source/lerobot-dataset-v3.mdx @@ -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 diff --git a/docs/source/training_dataset_streaming.mdx b/docs/source/training_dataset_streaming.mdx index 7cd053ce6..1eb0483aa 100644 --- a/docs/source/training_dataset_streaming.mdx +++ b/docs/source/training_dataset_streaming.mdx @@ -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 diff --git a/scripts/bench_episode_byte_cache.py b/scripts/bench_episode_byte_cache.py deleted file mode 100644 index 73ceb5468..000000000 --- a/scripts/bench_episode_byte_cache.py +++ /dev/null @@ -1,1373 +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 - -from __future__ import annotations - -import argparse -import json -import os -import random -import resource -import socket -import threading -import time -from collections.abc import Sequence -from concurrent.futures import ThreadPoolExecutor -from pathlib import Path - -import fsspec -import numpy as np -import pyarrow as pa -import pyarrow.compute as pc -import pyarrow.parquet as pq - -from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata -from lerobot.datasets.video_utils import VideoDecoderCache, decode_video_frames_torchcodec -from lerobot.streaming.episode_cache import EpisodeByteCache -from lerobot.streaming.episode_pool import ExactCoveragePool -from lerobot.streaming.manifest import EpisodeVideoManifest - - -def parse_args() -> argparse.Namespace: - 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"), - default="both", - help=argparse.SUPPRESS, - ) - parser.add_argument( - "--range-backend", - choices=("fsspec", "native-http"), - default="fsspec", - help="Range reader used by indexed/full episode-pool fetch tracks.", - ) - parser.add_argument("--num-episodes", type=int, default=512) - parser.add_argument( - "--manifest-episodes", - type=int, - default=None, - help="Limit manifest construction to the first N episodes for local smoke tests.", - ) - parser.add_argument("--pool-size", type=int, default=16) - parser.add_argument( - "--workers", - type=int, - default=8, - help="Concurrent camera-fetch jobs. Total connections ~= workers x range-subranges; " - "measure the host-specific saturation point before increasing this value.", - ) - parser.add_argument( - "--range-subranges", - type=int, - default=1, - help="Split each camera byte-range GET into N concurrent sub-range GETs (native-http only). " - "Divides per-episode latency by ~N under the per-host throughput ceiling.", - ) - parser.add_argument( - "--native-http-connections", - type=int, - default=None, - help="Max HTTP connections for --range-backend native-http. Defaults to --workers.", - ) - parser.add_argument( - "--native-http-retries", - type=int, - default=8, - help="Retries per native HTTP range request.", - ) - parser.add_argument( - "--native-http-timeout", - type=float, - default=120.0, - help="Timeout in seconds for native HTTP requests.", - ) - parser.add_argument( - "--progress-interval", - type=float, - default=10.0, - help="Print episode-pool fill progress every N seconds. Set 0 to disable.", - ) - parser.add_argument( - "--http-failure-log", - default=None, - help="Optional JSONL file for failed/retried HTTP range attempts.", - ) - parser.add_argument( - "--include-decode", - action="store_true", - help="Also run decoder-opening/frame-decode comparison tracks. Fetch-only is the default.", - ) - parser.add_argument("--include-pool-sampling", action="store_true") - parser.add_argument("--pool-random-samples", type=int, default=4096) - parser.add_argument("--batch-size", type=int, default=512) - parser.add_argument("--target-samples-s", type=float, default=500.0) - parser.add_argument("--stream-samples", type=int, default=4096) - parser.add_argument("--stream-prefetch-episodes", type=int, default=16) - parser.add_argument("--decode-workers", type=int, default=1) - parser.add_argument("--prefetch-ahead", type=int, default=8) - parser.add_argument("--frames-per-episode", type=int, default=16) - parser.add_argument("--max-probe-mb", type=int, default=64) - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--byte-budget-gb", type=float, default=80) - parser.add_argument("--distributed-shard-count", type=int, default=1) - parser.add_argument("--distributed-shard-index", type=int, default=0) - parser.add_argument("--summary-json", default=None) - return parser.parse_args() - - -def _episode_shard( - total: int, - requested: int, - seed: int, - *, - shard_count: int = 1, - shard_index: int = 0, -) -> list[int]: - rng = random.Random(seed) - upper = min(total, requested) - if shard_count < 1: - raise ValueError(f"distributed-shard-count must be >= 1, got {shard_count}") - if shard_index < 0 or shard_index >= shard_count: - raise ValueError(f"distributed-shard-index must be in [0, {shard_count}), got {shard_index}") - permutation = list(range(upper)) - rng.shuffle(permutation) - return permutation[shard_index::shard_count] - - -def _episode_pool( - total: int, - requested: int, - pool_size: int, - seed: int, - *, - shard_count: int = 1, - shard_index: int = 0, -) -> list[int]: - shard = _episode_shard( - total, - requested, - seed, - shard_count=shard_count, - shard_index=shard_index, - ) - if pool_size > len(shard): - raise ValueError( - f"pool-size={pool_size} exceeds shard episodes={len(shard)} for shard {shard_index}/{shard_count}" - ) - return shard[:pool_size] - - -def _timestamps(manifest: EpisodeVideoManifest, episodes: Sequence[int], frames_per_episode: int, seed: int): - rng = random.Random(seed) - out: dict[tuple[int, str], list[float]] = {} - for ep in episodes: - for camera_key in manifest.video_keys: - span = manifest.lookup(ep, camera_key) - lo = span.first_pts - hi = max(span.last_pts, lo) - out[(ep, camera_key)] = sorted(rng.uniform(lo, hi) for _ in range(frames_per_episode)) - return out - - -def _timestamps_from_meta( - meta: LeRobotDatasetMetadata, episodes: Sequence[int], frames_per_episode: int, seed: int -) -> dict[tuple[int, str], list[float]]: - rng = random.Random(seed) - out: dict[tuple[int, str], list[float]] = {} - for ep in episodes: - row = meta.episodes[ep] - for camera_key in meta.video_keys: - lo = float(row[f"videos/{camera_key}/from_timestamp"]) - hi = max(float(row[f"videos/{camera_key}/to_timestamp"]), lo) - out[(ep, camera_key)] = sorted(rng.uniform(lo, hi) for _ in range(frames_per_episode)) - return out - - -def _bytes_for(manifest: EpisodeVideoManifest, episodes: Sequence[int]) -> int: - total = 0 - for ep in episodes: - for camera_key in manifest.video_keys: - total += manifest.lookup(ep, camera_key).mdat_length - return total - - -def _random_training_samples(episodes: Sequence[int], count: int, seed: int) -> list[tuple[int, float]]: - rng = random.Random(seed) - out = [] - for _ in range(count): - ep = rng.choice(episodes) - out.append((ep, rng.random())) - return out - - -def _sampling_randomness(samples: Sequence[tuple[int, float]], *, batch_size: int) -> dict[str, float]: - if not samples: - return { - "sample_count": 0.0, - "unique_episodes": 0.0, - "unique_episode_fraction": 0.0, - "mean_samples_per_used_episode": 0.0, - "max_samples_per_episode": 0.0, - "mean_unique_episodes_per_batch": 0.0, - "min_unique_episodes_per_batch": 0.0, - } - counts: dict[int, int] = {} - for ep, _ts in samples: - counts[ep] = counts.get(ep, 0) + 1 - batch_uniques = [ - len({ep for ep, _ts in samples[idx : idx + batch_size]}) - for idx in range(0, len(samples), batch_size) - if samples[idx : idx + batch_size] - ] - return { - "sample_count": float(len(samples)), - "unique_episodes": float(len(counts)), - "unique_episode_fraction": len(counts) / len(samples), - "mean_samples_per_used_episode": len(samples) / len(counts), - "max_samples_per_episode": float(max(counts.values())), - "mean_unique_episodes_per_batch": float(np.mean(batch_uniques)), - "min_unique_episodes_per_batch": float(min(batch_uniques)), - } - - -def _decode_all( - cache: EpisodeByteCache, timestamps: dict[tuple[int, str], list[float]], *, decode_workers: int -) -> float: - start = time.perf_counter() - items = list(timestamps.items()) - if decode_workers <= 1: - for (ep, camera_key), ts in items: - cache.get_frames(ep, camera_key, ts) - else: - with ThreadPoolExecutor(max_workers=decode_workers) as pool: - futures = [pool.submit(cache.get_frames, ep, camera_key, ts) for (ep, camera_key), ts in items] - for future in futures: - future.result() - return time.perf_counter() - start - - -def _decoder_locks( - manifest: EpisodeVideoManifest, episodes: Sequence[int] -) -> dict[tuple[int, str], threading.Lock]: - return {(ep, camera_key): threading.Lock() for ep in episodes for camera_key in manifest.video_keys} - - -def _open_resident_decoders( - cache: EpisodeByteCache, episodes: Sequence[int], *, decode_workers: int -) -> tuple[float, int]: - items = [(ep, camera_key) for ep in episodes for camera_key in cache.manifest.video_keys] - start = time.perf_counter() - if decode_workers <= 1: - for ep, camera_key in items: - cache.get_decoder(ep, camera_key) - else: - with ThreadPoolExecutor(max_workers=decode_workers) as pool: - futures = [pool.submit(cache.get_decoder, ep, camera_key) for ep, camera_key in items] - for future in futures: - future.result() - return time.perf_counter() - start, len(items) - - -def _decode_training_sample( - cache: EpisodeByteCache, - episode_index: int, - relative_t: float, - locks: dict[tuple[int, str], threading.Lock], -) -> None: - for camera_key in cache.manifest.video_keys: - span = cache.manifest.lookup(episode_index, camera_key) - timestamp = span.first_pts + relative_t * max(span.last_pts - span.first_pts, 0.0) - with locks[(episode_index, camera_key)]: - cache.get_frames(episode_index, camera_key, [timestamp]) - - -def run_pool_random_decode( - cache: EpisodeByteCache, - episodes: Sequence[int], - *, - sample_count: int, - batch_size: int, - decode_workers: int, - seed: int, -) -> dict[str, float]: - samples = _random_training_samples(episodes, sample_count, seed) - touched_episodes = sorted({ep for ep, _ts in samples}) - decoder_open_s, decoder_count = _open_resident_decoders( - cache, touched_episodes, decode_workers=decode_workers - ) - locks = _decoder_locks(cache.manifest, touched_episodes) - - start = time.perf_counter() - if decode_workers <= 1: - for ep, ts in samples: - _decode_training_sample(cache, ep, ts, locks) - else: - with ThreadPoolExecutor(max_workers=decode_workers) as pool: - futures = [pool.submit(_decode_training_sample, cache, ep, ts, locks) for ep, ts in samples] - for future in futures: - future.result() - decode_s = time.perf_counter() - start - - randomness = _sampling_randomness(samples, batch_size=batch_size) - camera_frames = sample_count * len(cache.manifest.video_keys) - result = { - "decoder_open_s": decoder_open_s, - "decoder_count": float(decoder_count), - "decoder_open_ms": decoder_open_s * 1000 / max(decoder_count, 1), - "decode_s": decode_s, - "training_samples_s": sample_count / decode_s if decode_s > 0 else float("inf"), - "camera_frames_s": camera_frames / decode_s if decode_s > 0 else float("inf"), - "decode_ms_sample": decode_s * 1000 / max(sample_count, 1), - "decode_ms_camera_frame": decode_s * 1000 / max(camera_frames, 1), - } - result.update(randomness) - return result - - -def run_exact_coverage_stream( - cache: EpisodeByteCache, - shard_episodes: Sequence[int], - *, - pool_size: int, - sample_count: int, - target_samples_s: float, - prefetch_ahead: int, - batch_size: int, - decode_workers: int, - seed: int, -) -> dict[str, float]: - """Streaming keep-up with EXACT, exactly-once frame coverage (a real epoch). - - Drives the resident pool from ExactCoveragePool: every frame of every shard episode is decoded - exactly once, at most ``pool_size`` episodes resident, deterministic from ``seed``. Episodes are - prefetched ``prefetch_ahead`` beyond the sampling frontier so a freshly admitted episode's bytes - are already resident when it becomes eligible to be drawn. - """ - manifest = cache.manifest - cams = manifest.video_keys - frame_counts = [ - (int(ep), min(int(manifest.lookup(ep, c).frame_count) for c in cams)) for ep in shard_episodes - ] - total_frames = sum(n for _, n in frame_counts) - pool = ExactCoveragePool(frame_counts, pool_size, seed=seed) - counts = dict(frame_counts) - order = pool.admission_order - locks: dict[tuple[int, str], threading.Lock] = {} - - prefetch_frontier = 0 - - def prefetch_upto(idx: int) -> None: - nonlocal prefetch_frontier - limit = min(idx, len(order)) - while prefetch_frontier < limit: - cache.submit_prefetch(order[prefetch_frontier]) - prefetch_frontier += 1 - - refill_wait_s = 0.0 - - def make_ready(ep: int) -> None: - nonlocal refill_wait_s - wait_start = time.perf_counter() - cache.ensure_ready(ep) - refill_wait_s += time.perf_counter() - wait_start - _open_resident_decoders(cache, [ep], decode_workers=decode_workers) - for c in cams: - locks[(ep, c)] = threading.Lock() - - prefetch_upto(pool_size + prefetch_ahead) - for ep in pool.newly_admitted: - make_ready(ep) - pool.newly_admitted.clear() - - decode_pool = ThreadPoolExecutor(max_workers=decode_workers) if decode_workers > 1 else None - deadline_miss_s = 0.0 - samples_done = 0 - decoded_samples: list[tuple[int, float]] = [] - epoch_complete = False - start = time.perf_counter() - try: - while samples_done < sample_count and not epoch_complete: - batch_start = time.perf_counter() - current_batch_size = min(batch_size, sample_count - samples_done) - batch: list[tuple[int, float]] = [] - for _ in range(current_batch_size): - try: - ep, frame_index = next(pool) - except StopIteration: - epoch_complete = True - break - n = counts[ep] - batch.append((ep, frame_index / max(n - 1, 1))) - if pool.newly_admitted: - for new_ep in pool.newly_admitted: - make_ready(new_ep) - pool.newly_admitted.clear() - prefetch_upto(pool.admitted_count + prefetch_ahead) - if not batch: - break - if decode_pool is not None: - futures = [ - decode_pool.submit(_decode_training_sample, cache, ep, rel, locks) for ep, rel in batch - ] - for future in futures: - future.result() - else: - for ep, rel in batch: - _decode_training_sample(cache, ep, rel, locks) - decoded_samples.extend(batch) - samples_done += len(batch) - target_batch_s = len(batch) / target_samples_s if target_samples_s > 0 else 0.0 - batch_elapsed = time.perf_counter() - batch_start - if target_batch_s > 0 and batch_elapsed < target_batch_s: - time.sleep(target_batch_s - batch_elapsed) - elif target_batch_s > 0: - deadline_miss_s += batch_elapsed - target_batch_s - finally: - if decode_pool is not None: - decode_pool.shutdown(wait=True) - - elapsed = time.perf_counter() - start - replacements = max(0, pool.admitted_count - min(pool_size, len(order))) - result = { - "target_samples_s": target_samples_s, - "actual_samples_s": samples_done / elapsed if elapsed > 0 else float("inf"), - "stream_wall_s": elapsed, - "refill_wait_s": refill_wait_s, - "deadline_miss_s": deadline_miss_s, - "samples_done": float(samples_done), - "shard_total_frames": float(total_frames), - "epoch_complete": 1.0 if epoch_complete else 0.0, - "prefetch_ahead": float(prefetch_ahead), - "prefetch_episodes": float(prefetch_ahead), - "replacements": float(replacements), - "replacement_episodes_s": replacements / elapsed if elapsed > 0 else float("inf"), - "samples_per_episode": samples_done / replacements if replacements else float(samples_done), - "batch_size": float(batch_size), - "decode_workers": float(decode_workers), - "kept_up": 1.0 - if samples_done / elapsed >= target_samples_s * 0.98 and deadline_miss_s < elapsed * 0.02 - else 0.0, - } - result.update( - {f"stream_{k}": v for k, v in _sampling_randomness(decoded_samples, batch_size=batch_size).items()} - ) - return result - - -def _fill_cache( - cache: EpisodeByteCache, episodes: Sequence[int], *, progress_interval: float = 10.0 -) -> float: - start = time.perf_counter() - for ep in episodes: - cache.submit_prefetch(ep) - last_progress = start - for idx, ep in enumerate(episodes, start=1): - cache.ensure_ready(ep) - now = time.perf_counter() - if progress_interval > 0 and now - last_progress >= progress_interval: - timings = cache.timing_summary() - byte_count = timings.get("range_bytes", 0.0) - elapsed = max(now - start, 1e-9) - jobs = timings.get("jobs", 0.0) - total_jobs = len(episodes) * len(cache.manifest.video_keys) - _log( - "fill_progress: " - f"episodes_ready={idx}/{len(episodes)} " - f"camera_jobs={jobs:.0f}/{total_jobs} " - f"fetched={byte_count / 1024**3:.2f} GiB " - f"fetch={byte_count / elapsed / 1024**2:.1f} MiB/s " - f"elapsed={_format_duration(elapsed)}" - ) - last_progress = now - return time.perf_counter() - start - - -def _samples_per_s(elapsed_s: float, episodes: Sequence[int], frames_per_episode: int) -> float: - if elapsed_s <= 0: - return float("inf") - return len(episodes) * frames_per_episode / elapsed_s - - -def _log(message: str) -> None: - print(message, flush=True) - - -def _format_duration(seconds: float) -> str: - if seconds < 60: - return f"{seconds:.1f}s" - if seconds < 3600: - return f"{seconds / 60:.1f}m" - return f"{seconds / 3600:.1f}h" - - -def _current_rss_mib() -> float | None: - status_path = Path("/proc/self/status") - if not status_path.exists(): - return None - for line in status_path.read_text().splitlines(): - if line.startswith("VmRSS:"): - return float(line.split()[1]) / 1024 - return None - - -def _peak_rss_mib() -> float: - rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - # Linux reports KiB even for very large processes; macOS reports bytes. - if Path("/proc/self/status").exists(): - return rss / 1024 - return rss / 1024**2 - - -def _memory_snapshot() -> dict[str, float | None]: - return {"rss_mib": _current_rss_mib(), "peak_rss_mib": _peak_rss_mib()} - - -def _print_memory_summary(start: dict[str, float | None], end: dict[str, float | None]) -> None: - start_rss = start["rss_mib"] - end_rss = end["rss_mib"] - delta = None if start_rss is None or end_rss is None else end_rss - start_rss - print() - print("| Memory | MiB |") - print("|---|---:|") - if start_rss is not None: - print(f"| rss start | {start_rss:.1f} |") - if end_rss is not None: - print(f"| rss end | {end_rss:.1f} |") - if delta is not None: - print(f"| rss delta | {delta:.1f} |") - print(f"| peak rss | {end['peak_rss_mib']:.1f} |") - - -def _write_summary_json(path: str | None, payload: dict) -> None: - if path is None: - return - out = Path(path).expanduser() - out.parent.mkdir(parents=True, exist_ok=True) - tmp = out.with_suffix(out.suffix + ".tmp") - tmp.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") - tmp.replace(out) - print(f"summary_json: {out}") - - -def _root_join(data_root: str, relative_path: str) -> str: - if data_root.startswith("hf://"): - return f"{data_root.rstrip('/')}/{relative_path}" - return str(Path(data_root) / relative_path) - - -class EpisodeParquetReader: - def __init__(self, meta: LeRobotDatasetMetadata, data_root: str): - self.meta = meta - self.data_root = data_root - protocol = "hf" if data_root.startswith("hf://") else "file" - self.fs = fsspec.filesystem(protocol) - self._episode_row_groups = self._build_episode_row_groups() - self._table_cache: dict[str, pa.Table] = {} - self._cache_lock = threading.Lock() - - def read_episode(self, episode_index: int) -> None: - relative_path = str(self.meta.get_data_file_path(episode_index)) - table = self._read_table(relative_path) - table.filter(pc.equal(table["episode_index"], episode_index)) - - def _read_table(self, relative_path: str) -> pa.Table: - with self._cache_lock: - table = self._table_cache.get(relative_path) - if table is not None: - return table - with self.fs.open( - _root_join(self.data_root, relative_path), "rb", block_size=2**20, cache_type="none" - ) as f: - table = pq.ParquetFile(f).read() - with self._cache_lock: - return self._table_cache.setdefault(relative_path, table) - - def submit_read_episode(self, pool: ThreadPoolExecutor, episode_index: int): - return pool.submit(self.read_episode, episode_index) - - def read_episodes(self, episodes: Sequence[int], *, workers: int) -> float: - start = time.perf_counter() - if workers <= 1: - for ep in episodes: - self.read_episode(ep) - else: - with ThreadPoolExecutor(max_workers=workers) as pool: - futures = [pool.submit(self.read_episode, ep) for ep in episodes] - for future in futures: - future.result() - return time.perf_counter() - start - - def _build_episode_row_groups(self) -> dict[int, int]: - counts: dict[tuple[int, int], int] = {} - row_groups = {} - for ep_idx in range(int(self.meta.total_episodes)): - ep = self.meta.episodes[ep_idx] - key = (int(ep["data/chunk_index"]), int(ep["data/file_index"])) - row_groups[ep_idx] = counts.get(key, 0) - counts[key] = row_groups[ep_idx] + 1 - return row_groups - - -def run_fetch_pool( - manifest: EpisodeVideoManifest, - data_root: str, - episodes: Sequence[int], - dataset_episode_count: int, - benchmark_episode_count: int, - byte_budget: int, - workers: int, - range_backend: str, - args: argparse.Namespace, -) -> dict[str, float]: - with EpisodeByteCache( - manifest, - data_root, - byte_budget=byte_budget, - workers=workers, - range_backend=range_backend, - native_http_connections=args.native_http_connections, - native_http_timeout=args.native_http_timeout, - native_http_retries=args.native_http_retries, - native_http_subranges=args.range_subranges, - open_decoders=False, - ) as cache: - elapsed = _fill_cache(cache, episodes, progress_interval=args.progress_interval) - timings = cache.timing_summary() - random_decode = None - stream_sim = None - if args.include_pool_sampling: - _log("pool_sampling: warming resident decoders and decoding random samples") - random_decode = run_pool_random_decode( - cache, - episodes, - sample_count=args.pool_random_samples, - batch_size=args.batch_size, - decode_workers=args.decode_workers, - seed=args.seed + 3, - ) - _log( - f"pool_stream: exact coverage, consuming {args.target_samples_s:.1f} samples/s while prefetching ahead" - ) - shard_episodes = _episode_shard( - dataset_episode_count, - benchmark_episode_count, - args.seed, - shard_count=args.distributed_shard_count, - shard_index=args.distributed_shard_index, - ) - stream_sim = run_exact_coverage_stream( - cache, - shard_episodes, - pool_size=len(episodes), - sample_count=args.stream_samples, - target_samples_s=args.target_samples_s, - prefetch_ahead=args.stream_prefetch_episodes, - batch_size=args.batch_size, - decode_workers=args.decode_workers, - seed=args.seed + 5, - ) - byte_count = _bytes_for(manifest, episodes) - episode_mb = byte_count / len(episodes) / 1024**2 - job_count = max(timings["jobs"], 1.0) - result = { - "fetch_s": elapsed, - "fetch_mbps": byte_count / elapsed / 1024**2, - "fetch_episodes_s": len(episodes) / elapsed, - "episode_mb": episode_mb, - "avg_mb_miss": byte_count / (len(episodes) * len(manifest.video_keys)) / 1024**2, - "jobs": timings["jobs"], - "lookup_ms": timings["lookup_s"] * 1000 / job_count, - "range_fetch_ms": timings["fetch_s"] * 1000 / job_count, - "synthesize_ms": timings["synthesize_s"] * 1000 / job_count, - "store_ms": timings["store_s"] * 1000 / job_count, - } - result.update({key: value for key, value in timings.items() if key.startswith("range_")}) - if random_decode is not None: - result.update({f"pool_decode_{key}": value for key, value in random_decode.items()}) - if stream_sim is not None: - result.update({f"pool_stream_{key}": value for key, value in stream_sim.items()}) - return result - - -def run_parallel( - manifest: EpisodeVideoManifest, - data_root: str, - episodes: Sequence[int], - timestamps: dict[tuple[int, str], list[float]], - byte_budget: int, - workers: int, - decode_workers: int, - frames_per_episode: int, - parquet_reader: EpisodeParquetReader, - range_backend: str, -) -> dict[str, float]: - with EpisodeByteCache( - manifest, - data_root, - byte_budget=byte_budget, - workers=workers, - range_backend=range_backend, - open_decoders=False, - ) as cache: - parquet_s = parquet_reader.read_episodes(episodes, workers=workers) - fetch_s = _fill_cache(cache, episodes) - decoder_start = time.perf_counter() - for ep in episodes: - for camera_key in manifest.video_keys: - cache.get_decoder(ep, camera_key) - decoder_s = time.perf_counter() - decoder_start - decode_s = _decode_all(cache, timestamps, decode_workers=decode_workers) - byte_count = _bytes_for(manifest, episodes) - return { - "fetch_s": fetch_s, - "fetch_mbps": byte_count / fetch_s / 1024**2, - "fetch_episodes_s": len(episodes) / fetch_s, - "parquet_s": parquet_s, - "decoder_ms_miss": decoder_s * 1000 / (len(episodes) * len(manifest.video_keys)), - "decode_samples_s": _samples_per_s(decode_s, episodes, frames_per_episode), - } - - -def run_overlapped( - manifest: EpisodeVideoManifest, - data_root: str, - episodes: Sequence[int], - timestamps: dict[tuple[int, str], list[float]], - byte_budget: int, - workers: int, - decode_workers: int, - frames_per_episode: int, - prefetch_ahead: int, - parquet_reader: EpisodeParquetReader, - range_backend: str, -) -> dict[str, float]: - with EpisodeByteCache( - manifest, - data_root, - byte_budget=byte_budget, - workers=workers, - range_backend=range_backend, - open_decoders=True, - ) as cache: - start = time.perf_counter() - video_wait_decode_s = 0.0 - parquet_wait_s = 0.0 - parquet_pool = ThreadPoolExecutor(max_workers=max(1, min(workers, len(episodes)))) - parquet_futures = { - ep: parquet_reader.submit_read_episode(parquet_pool, ep) for ep in episodes[:prefetch_ahead] - } - for ep in episodes[:prefetch_ahead]: - cache.submit_prefetch(ep) - try: - for idx, ep in enumerate(episodes): - next_idx = idx + prefetch_ahead - if next_idx < len(episodes): - next_ep = episodes[next_idx] - cache.submit_prefetch(next_ep) - parquet_futures[next_ep] = parquet_reader.submit_read_episode(parquet_pool, next_ep) - - parquet_start = time.perf_counter() - parquet_futures.pop(ep).result() - parquet_wait_s += time.perf_counter() - parquet_start - - video_start = time.perf_counter() - cache.ensure_ready(ep) - if decode_workers <= 1: - for camera_key in manifest.video_keys: - cache.get_frames(ep, camera_key, timestamps[(ep, camera_key)]) - else: - with ThreadPoolExecutor(max_workers=decode_workers) as pool: - futures = [ - pool.submit(cache.get_frames, ep, camera_key, timestamps[(ep, camera_key)]) - for camera_key in manifest.video_keys - ] - for future in futures: - future.result() - video_wait_decode_s += time.perf_counter() - video_start - finally: - parquet_pool.shutdown(wait=True) - elapsed = time.perf_counter() - start - return { - "samples_s": _samples_per_s(elapsed, episodes, frames_per_episode), - "video_samples_s": _samples_per_s(video_wait_decode_s, episodes, frames_per_episode), - "parquet_samples_s": _samples_per_s(parquet_wait_s, episodes, frames_per_episode), - "wall_s": elapsed, - "video_wait_decode_s": video_wait_decode_s, - "parquet_wait_s": parquet_wait_s, - } - - -_remote_decoder_local = threading.local() - - -def _remote_decoder_cache() -> VideoDecoderCache: - cache = getattr(_remote_decoder_local, "cache", None) - if cache is None: - cache = VideoDecoderCache(max_size=None) - _remote_decoder_local.cache = cache - return cache - - -def _decode_remote_source( - meta: LeRobotDatasetMetadata, - data_root: str, - episode_index: int, - camera_key: str, - timestamps: list[float], -): - video_path = _root_join(data_root, str(meta.get_video_file_path(episode_index, camera_key))) - return decode_video_frames_torchcodec( - video_path, - timestamps, - tolerance_s=1.0 / float(meta.fps), - decoder_cache=_remote_decoder_cache(), - return_uint8=True, - ) - - -def run_remote_decoder( - meta: LeRobotDatasetMetadata, - data_root: str, - episodes: Sequence[int], - timestamps: dict[tuple[int, str], list[float]], - *, - frames_per_episode: int, - decode_workers: int, - parquet_reader: EpisodeParquetReader, -) -> dict[str, float]: - items = [ - (ep, camera_key, timestamps[(ep, camera_key)]) for ep in episodes for camera_key in meta.video_keys - ] - - start = time.perf_counter() - for ep, camera_key, ts in items: - if camera_key == meta.video_keys[0]: - parquet_reader.read_episode(ep) - _decode_remote_source(meta, data_root, ep, camera_key, ts) - sequential_s = time.perf_counter() - start - - start = time.perf_counter() - if decode_workers <= 1: - for ep, camera_key, ts in items: - if camera_key == meta.video_keys[0]: - parquet_reader.read_episode(ep) - _decode_remote_source(meta, data_root, ep, camera_key, ts) - else: - with ThreadPoolExecutor(max_workers=decode_workers) as pool: - parquet_futures = [pool.submit(parquet_reader.read_episode, ep) for ep in episodes] - futures = [ - pool.submit(_decode_remote_source, meta, data_root, ep, camera_key, ts) - for ep, camera_key, ts in items - ] - for future in parquet_futures: - future.result() - for future in futures: - future.result() - parallel_s = time.perf_counter() - start - - return { - "sequential_samples_s": _samples_per_s(sequential_s, episodes, frames_per_episode), - "parallel_samples_s": _samples_per_s(parallel_s, episodes, frames_per_episode), - } - - -def _print_range_timing_summary(fetch_pool: dict[str, float]) -> None: - range_jobs = fetch_pool.get("range_jobs", 0.0) - if range_jobs <= 0: - return - - print() - print("| Range Read Stage | avg ms/range |") - print("|---|---:|") - for key, label in ( - ("range_open_s", "fsspec handle open/lookup"), - ("range_seek_s", "fsspec seek"), - ("range_read_s", "fsspec read"), - ("range_resolve_s", "http URL resolve"), - ("range_header_s", "http response headers"), - ("range_first_byte_s", "http first body byte"), - ("range_body_s", "http body drain"), - ("range_chunk_gap_s", "http chunk wait"), - ("range_join_s", "join response chunks"), - ("range_failed_attempt_s", "http failed attempts"), - ("range_retry_sleep_s", "http retry sleep"), - ): - value = fetch_pool.get(key) - if value is not None: - print(f"| {label} | {value * 1000 / range_jobs:.3f} |") - if "range_retry_attempts" in fetch_pool: - print(f"| http retries | {fetch_pool['range_retry_attempts'] / range_jobs:.3f} |") - if "range_exception_attempts" in fetch_pool: - print(f"| http exceptions | {fetch_pool['range_exception_attempts'] / range_jobs:.3f} |") - if fetch_pool["range_exception_attempts"] > 0: - print( - f"| http failed attempt avg s | " - f"{fetch_pool.get('range_failed_attempt_s', 0.0) / fetch_pool['range_exception_attempts']:.1f} |" - ) - if fetch_pool.get("range_failed_requests"): - print(f"| http failed requests | {fetch_pool['range_failed_requests']:.0f} |") - exception_counts = { - key.removeprefix("range_exception_"): value - for key, value in fetch_pool.items() - if key.startswith("range_exception_") and key != "range_exception_attempts" - } - if exception_counts: - summary = ", ".join(f"{name}={count:.0f}" for name, count in sorted(exception_counts.items())) - print(f"| http exception counts | {summary} |") - failed_status_counts = { - key.removeprefix("range_failed_status_"): value - for key, value in fetch_pool.items() - if key.startswith("range_failed_status_") - } - if failed_status_counts: - summary = ", ".join(f"{status}={count:.0f}" for status, count in sorted(failed_status_counts.items())) - print(f"| http failed status counts | {summary} |") - status_counts = { - key.removeprefix("range_status_"): value - for key, value in fetch_pool.items() - if key.startswith("range_status_") - } - if status_counts: - summary = ", ".join(f"{status}={count:.0f}" for status, count in sorted(status_counts.items())) - print(f"| http status counts | {summary} |") - for method in ("head", "get"): - request_count = fetch_pool.get(f"range_hffs_{method}_requests", 0.0) - if request_count <= 0: - continue - print(f"| hffs {method.upper()} requests/range | {request_count / range_jobs:.3f} |") - print( - f"| hffs {method.upper()} total | {fetch_pool[f'range_hffs_{method}_s'] * 1000 / range_jobs:.3f} |" - ) - retries = fetch_pool.get(f"range_hffs_{method}_retries", 0.0) - exceptions = fetch_pool.get(f"range_hffs_{method}_exception_attempts", 0.0) - if retries: - print(f"| hffs {method.upper()} retries/range | {retries / range_jobs:.3f} |") - print( - f"| hffs {method.upper()} retry sleep | " - f"{fetch_pool.get(f'range_hffs_{method}_retry_sleep_s', 0.0) * 1000 / range_jobs:.3f} |" - ) - if exceptions: - print(f"| hffs {method.upper()} exceptions/range | {exceptions / range_jobs:.3f} |") - print( - f"| hffs {method.upper()} failed attempts | " - f"{fetch_pool.get(f'range_hffs_{method}_failed_attempt_s', 0.0) * 1000 / range_jobs:.3f} |" - ) - bytes_read = fetch_pool.get(f"range_hffs_{method}_bytes", 0.0) - total_s = fetch_pool.get(f"range_hffs_{method}_s", 0.0) - if bytes_read > 0 and total_s > 0: - print(f"| hffs {method.upper()} MiB/s | {bytes_read / total_s / 1024**2:.1f} |") - hffs_status_counts = { - key.removeprefix(f"range_hffs_{method}_status_"): value - for key, value in fetch_pool.items() - if key.startswith(f"range_hffs_{method}_status_") - } - if hffs_status_counts: - summary = ", ".join( - f"{status}={count:.0f}" for status, count in sorted(hffs_status_counts.items()) - ) - print(f"| hffs {method.upper()} status counts | {summary} |") - hffs_failed_status_counts = { - key.removeprefix(f"range_hffs_{method}_failed_status_"): value - for key, value in fetch_pool.items() - if key.startswith(f"range_hffs_{method}_failed_status_") - } - if hffs_failed_status_counts: - summary = ", ".join( - f"{status}={count:.0f}" for status, count in sorted(hffs_failed_status_counts.items()) - ) - print(f"| hffs {method.upper()} failed status counts | {summary} |") - hffs_exception_counts = { - key.removeprefix(f"range_hffs_{method}_exception_"): value - for key, value in fetch_pool.items() - if key.startswith(f"range_hffs_{method}_exception_") - and key != f"range_hffs_{method}_exception_attempts" - } - if hffs_exception_counts: - summary = ", ".join( - f"{name}={count:.0f}" for name, count in sorted(hffs_exception_counts.items()) - ) - print(f"| hffs {method.upper()} exception counts | {summary} |") - chunks = fetch_pool.get("range_chunks", 0.0) - if chunks > 0: - bytes_read = fetch_pool.get("range_bytes", 0.0) - body_s = fetch_pool.get("range_body_s", 0.0) - print(f"| http chunks/range | {chunks / range_jobs:.1f} |") - print(f"| http avg KiB/chunk | {bytes_read / chunks / 1024:.1f} |") - if body_s > 0: - print(f"| http body MiB/s | {bytes_read / body_s / 1024**2:.1f} |") - print(f"| range reads | {range_jobs:.0f} |") - print(f"| avg MiB/range | {fetch_pool.get('range_bytes', 0.0) / range_jobs / 1024**2:.1f} |") - - -def run_indexed_strategy( - meta: LeRobotDatasetMetadata, - data_root: str, - args: argparse.Namespace, - parquet_reader: EpisodeParquetReader, - *, - range_backend: str = "fsspec", - label: str = "indexed", - sidecar_path: str | None = None, -) -> None: - _log(f"starting_strategy: {label}") - memory_start = _memory_snapshot() - manifest_start = time.perf_counter() - dataset_episode_count = int(meta.total_episodes) - manifest_episode_count = args.manifest_episodes or dataset_episode_count - manifest_episode_count = min(manifest_episode_count, dataset_episode_count, args.num_episodes) - manifest = EpisodeVideoManifest.build( - meta, - data_root, - episode_indices=range(manifest_episode_count), - range_backend=range_backend, - workers=args.workers, - max_probe_bytes=args.max_probe_mb * 1024 * 1024, - sidecar_path=sidecar_path, - ) - manifest_s = time.perf_counter() - manifest_start - _log(f"{label}: manifest_build_s={manifest_s:.2f}") - - benchmark_episode_count = min(dataset_episode_count, args.num_episodes) - episodes = _episode_pool( - dataset_episode_count, - args.num_episodes, - args.pool_size, - args.seed, - shard_count=args.distributed_shard_count, - shard_index=args.distributed_shard_index, - ) - byte_budget = int(args.byte_budget_gb * 1024**3) - byte_count = _bytes_for(manifest, episodes) - _log( - f"{label}: planned_video_fetch={byte_count / 1024**3:.2f} GiB per fetch track " - f"({byte_count / len(episodes) / 1024**2:.1f} MiB/episode)" - ) - - _log(f"{label}: filling episode byte cache with {args.workers} workers") - fetch_pool = run_fetch_pool( - manifest, - data_root, - episodes, - dataset_episode_count, - benchmark_episode_count, - byte_budget, - args.workers, - range_backend, - args, - ) - estimated_dataset_s = dataset_episode_count / fetch_pool["fetch_episodes_s"] - estimated_benchmark_s = benchmark_episode_count / fetch_pool["fetch_episodes_s"] - - print(f"manifest_build_s: {manifest_s:.2f}") - print(f"strategy: {label}") - print(f"range_backend: {range_backend}") - print(f"mp4_sidecar: {sidecar_path or 'none'}") - print(f"data_root: {data_root}") - print(f"dataset_episodes: {dataset_episode_count}") - print(f"benchmark_episodes: {benchmark_episode_count}") - print(f"distributed_shard_count: {args.distributed_shard_count}") - print(f"distributed_shard_index: {args.distributed_shard_index}") - print(f"pool_episodes: {len(episodes)}") - print(f"sampled_episodes: {episodes}") - print(f"cameras: {manifest.video_keys}") - print() - print( - "| Track | fetch MB/s | fetch eps/s | wall s | est benchmark | est full dataset | avg MB/camera | notes |" - ) - print("|---|---:|---:|---:|---:|---:|---:|---|") - print( - f"| EPISODE POOL FETCH | {fetch_pool['fetch_mbps']:.1f} | " - f"{fetch_pool['fetch_episodes_s']:.2f} | {fetch_pool['fetch_s']:.2f} | " - f"{_format_duration(estimated_benchmark_s)} | {_format_duration(estimated_dataset_s)} | " - f"{fetch_pool['avg_mb_miss']:.1f} | {args.workers} workers, no decoder open/frame decode |" - ) - print() - print("| Camera Job Stage | avg ms/job |") - print("|---|---:|") - print(f"| manifest lookup | {fetch_pool['lookup_ms']:.3f} |") - print(f"| remote byte-range fetch | {fetch_pool['range_fetch_ms']:.3f} |") - print(f"| synthesize mini-MP4 | {fetch_pool['synthesize_ms']:.3f} |") - print(f"| store in shared cache | {fetch_pool['store_ms']:.3f} |") - print(f"| camera jobs | {fetch_pool['jobs']:.0f} |") - _print_range_timing_summary(fetch_pool) - if args.include_pool_sampling: - print() - print("| Resident Pool Decode | value |") - print("|---|---:|") - print(f"| random training samples | {fetch_pool['pool_decode_sample_count']:.0f} |") - print(f"| decoder opens | {fetch_pool['pool_decode_decoder_count']:.0f} |") - print(f"| decoder open ms/episode-camera | {fetch_pool['pool_decode_decoder_open_ms']:.3f} |") - print(f"| decode wall s | {fetch_pool['pool_decode_decode_s']:.3f} |") - print(f"| training samples/s | {fetch_pool['pool_decode_training_samples_s']:.1f} |") - print(f"| camera frames/s | {fetch_pool['pool_decode_camera_frames_s']:.1f} |") - print(f"| decode ms/training sample | {fetch_pool['pool_decode_decode_ms_sample']:.3f} |") - print(f"| decode ms/camera frame | {fetch_pool['pool_decode_decode_ms_camera_frame']:.3f} |") - print() - print("| Resident Pool Randomness | value |") - print("|---|---:|") - print(f"| pool episodes | {len(episodes)} |") - print(f"| batch size | {args.batch_size} |") - print(f"| unique episodes sampled | {fetch_pool['pool_decode_unique_episodes']:.0f} |") - print( - f"| mean unique episodes/batch | {fetch_pool['pool_decode_mean_unique_episodes_per_batch']:.1f} |" - ) - print( - f"| min unique episodes/batch | {fetch_pool['pool_decode_min_unique_episodes_per_batch']:.0f} |" - ) - print( - f"| mean samples/used episode | {fetch_pool['pool_decode_mean_samples_per_used_episode']:.2f} |" - ) - print(f"| max samples/episode | {fetch_pool['pool_decode_max_samples_per_episode']:.0f} |") - print() - print("| Streaming Keep-Up Simulation | value |") - print("|---|---:|") - print(f"| target samples/s | {fetch_pool['pool_stream_target_samples_s']:.1f} |") - print(f"| actual samples/s | {fetch_pool['pool_stream_actual_samples_s']:.1f} |") - print(f"| kept up | {'yes' if fetch_pool['pool_stream_kept_up'] else 'no'} |") - print(f"| batch size | {fetch_pool['pool_stream_batch_size']:.0f} |") - print(f"| decode workers | {fetch_pool['pool_stream_decode_workers']:.0f} |") - print(f"| stream wall s | {fetch_pool['pool_stream_stream_wall_s']:.3f} |") - print(f"| refill wait s | {fetch_pool['pool_stream_refill_wait_s']:.3f} |") - print(f"| deadline miss s | {fetch_pool['pool_stream_deadline_miss_s']:.3f} |") - print(f"| replacement episodes | {fetch_pool['pool_stream_replacements']:.0f} |") - print(f"| replacement episodes/s | {fetch_pool['pool_stream_replacement_episodes_s']:.2f} |") - print(f"| samples per replacement episode | {fetch_pool['pool_stream_samples_per_episode']:.0f} |") - print(f"| prefetch replacement episodes | {fetch_pool['pool_stream_prefetch_episodes']:.0f} |") - print( - f"| stream mean unique episodes/batch | " - f"{fetch_pool['pool_stream_stream_mean_unique_episodes_per_batch']:.1f} |" - ) - print( - f"| stream min unique episodes/batch | " - f"{fetch_pool['pool_stream_stream_min_unique_episodes_per_batch']:.0f} |" - ) - memory_end = _memory_snapshot() - _print_memory_summary(memory_start, memory_end) - summary = { - "hostname": socket.gethostname(), - "strategy": label, - "range_backend": range_backend, - "data_root": data_root, - "mp4_sidecar": sidecar_path, - "dataset_episodes": dataset_episode_count, - "benchmark_episodes": benchmark_episode_count, - "distributed_shard_count": args.distributed_shard_count, - "distributed_shard_index": args.distributed_shard_index, - "pool_episodes": len(episodes), - "sampled_episodes": episodes, - "workers": args.workers, - "decode_workers": args.decode_workers, - "manifest_build_s": manifest_s, - "fetch_bytes": byte_count, - "fetch_gib": byte_count / 1024**3, - "fetch_s": fetch_pool["fetch_s"], - "fetch_mib_s": fetch_pool["fetch_mbps"], - "fetch_episodes_s": fetch_pool["fetch_episodes_s"], - "avg_mb_camera": fetch_pool["avg_mb_miss"], - "range_reads": fetch_pool.get("range_jobs", 0.0), - "range_hffs_get_exceptions": fetch_pool.get("range_hffs_get_exception_attempts", 0.0), - "range_hffs_get_retries": fetch_pool.get("range_hffs_get_retries", 0.0), - "rss_start_mib": memory_start["rss_mib"], - "rss_end_mib": memory_end["rss_mib"], - "peak_rss_mib": memory_end["peak_rss_mib"], - } - for key, value in fetch_pool.items(): - if key.startswith("pool_decode_") or key.startswith("pool_stream_"): - summary[key] = value - _write_summary_json(args.summary_json, summary) - - if args.include_decode: - timestamps = _timestamps(manifest, episodes, args.frames_per_episode, args.seed + 1) - _log(f"{label}: running parallel video fetch + decode-only") - parallel = run_parallel( - manifest, - data_root, - episodes, - timestamps, - byte_budget, - args.workers, - args.decode_workers, - args.frames_per_episode, - parquet_reader, - range_backend, - ) - _log(f"{label}: running overlapped end-to-end") - overlapped = run_overlapped( - manifest, - data_root, - episodes, - timestamps, - byte_budget, - args.workers, - args.decode_workers, - args.frames_per_episode, - args.prefetch_ahead, - parquet_reader, - range_backend, - ) - print( - f"| DECODE COMPARISON | {parallel['fetch_mbps']:.1f} | {parallel['fetch_episodes_s']:.2f} | " - f"{parallel['fetch_s']:.2f} | " - f"{_format_duration(benchmark_episode_count / parallel['fetch_episodes_s'])} | " - f"{_format_duration(dataset_episode_count / parallel['fetch_episodes_s'])} | " - f"{fetch_pool['avg_mb_miss']:.1f} | " - f"decoder open {parallel['decoder_ms_miss']:.1f} ms/miss, " - f"decode {parallel['decode_samples_s']:.1f} samples/s, parquet {parallel['parquet_s']:.2f}s |" - ) - print( - f"| OVERLAPPED E2E | - | - | {overlapped['wall_s']:.2f} | - | - | " - f"{fetch_pool['avg_mb_miss']:.1f} | " - f"{overlapped['samples_s']:.1f} samples/s; video+decode " - f"{overlapped['video_wait_decode_s']:.2f}s, parquet {overlapped['parquet_wait_s']:.2f}s |" - ) - - -def run_remote_strategy( - meta: LeRobotDatasetMetadata, - data_root: str, - args: argparse.Namespace, - parquet_reader: EpisodeParquetReader, -) -> None: - _log("starting_strategy: remote-decoder") - episodes = _episode_pool( - int(meta.total_episodes), - args.num_episodes, - args.pool_size, - args.seed, - shard_count=args.distributed_shard_count, - shard_index=args.distributed_shard_index, - ) - timestamps = _timestamps_from_meta(meta, episodes, args.frames_per_episode, args.seed + 1) - _log("remote-decoder: running direct source MP4 decoder") - result = run_remote_decoder( - meta, - data_root, - episodes, - timestamps, - frames_per_episode=args.frames_per_episode, - decode_workers=args.decode_workers, - parquet_reader=parquet_reader, - ) - print("strategy: remote-decoder") - print(f"data_root: {data_root}") - print(f"episodes: {episodes}") - print(f"cameras: {list(meta.video_keys)}") - print() - print("| Track | samples/s | notes |") - print("|---|---:|---|") - print(f"| REMOTE SEQUENTIAL | {result['sequential_samples_s']:.1f} | direct source MP4 decoder |") - print( - f"| REMOTE PARALLEL | {result['parallel_samples_s']:.1f} | " - f"direct source MP4 decoder, {args.decode_workers} workers |" - ) - - -def main() -> None: - args = parse_args() - if args.strategy == "full": - args.strategy = "both" - if args.strategy == "native-http": - args.range_backend = "native-http" - if args.http_failure_log: - os.environ["LEROBOT_HTTP_FAILURE_LOG"] = args.http_failure_log - print(f"http_failure_log: {args.http_failure_log}") - data_root = args.data_root - 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 = 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}") - - if sidecar_path is not None and args.strategy == "both": - if args.include_decode: - run_remote_strategy(meta, data_root, args, parquet_reader) - print() - run_indexed_strategy( - meta, - data_root, - args, - parquet_reader, - range_backend=args.range_backend, - label=f"indexed-sidecar-{args.range_backend}", - sidecar_path=str(sidecar_path), - ) - return - if sidecar_path is not None and args.strategy == "indexed": - run_indexed_strategy( - meta, - data_root, - args, - parquet_reader, - range_backend=args.range_backend, - label=f"indexed-sidecar-{args.range_backend}", - sidecar_path=str(sidecar_path), - ) - return - if sidecar_path is not None and args.strategy == "native-http": - run_indexed_strategy( - meta, - data_root, - args, - parquet_reader, - range_backend="native-http", - label="indexed-sidecar-native-http", - sidecar_path=str(sidecar_path), - ) - return - if args.strategy == "both": - 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() - - if args.strategy in ("both", "indexed"): - run_indexed_strategy( - meta, - data_root, - args, - parquet_reader, - range_backend="fsspec", - label="indexed", - sidecar_path=None, - ) - if args.strategy == "both": - print() - if args.strategy == "remote-decoder" or (args.strategy == "both" and args.include_decode): - run_remote_strategy(meta, data_root, args, parquet_reader) - if args.strategy == "both" and args.include_decode: - print() - if args.strategy in ("both", "native-http"): - run_indexed_strategy( - meta, - data_root, - args, - parquet_reader, - range_backend="native-http", - label="indexed-native-http", - sidecar_path=None, - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/bench_streaming_dataset.py b/scripts/bench_streaming_dataset.py deleted file mode 100644 index c2aa0c821..000000000 --- a/scripts/bench_streaming_dataset.py +++ /dev/null @@ -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() diff --git a/scripts/summarize_episode_pool_bench.py b/scripts/summarize_episode_pool_bench.py deleted file mode 100644 index e9308c92f..000000000 --- a/scripts/summarize_episode_pool_bench.py +++ /dev/null @@ -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() diff --git a/src/lerobot/streaming/episode_video.py b/src/lerobot/streaming/episode_video.py deleted file mode 100644 index 206bc4592..000000000 --- a/src/lerobot/streaming/episode_video.py +++ /dev/null @@ -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", -] diff --git a/tests/test_streaming_core_imports.py b/tests/test_streaming_core_imports.py index 00536a010..071056349 100644 --- a/tests/test_streaming_core_imports.py +++ b/tests/test_streaming_core_imports.py @@ -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