From ee06d1005f7d1bb515fa974c8f6e86200f57a69c Mon Sep 17 00:00:00 2001 From: Pepijn Date: Thu, 23 Jul 2026 17:12:04 +0200 Subject: [PATCH] refactor(dataset): make streaming pool rank-owned --- docs/source/training_dataset_streaming.mdx | 47 +++--- examples/training/train_with_streaming.py | 6 +- scripts/bench_streaming_dataset.py | 25 +++- src/lerobot/configs/default.py | 4 +- src/lerobot/datasets/streaming_dataset.py | 134 ++++++++---------- src/lerobot/scripts/lerobot_train.py | 26 ++-- src/lerobot/streaming/episode_video.py | 132 ++++++++++++----- src/lerobot/streaming/mp4.py | 41 +++++- .../datasets/test_episode_video_streaming.py | 10 ++ tests/datasets/test_exact_coverage_pool.py | 49 +++++++ tests/datasets/test_streaming_production.py | 57 ++++++-- 11 files changed, 370 insertions(+), 161 deletions(-) diff --git a/docs/source/training_dataset_streaming.mdx b/docs/source/training_dataset_streaming.mdx index 52bc3b45c..477dcb354 100644 --- a/docs/source/training_dataset_streaming.mdx +++ b/docs/source/training_dataset_streaming.mdx @@ -17,11 +17,16 @@ during training. Recording and rollout encoders are not used by this training pa ## How an epoch is read -Each distributed rank and DataLoader worker owns a disjoint set of complete episodes. Within that -set, an exact-coverage sampler mixes a bounded pool of episodes while visiting every selected frame -once per epoch. Parquet columns and compressed MP4 byte ranges are prefetched from the same episode -admission frontier. Temporal history and future windows are resolved inside the complete episode, -including the same boundary padding masks as map-style loading. +Each distributed rank owns a deterministic, frame-balanced set of complete episodes. One logical +exact-coverage pool per rank mixes those episodes while visiting every selected frame once per +rank-local coverage epoch. Parquet columns and compressed MP4 byte ranges are prefetched from the +same byte-aware admission frontier. Temporal history and future windows are resolved inside the +complete episode, including the same boundary padding masks as map-style loading. + +When `--num_workers` is nonzero, training uses one dedicated DataLoader process per rank. Its +bounded result queue holds decoded batches while the policy trains. The configured worker count is +instead used as internal Parquet and byte-range fetch concurrency, so increasing it does not create +independent samplers or multiply the cache budget. The default map-style `LeRobotDataset` behavior is unchanged when `--dataset.streaming=false`. @@ -55,14 +60,15 @@ The production defaults are: | Option | Default | Meaning | | ----------------------------- | ------: | --------------------------------------------------- | -| `streaming_episode_pool_size` | 32 | Complete episodes mixed by each worker | +| `streaming_episode_pool_size` | 32 | Maximum complete episodes mixed by each rank | | `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool | -| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per worker | +| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per rank | | `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root | -Memory limits are per DataLoader worker. For example, four workers with the default byte budget can -use up to 32 GiB for compressed episode video bytes, plus Parquet batches, decoders, and framework -overhead. Start with fewer workers or a smaller budget on memory-constrained hosts: +The active episode set is capped by both episode count and the exact synthesized mini-MP4 sizes +computed from the sidecar. An episode larger than the complete rank budget fails before training +fetches its payload. The cache, decoder LRU, and decoded-batch queue remain independently bounded. +Start with a smaller pool or budget on memory-constrained hosts: ```bash lerobot-train \ @@ -71,7 +77,7 @@ lerobot-train \ --dataset.streaming_episode_pool_size=16 \ --dataset.streaming_prefetch_episodes=4 \ --dataset.streaming_byte_budget_gb=4 \ - --num_workers=2 \ + --num_workers=4 \ --policy.type=act \ --output_dir=outputs/train/act_streaming ``` @@ -85,8 +91,12 @@ If metadata remains in a dataset repository while payload files are mirrored els The earlier streaming reader used a bounded row shuffle buffer. The episode reader instead has deterministic exact-coverage ordering derived from the seed and epoch. Checkpoint resume restores the per-rank sample offset using the checkpoint batch size. Changing distributed world size or -DataLoader worker count changes episode ownership; changing batch size changes the batch boundaries. -For sample-exact comparisons, resume with the same world size, worker count, and batch size. +batch size changes ownership or batch boundaries. For sample-exact comparisons, resume with the +same world size and batch size. Keep the same internal fetch concurrency when comparing performance. + +Streaming checkpoints created by the earlier multi-worker sampler record a different ownership +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 @@ -97,7 +107,8 @@ uv run python scripts/bench_streaming_dataset.py \ --repo-id=OWNER/DATASET \ --revision=COMMIT_SHA \ --batch-size=16 \ - --num-workers=4 \ + --num-workers=1 \ + --fetch-workers=4 \ --summary-json=streaming-benchmark.json ``` @@ -117,10 +128,10 @@ from the integrated dataset benchmark and end-to-end training A/B. publish a validated complete sidecar explicitly. - **A sidecar lock times out:** another process may still be indexing the same revision. Confirm it is healthy before removing a stale lock. -- **A rank owns no data:** reduce the number of ranks or workers so every rank owns at least one - selected episode. -- **The byte budget is exceeded:** lower the episode pool, increase the per-worker byte budget, or - use fewer/larger episodes per worker. +- **A rank owns no data:** reduce the number of ranks so every rank owns at least one selected + episode. +- **The byte budget is exceeded:** lower the episode pool, increase the per-rank byte budget, or + use an explicitly prepared payload layout with smaller episode ranges. - **Remote reads cannot authenticate:** verify the normal Hugging Face token or fsspec credentials are available in every worker environment. Credentials are never embedded in the sidecar. - **Refill stalls are high:** compare p95/p99 batch wait, reduce network contention, raise prefetch diff --git a/examples/training/train_with_streaming.py b/examples/training/train_with_streaming.py index 9cf4cd764..b1696c25b 100644 --- a/examples/training/train_with_streaming.py +++ b/examples/training/train_with_streaming.py @@ -83,10 +83,12 @@ def main(): optimizer = torch.optim.Adam(policy.parameters(), lr=1e-4) dataloader = torch.utils.data.DataLoader( dataset, - num_workers=4, + # One worker owns the rank-level pool. Internal fetch concurrency is configured through + # StreamingLeRobotDataset.max_num_shards (the lerobot-train CLI derives it from num_workers). + num_workers=1, batch_size=16, pin_memory=device.type != "cpu", - prefetch_factor=2, # loads batches with multiprocessing while policy trains + prefetch_factor=2, # bounded decoded-batch queue while the policy trains persistent_workers=True, ) diff --git a/scripts/bench_streaming_dataset.py b/scripts/bench_streaming_dataset.py index dc14cf070..3f163c629 100644 --- a/scripts/bench_streaming_dataset.py +++ b/scripts/bench_streaming_dataset.py @@ -22,6 +22,7 @@ import statistics import subprocess import sys import time +from collections.abc import Sequence from pathlib import Path import torch @@ -37,7 +38,14 @@ def parse_args() -> argparse.Namespace: 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, default=4) + 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) @@ -48,7 +56,7 @@ def parse_args() -> argparse.Namespace: return parser.parse_args() -def percentile(values: list[float], quantile: float) -> float: +def percentile(values: Sequence[float], quantile: float) -> float: if not values: return 0.0 ordered = sorted(values) @@ -83,6 +91,8 @@ def child_process_max_rss_mb() -> float: 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() @@ -95,7 +105,7 @@ def main() -> None: episode_pool_size=args.episode_pool_size, prefetch_episodes=args.prefetch_episodes, byte_budget_gb=args.byte_budget_gb, - max_num_shards=max(1, args.num_workers), + max_num_shards=max(1, args.fetch_workers), return_uint8=True, ) dataset_init_s = time.perf_counter() - init_start @@ -112,6 +122,7 @@ def main() -> None: waits: list[float] = [] measured_samples = 0 measured_indices: list[int] = [] + unique_episodes_per_batch: list[int] = [] first_batch_s = 0.0 exhausted = False @@ -129,7 +140,9 @@ def main() -> None: 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) @@ -155,12 +168,18 @@ def main() -> None: "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, diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index da27610fa..cf3de3c85 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -46,11 +46,11 @@ class DatasetConfig: streaming: bool = False # Optional data-plane root used by training-time streaming. Metadata still resolves from repo_id/root. streaming_data_root: str | None = None - # Number of complete episodes mixed by the exact-coverage sampler in each DataLoader worker. + # Number of complete episodes mixed by the rank-level exact-coverage sampler. streaming_episode_pool_size: int = 32 # Complete episodes fetched ahead of the current admission frontier. streaming_prefetch_episodes: int = 8 - # Hard per-worker cap for synthesized episode-video bytes. + # Hard per-rank cap for synthesized episode-video bytes. streaming_byte_budget_gb: float = 8.0 # Fraction of episodes held out per task for offline evaluation (0.0 = disabled). eval_split: float = 0.0 diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index 2fb4c7b28..5250e34be 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -15,7 +15,7 @@ # limitations under the License. import io import os -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from concurrent.futures import Future, ThreadPoolExecutor from pathlib import Path @@ -41,12 +41,31 @@ from .utils import check_version_compatibility from .video_utils import decode_video_frames_pyav +def _balanced_episode_shards( + episode_indices: list[int], + episode_frame_counts: Mapping[int, int], + *, + world_size: int, +) -> list[list[int]]: + """Assign whole episodes deterministically with greedy frame-count balancing.""" + if world_size <= 0: + raise ValueError("world_size must be positive") + shards: list[list[int]] = [[] for _ in range(world_size)] + shard_frames = [0] * world_size + for episode in sorted(episode_indices, key=lambda item: (-episode_frame_counts[item], item)): + rank = min(range(world_size), key=lambda item: (shard_frames[item], item)) + shards[rank].append(episode) + shard_frames[rank] += episode_frame_counts[episode] + return shards + + class StreamingLeRobotDataset(torch.utils.data.IterableDataset): """Episode-scoped streaming reader for LeRobot datasets. - Metadata is cached locally, while each worker reads only the Parquet rows and MP4 byte ranges - needed for the complete episodes it owns. Episode ownership is disjoint across distributed - ranks and dataloader workers, and every selected frame is yielded exactly once per iteration. + Metadata is cached locally, while each rank reads only the Parquet rows and MP4 byte ranges + needed for the complete episodes it owns. One DataLoader worker per rank owns the logical pool; + its bounded result queue provides decoded-batch prefetch without creating independent samplers. + Episode ownership is disjoint and every selected frame is yielded exactly once per iteration. MP4 sidecars are resolved automatically and built in a revision-keyed local cache when absent. Example: @@ -123,8 +142,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): and fsspec URLs. episode_pool_size (int | None, optional): Number of complete episodes in the sampling pool. prefetch_episodes (int, optional): Episodes prefetched beyond the active pool. - byte_budget_gb (float, optional): Per-worker upper bound for synthesized episode video bytes. - repeat (bool, optional): Repeat worker-local exact-coverage epochs without yielding a + byte_budget_gb (float, optional): Per-rank upper bound for synthesized episode video bytes. + repeat (bool, optional): Repeat rank-local exact-coverage epochs without yielding a short final training batch. The training factory enables this; direct iteration is finite by default. """ @@ -266,39 +285,40 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): self._active_epoch = epoch worker = torch.utils.data.get_worker_info() - worker_index = worker.id if worker is not None else 0 - worker_count = worker.num_workers if worker is not None else 1 + if worker is not None and worker.num_workers > 1: + raise RuntimeError( + "StreamingLeRobotDataset uses one rank-level sampling pool and supports at most " + "one DataLoader worker per rank" + ) resume_offset = self._resume_offset self._resume_offset = 0 - worker_resume_offset, logical_worker_index = self._worker_resume_state( - resume_offset, - batch_size=self._resume_batch_size, - worker_index=worker_index, - worker_count=worker_count, - ) - consumer_episodes, _consumer_index, _consumer_count = self._consumer_episodes( - worker_id=logical_worker_index, - num_workers=worker_count, - ) + consumer_episodes, _rank, _world_size = self._rank_episodes() consumer_frame_count = sum(self._episode_frame_count(episode) for episode in consumer_episodes) if consumer_frame_count: - worker_epoch_delta, worker_resume_offset = divmod( - worker_resume_offset, - consumer_frame_count, - ) + worker_epoch_delta, resume_offset = divmod(resume_offset, consumer_frame_count) else: worker_epoch_delta = 0 epoch += worker_epoch_delta self._active_epoch = epoch if self.shuffle: self._next_epoch = epoch + 1 + + max_workers = min(self.max_num_shards, max(1, self.episode_pool_size + self.prefetch_episodes)) + video_cache = self._make_video_cache(consumer_episodes, max_workers) + episode_byte_sizes = ( + {episode: video_cache.manifest.episode_byte_size(episode) for episode in consumer_episodes} + if video_cache is not None + else None + ) planner = ExactCoveragePool( [(episode, self._episode_frame_count(episode)) for episode in consumer_episodes], pool_size=self.episode_pool_size, seed=self.seed, epoch=epoch, + episode_byte_sizes=episode_byte_sizes, + byte_budget=self.byte_budget if episode_byte_sizes is not None else None, ) - for _ in range(worker_resume_offset): + for _ in range(resume_offset): try: next(planner) except StopIteration: @@ -307,11 +327,9 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): planner.evicted.clear() parquet_reader = EpisodeParquetReader(self._data_root, columns=self._projected_columns) - max_workers = min(self.max_num_shards, max(1, self.episode_pool_size + self.prefetch_episodes)) executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="lerobot-parquet") episode_futures: dict[int, Future[datasets.Dataset]] = {} - scheduled_cursor = 0 - video_cache = self._make_video_cache(consumer_episodes, max_workers) + scheduled_episodes: set[int] = set() retained_video_episodes: set[int] = set() if video_cache is not None: for episode_index in planner.resident: @@ -326,17 +344,14 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): return future def schedule_frontier() -> None: - nonlocal scheduled_cursor - target = min( - len(planner.admission_order), - planner.admitted_count + self.prefetch_episodes, - ) - while scheduled_cursor < target: - episode_index = planner.admission_order[scheduled_cursor] + frontier = [*planner.resident, *planner.prefetch_candidates(self.prefetch_episodes)] + for episode_index in frontier: + if episode_index in scheduled_episodes: + continue submit(episode_index) if video_cache is not None: video_cache.submit_prefetch(episode_index) - scheduled_cursor += 1 + scheduled_episodes.add(episode_index) schedule_frontier() try: @@ -379,12 +394,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): if video_cache is not None: video_cache.close() - def _consumer_episodes( - self, - *, - worker_id: int | None = None, - num_workers: int | None = None, - ) -> tuple[list[int], int, int]: + def _rank_episodes(self) -> tuple[list[int], int, int]: if torch.distributed.is_available() and torch.distributed.is_initialized(): rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() @@ -393,49 +403,23 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): world_size = int(os.environ.get("WORLD_SIZE", "1")) if world_size <= 0 or rank < 0 or rank >= world_size: raise ValueError(f"Invalid distributed rank/world size: rank={rank}, world_size={world_size}") - - worker = torch.utils.data.get_worker_info() - if worker_id is None: - worker_id = worker.id if worker is not None else 0 - if num_workers is None: - num_workers = worker.num_workers if worker is not None else 1 - consumer_index = rank * num_workers + worker_id - consumer_count = world_size * num_workers - return self._selected_episodes[consumer_index::consumer_count], consumer_index, consumer_count - - @staticmethod - def _worker_resume_state( - offset: int, - *, - batch_size: int, - worker_index: int, - worker_count: int, - ) -> tuple[int, int]: - completed_batches, partial_batch = divmod(offset, batch_size) - completed_cycles, next_logical_worker = divmod(completed_batches, worker_count) - logical_worker = (worker_index + next_logical_worker) % worker_count - worker_batches = completed_cycles + int(logical_worker < next_logical_worker) - worker_offset = worker_batches * batch_size - if partial_batch and logical_worker == next_logical_worker: - worker_offset += partial_batch - return worker_offset, logical_worker + counts = {episode: self._episode_frame_count(episode) for episode in self._selected_episodes} + shards = _balanced_episode_shards(self._selected_episodes, counts, world_size=world_size) + return shards[rank], rank, world_size def _episode_frame_count(self, episode_index: int) -> int: episode = self.meta.episodes[episode_index] return int(episode["dataset_to_index"] - episode["dataset_from_index"]) def num_frames_for_rank(self, rank: int, world_size: int, num_workers: int) -> int: - """Return frames owned by one training rank under episode-scoped worker sharding.""" + """Return frames owned by one training rank under balanced whole-episode sharding.""" if world_size <= 0 or rank < 0 or rank >= world_size: raise ValueError(f"Invalid distributed rank/world size: rank={rank}, world_size={world_size}") - workers = max(1, num_workers) - consumer_count = world_size * workers - consumer_indices = set(range(rank * workers, (rank + 1) * workers)) - return sum( - self._episode_frame_count(episode) - for position, episode in enumerate(self._selected_episodes) - if position % consumer_count in consumer_indices - ) + if num_workers > 1: + raise ValueError("Rank-level streaming supports at most one DataLoader worker per rank") + counts = {episode: self._episode_frame_count(episode) for episode in self._selected_episodes} + shards = _balanced_episode_shards(self._selected_episodes, counts, world_size=world_size) + return sum(counts[episode] for episode in shards[rank]) def _load_episode_dataset( self, diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 0c96ac38f..fe89d2564 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -460,22 +460,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): else: shuffle = True sampler = None - train_num_workers = cfg.num_workers - if train_num_workers > 0: - max_nonempty_workers = dataset.num_episodes // accelerator.num_processes - if max_nonempty_workers == 0: - raise ValueError( - "At least one distributed rank owns no streaming episode. " - "Reduce the distributed world size." - ) - train_num_workers = min(train_num_workers, max_nonempty_workers) - if train_num_workers != cfg.num_workers and is_main_process: - logging.warning( - "Reducing streaming DataLoader workers from %d to %d so every rank/worker " - "owns at least one complete episode.", - cfg.num_workers, - train_num_workers, - ) + # One dedicated DataLoader process owns the rank-level planner/cache. Its bounded result + # queue provides decoded-batch prefetch; max_num_shards controls the internal Parquet/range + # fetch executors independently. + train_num_workers = min(cfg.num_workers, 1) + if cfg.num_workers > 1 and is_main_process: + logging.info( + "Using one streaming DataLoader worker per rank; %d configured workers remain " + "available as the dataset's internal fetch concurrency.", + cfg.num_workers, + ) rank_frames = dataset.num_frames_for_rank( accelerator.process_index, accelerator.num_processes, diff --git a/src/lerobot/streaming/episode_video.py b/src/lerobot/streaming/episode_video.py index aed528447..029a07227 100644 --- a/src/lerobot/streaming/episode_video.py +++ b/src/lerobot/streaming/episode_video.py @@ -17,7 +17,7 @@ import posixpath import threading import time from collections import OrderedDict -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from concurrent.futures import Future, ThreadPoolExecutor, as_completed from dataclasses import dataclass from datetime import UTC, datetime @@ -33,7 +33,13 @@ import numpy as np from huggingface_hub import HfApi, HfFileSystem, constants from huggingface_hub.utils import get_session, hf_raise_for_status -from lerobot.streaming.mp4 import Mp4Index, Mp4SampleSlice, fetch_mp4_index, synthesize_mp4 +from lerobot.streaming.mp4 import ( + Mp4Index, + Mp4SampleSlice, + fetch_mp4_index, + synthesize_mp4, + synthesized_mp4_size, +) if TYPE_CHECKING: from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata @@ -1041,6 +1047,16 @@ class EpisodeVideoManifest: source_start_pts=span.source_start_pts, ) + def episode_byte_size(self, episode_index: int) -> int: + """Exact synthesized video bytes retained while an episode is active.""" + return sum( + synthesized_mp4_size( + self.mp4_index(episode_index, camera_key), + self.sample_slice(episode_index, camera_key), + ) + for camera_key in self.video_keys + ) + class ExactCoveragePool: """Deterministic, exactly-once frame coverage over a byte-cache episode pool. @@ -1052,7 +1068,8 @@ class ExactCoveragePool: 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. + - Episodes are admitted in a seeded global permutation until either ``pool_size`` or the + optional indexed-byte budget is reached. - 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 @@ -1065,9 +1082,9 @@ class ExactCoveragePool: 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). + Determinism: the order is a pure function of ``(seed, epoch)``, the episode frame counts, and + optional byte sizes/budget. Resume is a deterministic fast-forward: re-instantiate with the + same inputs and skip ``n`` samples (tabular only, no decode). """ def __init__( @@ -1077,34 +1094,71 @@ class ExactCoveragePool: *, seed: int, epoch: int = 0, + episode_byte_sizes: Mapping[int, int] | None = None, + byte_budget: int | None = None, ): 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.pool_size = max(1, pool_size) + self._byte_budget = byte_budget + if byte_budget is not None and byte_budget <= 0: + raise ValueError("byte_budget must be positive") + if byte_budget is not None and episode_byte_sizes is None: + raise ValueError("episode_byte_sizes are required when byte_budget is set") + self._byte_sizes = { + episode: int(episode_byte_sizes[episode]) if episode_byte_sizes is not None else 0 + for episode in self._counts + } + if any(size < 0 for size in self._byte_sizes.values()): + raise ValueError("episode byte sizes must be non-negative") + if byte_budget is not None: + oversized = next( + ((episode, size) for episode, size in self._byte_sizes.items() if size > byte_budget), + None, + ) + if oversized is not None: + episode, size = oversized + raise ValueError( + f"Episode {episode} requires {size} bytes, exceeding the byte budget {byte_budget}" + ) + + # Preserve the full seeded order for benchmark/tooling compatibility. Byte-aware admission + # may temporarily skip an entry, but every episode remains in this deterministic frontier. self.admission_order: list[int] = order.tolist() - self._admit_cursor = 0 - self._remaining: dict[int, list[int]] = {} + self._pending: list[int] = list(self.admission_order) + self._admitted_count = 0 + self._remaining: dict[int, tuple[np.ndarray, int]] = {} self._remaining_total = 0 + self._resident_bytes = 0 self.newly_admitted: list[int] = [] self.evicted: list[int] = [] - for _ in range(max(1, pool_size)): - self._admit_one() + self._admit_available() - 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) + def _admit_available(self) -> None: + while len(self._remaining) < self.pool_size and self._pending: + available_bytes = None if self._byte_budget is None else self._byte_budget - self._resident_bytes + pending_index = next( + ( + index + for index, episode in enumerate(self._pending) + if available_bytes is None or self._byte_sizes[episode] <= available_bytes + ), + None, + ) + if pending_index is None: + return + + episode = self._pending.pop(pending_index) + frame_count = self._counts[episode] + frames = np.arange(frame_count, dtype=np.int64) + self._rng.shuffle(frames) + self._remaining[episode] = (frames, frame_count) + self._remaining_total += frame_count + self._resident_bytes += self._byte_sizes[episode] + self._admitted_count += 1 + self.newly_admitted.append(episode) @property def remaining_total(self) -> int: @@ -1113,12 +1167,22 @@ class ExactCoveragePool: @property def admitted_count(self) -> int: """Number of episodes pulled from the admission order so far (pool fills + rotations).""" - return self._admit_cursor + return self._admitted_count @property def resident(self) -> list[int]: return list(self._remaining) + @property + def resident_bytes(self) -> int: + return self._resident_bytes + + def prefetch_candidates(self, count: int) -> list[int]: + """Return the next deterministic pending frontier without admitting it.""" + if count <= 0: + return [] + return self._pending[:count] + def __iter__(self) -> ExactCoveragePool: return self @@ -1129,18 +1193,22 @@ class ExactCoveragePool: # 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): + for ep, (_frames, remaining) in self._remaining.items(): + if target < remaining: chosen = ep break - target -= len(frames) - frames = self._remaining[chosen] - frame_index = frames.pop() + target -= remaining + frames, remaining = self._remaining[chosen] + remaining -= 1 + frame_index = int(frames[remaining]) self._remaining_total -= 1 - if not frames: + if remaining == 0: del self._remaining[chosen] self.evicted.append(chosen) - self._admit_one() + self._resident_bytes -= self._byte_sizes[chosen] + self._admit_available() + else: + self._remaining[chosen] = (frames, remaining) return chosen, frame_index diff --git a/src/lerobot/streaming/mp4.py b/src/lerobot/streaming/mp4.py index 4baf513c6..4b58d20bc 100644 --- a/src/lerobot/streaming/mp4.py +++ b/src/lerobot/streaming/mp4.py @@ -329,10 +329,49 @@ def synthesize_mp4(index: Mp4Index, sample_slice: Mp4SampleSlice, mdat_payload: sync = index.sync_samples[(index.sync_samples >= lo) & (index.sync_samples < hi)] - lo + 1 moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=0) header_size = len(index.ftyp) + len(moov) - moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=header_size + 8) + mdat_header_size = 8 if len(mdat_payload) + 8 <= 0xFFFFFFFF else 16 + moov = _make_moov( + index, + durations, + sizes, + rel_offsets, + sync, + mdat_data_offset=header_size + mdat_header_size, + ) return index.ftyp + moov + _box(b"mdat", mdat_payload) +def synthesized_mp4_size(index: Mp4Index, sample_slice: Mp4SampleSlice) -> int: + """Return the exact synthesized mini-MP4 size without fetching its media payload.""" + lo = sample_slice.sample_lo + hi = sample_slice.sample_hi + 1 + if lo < 0 or hi > len(index.sample_sizes) or lo >= hi: + raise ValueError(f"Invalid sample range [{lo}, {hi}) for {index.file_path}") + + offsets = index.sample_offsets[lo:hi] + sizes = index.sample_sizes[lo:hi] + rel_offsets = offsets - sample_slice.byte_offset + if int(rel_offsets.min()) != 0: + raise ValueError("Sample slice must start at the minimum referenced sample offset") + if int((rel_offsets + sizes).max()) > sample_slice.byte_length: + raise ValueError("Sample slice does not cover all referenced samples") + + durations = index.sample_durations[lo:hi] + sync = index.sync_samples[(index.sync_samples >= lo) & (index.sync_samples < hi)] - lo + 1 + moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=0) + header_size = len(index.ftyp) + len(moov) + mdat_header_size = 8 if sample_slice.byte_length + 8 <= 0xFFFFFFFF else 16 + moov = _make_moov( + index, + durations, + sizes, + rel_offsets, + sync, + mdat_data_offset=header_size + mdat_header_size, + ) + return len(index.ftyp) + len(moov) + mdat_header_size + sample_slice.byte_length + + def iter_boxes( data: bytes, start: int, diff --git a/tests/datasets/test_episode_video_streaming.py b/tests/datasets/test_episode_video_streaming.py index eedb1fdc6..7fe78dc8c 100644 --- a/tests/datasets/test_episode_video_streaming.py +++ b/tests/datasets/test_episode_video_streaming.py @@ -38,6 +38,7 @@ from lerobot.streaming.mp4 import ( _vmhd, parse_mp4_index, synthesize_mp4, + synthesized_mp4_size, ) @@ -90,6 +91,15 @@ def test_synthesized_mp4_rebases_one_chunk_per_sample_offsets(): np.testing.assert_array_equal(mini_index.sample_sizes, np.array([10, 10, 10])) +def test_synthesized_mp4_size_matches_materialized_bytes(): + mp4 = parse_mp4_index("test.mp4", _minimal_mp4([10_000, 10_050, 10_025])) + sample_slice = mp4.sample_slice(0.0, 2.0, keyframe_pad_s=0, keyframe_pad_fraction=0) + + mini = synthesize_mp4(mp4, sample_slice, b"x" * sample_slice.byte_length) + + assert synthesized_mp4_size(mp4, sample_slice) == len(mini) + + def test_parser_accepts_co64_chunk_offsets(): mp4 = parse_mp4_index("test.mp4", _minimal_mp4([10_000, 10_050, 10_025], use_co64=True)) diff --git a/tests/datasets/test_exact_coverage_pool.py b/tests/datasets/test_exact_coverage_pool.py index e31a7354d..81002c1ae 100644 --- a/tests/datasets/test_exact_coverage_pool.py +++ b/tests/datasets/test_exact_coverage_pool.py @@ -16,6 +16,8 @@ from collections import Counter +import pytest + from lerobot.streaming.episode_video import ExactCoveragePool EPISODES = [(0, 5), (1, 3), (2, 8), (3, 1), (4, 6), (5, 4), (6, 7), (7, 2)] @@ -95,3 +97,50 @@ 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}) + + +def test_byte_aware_admission_never_exceeds_budget(): + sizes = {0: 7, 1: 6, 2: 4, 3: 3, 4: 2} + pool = ExactCoveragePool( + EPISODES[:5], + pool_size=4, + seed=9, + episode_byte_sizes=sizes, + byte_budget=10, + ) + out = [] + max_resident_bytes = pool.resident_bytes + while pool.remaining_total: + out.append(next(pool)) + max_resident_bytes = max(max_resident_bytes, pool.resident_bytes) + + assert Counter(out) == Counter((ep, frame) for ep, count in EPISODES[:5] for frame in range(count)) + assert max_resident_bytes <= 10 + assert len(pool.admission_order) == len(EPISODES[:5]) + + +def test_byte_aware_admission_rejects_one_oversized_episode(): + with pytest.raises(ValueError, match="Episode 1.*byte budget"): + ExactCoveragePool( + [(0, 2), (1, 3)], + pool_size=2, + seed=0, + episode_byte_sizes={0: 4, 1: 11}, + byte_budget=10, + ) + + +def test_prefetch_candidates_follow_deterministic_pending_frontier(): + pool = ExactCoveragePool( + EPISODES, + pool_size=2, + seed=17, + episode_byte_sizes={episode: 1 for episode, _ in EPISODES}, + byte_budget=2, + ) + + candidates = pool.prefetch_candidates(3) + + assert len(candidates) == 3 + assert not set(candidates) & set(pool.resident) + assert candidates == pool.prefetch_candidates(3) diff --git a/tests/datasets/test_streaming_production.py b/tests/datasets/test_streaming_production.py index 1db2c6751..197e16916 100644 --- a/tests/datasets/test_streaming_production.py +++ b/tests/datasets/test_streaming_production.py @@ -19,7 +19,7 @@ import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") -from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset +from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset, _balanced_episode_shards from lerobot.utils.utils import cycle from tests.fixtures.constants import DUMMY_REPO_ID @@ -209,6 +209,28 @@ def test_streaming_reads_video_bytes_from_configured_fsspec_root( assert torch.equal(sample[camera_key], reference[camera_key]) +def test_streaming_rejects_episode_larger_than_rank_byte_budget( + tmp_path: Path, lerobot_dataset_factory +) -> None: + root = tmp_path / "dataset" + lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=2, + total_frames=10, + ) + streaming = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + shuffle=False, + buffer_size=2, + byte_budget_gb=1 / 1024**3, + ) + + with pytest.raises(ValueError, match="Episode .*byte budget"): + next(iter(streaming)) + + def test_streaming_rank_shards_are_disjoint(tmp_path: Path, lerobot_dataset_factory, monkeypatch) -> None: root = tmp_path / "dataset" map_dataset = lerobot_dataset_factory( @@ -239,9 +261,22 @@ def test_streaming_rank_shards_are_disjoint(tmp_path: Path, lerobot_dataset_fact assert per_rank[0] | per_rank[1] == set(range(len(map_dataset))) -def test_streaming_workers_do_not_duplicate_frames(tmp_path: Path, lerobot_dataset_factory) -> None: +def test_rank_shards_are_greedily_balanced_by_frame_count() -> None: + shards = _balanced_episode_shards( + [0, 1, 2, 3, 4], + {0: 100, 1: 90, 2: 20, 3: 10, 4: 5}, + world_size=2, + ) + + assert {episode for shard in shards for episode in shard} == {0, 1, 2, 3, 4} + assert set(shards[0]).isdisjoint(shards[1]) + totals = [sum({0: 100, 1: 90, 2: 20, 3: 10, 4: 5}[episode] for episode in shard) for shard in shards] + assert max(totals) - min(totals) <= 15 + + +def test_streaming_rejects_multiple_sampling_workers(tmp_path: Path, lerobot_dataset_factory) -> None: root = tmp_path / "dataset" - map_dataset = lerobot_dataset_factory( + lerobot_dataset_factory( root=root, repo_id=DUMMY_REPO_ID, total_episodes=8, @@ -256,10 +291,8 @@ def test_streaming_workers_do_not_duplicate_frames(tmp_path: Path, lerobot_datas ) loader = torch.utils.data.DataLoader(streaming, batch_size=None, num_workers=2) - indices = [int(item["index"]) for item in loader] - - assert len(indices) == len(map_dataset) - assert set(indices) == set(range(len(map_dataset))) + with pytest.raises(RuntimeError, match="one DataLoader worker per rank"): + list(loader) def test_streaming_persistent_workers_advance_epochs(tmp_path: Path, lerobot_dataset_factory) -> None: @@ -281,7 +314,7 @@ def test_streaming_persistent_workers_advance_epochs(tmp_path: Path, lerobot_dat loader = torch.utils.data.DataLoader( streaming, batch_size=None, - num_workers=2, + num_workers=1, persistent_workers=True, ) try: @@ -317,7 +350,7 @@ def test_streaming_worker_exception_propagates_and_workers_stop( loader = torch.utils.data.DataLoader( streaming, batch_size=None, - num_workers=2, + num_workers=1, persistent_workers=True, ) try: @@ -376,7 +409,7 @@ def test_streaming_worker_resume_reproduces_remaining_stream( ) def load(dataset: StreamingLeRobotDataset) -> list[int]: - loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=2) + loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=1) if batch_size is None: return [int(item["index"]) for item in loader] return [int(index) for batch in loader for index in batch["index"]] @@ -450,7 +483,7 @@ def test_streaming_worker_resume_after_epoch_boundary(tmp_path: Path, lerobot_da loader = torch.utils.data.DataLoader( dataset, batch_size=4, - num_workers=2, + num_workers=1, persistent_workers=True, ) try: @@ -502,7 +535,7 @@ def test_streaming_local_training_step_smoke(tmp_path: Path, lerobot_dataset_fac buffer_size=2, repeat=True, ) - loader = torch.utils.data.DataLoader(dataset, batch_size=4, num_workers=2) + loader = torch.utils.data.DataLoader(dataset, batch_size=4, num_workers=1) iterator = iter(loader) try: batch = next(iterator)