diff --git a/scripts/bench_episode_byte_cache.py b/scripts/bench_episode_byte_cache.py index c8dc65e74..2b777da54 100644 --- a/scripts/bench_episode_byte_cache.py +++ b/scripts/bench_episode_byte_cache.py @@ -32,6 +32,7 @@ import pyarrow.parquet as pq from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata from lerobot.datasets.episode_video_streaming import ( EpisodeByteCache, + ExactCoveragePool, EpisodeVideoManifest, NativeHTTPRangeFetcher, assert_hf_hub_range_cache_branch, @@ -77,6 +78,13 @@ def parse_args() -> argparse.Namespace: help="Concurrent camera-fetch jobs. Total connections ~= workers x range-subranges; " "the HF bucket path saturates around 64 connections per host, so keep the product near 64.", ) + parser.add_argument( + "--coverage", + choices=["sampled", "exact"], + default="sampled", + help="sampled: with-replacement random draws (no epoch guarantee). exact: every frame of " + "every shard episode decoded exactly once per epoch (deterministic, pool-bounded).", + ) parser.add_argument( "--range-subranges", type=int, @@ -350,6 +358,133 @@ def run_pool_random_decode( 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). + + Unlike run_pool_stream_simulation (with-replacement sampling, no coverage guarantee), this + 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 + result = { + "coverage_mode": "exact", + "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), + "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 run_pool_stream_simulation( cache: EpisodeByteCache, resident_episodes: Sequence[int], @@ -765,22 +900,43 @@ def run_fetch_pool( _log( f"pool_stream: consuming {args.target_samples_s:.1f} samples/s while prefetching replacements" ) - stream_sim = run_pool_stream_simulation( - cache, - episodes, - dataset_episode_count=dataset_episode_count, - num_episodes=benchmark_episode_count, - sample_count=args.stream_samples, - target_samples_s=args.target_samples_s, - samples_per_episode=args.pool_samples_per_episode, - prefetch_episodes=args.stream_prefetch_episodes, - shard_count=args.distributed_shard_count, - shard_index=args.distributed_shard_index, - shard_seed=args.seed, - batch_size=args.batch_size, - decode_workers=args.decode_workers, - seed=args.seed + 4, - ) + if args.coverage == "exact": + _log("pool_stream: EXACT coverage (every frame once) 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, + ) + else: + stream_sim = run_pool_stream_simulation( + cache, + episodes, + dataset_episode_count=dataset_episode_count, + num_episodes=benchmark_episode_count, + sample_count=args.stream_samples, + target_samples_s=args.target_samples_s, + samples_per_episode=args.pool_samples_per_episode, + prefetch_episodes=args.stream_prefetch_episodes, + shard_count=args.distributed_shard_count, + shard_index=args.distributed_shard_index, + shard_seed=args.seed, + batch_size=args.batch_size, + decode_workers=args.decode_workers, + seed=args.seed + 4, + ) byte_count = _bytes_for(manifest, episodes) episode_mb = byte_count / len(episodes) / 1024**2 job_count = max(timings["jobs"], 1.0) diff --git a/src/lerobot/datasets/episode_video_streaming.py b/src/lerobot/datasets/episode_video_streaming.py index f48d578ec..10d5ce4bd 100644 --- a/src/lerobot/datasets/episode_video_streaming.py +++ b/src/lerobot/datasets/episode_video_streaming.py @@ -981,6 +981,108 @@ class EpisodeVideoManifest: ) +class ExactCoveragePool: + """Deterministic, exactly-once frame coverage over a byte-cache episode pool. + + The sampled/with-replacement pool (``run_pool_stream_simulation``) never guarantees a full + epoch: frames are drawn randomly and episodes rotate on a fixed cadence. This planner instead + enumerates *every frame of every episode exactly once per epoch* while keeping at most + ``pool_size`` episodes resident, so batch mixing stays high but coverage is complete and + reproducible. + + Mechanics (this is the "evict only when all frames sampled" model): + - Episodes are admitted in a seeded global permutation; the first ``pool_size`` fill the pool. + - Each resident episode carries a seeded shuffle of its own frame indices. + - Each draw picks a resident episode with probability proportional to its *remaining* frames + (i.e. a uniform draw over all remaining frames in the pool, the map-style ideal) and pops + one frame. + - An episode is evicted only when its last frame is emitted; a new episode is then admitted. + - The epoch ends when the admission order is exhausted and every resident episode is drained. + + Newly admitted episodes are surfaced via :attr:`newly_admitted` (drain it to drive prefetch) + and evictions via :attr:`evicted` (drain to release cache bytes). The planner does no I/O and + is fully unit-testable. It yields ``(episode_index, frame_index)``; map to a decode timestamp + with ``frame_index / max(frame_count - 1, 1)``. + + Determinism: the order is a pure function of ``(seed, epoch)`` and the episode->frame_count + map. Resume is a deterministic fast-forward: re-instantiate with the same seed/epoch and skip + ``n`` samples (tabular only, no decode). + """ + + def __init__( + self, + episode_frame_counts: Sequence[tuple[int, int]], + pool_size: int, + *, + seed: int, + epoch: int = 0, + ): + self._counts = {int(ep): int(n) for ep, n in episode_frame_counts if int(n) > 0} + self._rng = np.random.default_rng([seed, epoch]) + order = np.array(sorted(self._counts), dtype=np.int64) + self._rng.shuffle(order) + # Full admission order is exposed so a caller can prefetch episodes ahead of when they + # enter the sampling pool (a freshly admitted episode is immediately eligible to be drawn, + # so its bytes must already be resident). + self.admission_order: list[int] = order.tolist() + self._admit_cursor = 0 + self._remaining: dict[int, list[int]] = {} + self._remaining_total = 0 + self.newly_admitted: list[int] = [] + self.evicted: list[int] = [] + for _ in range(max(1, pool_size)): + self._admit_one() + + def _admit_one(self) -> None: + if self._admit_cursor >= len(self.admission_order): + return + ep = self.admission_order[self._admit_cursor] + self._admit_cursor += 1 + n = self._counts[ep] + frames = np.arange(n, dtype=np.int64) + self._rng.shuffle(frames) + self._remaining[ep] = frames.tolist() + self._remaining_total += n + self.newly_admitted.append(ep) + + @property + def remaining_total(self) -> int: + return self._remaining_total + + @property + def admitted_count(self) -> int: + """Number of episodes pulled from the admission order so far (pool fills + rotations).""" + return self._admit_cursor + + @property + def resident(self) -> list[int]: + return list(self._remaining) + + def __iter__(self) -> "ExactCoveragePool": + return self + + def __next__(self) -> tuple[int, int]: + if self._remaining_total == 0: + raise StopIteration + # Uniform draw over all remaining frames in the pool: walk the residents by cumulative + # remaining count. O(pool_size) per draw (~1024) -> negligible next to decode. + target = int(self._rng.integers(self._remaining_total)) + chosen = None + for ep, frames in self._remaining.items(): + if target < len(frames): + chosen = ep + break + target -= len(frames) + frames = self._remaining[chosen] + frame_index = frames.pop() + self._remaining_total -= 1 + if not frames: + del self._remaining[chosen] + self.evicted.append(chosen) + self._admit_one() + return chosen, frame_index + + class EpisodeByteCache: def __init__( self, diff --git a/tests/datasets/test_exact_coverage_pool.py b/tests/datasets/test_exact_coverage_pool.py new file mode 100644 index 000000000..476c17f76 --- /dev/null +++ b/tests/datasets/test_exact_coverage_pool.py @@ -0,0 +1,97 @@ +# Copyright 2025 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 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""ExactCoveragePool: exactly-once frame coverage over a bounded episode pool.""" + +from collections import Counter + +from lerobot.datasets.episode_video_streaming import ExactCoveragePool + +EPISODES = [(0, 5), (1, 3), (2, 8), (3, 1), (4, 6), (5, 4), (6, 7), (7, 2)] +TOTAL = sum(n for _, n in EPISODES) +EXPECTED = Counter((ep, i) for ep, n in EPISODES for i in range(n)) + + +def _drain(pool): + out, max_resident = [], 0 + while True: + try: + out.append(next(pool)) + except StopIteration: + break + max_resident = max(max_resident, len(pool.resident)) + return out, max_resident + + +def test_exact_once_coverage(): + out, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=42)) + assert len(out) == TOTAL + assert Counter(out) == EXPECTED # every (episode, frame) exactly once, no dups/misses + + +def test_pool_never_exceeds_size(): + _, max_resident = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=42)) + assert max_resident <= 3 + + +def test_deterministic_per_seed_and_epoch(): + a, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=7)) + b, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=7)) + c, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=8)) + d, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=7, epoch=1)) + assert a == b + assert a != c and a != d # seed and epoch both change the order + assert Counter(c) == EXPECTED and Counter(d) == EXPECTED # ... but coverage is preserved + + +def test_admission_and_eviction_events(): + pool = ExactCoveragePool(EPISODES, pool_size=3, seed=0) + admitted_ever, evicted_ever = set(), set() + # first three episodes admitted at construction + admitted_ever.update(pool.newly_admitted) + assert len(admitted_ever) == 3 + while True: + pool.newly_admitted.clear() + pool.evicted.clear() + try: + next(pool) + except StopIteration: + break + admitted_ever.update(pool.newly_admitted) + evicted_ever.update(pool.evicted) + assert admitted_ever == {ep for ep, _ in EPISODES} # every episode admitted exactly once + # every episode except the pool_size still resident at the end is evicted on exhaustion + assert len(evicted_ever) >= len(EPISODES) - 3 + + +def test_uniform_mixing_matches_coupon_collector(): + # 64 equal episodes, pool 64, first 64 draws -> ~64*(1-(1-1/64)^64) ~= 41 distinct + big = [(e, 100) for e in range(64)] + pool = ExactCoveragePool(big, pool_size=64, seed=0) + head = [next(pool)[0] for _ in range(64)] + assert len(set(head)) >= 30 # far above sequential (=1); ~41 expected + + +def test_large_epoch_bounded_and_complete(): + big = [(e, 90) for e in range(500)] + out, max_resident = _drain(ExactCoveragePool(big, pool_size=64, seed=3)) + assert len(out) == 500 * 90 + assert len(set(out)) == 500 * 90 # exactly once + assert max_resident <= 64 + + +def test_zero_length_episodes_skipped(): + pool = ExactCoveragePool([(0, 3), (1, 0), (2, 2)], pool_size=8, seed=0) + out, _ = _drain(pool) + assert Counter(out) == Counter({(0, 0): 1, (0, 1): 1, (0, 2): 1, (2, 0): 1, (2, 1): 1})