diff --git a/scripts/bench_episode_byte_cache.py b/scripts/bench_episode_byte_cache.py index 611f24f00..16e5098e0 100644 --- a/scripts/bench_episode_byte_cache.py +++ b/scripts/bench_episode_byte_cache.py @@ -12,6 +12,7 @@ from __future__ import annotations import argparse import random +import resource import tempfile import threading import time @@ -310,6 +311,44 @@ def _format_duration(seconds: float) -> str: 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; macOS reports bytes. + if rss > 10**8: + return rss / 1024**2 + return rss / 1024 + + +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 _root_join(data_root: str, relative_path: str) -> str: if data_root.startswith("hf://"): return f"{data_root.rstrip('/')}/{relative_path}" @@ -665,6 +704,7 @@ def run_indexed_strategy( 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 @@ -724,6 +764,7 @@ def run_indexed_strategy( 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_memory_summary(memory_start, _memory_snapshot()) if args.include_decode: timestamps = _timestamps(manifest, episodes, args.frames_per_episode, args.seed + 1) @@ -783,6 +824,7 @@ def run_random_frame_strategy( if args.frame_pool_size <= 0: raise ValueError(f"frame-pool-size must be > 0, got {args.frame_pool_size}") _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 @@ -869,6 +911,7 @@ def run_random_frame_strategy( print(f"| raw byte windows | {len(windows)} |") print(f"| coalesced byte ranges | {len(coalesced)} |") print(f"| useful sample MiB | {useful_bytes / 1024**2:.1f} |") + _print_memory_summary(memory_start, _memory_snapshot()) def run_remote_strategy(