From 834c282631c3bab7eda093080c6ab4813c9b731c Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 17 Jun 2026 20:16:30 +0200 Subject: [PATCH] Make episode cache benchmark fetch-only by default --- scripts/bench_episode_byte_cache.py | 114 +++++++++++++++------------- 1 file changed, 60 insertions(+), 54 deletions(-) diff --git a/scripts/bench_episode_byte_cache.py b/scripts/bench_episode_byte_cache.py index 34a97efda..5040abd51 100644 --- a/scripts/bench_episode_byte_cache.py +++ b/scripts/bench_episode_byte_cache.py @@ -60,6 +60,11 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("--pool-size", type=int, default=16) parser.add_argument("--workers", type=int, default=8) + 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("--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) @@ -280,23 +285,22 @@ class EpisodeParquetReader: return row_groups -def run_sequential( +def run_fetch_pool( manifest: EpisodeVideoManifest, data_root: str, episodes: Sequence[int], byte_budget: int, - parquet_reader: EpisodeParquetReader, + workers: int, range_backend: str, ) -> dict[str, float]: with EpisodeByteCache( manifest, data_root, byte_budget=byte_budget, - workers=1, + workers=workers, range_backend=range_backend, open_decoders=False, ) as cache: - parquet_s = parquet_reader.read_episodes(episodes, workers=1) elapsed = _fill_cache(cache, episodes) byte_count = _bytes_for(manifest, episodes) episode_mb = byte_count / len(episodes) / 1024**2 @@ -305,7 +309,6 @@ def run_sequential( "fetch_mbps": byte_count / elapsed / 1024**2, "fetch_episodes_s": len(episodes) / elapsed, "episode_mb": episode_mb, - "parquet_s": parquet_s, "avg_mb_miss": byte_count / (len(episodes) * len(manifest.video_keys)) / 1024**2, } @@ -519,7 +522,6 @@ def run_indexed_strategy( _log(f"{label}: manifest_build_s={manifest_s:.2f}") episodes = _episode_pool(int(meta.total_episodes), args.num_episodes, args.pool_size, args.seed) - timestamps = _timestamps(manifest, episodes, args.frames_per_episode, args.seed + 1) byte_budget = int(args.byte_budget_gb * 1024**3) byte_count = _bytes_for(manifest, episodes) _log( @@ -527,35 +529,8 @@ def run_indexed_strategy( f"({byte_count / len(episodes) / 1024**2:.1f} MiB/episode)" ) - _log(f"{label}: running sequential video fetch") - sequential = run_sequential(manifest, data_root, episodes, byte_budget, parquet_reader, range_backend) - _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, - ) + _log(f"{label}: filling episode byte cache with {args.workers} workers") + fetch_pool = run_fetch_pool(manifest, data_root, episodes, byte_budget, args.workers, range_backend) print(f"manifest_build_s: {manifest_s:.2f}") print(f"strategy: {label}") @@ -565,25 +540,55 @@ def run_indexed_strategy( print(f"episodes: {episodes}") print(f"cameras: {manifest.video_keys}") print() - print("| Track | fetch MB/s | fetch eps/s | samples/s | avg MB/miss | notes |") + print("| Track | fetch MB/s | fetch eps/s | wall s | avg MB/camera | notes |") print("|---|---:|---:|---:|---:|---|") print( - f"| SEQUENTIAL | {sequential['fetch_mbps']:.1f} | {sequential['fetch_episodes_s']:.2f} | - | " - f"{sequential['avg_mb_miss']:.1f} | 1 worker video fetch, parquet {sequential['parquet_s']:.2f}s |" - ) - print( - f"| PARALLEL | {parallel['fetch_mbps']:.1f} | {parallel['fetch_episodes_s']:.2f} | " - f"{parallel['decode_samples_s']:.1f} | " - f"{sequential['avg_mb_miss']:.1f} | decode-only, decoder open " - f"{parallel['decoder_ms_miss']:.1f} ms/miss, parquet {parallel['parquet_s']:.2f}s |" - ) - print( - f"| OVERLAPPED | - | - | {overlapped['samples_s']:.1f} | {sequential['avg_mb_miss']:.1f} | " - f"end-to-end; video {overlapped['video_samples_s']:.1f} samples/s " - f"({overlapped['video_wait_decode_s']:.2f}s), parquet {overlapped['parquet_samples_s']:.1f} " - f"samples/s ({overlapped['parquet_wait_s']:.2f}s) |" + f"| EPISODE POOL FETCH | {fetch_pool['fetch_mbps']:.1f} | " + f"{fetch_pool['fetch_episodes_s']:.2f} | {fetch_pool['fetch_s']:.2f} | " + f"{fetch_pool['avg_mb_miss']:.1f} | {args.workers} workers, no decoder open/frame decode |" ) + 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} | {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} | {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, @@ -635,8 +640,9 @@ def main() -> None: print(f"using_mp4_sidecar: {sidecar_path}") if sidecar_path is not None and args.strategy == "both": - run_remote_strategy(meta, data_root, args, parquet_reader) - print() + if args.include_decode: + run_remote_strategy(meta, data_root, args, parquet_reader) + print() run_indexed_strategy( meta, data_root, @@ -704,9 +710,9 @@ def main() -> None: ) if args.strategy == "both": print() - if args.strategy in ("both", "remote-decoder"): + 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": + if args.strategy == "both" and args.include_decode: print() if args.strategy in ("both", "native-http"): run_indexed_strategy(