diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 7f7a34e6a..9651266cd 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -37,6 +37,8 @@ - sections: - local: lerobot-dataset-v3 title: Using LeRobotDataset + - local: training_dataset_streaming + title: Training Dataset Streaming - local: porting_datasets_v3 title: Porting Large Datasets - local: using_dataset_tools diff --git a/docs/source/lerobot-dataset-v3.mdx b/docs/source/lerobot-dataset-v3.mdx index 0647af0b0..0f672c255 100644 --- a/docs/source/lerobot-dataset-v3.mdx +++ b/docs/source/lerobot-dataset-v3.mdx @@ -131,15 +131,68 @@ for batch in data_loader: # model.forward(batch) ``` -## Stream a dataset (no downloads) +## Stream a dataset during training -Use `StreamingLeRobotDataset` to iterate directly from the Hub without local copies. This allows to stream large datasets without the need to downloading them onto disk or loading them onto memory, and is a key feature of the new dataset format. +Enable training-time streaming with the public `--dataset.streaming=true` flag: + +```bash +lerobot-train \ + --dataset.repo_id=yaak-ai/L2D-v3 \ + --dataset.streaming=true \ + --policy.type=act \ + --output_dir=outputs/train/act_streaming +``` + +This is separate from `--dataset.streaming_encoding=true`, which controls video encoding while +recording. Training-time streaming leaves the map-style `LeRobotDataset` path unchanged. + +`StreamingLeRobotDataset` assigns complete episodes disjointly across distributed ranks and +DataLoader workers, reads only the selected episode rows from Parquet, and keeps a bounded pool of +episode video bytes. Every selected frame is visited exactly once per streaming epoch. + +On first use, LeRobot looks for a revision-matched MP4 index sidecar. If none is published with the +dataset, it builds one in the revision-keyed local LeRobot cache under a process lock and installs it +atomically. Training never uploads this sidecar. Dataset maintainers can build and publish one +explicitly with `scripts/build_mp4_sidecar.py --push`. + +The default memory cap is 8 GiB per DataLoader worker. It can be adjusted along with episode mixing +and prefetch: + +```bash +lerobot-train \ + --dataset.repo_id=yaak-ai/L2D-v3 \ + --dataset.streaming=true \ + --dataset.streaming_byte_budget_gb=4 \ + --dataset.streaming_episode_pool_size=16 \ + --dataset.streaming_prefetch_episodes=4 \ + --policy.type=act \ + --output_dir=outputs/train/act_streaming +``` + +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=4 \ + --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 from lerobot.datasets import StreamingLeRobotDataset repo_id = "yaak-ai/L2D-v3" -dataset = StreamingLeRobotDataset(repo_id) # streams directly from the Hub +dataset = StreamingLeRobotDataset(repo_id) ```
diff --git a/docs/source/training_dataset_streaming.mdx b/docs/source/training_dataset_streaming.mdx new file mode 100644 index 000000000..52bc3b45c --- /dev/null +++ b/docs/source/training_dataset_streaming.mdx @@ -0,0 +1,127 @@ +# Training Dataset Streaming + +Training-time dataset streaming lets `lerobot-train` consume a LeRobotDataset v3 without first +downloading its complete Parquet and video payload. Enable it through the existing public switch: + +```bash +lerobot-train \ + --dataset.repo_id=OWNER/DATASET \ + --dataset.streaming=true \ + --policy.type=act \ + --output_dir=outputs/train/act_streaming +``` + +This feature is independent of `--dataset.streaming_encoding=true`. `streaming_encoding` controls +how videos are written while recording; `dataset.streaming` controls how an existing dataset is read +during training. Recording and rollout encoders are not used by this training path. + +## 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. + +The default map-style `LeRobotDataset` behavior is unchanged when `--dataset.streaming=false`. + +## MP4 sidecars and the first run + +Video streaming uses a small MP4 index sidecar. Dataset initialization first checks the +revision-keyed local cache, then looks for a valid published sidecar. If neither is available, +LeRobot builds the sidecar locally while holding a process lock and installs it atomically. A +failed or interrupted build does not replace the previous valid file. + +Training is read-only: it never uploads a sidecar or modifies the dataset repository. On a cluster +with node-local caches, the first job may build once per node. A shared LeRobot cache avoids that +duplication. + +Dataset maintainers can build a sidecar ahead of time: + +```bash +uv run python scripts/build_mp4_sidecar.py \ + --repo-id=OWNER/DATASET \ + --revision=COMMIT_SHA \ + --data-root=hf://datasets/OWNER/DATASET@COMMIT_SHA \ + --output=/tmp/dataset-mp4-sidecar.npz +``` + +Publication is always explicit. Add `--push` only after validating the complete-dataset sidecar. +Subset sidecars cannot be published. + +## Configuration + +The production defaults are: + +| Option | Default | Meaning | +| ----------------------------- | ------: | --------------------------------------------------- | +| `streaming_episode_pool_size` | 32 | Complete episodes mixed by each worker | +| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool | +| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per worker | +| `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: + +```bash +lerobot-train \ + --dataset.repo_id=OWNER/DATASET \ + --dataset.streaming=true \ + --dataset.streaming_episode_pool_size=16 \ + --dataset.streaming_prefetch_episodes=4 \ + --dataset.streaming_byte_budget_gb=4 \ + --num_workers=2 \ + --policy.type=act \ + --output_dir=outputs/train/act_streaming +``` + +If metadata remains in a dataset repository while payload files are mirrored elsewhere, set +`--dataset.streaming_data_root`. Supported values include local paths, revision-qualified +`hf://datasets/...` roots, `hf://buckets/...` roots, and fsspec URLs. + +## Resume and shuffle migration + +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. + +## 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=4 \ + --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 + 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. +- **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 + gradually, and verify that the dataset sidecar matches the exact revision. diff --git a/examples/training/train_with_streaming.py b/examples/training/train_with_streaming.py index 4629bcda6..9cf4cd764 100644 --- a/examples/training/train_with_streaming.py +++ b/examples/training/train_with_streaming.py @@ -12,8 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""This script demonstrates how to train a Diffusion Policy on the PushT environment, -using a dataset processed in streaming mode.""" +"""Train directly from episode-scoped Parquet and MP4 streams. + +For normal training, prefer ``lerobot-train --dataset.streaming=true`` so distributed sharding, +checkpoint resume, and device placement are configured by the training pipeline. This lower-level +example shows the underlying Python API. +""" from pathlib import Path @@ -64,9 +68,17 @@ def main(): ACTION: [t / dataset_metadata.fps for t in range(cfg.n_action_steps)], } - # Instantiating the training dataset in streaming mode allows to not consume up memory as the data is fetched - # iteratively rather than being load into memory all at once. Retrieved frames are shuffled across epochs - dataset = StreamingLeRobotDataset(dataset_id, delta_timestamps=delta_timestamps, tolerance_s=1e-3) + # The first run resolves or locally builds a revision-safe MP4 index sidecar. It is never + # uploaded implicitly. Episode rows and video byte ranges are then prefetched together. + dataset = StreamingLeRobotDataset( + dataset_id, + delta_timestamps=delta_timestamps, + tolerance_s=1e-3, + episode_pool_size=16, + prefetch_episodes=4, + byte_budget_gb=4, + repeat=True, + ) optimizer = torch.optim.Adam(policy.parameters(), lr=1e-4) dataloader = torch.utils.data.DataLoader( @@ -74,8 +86,8 @@ def main(): num_workers=4, batch_size=16, pin_memory=device.type != "cpu", - drop_last=True, prefetch_factor=2, # loads batches with multiprocessing while policy trains + persistent_workers=True, ) # Run training loop. diff --git a/pyproject.toml b/pyproject.toml index 937f48a86..345e0b544 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,9 @@ dependencies = [ # Config & Hub "draccus==0.10.0", # TODO: Relax version constraint "huggingface-hub>=1.0.0,<2.0.0", + "filelock>=3.12.0,<4.0.0", + "fsspec>=2023.5.0,<2027.0.0", + "httpx>=0.27.0,<1.0.0", "requests>=2.32.0,<3.0.0", # Environments diff --git a/scripts/bench_episode_byte_cache.py b/scripts/bench_episode_byte_cache.py index 41b27e551..45b07e478 100644 --- a/scripts/bench_episode_byte_cache.py +++ b/scripts/bench_episode_byte_cache.py @@ -16,7 +16,6 @@ import os import random import resource import socket -import tempfile import threading import time from collections.abc import Sequence @@ -30,27 +29,24 @@ import pyarrow.compute as pc 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, -) from lerobot.datasets.video_utils import VideoDecoderCache, decode_video_frames_torchcodec - -DEFAULT_REPO = "allenai/MolmoAct2-BimanualYAM-Dataset" -DEFAULT_REVISION = "e9f21ae15074330839f2ac25ed4b49d76dfa1f9c" -DEFAULT_DATA_ROOT = "hf://buckets/pepijn223/MolmoAct2-BimanualYAM-Dataset-bucket" -SIDECAR_CACHE_DIR = Path(tempfile.gettempdir()) / "lerobot-sidecars" -FULL_SIDECAR_NAME = "molmoact2-full.npz" +from lerobot.streaming.episode_video import ( + EpisodeByteCache, + EpisodeVideoManifest, + ExactCoveragePool, +) def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Benchmark episode-level streaming mini-MP4 cache.") - parser.add_argument("--repo-id", default=DEFAULT_REPO) - parser.add_argument("--revision", default=DEFAULT_REVISION) - parser.add_argument("--data-root", default=DEFAULT_DATA_ROOT) + 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"), @@ -144,7 +140,6 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--in-memory", action="store_true", help="Accepted for compatibility; manifest is always in memory." ) - parser.add_argument("--no-hub-branch-assert", action="store_true") return parser.parse_args() @@ -461,7 +456,6 @@ def run_exact_coverage_stream( 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, @@ -584,74 +578,6 @@ def _root_join(data_root: str, relative_path: str) -> str: return str(Path(data_root) / relative_path) -def _find_or_download_sidecar(data_root: str, manifest_episode_count: int) -> Path | None: - _ = manifest_episode_count - local = SIDECAR_CACHE_DIR / FULL_SIDECAR_NAME - if _valid_sidecar(local): - return local - if local.exists(): - print(f"mp4_sidecar_invalid_local: {local}") - local.unlink() - remote_relative = f"meta/mp4-sidecars/{FULL_SIDECAR_NAME}" - remote = _root_join(data_root, remote_relative) - protocol = "hf" if data_root.startswith("hf://") else "file" - fs = fsspec.filesystem(protocol) - if not fs.exists(remote): - return None - local.parent.mkdir(parents=True, exist_ok=True) - print(f"downloading_mp4_sidecar: {remote} -> {local}") - if data_root.startswith("hf://"): - _download_sidecar_native_http(data_root, remote_relative, local) - else: - fs.get(remote, str(local)) - return local - - -def _valid_sidecar(path: Path) -> bool: - if not path.exists(): - return False - try: - with np.load(path, allow_pickle=False) as data: - return "manifest_json" in data - except Exception: - return False - - -def _download_sidecar_native_http(data_root: str, relative_path: str, local: Path) -> None: - fetcher = NativeHTTPRangeFetcher(data_root, max_connections=16) - tmp = local.with_suffix(local.suffix + ".tmp") - try: - size = fetcher.info_size(relative_path) - chunk_size = 16 * 1024 * 1024 - ranges = [(offset, min(chunk_size, size - offset)) for offset in range(0, size, chunk_size)] - with tmp.open("wb") as out_file: - out_file.truncate(size) - - def read_chunk(offset_length: tuple[int, int]) -> tuple[int, bytes]: - offset, length = offset_length - return offset, fetcher.read_range(relative_path, offset, length) - - start = time.perf_counter() - done = 0 - with ThreadPoolExecutor(max_workers=8) as pool: - futures = [pool.submit(read_chunk, item) for item in ranges] - with tmp.open("r+b") as rw_file: - for future in futures: - offset, data = future.result() - rw_file.seek(offset) - rw_file.write(data) - done += len(data) - elapsed = max(time.perf_counter() - start, 1e-9) - print( - f"sidecar_download: {done / 1024**2:.1f}/{size / 1024**2:.1f} MiB " - f"({done / elapsed / 1024**2:.1f} MiB/s)", - flush=True, - ) - tmp.replace(local) - finally: - fetcher.close() - - class EpisodeParquetReader: def __init__(self, meta: LeRobotDatasetMetadata, data_root: str): self.meta = meta @@ -1369,15 +1295,14 @@ def main() -> None: os.environ["LEROBOT_HTTP_FAILURE_LOG"] = args.http_failure_log print(f"http_failure_log: {args.http_failure_log}") data_root = args.data_root - if data_root.startswith("hf://") and not args.no_hub_branch_assert: - assert_hf_hub_range_cache_branch() - 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 = _find_or_download_sidecar(data_root, manifest_episode_count) + 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}") @@ -1419,15 +1344,7 @@ def main() -> None: ) return if args.strategy == "both": - expected_sidecar = SIDECAR_CACHE_DIR / FULL_SIDECAR_NAME - expected_remote = _root_join(data_root, f"meta/mp4-sidecars/{FULL_SIDECAR_NAME}") - print(f"mp4_sidecar_missing_local: {expected_sidecar}") - print(f"mp4_sidecar_missing_remote: {expected_remote}") - print( - "build_mp4_sidecar: " - "uv run --no-sync python scripts/build_mp4_sidecar.py " - f"--workers {args.workers} --range-backend native-http --output {expected_sidecar}" - ) + 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() diff --git a/scripts/bench_streaming_dataset.py b/scripts/bench_streaming_dataset.py new file mode 100644 index 000000000..dc14cf070 --- /dev/null +++ b/scripts/bench_streaming_dataset.py @@ -0,0 +1,179 @@ +#!/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 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, 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("--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: list[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() + 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, + max_num_shards=max(1, args.num_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] = [] + 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() + measured_indices.extend(int(index) for index in 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)), + "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, + "prefetch_factor": args.prefetch_factor, + "episode_pool_size": args.episode_pool_size, + "prefetch_episodes": args.prefetch_episodes, + "byte_budget_gb": args.byte_budget_gb, + "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/build_mp4_sidecar.py b/scripts/build_mp4_sidecar.py index 3fcb9ed8f..783daca00 100644 --- a/scripts/build_mp4_sidecar.py +++ b/scripts/build_mp4_sidecar.py @@ -17,76 +17,78 @@ from pathlib import Path import fsspec from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata -from lerobot.datasets.episode_video_streaming import EpisodeVideoManifest, assert_hf_hub_range_cache_branch - -DEFAULT_REPO = "allenai/MolmoAct2-BimanualYAM-Dataset" -DEFAULT_REVISION = "e9f21ae15074330839f2ac25ed4b49d76dfa1f9c" -DEFAULT_DATA_ROOT = "hf://buckets/pepijn223/MolmoAct2-BimanualYAM-Dataset-bucket" +from lerobot.datasets.streaming_sidecar import ( + build_mp4_sidecar, + make_sidecar_spec, + published_sidecar_url, + range_backend_for_root, +) +from lerobot.streaming.sidecar import SidecarSpec def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Build a reusable MP4 byte-index sidecar for streaming.") - parser.add_argument("--repo-id", default=DEFAULT_REPO) - parser.add_argument("--revision", default=DEFAULT_REVISION) - parser.add_argument("--data-root", default=DEFAULT_DATA_ROOT) + parser.add_argument("--repo-id", required=True) + parser.add_argument("--revision", default=None) + parser.add_argument("--data-root", required=True) parser.add_argument("--output", required=True) parser.add_argument("--episodes", type=int, default=None) parser.add_argument("--workers", type=int, default=8) - parser.add_argument("--range-backend", choices=("fsspec", "native-http"), default="native-http") + parser.add_argument("--range-backend", choices=("fsspec", "native-http"), default=None) parser.add_argument("--max-probe-mb", type=int, default=64) - parser.add_argument( - "--no-push", action="store_true", help="Do not upload the sidecar to data_root/meta/mp4-sidecars." - ) - parser.add_argument("--no-hub-branch-assert", action="store_true") + parser.add_argument("--push", action="store_true", help="Explicitly publish the sidecar to data_root.") return parser.parse_args() -def push_sidecar(local_path: str, data_root: str) -> list[str]: - if not data_root.startswith("hf://"): - return [] +def push_sidecar(local_path: str, spec: SidecarSpec) -> list[str]: + if not spec.data_root.startswith("hf://"): + raise ValueError("--push currently supports only hf:// data roots") - local = Path(local_path) fs = fsspec.filesystem("hf") - remote_dir = f"{data_root.rstrip('/')}/meta/mp4-sidecars" - remote_paths = [f"{remote_dir}/{local.name}"] - - for remote in remote_paths: - fs.put(str(local), remote) - return remote_paths + remote = published_sidecar_url(spec) + fs.put(str(Path(local_path)), remote) + return [remote] def main() -> None: args = parse_args() - if args.data_root.startswith("hf://") and not args.no_hub_branch_assert: - assert_hf_hub_range_cache_branch() meta = LeRobotDatasetMetadata(args.repo_id, revision=args.revision) meta.ensure_readable() total = ( int(meta.total_episodes) if args.episodes is None else min(args.episodes, int(meta.total_episodes)) ) - rel_paths = sorted( - {str(meta.get_video_file_path(ep_idx, key)) for ep_idx in range(total) for key in meta.video_keys} - ) + spec = make_sidecar_spec(meta, args.data_root) + if total != int(meta.total_episodes): + selected_paths = { + str(meta.get_video_file_path(ep_idx, key)) for ep_idx in range(total) for key in meta.video_keys + } + spec = SidecarSpec( + repo_id=spec.repo_id, + revision=spec.revision, + data_root=spec.data_root, + source_files=tuple(item for item in spec.source_files if item[0] in selected_paths), + ) start = time.perf_counter() - EpisodeVideoManifest.write_file_sidecar( + build_mp4_sidecar( args.output, - rel_paths, - args.data_root, - range_backend=args.range_backend, + spec, + range_backend=args.range_backend or range_backend_for_root(args.data_root), workers=args.workers, max_probe_bytes=args.max_probe_mb * 1024 * 1024, ) elapsed = time.perf_counter() - start print(f"wrote {args.output}") - print(f"episodes={total} files={len(rel_paths)} elapsed_s={elapsed:.2f}") - if args.no_push: - print("push_skipped: --no-push") - else: - pushed = push_sidecar(args.output, args.data_root) + print(f"episodes={total} files={len(spec.source_files)} elapsed_s={elapsed:.2f}") + if args.push: + if total != int(meta.total_episodes): + raise ValueError("Only a complete dataset sidecar can be published") + pushed = push_sidecar(args.output, spec) for remote in pushed: print(f"pushed {remote}") + else: + print("push_skipped: pass --push for explicit publication") if __name__ == "__main__": diff --git a/src/lerobot/common/train_utils.py b/src/lerobot/common/train_utils.py index b26196f14..4213e18c0 100644 --- a/src/lerobot/common/train_utils.py +++ b/src/lerobot/common/train_utils.py @@ -53,7 +53,11 @@ def get_step_checkpoint_dir(output_dir: Path, total_steps: int, step: int) -> Pa def save_training_step( - step: int, save_dir: Path, num_processes: int | None = None, batch_size: int | None = None + step: int, + save_dir: Path, + num_processes: int | None = None, + batch_size: int | None = None, + num_workers: int | None = None, ) -> None: state: dict = {"step": step} # num_processes and batch_size are recorded so a resumed run can detect a changed world size or @@ -64,6 +68,8 @@ def save_training_step( state["num_processes"] = num_processes if batch_size is not None: state["batch_size"] = batch_size + if num_workers is not None: + state["num_workers"] = num_workers write_json(state, save_dir / TRAINING_STEP) @@ -82,6 +88,11 @@ def load_training_batch_size(checkpoint_dir: Path) -> int | None: return load_json(checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP).get("batch_size") +def load_training_num_workers(checkpoint_dir: Path) -> int | None: + """DataLoader worker count recorded at checkpoint time, or None for older checkpoints.""" + return load_json(checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP).get("num_workers") + + def update_last_checkpoint(checkpoint_dir: Path) -> Path: last_checkpoint_dir = checkpoint_dir.parent / LAST_CHECKPOINT_LINK if last_checkpoint_dir.is_symlink(): @@ -101,6 +112,7 @@ def save_checkpoint( postprocessor: PolicyProcessorPipeline | None = None, num_processes: int | None = None, batch_size: int | None = None, + num_workers: int | None = None, model_state_dict: dict | None = None, optim_state_dict: dict | None = None, ) -> None: @@ -132,6 +144,8 @@ def save_checkpoint( resume. Defaults to None (not recorded). batch_size (int | None, optional): Per-process batch size to record for sample-exact resume. Defaults to None (not recorded). + num_workers (int | None, optional): Per-process DataLoader worker count to record for + sample-exact streaming resume. Defaults to None (not recorded). model_state_dict: Pre-gathered full (unsharded) model state dict. Required under FSDP, where `policy.state_dict()` would return sharded tensors; the caller gathers it via a cross-rank collective and passes it here so rank 0 can write it directly. It holds @@ -160,6 +174,7 @@ def save_checkpoint( scheduler, num_processes=num_processes, batch_size=batch_size, + num_workers=num_workers, optim_state_dict=optim_state_dict, ) @@ -171,6 +186,7 @@ def save_training_state( scheduler: LRScheduler | None = None, num_processes: int | None = None, batch_size: int | None = None, + num_workers: int | None = None, optim_state_dict: dict | None = None, ) -> None: """ @@ -185,12 +201,20 @@ def save_training_state( Defaults to None. num_processes (int | None, optional): Distributed world size to record. Defaults to None. batch_size (int | None, optional): Per-process batch size to record. Defaults to None. + num_workers (int | None, optional): Per-process DataLoader worker count to record. + Defaults to None. optim_state_dict: Pre-gathered full optimizer state dict (for FSDP). Saved instead of `optimizer.state_dict()` when provided. Defaults to None. """ save_dir = checkpoint_dir / TRAINING_STATE_DIR save_dir.mkdir(parents=True, exist_ok=True) - save_training_step(train_step, save_dir, num_processes=num_processes, batch_size=batch_size) + save_training_step( + train_step, + save_dir, + num_processes=num_processes, + batch_size=batch_size, + num_workers=num_workers, + ) save_rng_state(save_dir) if optimizer is not None: save_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict) diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index 38991a665..da27610fa 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -44,6 +44,14 @@ class DatasetConfig: # Has no effect on datasets without depth cameras. depth_output_unit: str = DEFAULT_DEPTH_UNIT 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. + 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. + 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 @@ -54,6 +62,12 @@ class DatasetConfig: ) if not (0.0 <= self.eval_split < 1.0): raise ValueError(f"eval_split must be in [0.0, 1.0), got {self.eval_split}") + if self.streaming_episode_pool_size <= 0: + raise ValueError("streaming_episode_pool_size must be positive") + if self.streaming_prefetch_episodes < 0: + raise ValueError("streaming_prefetch_episodes must be non-negative") + if self.streaming_byte_budget_gb <= 0: + raise ValueError("streaming_byte_budget_gb must be positive") if self.episodes is not None: if any(ep < 0 for ep in self.episodes): raise ValueError( diff --git a/src/lerobot/datasets/episode_parquet.py b/src/lerobot/datasets/episode_parquet.py new file mode 100644 index 000000000..fc24798ef --- /dev/null +++ b/src/lerobot/datasets/episode_parquet.py @@ -0,0 +1,106 @@ +# 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 + +"""Episode-scoped Parquet reads for training-time dataset streaming.""" + +from __future__ import annotations + +import posixpath +from collections.abc import Sequence +from pathlib import Path + +import fsspec +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq + + +class EpisodeParquetReader: + """Read complete episodes with column projection from local or fsspec roots.""" + + def __init__(self, data_root: str | Path, *, columns: Sequence[str]): + if not columns: + raise ValueError("EpisodeParquetReader requires at least one projected column") + self.columns = tuple(dict.fromkeys(columns)) + self._read_columns = ( + self.columns if "episode_index" in self.columns else (*self.columns, "episode_index") + ) + self._filesystem, self._root_path = fsspec.core.url_to_fs(str(data_root)) + + def read_episode( + self, + relative_path: str | Path, + *, + episode_index: int, + expected_rows: int, + ) -> pa.Table: + if expected_rows <= 0: + raise ValueError(f"Episode {episode_index} must contain at least one row") + + path = posixpath.join(self._root_path.rstrip("/"), str(relative_path).lstrip("/")) + with self._filesystem.open(path, "rb") as source: + parquet = pq.ParquetFile(source) + available = set(parquet.schema_arrow.names) + missing = sorted(set(self._read_columns) - available) + if missing: + raise ValueError(f"Parquet file {relative_path} is missing projected columns: {missing}") + row_group = self._matching_row_group(parquet, episode_index) + table = ( + parquet.read_row_group(row_group, columns=list(self._read_columns)) + if row_group is not None + else parquet.read(columns=list(self._read_columns)) + ) + + if row_group is None: + table = table.filter(pc.equal(table.column("episode_index"), episode_index)) + self._validate_complete_episode(table, episode_index, expected_rows, relative_path) + if "episode_index" not in self.columns: + table = table.drop_columns(["episode_index"]) + return table + + @staticmethod + def _matching_row_group(parquet: pq.ParquetFile, episode_index: int) -> int | None: + episode_column = next( + index + for index in range(parquet.metadata.num_columns) + if parquet.metadata.schema.column(index).path == "episode_index" + ) + matches = [] + for row_group in range(parquet.metadata.num_row_groups): + statistics = parquet.metadata.row_group(row_group).column(episode_column).statistics + if ( + statistics is not None + and statistics.has_min_max + and int(statistics.min) == episode_index + and int(statistics.max) == episode_index + ): + matches.append(row_group) + return matches[0] if len(matches) == 1 else None + + @staticmethod + def _validate_complete_episode( + table: pa.Table, + episode_index: int, + expected_rows: int, + relative_path: str | Path, + ) -> None: + actual_rows = len(table) + if actual_rows != expected_rows: + raise ValueError( + f"Parquet episode {episode_index} in {relative_path}: " + f"expected {expected_rows} rows, found {actual_rows}" + ) + episodes = table.column("episode_index").to_pylist() + if any(int(value) != episode_index for value in episodes): + raise ValueError(f"Parquet file {relative_path} returned rows outside episode {episode_index}") + if "frame_index" in table.column_names: + frame_indices = [int(value) for value in table.column("frame_index").to_pylist()] + if frame_indices != list(range(expected_rows)): + raise ValueError( + f"Parquet episode {episode_index} in {relative_path} has non-contiguous frame indices" + ) diff --git a/src/lerobot/datasets/factory.py b/src/lerobot/datasets/factory.py index da7b4365a..b5a3c8db6 100644 --- a/src/lerobot/datasets/factory.py +++ b/src/lerobot/datasets/factory.py @@ -66,7 +66,9 @@ def resolve_delta_timestamps( return delta_timestamps -def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset: +def make_dataset( + cfg: TrainPipelineConfig, +) -> LeRobotDataset | StreamingLeRobotDataset | MultiLeRobotDataset: """Handles the logic of setting up delta timestamps and image transforms before creating a dataset. Args: @@ -108,9 +110,15 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas delta_timestamps=delta_timestamps, image_transforms=image_transforms, revision=cfg.dataset.revision, - max_num_shards=cfg.num_workers, + max_num_shards=max(1, cfg.num_workers), tolerance_s=cfg.tolerance_s, return_uint8=True, + depth_output_unit=cfg.dataset.depth_output_unit, + data_root=cfg.dataset.streaming_data_root, + episode_pool_size=cfg.dataset.streaming_episode_pool_size, + prefetch_episodes=cfg.dataset.streaming_prefetch_episodes, + byte_budget_gb=cfg.dataset.streaming_byte_budget_gb, + repeat=True, ) else: raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.") @@ -138,7 +146,10 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas def make_train_eval_datasets( cfg: TrainPipelineConfig, -) -> tuple[LeRobotDataset | MultiLeRobotDataset, LeRobotDataset | None]: +) -> tuple[ + LeRobotDataset | StreamingLeRobotDataset | MultiLeRobotDataset, + LeRobotDataset | None, +]: """Create train and optional eval datasets by splitting episodes based on eval_split. The last ceil(n_episodes * eval_split) episodes per task are held out for evaluation. @@ -181,17 +192,37 @@ def make_train_eval_datasets( ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None ) - train_dataset = LeRobotDataset( - cfg.dataset.repo_id, - root=cfg.dataset.root, - episodes=train_episodes, - delta_timestamps=delta_timestamps, - image_transforms=train_image_transforms, - revision=cfg.dataset.revision, - video_backend=cfg.dataset.video_backend, - return_uint8=True, - tolerance_s=cfg.tolerance_s, - ) + if cfg.dataset.streaming: + train_dataset = StreamingLeRobotDataset( + cfg.dataset.repo_id, + root=cfg.dataset.root, + episodes=train_episodes, + delta_timestamps=delta_timestamps, + image_transforms=train_image_transforms, + revision=cfg.dataset.revision, + max_num_shards=max(1, cfg.num_workers), + tolerance_s=cfg.tolerance_s, + return_uint8=True, + depth_output_unit=cfg.dataset.depth_output_unit, + data_root=cfg.dataset.streaming_data_root, + episode_pool_size=cfg.dataset.streaming_episode_pool_size, + prefetch_episodes=cfg.dataset.streaming_prefetch_episodes, + byte_budget_gb=cfg.dataset.streaming_byte_budget_gb, + repeat=True, + ) + else: + train_dataset = LeRobotDataset( + cfg.dataset.repo_id, + root=cfg.dataset.root, + episodes=train_episodes, + delta_timestamps=delta_timestamps, + image_transforms=train_image_transforms, + revision=cfg.dataset.revision, + video_backend=cfg.dataset.video_backend, + return_uint8=True, + depth_output_unit=cfg.dataset.depth_output_unit, + tolerance_s=cfg.tolerance_s, + ) eval_dataset = LeRobotDataset( cfg.dataset.repo_id, @@ -202,6 +233,7 @@ def make_train_eval_datasets( revision=cfg.dataset.revision, video_backend=cfg.dataset.video_backend, return_uint8=True, + depth_output_unit=cfg.dataset.depth_output_unit, tolerance_s=cfg.tolerance_s, ) diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index 14d4a52a4..2fb4c7b28 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -13,208 +13,46 @@ # 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. -from collections import deque -from collections.abc import Callable, Generator, Iterable, Iterator +import io +import os +from collections.abc import Callable, Iterator +from concurrent.futures import Future, ThreadPoolExecutor from pathlib import Path import datasets import numpy as np import torch -from datasets import load_dataset from lerobot.configs import DEFAULT_DEPTH_UNIT, DEPTH_METER_UNIT, DepthEncoderConfig -from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE +from lerobot.streaming.episode_video import EpisodeByteCache, EpisodeVideoManifest, ExactCoveragePool +from lerobot.utils.constants import HF_LEROBOT_HOME from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata from .depth_utils import MM_PER_METRE, dequantize_depth -from .feature_utils import get_delta_indices -from .io_utils import item_to_torch -from .utils import ( - check_version_compatibility, - find_float_index, - is_float_in_list, - safe_shard, +from .episode_parquet import EpisodeParquetReader +from .feature_utils import check_delta_timestamps, get_delta_indices, get_hf_features_from_features +from .io_utils import hf_transform_to_torch +from .streaming_sidecar import ( + ensure_dataset_mp4_sidecar, + range_backend_for_root, + streaming_data_root, ) -from .video_utils import ( - VideoDecoderCache, - decode_video_frames, - decode_video_frames_torchcodec, -) - - -class LookBackError(Exception): - """ - Exception raised when trying to look back in the history of a Backtrackable object. - """ - - pass - - -class LookAheadError(Exception): - """ - Exception raised when trying to look ahead in the future of a Backtrackable object. - """ - - pass - - -class Backtrackable[T]: - """ - Wrap any iterator/iterable so you can step back up to `history` items - and look ahead up to `lookahead` items. - - This is useful for streaming datasets where you need to access previous and future items - but can't load the entire dataset into memory. - - Example: - ------- - ```python - ds = load_dataset("c4", "en", streaming=True, split="train") - rev = Backtrackable(ds, history=3, lookahead=2) - - x0 = next(rev) # forward - x1 = next(rev) - x2 = next(rev) - - # Look ahead - x3_peek = rev.peek_ahead(1) # next item without moving cursor - x4_peek = rev.peek_ahead(2) # two items ahead - - # Look back - x1_again = rev.peek_back(1) # previous item without moving cursor - x0_again = rev.peek_back(2) # two items back - - # Move backward - x1_back = rev.prev() # back one step - next(rev) # returns x2, continues forward from where we were - ``` - """ - - __slots__ = ("_source", "_back_buf", "_ahead_buf", "_cursor", "_history", "_lookahead") - - def __init__(self, iterable: Iterable[T], *, history: int = 1, lookahead: int = 0): - if history < 1: - raise ValueError("history must be >= 1") - if lookahead <= 0: - raise ValueError("lookahead must be > 0") - - self._source: Iterator[T] = iter(iterable) - self._back_buf: deque[T] = deque(maxlen=history) - self._ahead_buf: deque[T] = deque(maxlen=lookahead) if lookahead > 0 else deque() - self._cursor: int = 0 - self._history = history - self._lookahead = lookahead - - def __iter__(self) -> "Backtrackable[T]": - return self - - def __next__(self) -> T: - # If we've stepped back, consume from back buffer first - if self._cursor < 0: # -1 means "last item", etc. - self._cursor += 1 - return self._back_buf[self._cursor] - - # If we have items in the ahead buffer, use them first - item = self._ahead_buf.popleft() if self._ahead_buf else next(self._source) - - # Add current item to back buffer and reset cursor - self._back_buf.append(item) - self._cursor = 0 - return item - - def prev(self) -> T: - """ - Step one item back in history and return it. - Raises IndexError if already at the oldest buffered item. - """ - if len(self._back_buf) + self._cursor <= 1: - raise LookBackError("At start of history") - - self._cursor -= 1 - return self._back_buf[self._cursor] - - def peek_back(self, n: int = 1) -> T: - """ - Look `n` items back (n=1 == previous item) without moving the cursor. - """ - if n < 0 or n + 1 > len(self._back_buf) + self._cursor: - raise LookBackError("peek_back distance out of range") - - return self._back_buf[self._cursor - (n + 1)] - - def peek_ahead(self, n: int = 1) -> T: - """ - Look `n` items ahead (n=1 == next item) without moving the cursor. - Fills the ahead buffer if necessary. - """ - if n < 1: - raise LookAheadError("peek_ahead distance must be 1 or more") - elif n > self._lookahead: - raise LookAheadError("peek_ahead distance exceeds lookahead limit") - - # Fill ahead buffer if we don't have enough items - while len(self._ahead_buf) < n: - try: - item = next(self._source) - self._ahead_buf.append(item) - - except StopIteration as err: - raise LookAheadError("peek_ahead: not enough items in source") from err - - return self._ahead_buf[n - 1] - - def history(self) -> list[T]: - """ - Return a copy of the buffered history (most recent last). - The list length ≤ `history` argument passed at construction. - """ - if self._cursor == 0: - return list(self._back_buf) - - # When cursor<0, slice so the order remains chronological - return list(self._back_buf)[: self._cursor or None] - - def can_peek_back(self, steps: int = 1) -> bool: - """ - Check if we can go back `steps` items without raising an IndexError. - """ - return steps <= len(self._back_buf) + self._cursor - - def can_peek_ahead(self, steps: int = 1) -> bool: - """ - Check if we can peek ahead `steps` items. - This may involve trying to fill the ahead buffer. - """ - if self._lookahead > 0 and steps > self._lookahead: - return False - - # Try to fill ahead buffer to check if we can peek that far - try: - while len(self._ahead_buf) < steps: - if self._lookahead > 0 and len(self._ahead_buf) >= self._lookahead: - return False - item = next(self._source) - self._ahead_buf.append(item) - return True - except StopIteration: - return False +from .utils import check_version_compatibility +from .video_utils import decode_video_frames_pyav class StreamingLeRobotDataset(torch.utils.data.IterableDataset): - """LeRobotDataset with streaming capabilities. + """Episode-scoped streaming reader for LeRobot datasets. - This class extends LeRobotDataset to add streaming functionality, allowing data to be streamed - rather than loaded entirely into memory. This is especially useful for large datasets that may - not fit in memory or when you want to quickly explore a dataset without downloading it completely. - - The key innovation is using a Backtrackable iterator that maintains a bounded buffer of recent - items, allowing us to access previous frames for delta timestamps without loading the entire - dataset into memory. + 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. + MP4 sidecars are resolved automatically and built in a revision-keyed local cache when absent. Example: Basic usage: ```python - from lerobot.common.datasets.streaming_dataset import StreamingLeRobotDataset + from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset # Create a streaming dataset with delta timestamps delta_timestamps = { @@ -225,8 +63,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): dataset = StreamingLeRobotDataset( repo_id="your-dataset-repo-id", delta_timestamps=delta_timestamps, - streaming=True, - buffer_size=1000, ) # Iterate over the dataset @@ -244,7 +80,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): root: str | Path | None = None, episodes: list[int] | None = None, image_transforms: Callable | None = None, - delta_timestamps: dict[list[float]] | None = None, + delta_timestamps: dict[str, list[float]] | None = None, tolerance_s: float = 1e-4, revision: str | None = None, force_cache_sync: bool = False, @@ -256,6 +92,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): shuffle: bool = True, return_uint8: bool = False, depth_output_unit: str = DEFAULT_DEPTH_UNIT, + data_root: str | Path | None = None, + episode_pool_size: int | None = None, + prefetch_episodes: int = 8, + byte_budget_gb: float = 8.0, + repeat: bool = False, ): """Initialize a StreamingLeRobotDataset. @@ -270,14 +111,22 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): tolerance_s (float, optional): Tolerance in seconds for timestamp matching. revision (str, optional): Git revision id (branch name, tag, or commit hash). force_cache_sync (bool, optional): Flag to sync and refresh local files first. - streaming (bool, optional): Whether to stream the dataset or load it all. Defaults to True. - buffer_size (int, optional): Buffer size for shuffling when streaming. Defaults to 1000. - max_num_shards (int, optional): Number of shards to re-shard the input dataset into. Defaults to 16. + streaming (bool, optional): Compatibility flag retained by the public API. + buffer_size (int, optional): Compatibility setting used to derive the default episode pool size. + max_num_shards (int, optional): Maximum number of concurrent episode loading workers. seed (int, optional): Reproducibility random seed. rng (np.random.Generator | None, optional): Random number generator. shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True. depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm"). Defaults to "mm". + data_root (str | Path | None, optional): Dataset payload root. Supports local paths, ``hf://``, + 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 + short final training batch. The training factory enables this; direct iteration is + finite by default. """ super().__init__() self.repo_id = repo_id @@ -295,11 +144,28 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): self.streaming = streaming self.buffer_size = buffer_size + self.max_num_shards = max_num_shards self._return_uint8 = return_uint8 self._depth_output_unit = depth_output_unit - - # We cache the video decoders to avoid re-initializing them at each frame (avoiding a ~10x slowdown) - self.video_decoder_cache = None + if buffer_size <= 0: + raise ValueError("buffer_size must be positive") + if max_num_shards <= 0: + raise ValueError("max_num_shards must be positive") + if episode_pool_size is not None and episode_pool_size <= 0: + raise ValueError("episode_pool_size must be positive") + if prefetch_episodes < 0: + raise ValueError("prefetch_episodes must be non-negative") + if byte_budget_gb <= 0: + raise ValueError("byte_budget_gb must be positive") + self.episode_pool_size = episode_pool_size or min(buffer_size, 32) + self.prefetch_episodes = prefetch_episodes + self.byte_budget = int(byte_budget_gb * 1024**3) + self.repeat = repeat + self._next_epoch = 0 + self._active_epoch = 0 + self._resume_offset = 0 + self._resume_batch_size = 1 + self._state_offset = 0 if self._requested_root is not None: self.root.mkdir(exist_ok=True, parents=True) @@ -326,34 +192,53 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): if key in self.meta.image_keys } - self.delta_timestamps = None - self.delta_indices = None + selected_episodes = list(range(self.meta.total_episodes)) if episodes is None else list(episodes) + if len(set(selected_episodes)) != len(selected_episodes): + raise ValueError("episodes must not contain duplicates") + invalid_episodes = [ + episode for episode in selected_episodes if episode < 0 or episode >= self.meta.total_episodes + ] + if invalid_episodes: + raise ValueError( + f"Episode indices out of range for dataset with {self.meta.total_episodes} episodes: " + f"{invalid_episodes}" + ) + self._selected_episodes = selected_episodes + + self.delta_timestamps: dict[str, list[float]] | None = None + self.delta_indices: dict[str, list[int]] | None = None if delta_timestamps is not None: - self._validate_delta_timestamp_keys(delta_timestamps) # raises ValueError if invalid + check_delta_timestamps(delta_timestamps, self.fps, tolerance_s) self.delta_timestamps = delta_timestamps self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps) - self.hf_dataset: datasets.IterableDataset = load_dataset( - self.repo_id if not self.streaming_from_local else str(self.root), - split="train", - streaming=self.streaming, - data_files="data/*/*.parquet", - revision=self.revision, + self._data_root = streaming_data_root( + self.meta, + requested_root=self._requested_root, + configured_data_root=str(data_root) if data_root is not None else None, ) - - self.num_shards = min(self.hf_dataset.num_shards, max_num_shards) + sidecar_backend = range_backend_for_root(self._data_root) + self._sidecar_path = ensure_dataset_mp4_sidecar( + self.meta, + self._data_root, + workers=max_num_shards, + range_backend=sidecar_backend, + ) + self._hf_features = get_hf_features_from_features(self.meta.features) + self._projected_columns = tuple(self._hf_features) + self.num_shards = min(max_num_shards, max(1, len(self._selected_episodes))) @property - def num_frames(self): - return self.meta.total_frames + def num_frames(self) -> int: + return sum(self._episode_frame_count(episode) for episode in self._selected_episodes) @property - def num_episodes(self): - return self.meta.total_episodes + def num_episodes(self) -> int: + return len(self._selected_episodes) @property - def fps(self): + def fps(self) -> int: return self.meta.fps @property @@ -361,397 +246,354 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): """Physical unit (``"m"`` or ``"mm"``) depth maps are returned in on read.""" return self._depth_output_unit - @staticmethod - def _iter_random_indices( - rng: np.random.Generator, buffer_size: int, random_batch_size=100 - ) -> Iterator[int]: - while True: - yield from (int(i) for i in rng.integers(0, buffer_size, size=random_batch_size)) - - @staticmethod - def _infinite_generator_over_elements(rng: np.random.Generator, elements: list[int]) -> Iterator[int]: - while True: - yield rng.choice(elements) - - # TODO(fracapuano): Implement multi-threaded prefetching to accelerate data loading. - # The current sequential iteration is a bottleneck. A producer-consumer pattern - # could be used with a ThreadPoolExecutor to run `make_frame` (especially video decoding) - # in parallel, feeding a queue from which this iterator will yield processed items. def __iter__(self) -> Iterator[dict[str, torch.Tensor]]: - if self.video_decoder_cache is None: - self.video_decoder_cache = VideoDecoderCache() - - # keep the same seed across exhaustions if shuffle is False, otherwise shuffle data across exhaustions - rng = np.random.default_rng(self.seed) if not self.shuffle else self.rng - - buffer_indices_generator = self._iter_random_indices(rng, self.buffer_size) - - idx_to_backtrack_dataset = { - idx: self._make_backtrackable_dataset(safe_shard(self.hf_dataset, idx, self.num_shards)) - for idx in range(self.num_shards) - } - - # This buffer is populated while iterating on the dataset's shards - # the logic is to add 2 levels of randomness: - # (1) sample one shard at random from the ones available, and - # (2) sample one frame from the shard sampled at (1) - frames_buffer = [] - while available_shards := list(idx_to_backtrack_dataset.keys()): - shard_key = next(self._infinite_generator_over_elements(rng, available_shards)) - backtrack_dataset = idx_to_backtrack_dataset[shard_key] # selects which shard to iterate on + if self.repeat: + return self._repeat_iterator() + return self._iter_once() + def _repeat_iterator(self) -> Iterator[dict[str, torch.Tensor]]: + while True: + iterator = self._iter_once() try: - for frame in self.make_frame(backtrack_dataset): - if len(frames_buffer) == self.buffer_size: - i = next(buffer_indices_generator) # samples a element from the buffer - yield frames_buffer[i] - frames_buffer[i] = frame - else: - frames_buffer.append(frame) - break # random shard sampled, switch shard - except ( - RuntimeError, - StopIteration, - ): # NOTE: StopIteration inside a generator throws a RuntimeError since python 3.7 - del idx_to_backtrack_dataset[shard_key] # Remove exhausted shard, onto another shard + first = next(iterator) + except StopIteration: + return + yield first + yield from iterator - # Once shards are all exhausted, shuffle the buffer and yield the remaining frames - rng.shuffle(frames_buffer) - yield from frames_buffer + def _iter_once(self) -> Iterator[dict[str, torch.Tensor]]: + epoch = self._next_epoch if self.shuffle else 0 + self._active_epoch = epoch - def _get_window_steps( - self, delta_timestamps: dict[str, list[float]] | None = None, dynamic_bounds: bool = False + 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 + 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_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, + ) + else: + worker_epoch_delta = 0 + epoch += worker_epoch_delta + self._active_epoch = epoch + if self.shuffle: + self._next_epoch = epoch + 1 + planner = ExactCoveragePool( + [(episode, self._episode_frame_count(episode)) for episode in consumer_episodes], + pool_size=self.episode_pool_size, + seed=self.seed, + epoch=epoch, + ) + for _ in range(worker_resume_offset): + try: + next(planner) + except StopIteration: + return + planner.newly_admitted.clear() + 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) + retained_video_episodes: set[int] = set() + if video_cache is not None: + for episode_index in planner.resident: + video_cache.retain_episode(episode_index) + retained_video_episodes.add(episode_index) + + def submit(episode_index: int) -> Future[datasets.Dataset]: + future = episode_futures.get(episode_index) + if future is None: + future = executor.submit(self._load_episode_dataset, parquet_reader, episode_index) + episode_futures[episode_index] = future + 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] + submit(episode_index) + if video_cache is not None: + video_cache.submit_prefetch(episode_index) + scheduled_cursor += 1 + + schedule_frontier() + try: + while True: + try: + episode_index, frame_index = next(planner) + except StopIteration: + self._active_epoch = epoch + 1 if self.shuffle else 0 + self._state_offset = 0 + break + + episode_dataset = submit(episode_index).result() + item = self._make_episode_item( + episode_dataset, + episode_index, + frame_index, + video_cache=video_cache, + ) + + for evicted_episode in planner.evicted: + episode_futures.pop(evicted_episode, None) + if video_cache is not None and evicted_episode in retained_video_episodes: + video_cache.release_episode(evicted_episode) + retained_video_episodes.remove(evicted_episode) + if video_cache is not None: + for admitted_episode in planner.newly_admitted: + if admitted_episode not in retained_video_episodes: + video_cache.retain_episode(admitted_episode) + retained_video_episodes.add(admitted_episode) + planner.evicted.clear() + planner.newly_admitted.clear() + schedule_frontier() + + self._state_offset += 1 + yield item + finally: + for future in episode_futures.values(): + future.cancel() + executor.shutdown(wait=True, cancel_futures=True) + 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]: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + else: + rank = int(os.environ.get("RANK", "0")) + 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]: - if delta_timestamps is None: - return 1, 1 + 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 - if not dynamic_bounds: - # Fix the windows - lookback = LOOKBACK_BACKTRACKTABLE - lookahead = LOOKAHEAD_BACKTRACKTABLE - else: - # Dynamically adjust the windows based on the given delta_timesteps - all_timestamps = sum(delta_timestamps.values(), []) - lookback = min(all_timestamps) * self.fps - lookahead = max(all_timestamps) * self.fps + 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"]) - # When lookback is >=0 it means no negative timesteps have been provided - lookback = 0 if lookback >= 0 else (lookback * -1) + 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.""" + 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 + ) - return lookback, lookahead - - def _make_backtrackable_dataset(self, dataset: datasets.IterableDataset) -> Backtrackable: - lookback, lookahead = self._get_window_steps(self.delta_timestamps) - return Backtrackable(dataset, history=lookback, lookahead=lookahead) - - def _make_timestamps_from_indices( - self, start_ts: float, indices: dict[str, list[int]] | None = None - ) -> dict[str, list[float]]: - if indices is not None: - return { - key: ( - start_ts + torch.tensor(indices[key]) / self.fps - ).tolist() # NOTE: why not delta_timestamps directly? - for key in self.delta_timestamps - } - else: - return dict.fromkeys(self.meta.video_keys, [start_ts]) - - def _make_padding_camera_frame(self, camera_key: str): - """Variable-shape padding frame for given camera keys, given in (H, W, C)""" - return torch.zeros(self.meta.info.features[camera_key]["shape"]).permute(-1, 0, 1) - - def _get_video_frame_padding_mask( + def _load_episode_dataset( self, - video_frames: dict[str, torch.Tensor], - query_timestamps: dict[str, list[float]], - original_timestamps: dict[str, list[float]], - ) -> dict[str, torch.BoolTensor]: - padding_mask = {} + reader: EpisodeParquetReader, + episode_index: int, + ) -> datasets.Dataset: + table = reader.read_episode( + self.meta.get_data_file_path(episode_index), + episode_index=episode_index, + expected_rows=self._episode_frame_count(episode_index), + ) + # from_dict applies the declared HF feature encoders (images, nested language/JSON fields) + # while retaining the episode-sized memory bound. + dataset = datasets.Dataset.from_dict(table.to_pydict(), features=self._hf_features) + dataset.set_transform(hf_transform_to_torch) + return dataset - for video_key, timestamps in original_timestamps.items(): - if video_key not in video_frames: - continue # only padding on video keys that are available - frames = [] - mask = [] - padding_frame = self._make_padding_camera_frame(video_key) - for ts in timestamps: - if is_float_in_list(ts, query_timestamps[video_key]): - idx = find_float_index(ts, query_timestamps[video_key]) - frames.append(video_frames[video_key][idx, :]) - mask.append(False) - else: - frames.append(padding_frame) - mask.append(True) + def _make_video_cache( + self, + episode_indices: list[int], + workers: int, + ) -> EpisodeByteCache | None: + if self._sidecar_path is None or not episode_indices: + return None + range_backend = range_backend_for_root(self._data_root) + manifest = EpisodeVideoManifest.build( + self.meta, + self._data_root, + episode_indices=episode_indices, + range_backend=range_backend, + workers=workers, + sidecar_path=self._sidecar_path, + ) + decoder_limit = max(1, min(64, self.episode_pool_size * max(1, len(self.meta.video_keys)))) + return EpisodeByteCache( + manifest, + self._data_root, + byte_budget=self.byte_budget, + workers=workers, + range_backend=range_backend, + max_open_decoders=decoder_limit, + ) - padding_mask[f"{video_key}_is_pad"] = torch.BoolTensor(mask) + def _make_episode_item( + self, + episode_dataset: datasets.Dataset, + episode_index: int, + frame_index: int, + *, + video_cache: EpisodeByteCache | None, + ) -> dict: + item = episode_dataset[frame_index] + episode = self.meta.episodes[episode_index] + episode_start = int(episode["dataset_from_index"]) - return padding_mask - - def make_frame(self, dataset_iterator: Backtrackable) -> Generator: - """Makes a frame starting from a dataset iterator""" - item = next(dataset_iterator) - item = item_to_torch(item) - - updates = [] # list of "updates" to apply to the item retrieved from hf_dataset (w/o camera features) - - # Get episode index from the item - ep_idx = item["episode_index"] - - # "timestamp" restarts from 0 for each episode, whereas we need a global timestep within the single .mp4 file (given by index/fps) - current_ts = item["index"] / self.fps - - episode_boundaries_ts = { - key: ( - self.meta.episodes[ep_idx][f"videos/{key}/from_timestamp"], - self.meta.episodes[ep_idx][f"videos/{key}/to_timestamp"], - ) - for key in self.meta.video_keys - } - - # Apply delta querying logic if necessary if self.delta_indices is not None: - query_result, padding = self._get_delta_frames(dataset_iterator, item) - updates.append(query_result) - updates.append(padding) - - # Load video frames, when needed - if len(self.meta.video_keys) > 0: - original_timestamps = self._make_timestamps_from_indices(current_ts, self.delta_indices) - - # Some timestamps might not result available considering the episode's boundaries - query_timestamps = self._get_query_timestamps( - current_ts, self.delta_indices, episode_boundaries_ts - ) - video_frames = self._query_videos(query_timestamps, ep_idx) - - if self.image_transforms is not None: - image_keys = self.meta.camera_keys - for cam in image_keys: - video_frames[cam] = self.image_transforms(video_frames[cam]) - - updates.append(video_frames) - - if self.delta_indices is not None: - # We always return the same number of frames. Unavailable frames are padded. - padding_mask = self._get_video_frame_padding_mask( - video_frames, query_timestamps, original_timestamps + for key, delta_indices in self.delta_indices.items(): + target_indices = [ + max(0, min(len(episode_dataset) - 1, frame_index + delta)) for delta in delta_indices + ] + item[f"{key}_is_pad"] = torch.BoolTensor( + [ + frame_index + delta < 0 or frame_index + delta >= len(episode_dataset) + for delta in delta_indices + ] ) - updates.append(padding_mask) + if key not in self.meta.video_keys: + item[key] = torch.stack(episode_dataset[target_indices][key]) - result = item.copy() - for update in updates: - result.update(update) + if self.meta.video_keys: + if video_cache is None: + raise RuntimeError("Video dataset streaming requires an episode byte cache") + episode_metadata = self.meta.episodes[episode_index] + for video_key in self.meta.video_keys: + if self.delta_indices is not None and video_key in self.delta_indices: + target_indices = [ + max(0, min(len(episode_dataset) - 1, frame_index + delta)) + for delta in self.delta_indices[video_key] + ] + else: + target_indices = [frame_index] + local_timestamps = [ + float(episode_dataset[index]["timestamp"].item()) for index in target_indices + ] + from_timestamp = float(episode_metadata[f"videos/{video_key}/from_timestamp"]) + query_timestamps = [from_timestamp + timestamp for timestamp in local_timestamps] + if video_key in self.meta.depth_keys: + source_start = video_cache.manifest.lookup(episode_index, video_key).source_start_pts + frames = decode_video_frames_pyav( + io.BytesIO(video_cache.get_bytes(episode_index, video_key)), + [timestamp - source_start for timestamp in query_timestamps], + self.tolerance_s, + return_uint8=False, + is_depth=True, + ) + depth_encoder = self._depth_encoder_configs[video_key] + frames = dequantize_depth( + frames, + depth_min=depth_encoder.depth_min, + depth_max=depth_encoder.depth_max, + shift=depth_encoder.shift, + use_log=depth_encoder.use_log, + output_unit=self._depth_output_unit, + ) + else: + frames = video_cache.get_frames(episode_index, video_key, query_timestamps) + if not self._return_uint8: + frames = frames.to(torch.float32) / 255.0 + item[video_key] = frames.squeeze(0) + + if self.image_transforms is not None: + for camera_key in self.meta.camera_keys: + if camera_key in self.meta.depth_keys: + continue + item[camera_key] = self.image_transforms(item[camera_key]) - # Convert raw-image depth features to the output unit (video depth is already converted). for key, stored_unit in self._image_depth_units.items(): - if key in result and stored_unit is not None and stored_unit != self._depth_output_unit: - result[key] = ( - result[key] * MM_PER_METRE - if stored_unit == DEPTH_METER_UNIT - else result[key] / MM_PER_METRE + if key in item and stored_unit is not None and stored_unit != self._depth_output_unit: + item[key] = ( + item[key] * MM_PER_METRE if stored_unit == DEPTH_METER_UNIT else item[key] / MM_PER_METRE ) - result["task"] = self.meta.tasks.iloc[item["task_index"]].name - - yield result - - def _get_query_timestamps( - self, - current_ts: float, - query_indices: dict[str, list[int]] | None = None, - episode_boundaries_ts: dict[str, tuple[float, float]] | None = None, - ) -> dict[str, list[float]]: - query_timestamps = {} - keys_to_timestamps = self._make_timestamps_from_indices(current_ts, query_indices) - for key in self.meta.video_keys: - if query_indices is not None and key in query_indices: - timestamps = keys_to_timestamps[key] - # Clamp out timesteps outside of episode boundaries - query_timestamps[key] = torch.clamp( - torch.tensor(timestamps), *episode_boundaries_ts[key] - ).tolist() - - else: - query_timestamps[key] = [current_ts] - - return query_timestamps - - def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict: - """Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function - in the main process (e.g. by using a second Dataloader with num_workers=0). It will result in a - Segmentation Fault. This probably happens because a memory reference to the video loader is created in - the main process and a subprocess fails to access it. - """ - - item = {} - for video_key, query_ts in query_timestamps.items(): - root = self.meta.url_root if self.streaming and not self.streaming_from_local else self.root - video_path = f"{root}/{self.meta.get_video_file_path(ep_idx, video_key)}" - if video_key in self.meta.depth_keys: - # Depth maps are 12-bit quantized and only decodable via pyav; dequantize back - # to physical units to match the non-streaming reader. - frames = decode_video_frames( - video_path, - query_ts, - self.tolerance_s, - backend="pyav", - return_uint8=False, - is_depth=True, - ) - depth_encoder = self._depth_encoder_configs[video_key] - frames = dequantize_depth( - frames, - depth_min=depth_encoder.depth_min, - depth_max=depth_encoder.depth_max, - shift=depth_encoder.shift, - use_log=depth_encoder.use_log, - output_unit=self._depth_output_unit, - ) - else: - frames = decode_video_frames_torchcodec( - video_path, - query_ts, - self.tolerance_s, - decoder_cache=self.video_decoder_cache, - return_uint8=self._return_uint8, - ) - - item[video_key] = frames.squeeze(0) if len(query_ts) == 1 else frames - + task_index = int(item["task_index"].item()) + item["task"] = self.meta.tasks.iloc[task_index].name + if int(item["episode_index"].item()) != episode_index: + raise RuntimeError(f"Episode reader returned episode {item['episode_index']} for {episode_index}") + if int(item["index"].item()) != episode_start + frame_index: + raise RuntimeError( + f"Episode {episode_index} frame {frame_index} has unexpected absolute index {item['index']}" + ) return item - def _get_delta_frames(self, dataset_iterator: Backtrackable, current_item: dict): - # TODO(fracapuano): Modularize this function, refactor the code - """Get frames with delta offsets using the backtrackable iterator. + def state_dict(self) -> dict[str, int]: + return { + "epoch": self._active_epoch, + "offset": self._state_offset, + "batch_size": self._resume_batch_size, + } - Args: - current_item (dict): Current item from the iterator. - ep_idx (int): Episode index. - - Returns: - tuple: (query_result, padding) - frames at delta offsets and padding info. - """ - current_episode_idx = current_item["episode_index"] - - # Prepare results - query_result = {} - padding = {} - - for key, delta_indices in self.delta_indices.items(): - if key in self.meta.video_keys: - continue # visual frames are decoded separately - - target_frames = [] - is_pad = [] - - # Create a results dictionary to store frames in processing order, then reconstruct original order for stacking - delta_results = {} - - # Separate and sort deltas by difficulty (easier operations first) - negative_deltas = sorted([d for d in delta_indices if d < 0], reverse=True) # [-1, -2, -3, ...] - positive_deltas = sorted([d for d in delta_indices if d > 0]) # [1, 2, 3, ...] - zero_deltas = [d for d in delta_indices if d == 0] - - # Process zero deltas (current frame) - for delta in zero_deltas: - delta_results[delta] = ( - current_item[key], - False, - ) - - # Process negative deltas in order of increasing difficulty - lookback_failed = False - - last_successful_frame = current_item[key] - - for delta in negative_deltas: - if lookback_failed: - delta_results[delta] = (last_successful_frame, True) - continue - - try: - steps_back = abs(delta) - if dataset_iterator.can_peek_back(steps_back): - past_item = dataset_iterator.peek_back(steps_back) - past_item = item_to_torch(past_item) - - if past_item["episode_index"] == current_episode_idx: - delta_results[delta] = (past_item[key], False) - last_successful_frame = past_item[key] - - else: - raise LookBackError("Retrieved frame is from different episode!") - else: - raise LookBackError("Cannot go back further than the history buffer!") - - except LookBackError: - delta_results[delta] = (last_successful_frame, True) - lookback_failed = True # All subsequent negative deltas will also fail - - # Process positive deltas in order of increasing difficulty - lookahead_failed = False - last_successful_frame = current_item[key] - - for delta in positive_deltas: - if lookahead_failed: - delta_results[delta] = (last_successful_frame, True) - continue - - try: - if dataset_iterator.can_peek_ahead(delta): - future_item = dataset_iterator.peek_ahead(delta) - future_item = item_to_torch(future_item) - - if future_item["episode_index"] == current_episode_idx: - delta_results[delta] = (future_item[key], False) - last_successful_frame = future_item[key] - - else: - raise LookAheadError("Retrieved frame is from different episode!") - else: - raise LookAheadError("Cannot go ahead further than the lookahead buffer!") - - except LookAheadError: - delta_results[delta] = (last_successful_frame, True) - lookahead_failed = True # All subsequent positive deltas will also fail - - # Reconstruct original order for stacking - for delta in delta_indices: - frame, is_padded = delta_results[delta] - - # add batch dimension for stacking - target_frames.append(frame) # frame.unsqueeze(0)) - is_pad.append(is_padded) - - # Stack frames and add to results - if target_frames: - query_result[key] = torch.stack(target_frames) - padding[f"{key}_is_pad"] = torch.BoolTensor(is_pad) - - return query_result, padding - - def _validate_delta_timestamp_keys(self, delta_timestamps: dict[list[float]]) -> None: - """ - Validate that all keys in delta_timestamps correspond to actual features in the dataset. - - Raises: - ValueError: If any delta timestamp key doesn't correspond to a dataset feature. - """ - if delta_timestamps is None: - return - - # Get all available feature keys from the dataset metadata - available_features = set(self.meta.features.keys()) - - # Get all keys from delta_timestamps - delta_keys = set(delta_timestamps.keys()) - - # Find any keys that don't correspond to features - invalid_keys = delta_keys - available_features - - if invalid_keys: + def load_state_dict(self, state: dict[str, int]) -> None: + epoch = int(state.get("epoch", 0)) + offset = int(state.get("offset", 0)) + batch_size = int(state.get("batch_size", 1)) + if epoch < 0 or offset < 0 or batch_size <= 0: raise ValueError( - f"The following delta_timestamp keys do not correspond to dataset features: {invalid_keys}. " - f"Available features are: {sorted(available_features)}" + "Streaming dataset epoch/offset must be non-negative and batch_size must be positive" ) + self._next_epoch = epoch + self._active_epoch = epoch + self._resume_offset = offset + self._resume_batch_size = batch_size + self._state_offset = offset + + def set_epoch(self, epoch: int) -> None: + if epoch < 0: + raise ValueError("epoch must be non-negative") + self._next_epoch = epoch + self._active_epoch = epoch + self._resume_offset = 0 + self._resume_batch_size = 1 + self._state_offset = 0 diff --git a/src/lerobot/datasets/streaming_sidecar.py b/src/lerobot/datasets/streaming_sidecar.py new file mode 100644 index 000000000..d597da374 --- /dev/null +++ b/src/lerobot/datasets/streaming_sidecar.py @@ -0,0 +1,141 @@ +# 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 + +"""LeRobot metadata adapter for automatic MP4 sidecar resolution.""" + +from __future__ import annotations + +import logging +import shutil +from pathlib import Path + +import fsspec + +from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata +from lerobot.streaming.episode_video import EpisodeVideoManifest +from lerobot.streaming.sidecar import SidecarSpec, ensure_mp4_sidecar, sidecar_cache_path +from lerobot.utils.constants import HF_LEROBOT_HOME + +DEFAULT_SIDECAR_CACHE = HF_LEROBOT_HOME / "streaming-sidecars" + + +def range_backend_for_root(data_root: str) -> str: + """Use direct HTTP only for HF roots; all local and other fsspec protocols stay generic.""" + return "native-http" if data_root.startswith("hf://") else "fsspec" + + +def streaming_data_root( + meta: LeRobotDatasetMetadata, + *, + requested_root: str | Path | None, + configured_data_root: str | None, +) -> str: + if configured_data_root is not None: + return configured_data_root.rstrip("/") + if requested_root is not None: + return str(Path(requested_root).expanduser()) + return f"hf://datasets/{meta.repo_id}@{meta.revision}" + + +def make_sidecar_spec(meta: LeRobotDatasetMetadata, data_root: str) -> SidecarSpec: + relative_paths = sorted( + { + str(meta.get_video_file_path(episode_index, video_key)) + for episode_index in range(int(meta.total_episodes)) + for video_key in meta.video_keys + } + ) + root = Path(data_root).expanduser() + source_files: tuple[tuple[str, int | None], ...] + if root.is_dir(): + source_files = tuple((path, (root / path).stat().st_size) for path in relative_paths) + else: + source_files = tuple((path, None) for path in relative_paths) + return SidecarSpec( + repo_id=meta.repo_id, + revision=str(meta.revision), + data_root=data_root.rstrip("/"), + source_files=source_files, + ) + + +def build_mp4_sidecar( + destination: str | Path, + spec: SidecarSpec, + *, + workers: int = 8, + range_backend: str = "native-http", + max_probe_bytes: int = 64 * 1024 * 1024, +) -> None: + EpisodeVideoManifest.write_file_sidecar( + destination, + [path for path, _size in spec.source_files], + spec.data_root, + spec=spec, + range_backend=range_backend, + workers=workers, + max_probe_bytes=max_probe_bytes, + ) + + +def published_sidecar_url(spec: SidecarSpec, cache_root: str | Path = DEFAULT_SIDECAR_CACHE) -> str: + name = sidecar_cache_path(cache_root, spec).name + return f"{spec.data_root}/meta/mp4-sidecars/{name}" + + +def download_published_sidecar( + destination: Path, + spec: SidecarSpec, + *, + cache_root: str | Path = DEFAULT_SIDECAR_CACHE, +) -> bool: + if Path(spec.data_root).expanduser().is_dir(): + return False + filesystem, source = fsspec.core.url_to_fs(published_sidecar_url(spec, cache_root)) + if not filesystem.exists(source): + return False + with filesystem.open(source, "rb") as remote, destination.open("wb") as local: + shutil.copyfileobj(remote, local) + return True + + +def ensure_dataset_mp4_sidecar( + meta: LeRobotDatasetMetadata, + data_root: str, + *, + cache_root: str | Path = DEFAULT_SIDECAR_CACHE, + workers: int = 8, + range_backend: str = "native-http", + lock_timeout_s: float = 30 * 60, +) -> Path | None: + if not meta.video_keys: + return None + + spec = make_sidecar_spec(meta, data_root) + logging.info( + "Resolving training-time MP4 sidecar for %s@%s (%d files)", + spec.repo_id, + spec.revision, + len(spec.source_files), + ) + return ensure_mp4_sidecar( + spec, + cache_root, + build=lambda path, target_spec: build_mp4_sidecar( + path, + target_spec, + workers=workers, + range_backend=range_backend, + ), + download=lambda path, target_spec: download_published_sidecar( + path, + target_spec, + cache_root=cache_root, + ), + lock_timeout_s=lock_timeout_s, + ) diff --git a/src/lerobot/datasets/video_utils.py b/src/lerobot/datasets/video_utils.py index ef3005dd8..da5905831 100644 --- a/src/lerobot/datasets/video_utils.py +++ b/src/lerobot/datasets/video_utils.py @@ -28,7 +28,7 @@ from dataclasses import asdict, dataclass, field from fractions import Fraction from pathlib import Path from threading import Lock -from typing import Any, ClassVar +from typing import Any, BinaryIO, ClassVar import av import fsspec @@ -105,7 +105,7 @@ def decode_video_frames( def decode_video_frames_pyav( - video_path: Path | str, + video_path: Path | str | BinaryIO, timestamps: list[float], tolerance_s: float, log_loaded_timestamps: bool = False, @@ -124,7 +124,7 @@ def decode_video_frames_pyav( video can be adjusted at encoding time to trade off decoding speed against file size. Args: - video_path: Path to the video file. + video_path: Path to the video file or a seekable binary file object. timestamps: List of timestamps (in seconds) to extract frames for. tolerance_s: Allowed deviation in seconds between a queried timestamp and the closest decoded frame. @@ -137,7 +137,8 @@ def decode_video_frames_pyav( torch.Tensor of shape (len(timestamps), C, H, W). """ # TODO(rcadene): also load audio stream at the same time - video_path = str(video_path) + video_source = str(video_path) if isinstance(video_path, (Path, str)) else video_path + video_label = str(video_path) if isinstance(video_path, (Path, str)) else "" # set the first and last requested timestamps # Note: previous timestamps are usually loaded, since we need to access the previous key frame @@ -151,7 +152,7 @@ def decode_video_frames_pyav( # av.time_base units (microseconds). `backward=True` lands us on the nearest keyframe at or # before `first_ts`, so we can then decode forward until we cover `last_ts`. See: # https://pyav.basswood-io.com/docs/stable/api/container.html#av.container.InputContainer.seek - with av.open(video_path) as container: + with av.open(video_source) as container: stream = container.streams.video[0] # Seek to the nearest keyframe at or before `first_ts` with a 1 frame margin container.seek( @@ -180,7 +181,7 @@ def decode_video_frames_pyav( if not loaded_frames: raise FrameTimestampError( - f"No frames could be decoded from {video_path} in the timestamp range [{first_ts}, {last_ts}]." + f"No frames could be decoded from {video_label} in the timestamp range [{first_ts}, {last_ts}]." ) query_ts = torch.tensor(timestamps) @@ -199,7 +200,7 @@ def decode_video_frames_pyav( " To be safe, we advise to ignore this item during training." f"\nqueried timestamps: {query_ts}" f"\nloaded timestamps: {loaded_ts_t}" - f"\nvideo: {video_path}" + f"\nvideo: {video_label}" f"\nbackend: pyav" ) diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index ec5565cf4..0c96ac38f 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -41,6 +41,7 @@ from lerobot.common.train_utils import ( load_fsdp_optimizer_state, load_training_batch_size, load_training_num_processes, + load_training_num_workers, load_training_state, push_checkpoint_to_hub, save_checkpoint, @@ -201,7 +202,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): require_package("accelerate", extra="training") from accelerate import Accelerator - from accelerate.utils import DistributedDataParallelKwargs, DistributedType + from accelerate.utils import DistributedDataParallelKwargs, DistributedType, send_to_device cfg.validate() @@ -459,6 +460,59 @@ 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, + ) + rank_frames = dataset.num_frames_for_rank( + accelerator.process_index, + accelerator.num_processes, + train_num_workers, + ) + if rank_frames == 0: + raise ValueError("This rank owns no streaming episodes. Reduce the distributed world size.") + if cfg.resume and step > 0: + saved_num_processes = load_training_num_processes(cfg.checkpoint_path) + saved_batch_size = load_training_batch_size(cfg.checkpoint_path) + saved_num_workers = load_training_num_workers(cfg.checkpoint_path) + if saved_num_processes not in (None, accelerator.num_processes): + raise ValueError( + "Sample-exact streaming resume requires the checkpoint world size " + f"({saved_num_processes}) to match the current world size " + f"({accelerator.num_processes})." + ) + if saved_batch_size not in (None, cfg.batch_size): + raise ValueError( + "Sample-exact streaming resume requires the checkpoint batch size " + f"({saved_batch_size}) to match the current batch size ({cfg.batch_size})." + ) + if saved_num_workers not in (None, train_num_workers): + raise ValueError( + "Sample-exact streaming resume requires the checkpoint DataLoader worker count " + f"({saved_num_workers}) to match the current count ({train_num_workers})." + ) + stream_offset = step * (saved_batch_size or cfg.batch_size) + dataset.load_state_dict( + { + "epoch": 0, + "offset": stream_offset, + "batch_size": cfg.batch_size, + } + ) + if is_main_process: + logging.info(f"Resuming streaming data order at local sample {stream_offset}") # Only swap in the language-aware collate when the dataset actually # declares language columns; otherwise stay on PyTorch's default @@ -466,15 +520,20 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None dataloader = torch.utils.data.DataLoader( dataset, - num_workers=cfg.num_workers, + num_workers=train_num_workers if cfg.dataset.streaming else cfg.num_workers, batch_size=cfg.batch_size, shuffle=shuffle and not cfg.dataset.streaming, sampler=sampler, pin_memory=device.type == "cuda", drop_last=False, collate_fn=collate_fn, - prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None, - persistent_workers=cfg.persistent_workers and cfg.num_workers > 0, + prefetch_factor=( + cfg.prefetch_factor + if (train_num_workers if cfg.dataset.streaming else cfg.num_workers) > 0 + else None + ), + persistent_workers=cfg.persistent_workers + and (train_num_workers if cfg.dataset.streaming else cfg.num_workers) > 0, ) # Build eval dataloader if a held-out split exists @@ -506,7 +565,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): # Prepare everything with accelerator accelerator.wait_for_everyone() - if eval_dataloader is not None: + if cfg.dataset.streaming: + # StreamingLeRobotDataset already assigns complete episodes disjointly to ranks and workers. + # Preparing this loader would make Accelerate shard its batches a second time. + if eval_dataloader is not None: + policy, optimizer, lr_scheduler, eval_dataloader = accelerator.prepare( + policy, optimizer, lr_scheduler, eval_dataloader + ) + else: + policy, optimizer, lr_scheduler = accelerator.prepare(policy, optimizer, lr_scheduler) + elif eval_dataloader is not None: policy, optimizer, dataloader, lr_scheduler, eval_dataloader = accelerator.prepare( policy, optimizer, dataloader, lr_scheduler, eval_dataloader ) @@ -569,6 +637,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): for _ in range(step, cfg.steps): start_time = time.perf_counter() batch = next(dl_iter) + if cfg.dataset.streaming: + batch = send_to_device(batch, device, non_blocking=device.type == "cuda") for cam_key in dataset.meta.camera_keys: if cam_key in batch and batch[cam_key].dtype == torch.uint8: batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0 @@ -665,6 +735,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): postprocessor=postprocessor, num_processes=accelerator.num_processes, batch_size=cfg.batch_size, + num_workers=train_num_workers if cfg.dataset.streaming else cfg.num_workers, model_state_dict=model_state_dict, optim_state_dict=optim_state_dict, ) diff --git a/src/lerobot/streaming/__init__.py b/src/lerobot/streaming/__init__.py new file mode 100644 index 000000000..5f70239b5 --- /dev/null +++ b/src/lerobot/streaming/__init__.py @@ -0,0 +1,9 @@ +# 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 + +"""Low-level primitives for training-time dataset streaming.""" diff --git a/src/lerobot/datasets/episode_video_streaming.py b/src/lerobot/streaming/episode_video.py similarity index 87% rename from src/lerobot/datasets/episode_video_streaming.py rename to src/lerobot/streaming/episode_video.py index ebb3e7269..aed528447 100644 --- a/src/lerobot/datasets/episode_video_streaming.py +++ b/src/lerobot/streaming/episode_video.py @@ -11,17 +11,19 @@ from __future__ import annotations import contextlib import io import json +import logging import os +import posixpath import threading import time from collections import OrderedDict -from concurrent.futures import Future, ThreadPoolExecutor +from collections.abc import Sequence +from concurrent.futures import Future, ThreadPoolExecutor, as_completed from dataclasses import dataclass from datetime import UTC, datetime -from importlib import metadata from pathlib import Path from types import MethodType -from typing import Any +from typing import TYPE_CHECKING, Any from urllib.parse import quote, urljoin, urlparse from uuid import uuid4 @@ -31,8 +33,11 @@ import numpy as np from huggingface_hub import HfApi, HfFileSystem, constants from huggingface_hub.utils import get_session, hf_raise_for_status -from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata -from lerobot.datasets.mp4 import Mp4Index, Mp4SampleSlice, fetch_mp4_index, synthesize_mp4 +from lerobot.streaming.mp4 import Mp4Index, Mp4SampleSlice, fetch_mp4_index, synthesize_mp4 + +if TYPE_CHECKING: + from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata + from lerobot.streaming.sidecar import SidecarSpec _HTTP_FAILURE_LOG_LOCK = threading.Lock() @@ -134,11 +139,15 @@ class ThreadLocalRangeFetcher: def __init__(self, data_root: str | Path, *, block_size: int = 2**20, cache_type: str = "none"): self.data_root = str(data_root).rstrip("/") - protocol = "hf" if self.data_root.startswith("hf://") else "file" - self.fs = fsspec.filesystem(protocol) + self.fs, self._root_path = fsspec.core.url_to_fs(self.data_root) + self._is_local = self.fs.protocol in ("file", "local") or ( + isinstance(self.fs.protocol, tuple) and "file" in self.fs.protocol + ) self.block_size = block_size self.cache_type = cache_type self._local = threading.local() + self._handles_lock = threading.Lock() + self._all_handles: dict[int, Any] = {} self._timing_lock = threading.Lock() self._timing_totals = { "range_jobs": 0.0, @@ -149,9 +158,9 @@ class ThreadLocalRangeFetcher: } def _url(self, relative_path: str) -> str: - if self.data_root.startswith("hf://"): - return f"{self.data_root}/{relative_path}" - return str(Path(self.data_root) / relative_path) + if self._is_local: + return str(Path(self._root_path) / relative_path) + return posixpath.join(self._root_path.rstrip("/"), relative_path.lstrip("/")) def _handle(self, relative_path: str): handles = getattr(self._local, "handles", None) @@ -165,6 +174,8 @@ class ThreadLocalRangeFetcher: ) self._instrument_hf_handle(handle) handles[relative_path] = handle + with self._handles_lock: + self._all_handles[id(handle)] = handle return handle def info_size(self, relative_path: str) -> int: @@ -332,13 +343,15 @@ class ThreadLocalRangeFetcher: return dict(self._timing_totals) def close(self) -> None: - handles = getattr(self._local, "handles", None) - if handles is None: - return - for handle in handles.values(): + with self._handles_lock: + handles = list(self._all_handles.values()) + self._all_handles.clear() + for handle in handles: with contextlib.suppress(Exception): handle.close() - handles.clear() + local_handles = getattr(self._local, "handles", None) + if local_handles is not None: + local_handles.clear() class NativeHTTPRangeFetcher: @@ -751,7 +764,7 @@ def make_range_fetcher( class EpisodeVideoManifest: - _FILE_SIDECAR_CACHE: dict[str, dict[str, VideoFileRecord]] = {} + _FILE_SIDECAR_CACHE: dict[str, tuple[tuple[int, int], dict[str, VideoFileRecord]]] = {} _FILE_SIDECAR_CACHE_LOCK = threading.Lock() def __init__( @@ -873,7 +886,14 @@ class EpisodeVideoManifest: try: with ThreadPoolExecutor(max_workers=workers) as pool: - return list(pool.map(build_file, rel_paths)) + futures = {pool.submit(build_file, path): path for path in rel_paths} + records = [] + progress_interval = max(1, len(futures) // 20) + for completed, future in enumerate(as_completed(futures), start=1): + records.append(future.result()) + if completed == len(futures) or completed % progress_interval == 0: + logging.info("Indexed %d/%d MP4 files for streaming sidecar", completed, len(futures)) + return sorted(records, key=lambda record: record.file_path) finally: fetcher.close() @@ -884,6 +904,7 @@ class EpisodeVideoManifest: rel_paths: list[str], data_root: str | Path, *, + spec: SidecarSpec, range_backend: str = "native-http", workers: int = 8, header_probe_bytes: int = 4 * 1024 * 1024, @@ -897,14 +918,22 @@ class EpisodeVideoManifest: header_probe_bytes=header_probe_bytes, max_probe_bytes=max_probe_bytes, ) - cls.save_file_sidecar(sidecar_path, records) + cls.save_file_sidecar(sidecar_path, records, spec=spec) @staticmethod - def save_file_sidecar(sidecar_path: str | Path, records: list[VideoFileRecord]) -> None: + def save_file_sidecar( + sidecar_path: str | Path, + records: list[VideoFileRecord], + *, + spec: SidecarSpec, + ) -> None: sidecar_path = Path(sidecar_path) sidecar_path.parent.mkdir(parents=True, exist_ok=True) payload = { - "version": 1, + "version": 2, + "sidecar": spec.with_source_files( + tuple((record.file_path, record.file_size) for record in records) + ).to_dict(), "files": [ {"file_path": record.file_path, "file_size": record.file_size, "mp4": record.mp4.to_dict()} for record in records @@ -918,17 +947,49 @@ class EpisodeVideoManifest: arrays[f"{file_idx}/sample_offsets"] = record.mp4.sample_offsets arrays[f"{file_idx}/sync_samples"] = record.mp4.sync_samples np.savez_compressed(sidecar_path, manifest_json=json.dumps(payload).encode("utf-8"), **arrays) + cache_key = str(sidecar_path.expanduser()) + with EpisodeVideoManifest._FILE_SIDECAR_CACHE_LOCK: + EpisodeVideoManifest._FILE_SIDECAR_CACHE.pop(cache_key, None) + + @staticmethod + def load_file_sidecar_metadata(sidecar_path: str | Path) -> dict[str, Any]: + with np.load(sidecar_path, allow_pickle=False) as data: + payload = json.loads(bytes(data["manifest_json"]).decode("utf-8")) + if payload.get("version") != 2 or not isinstance(payload.get("sidecar"), dict): + raise ValueError(f"Unsupported MP4 sidecar schema in {sidecar_path}") + return payload["sidecar"] + + @staticmethod + def validate_file_sidecar(sidecar_path: str | Path, spec: SidecarSpec) -> bool: + try: + from lerobot.streaming.sidecar import SidecarSpec + + candidate = SidecarSpec.from_dict(EpisodeVideoManifest.load_file_sidecar_metadata(sidecar_path)) + if not spec.matches(candidate): + return False + records = EpisodeVideoManifest.load_file_sidecar(sidecar_path) + except (OSError, ValueError, KeyError, TypeError, json.JSONDecodeError): + return False + + expected = dict(candidate.source_files) + actual = {path: record.file_size for path, record in records.items()} + return actual == expected @staticmethod def load_file_sidecar(sidecar_path: str | Path) -> dict[str, VideoFileRecord]: - cache_key = str(Path(sidecar_path).expanduser()) + path = Path(sidecar_path).expanduser() + cache_key = str(path) + stat = path.stat() + signature = (stat.st_mtime_ns, stat.st_size) with EpisodeVideoManifest._FILE_SIDECAR_CACHE_LOCK: cached = EpisodeVideoManifest._FILE_SIDECAR_CACHE.get(cache_key) - if cached is not None: - return cached + if cached is not None and cached[0] == signature: + return cached[1] - with np.load(sidecar_path, allow_pickle=False) as data: + with np.load(path, allow_pickle=False) as data: payload = json.loads(bytes(data["manifest_json"]).decode("utf-8")) + if payload.get("version") != 2: + raise ValueError(f"Unsupported MP4 sidecar schema in {path}") records = {} for file_idx, item in enumerate(payload["files"]): arrays = { @@ -944,7 +1005,7 @@ class EpisodeVideoManifest: mp4 = Mp4Index.from_dict(item["mp4"], arrays) records[item["file_path"]] = VideoFileRecord(item["file_path"], int(item["file_size"]), mp4) with EpisodeVideoManifest._FILE_SIDECAR_CACHE_LOCK: - EpisodeVideoManifest._FILE_SIDECAR_CACHE[cache_key] = records + EpisodeVideoManifest._FILE_SIDECAR_CACHE[cache_key] = (signature, records) return records def camera_id(self, camera_key: str) -> int: @@ -1058,7 +1119,7 @@ class ExactCoveragePool: def resident(self) -> list[int]: return list(self._remaining) - def __iter__(self) -> "ExactCoveragePool": + def __iter__(self) -> ExactCoveragePool: return self def __next__(self) -> tuple[int, int]: @@ -1097,7 +1158,12 @@ class EpisodeByteCache: native_http_retries: int = 4, native_http_subranges: int = 1, open_decoders: bool = True, + max_open_decoders: int = 64, ): + if byte_budget <= 0: + raise ValueError("byte_budget must be positive") + if max_open_decoders <= 0: + raise ValueError("max_open_decoders must be positive") self.manifest = manifest self.fetcher = make_range_fetcher( data_root, @@ -1110,9 +1176,12 @@ class EpisodeByteCache: ) self.byte_budget = byte_budget self.open_decoders = open_decoders + self.max_open_decoders = max_open_decoders self._pool = ThreadPoolExecutor(max_workers=workers) self._cache: OrderedDict[tuple[int, str], dict[str, Any]] = OrderedDict() + self._decoders: OrderedDict[tuple[int, str], Any] = OrderedDict() self._futures: dict[tuple[int, str], Future[dict[str, Any]]] = {} + self._retained_episodes: dict[int, int] = {} self._bytes = 0 self._lock = threading.Lock() self._timing_totals = { @@ -1124,10 +1193,12 @@ class EpisodeByteCache: } def close(self) -> None: - self._pool.shutdown(wait=True) + self._pool.shutdown(wait=True, cancel_futures=True) with self._lock: self._cache.clear() + self._decoders.clear() self._futures.clear() + self._retained_episodes.clear() self._bytes = 0 self.fetcher.close() @@ -1141,9 +1212,34 @@ class EpisodeByteCache: for camera_key in self.manifest.video_keys: self._submit(episode_index, camera_key) + def retain_episode(self, episode_index: int) -> None: + with self._lock: + self._retained_episodes[episode_index] = self._retained_episodes.get(episode_index, 0) + 1 + + def release_episode(self, episode_index: int) -> None: + with self._lock: + count = self._retained_episodes.get(episode_index, 0) + if count <= 1: + self._retained_episodes.pop(episode_index, None) + else: + self._retained_episodes[episode_index] = count - 1 + self._evict_locked() + + @property + def resident_bytes(self) -> int: + with self._lock: + return self._bytes + + @property + def open_decoder_count(self) -> int: + with self._lock: + return len(self._decoders) + def ensure_ready(self, episode_index: int) -> None: for camera_key in self.manifest.video_keys: self.get_bytes(episode_index, camera_key) + if self.open_decoders: + self.get_decoder(episode_index, camera_key) def is_ready(self, episode_index: int) -> bool: """Non-blocking: True when every camera of the episode is fetched (cached or future done). @@ -1164,20 +1260,30 @@ class EpisodeByteCache: def get_bytes(self, episode_index: int, camera_key: str) -> bytes: return self._get_entry(episode_index, camera_key)["bytes"] - def get_decoder(self, episode_index: int, camera_key: str): + def get_decoder(self, episode_index: int, camera_key: str) -> Any: + key = (episode_index, camera_key) entry = self._get_entry(episode_index, camera_key) - decoder = entry.get("decoder") - if decoder is None: - decoder = open_video_decoder(io.BytesIO(entry["bytes"])) - entry["decoder"] = decoder + with self._lock: + decoder = self._decoders.get(key) + if decoder is not None: + self._decoders.move_to_end(key) + return decoder + + decoder = open_video_decoder(io.BytesIO(entry["bytes"])) + with self._lock: + existing = self._decoders.get(key) + if existing is not None: + self._decoders.move_to_end(key) + return existing + self._decoders[key] = decoder + while len(self._decoders) > self.max_open_decoders: + self._decoders.popitem(last=False) return decoder def get_frames(self, episode_index: int, camera_key: str, timestamps: list[float]): span = self.manifest.lookup(episode_index, camera_key) local_ts = [ts - span.source_start_pts for ts in timestamps] decoder = self.get_decoder(episode_index, camera_key) - if hasattr(decoder, "get_frames_played_at"): - return decoder.get_frames_played_at(local_ts).data metadata = decoder.metadata fps = getattr(metadata, "average_fps", None) if fps is None: @@ -1224,7 +1330,13 @@ class EpisodeByteCache: return existing self._cache[key] = entry self._bytes += len(entry["bytes"]) - self._evict_locked() + try: + self._evict_locked() + except MemoryError: + failed_entry = self._cache.pop(key, None) + if failed_entry is not None: + self._bytes -= len(failed_entry["bytes"]) + raise timings = entry.pop("_timings", None) if timings is not None: self._timing_totals["lookup_s"] += timings["lookup_s"] @@ -1235,9 +1347,18 @@ class EpisodeByteCache: return entry def _evict_locked(self) -> None: - while self._bytes > self.byte_budget and self._cache: - _key, entry = self._cache.popitem(last=False) + while self._bytes > self.byte_budget: + key = next( + (candidate for candidate in self._cache if candidate[0] not in self._retained_episodes), + None, + ) + if key is None: + raise MemoryError( + f"Retained episode bytes exceed byte budget ({self._bytes} > {self.byte_budget})" + ) + entry = self._cache.pop(key) self._bytes -= len(entry["bytes"]) + self._decoders.pop(key, None) def _fetch_and_synthesize(self, episode_index: int, camera_key: str) -> dict[str, Any]: lookup_start = time.perf_counter() @@ -1263,15 +1384,12 @@ class EpisodeByteCache: synthesize_s = time.perf_counter() - synthesize_start entry: dict[str, Any] = { "bytes": mp4_bytes, - "decoder": None, "_timings": { "lookup_s": lookup_s, "fetch_s": fetch_s, "synthesize_s": synthesize_s, }, } - if self.open_decoders: - entry["decoder"] = open_video_decoder(io.BytesIO(mp4_bytes)) return entry @@ -1283,32 +1401,6 @@ def open_video_decoder(file_like_or_bytesio, frame_mappings=None): return VideoDecoder(file_like_or_bytesio, seek_mode="approximate") -def assert_hf_hub_range_cache_branch() -> None: - """Fail unless huggingface_hub was installed from the required range-cache branch.""" - - try: - dist = metadata.distribution("huggingface_hub") - except metadata.PackageNotFoundError as exc: - raise AssertionError("huggingface_hub is not installed") from exc - - candidates = [] - direct_url = dist.read_text("direct_url.json") - if direct_url: - candidates.append(direct_url) - with contextlib.suppress(json.JSONDecodeError): - parsed = json.loads(direct_url) - candidates.append(str(parsed.get("url", ""))) - candidates.append(str(parsed.get("vcs_info", {}).get("requested_revision", ""))) - candidates.append(str(parsed.get("vcs_info", {}).get("commit_id", ""))) - - text = "\n".join(candidates) - if "feat/hffs-cache-cdn-range-reads" not in text: - raise AssertionError( - "huggingface_hub must be installed from " - "git+https://github.com/huggingface/huggingface_hub.git@feat/hffs-cache-cdn-range-reads" - ) - - @dataclass class StageTimer: fetch_ms: float = 0.0 diff --git a/src/lerobot/datasets/mp4.py b/src/lerobot/streaming/mp4.py similarity index 99% rename from src/lerobot/datasets/mp4.py rename to src/lerobot/streaming/mp4.py index 5d16908d1..4baf513c6 100644 --- a/src/lerobot/datasets/mp4.py +++ b/src/lerobot/streaming/mp4.py @@ -6,6 +6,8 @@ # # http://www.apache.org/licenses/LICENSE-2.0 +"""MP4 indexing and in-memory episode synthesis primitives.""" + from __future__ import annotations import struct diff --git a/src/lerobot/streaming/sidecar.py b/src/lerobot/streaming/sidecar.py new file mode 100644 index 000000000..312662ed1 --- /dev/null +++ b/src/lerobot/streaming/sidecar.py @@ -0,0 +1,176 @@ +# 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 + +"""Revision-safe lifecycle for locally cached MP4 byte-index sidecars.""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +import re +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from uuid import uuid4 + +from filelock import FileLock, Timeout + +from lerobot.streaming.episode_video import EpisodeVideoManifest + +SIDECAR_SCHEMA_VERSION = 2 + + +class SidecarLockTimeoutError(TimeoutError): + """Raised when another process does not finish a sidecar build in time.""" + + +@dataclass(frozen=True) +class SidecarSpec: + repo_id: str + revision: str + data_root: str + source_files: tuple[tuple[str, int | None], ...] + schema_version: int = SIDECAR_SCHEMA_VERSION + + def __post_init__(self) -> None: + if not self.repo_id: + raise ValueError("repo_id must not be empty") + if not self.revision: + raise ValueError("revision must not be empty") + normalized = tuple( + sorted((str(path), None if size is None else int(size)) for path, size in self.source_files) + ) + if any(not path or size is not None and size < 0 for path, size in normalized): + raise ValueError("source file paths must be non-empty and sizes must be non-negative") + object.__setattr__(self, "source_files", normalized) + + def to_dict(self) -> dict[str, object]: + return { + "schema_version": self.schema_version, + "repo_id": self.repo_id, + "revision": self.revision, + "data_root": self.data_root, + "source_files": [{"path": path, "size": size} for path, size in self.source_files], + } + + @classmethod + def from_dict(cls, data: dict[str, object]) -> SidecarSpec: + source_files = data.get("source_files") + if not isinstance(source_files, list): + raise ValueError("MP4 sidecar source_files must be a list") + parsed_files: list[tuple[str, int | None]] = [] + for item in source_files: + if not isinstance(item, dict) or not isinstance(item.get("path"), str): + raise ValueError("Invalid MP4 sidecar source file entry") + size = item.get("size") + parsed_files.append((item["path"], None if size is None else int(size))) + return cls( + schema_version=int(data["schema_version"]), + repo_id=str(data["repo_id"]), + revision=str(data["revision"]), + data_root=str(data["data_root"]), + source_files=tuple(parsed_files), + ) + + def with_source_files(self, source_files: tuple[tuple[str, int], ...]) -> SidecarSpec: + return SidecarSpec( + repo_id=self.repo_id, + revision=self.revision, + data_root=self.data_root, + source_files=source_files, + schema_version=self.schema_version, + ) + + def matches(self, candidate: SidecarSpec) -> bool: + if ( + self.schema_version != candidate.schema_version + or self.repo_id != candidate.repo_id + or self.revision != candidate.revision + or self.data_root != candidate.data_root + ): + return False + expected = dict(self.source_files) + actual = dict(candidate.source_files) + if expected.keys() != actual.keys(): + return False + return all(size is None or actual[path] == size for path, size in expected.items()) + + +SidecarBuilder = Callable[[Path, SidecarSpec], None] +SidecarDownloader = Callable[[Path, SidecarSpec], bool] + + +def sidecar_cache_path(cache_root: str | Path, spec: SidecarSpec) -> Path: + identity = json.dumps( + { + "schema_version": spec.schema_version, + "repo_id": spec.repo_id, + "revision": spec.revision, + "data_root": spec.data_root, + }, + sort_keys=True, + separators=(",", ":"), + ) + digest = hashlib.sha256(identity.encode()).hexdigest()[:16] + repo_slug = re.sub(r"[^A-Za-z0-9_.-]+", "--", spec.repo_id).strip("-") or "dataset" + revision_slug = re.sub(r"[^A-Za-z0-9_.-]+", "-", spec.revision).strip("-")[:32] or "revision" + return ( + Path(cache_root).expanduser() / repo_slug / f"mp4-v{spec.schema_version}-{revision_slug}-{digest}.npz" + ) + + +def ensure_mp4_sidecar( + spec: SidecarSpec, + cache_root: str | Path, + *, + build: SidecarBuilder, + download: SidecarDownloader | None = None, + lock_timeout_s: float = 30 * 60, +) -> Path: + """Return a valid local sidecar, downloading or building it exactly once. + + This function never uploads. ``download`` and ``build`` must write only to the temporary path + provided to them; a validated file becomes visible at the cache path through ``os.replace``. + """ + + destination = sidecar_cache_path(cache_root, spec) + if EpisodeVideoManifest.validate_file_sidecar(destination, spec): + return destination + + destination.parent.mkdir(parents=True, exist_ok=True) + lock_path = destination.with_suffix(f"{destination.suffix}.lock") + try: + with FileLock(lock_path, timeout=lock_timeout_s): + if EpisodeVideoManifest.validate_file_sidecar(destination, spec): + return destination + + temporary = destination.parent / f".{destination.name}.{uuid4().hex}.tmp.npz" + try: + if download is not None: + logging.info("Looking for published MP4 sidecar for %s@%s", spec.repo_id, spec.revision) + if download(temporary, spec) and EpisodeVideoManifest.validate_file_sidecar( + temporary, spec + ): + os.replace(temporary, destination) + return destination + temporary.unlink(missing_ok=True) + + logging.info("Building MP4 sidecar for %s@%s", spec.repo_id, spec.revision) + build(temporary, spec) + if not EpisodeVideoManifest.validate_file_sidecar(temporary, spec): + raise ValueError("Built MP4 sidecar failed revision and source validation") + os.replace(temporary, destination) + return destination + finally: + temporary.unlink(missing_ok=True) + except Timeout as exc: + raise SidecarLockTimeoutError( + f"Timed out waiting {lock_timeout_s:g}s for MP4 sidecar lock {lock_path}" + ) from exc diff --git a/tests/configs/test_default.py b/tests/configs/test_default.py index 238b8bacd..57945b75e 100644 --- a/tests/configs/test_default.py +++ b/tests/configs/test_default.py @@ -36,3 +36,16 @@ def test_dataset_config_none_episodes_ok(): def test_dataset_config_empty_episodes_ok(): DatasetConfig(repo_id="user/repo", episodes=[]) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("streaming_episode_pool_size", 0, "episode_pool_size"), + ("streaming_prefetch_episodes", -1, "prefetch_episodes"), + ("streaming_byte_budget_gb", 0, "byte_budget_gb"), + ], +) +def test_dataset_config_rejects_invalid_streaming_resource_limits(field, value, message): + with pytest.raises(ValueError, match=message): + DatasetConfig(repo_id="user/repo", **{field: value}) diff --git a/tests/datasets/test_depth.py b/tests/datasets/test_depth.py index 5391cc558..ac9b51d43 100644 --- a/tests/datasets/test_depth.py +++ b/tests/datasets/test_depth.py @@ -323,14 +323,13 @@ class TestDepthUnitMetadata: np.testing.assert_allclose(float(np.asarray(stats["mean"]).reshape(-1)[0]), expected, rtol=0.05) np.testing.assert_allclose(float(np.asarray(stats["count"]).reshape(-1)[0]), count) - if not use_videos: - depth = read_dataset[0][DEPTH_KEY] - assert torch.allclose(depth, torch.full_like(depth, expected)) + from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset - from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset - - stream_dataset = StreamingLeRobotDataset( - repo_id=DUMMY_REPO_ID, root=tmp_path / "ds", depth_output_unit=output_unit - ) - stream_depth = next(iter(stream_dataset))[DEPTH_KEY] - assert torch.allclose(stream_depth, torch.full_like(stream_depth, expected)) + stream_dataset = StreamingLeRobotDataset( + repo_id=DUMMY_REPO_ID, root=tmp_path / "ds", depth_output_unit=output_unit + ) + stream_item = next(iter(stream_dataset)) + stream_depth = stream_item[DEPTH_KEY] + reference_depth = read_dataset[int(stream_item["index"])][DEPTH_KEY] + assert torch.allclose(stream_depth, reference_depth) + assert torch.allclose(stream_depth, torch.full_like(stream_depth, expected), rtol=0.05) diff --git a/tests/datasets/test_episode_parquet_reader.py b/tests/datasets/test_episode_parquet_reader.py new file mode 100644 index 000000000..80ab1ffef --- /dev/null +++ b/tests/datasets/test_episode_parquet_reader.py @@ -0,0 +1,110 @@ +#!/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 + +from pathlib import Path + +import fsspec +import pytest + +pytest.importorskip("pyarrow", reason="pyarrow is required (install lerobot[dataset])") + +import pyarrow as pa +import pyarrow.parquet as pq + +from lerobot.datasets.episode_parquet import EpisodeParquetReader + + +def _table(episodes: list[int]) -> pa.Table: + frame_counts: dict[int, int] = {} + frame_indices = [] + values = [] + ignored = [] + for episode in episodes: + frame_index = frame_counts.get(episode, 0) + frame_counts[episode] = frame_index + 1 + frame_indices.append(frame_index) + values.append(episode * 10 + frame_index) + ignored.append(f"ignored-{episode}-{frame_index}") + return pa.table( + { + "episode_index": episodes, + "frame_index": frame_indices, + "value": values, + "ignored": ignored, + } + ) + + +def _write_episode_row_groups(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + writer = pq.ParquetWriter(path, _table([0]).schema) + try: + writer.write_table(_table([0, 0])) + writer.write_table(_table([1, 1, 1])) + finally: + writer.close() + + +def test_reader_projects_columns_and_reads_matching_row_group(tmp_path: Path) -> None: + path = tmp_path / "data/chunk-000/file-000.parquet" + _write_episode_row_groups(path) + reader = EpisodeParquetReader(tmp_path, columns=("episode_index", "frame_index", "value")) + + table = reader.read_episode(path.relative_to(tmp_path), episode_index=1, expected_rows=3) + + assert table.column_names == ["episode_index", "frame_index", "value"] + assert table.column("value").to_pylist() == [10, 11, 12] + + +def test_reader_filters_legacy_mixed_row_group(tmp_path: Path) -> None: + path = tmp_path / "data/chunk-000/file-000.parquet" + path.parent.mkdir(parents=True) + pq.write_table(_table([0, 0, 1, 1, 1]), path) + reader = EpisodeParquetReader(tmp_path, columns=("episode_index", "frame_index", "value")) + + table = reader.read_episode(path.relative_to(tmp_path), episode_index=1, expected_rows=3) + + assert table.column("episode_index").to_pylist() == [1, 1, 1] + assert table.column("frame_index").to_pylist() == [0, 1, 2] + + +def test_reader_rejects_partial_episode(tmp_path: Path) -> None: + path = tmp_path / "data/chunk-000/file-000.parquet" + path.parent.mkdir(parents=True) + pq.write_table(_table([2, 2]), path) + reader = EpisodeParquetReader(tmp_path, columns=("episode_index", "frame_index")) + + with pytest.raises(ValueError, match="expected 3 rows, found 2"): + reader.read_episode(path.relative_to(tmp_path), episode_index=2, expected_rows=3) + + +def test_reader_rejects_missing_episode(tmp_path: Path) -> None: + path = tmp_path / "data/chunk-000/file-000.parquet" + path.parent.mkdir(parents=True) + pq.write_table(_table([0, 0]), path) + reader = EpisodeParquetReader(tmp_path, columns=("episode_index", "frame_index")) + + with pytest.raises(ValueError, match="episode 4"): + reader.read_episode(path.relative_to(tmp_path), episode_index=4, expected_rows=1) + + +def test_reader_supports_fsspec_remote_root() -> None: + filesystem = fsspec.filesystem("memory") + root = "memory://episode-reader" + path = "episode-reader/data/chunk-000/file-000.parquet" + with filesystem.open(path, "wb") as output: + pq.write_table(_table([0, 0, 0]), output) + reader = EpisodeParquetReader(root, columns=("episode_index", "value")) + + table = reader.read_episode("data/chunk-000/file-000.parquet", episode_index=0, expected_rows=3) + + assert table.column("value").to_pylist() == [0, 1, 2] diff --git a/tests/datasets/test_episode_video_streaming.py b/tests/datasets/test_episode_video_streaming.py index 226b04c9b..eedb1fdc6 100644 --- a/tests/datasets/test_episode_video_streaming.py +++ b/tests/datasets/test_episode_video_streaming.py @@ -10,12 +10,19 @@ import json import struct +import threading +from concurrent.futures import ThreadPoolExecutor import numpy as np import pytest -from lerobot.datasets.episode_video_streaming import assert_hf_hub_range_cache_branch -from lerobot.datasets.mp4 import ( +from lerobot.streaming.episode_video import ( + EpisodeByteCache, + EpisodeVideoManifest, + ThreadLocalRangeFetcher, + _log_http_failure, +) +from lerobot.streaming.mp4 import ( _box, _co64, _dinf, @@ -89,33 +96,115 @@ def test_parser_accepts_co64_chunk_offsets(): np.testing.assert_array_equal(mp4.sample_offsets, np.array([10_000, 10_050, 10_025])) -def test_hf_hub_branch_assertion_accepts_requested_revision(monkeypatch): - class FakeDist: - def read_text(self, name): - assert name == "direct_url.json" - return json.dumps( - { - "url": "https://github.com/huggingface/huggingface_hub.git", - "vcs_info": {"requested_revision": "feat/hffs-cache-cdn-range-reads"}, - } - ) - +def _fake_cache(monkeypatch, tmp_path, *, byte_budget=8, max_open_decoders=1): + manifest = EpisodeVideoManifest(video_keys=["camera"], files=[], spans={}) + cache = EpisodeByteCache( + manifest, + tmp_path, + byte_budget=byte_budget, + workers=1, + open_decoders=False, + max_open_decoders=max_open_decoders, + ) monkeypatch.setattr( - "lerobot.datasets.episode_video_streaming.metadata.distribution", lambda _: FakeDist() + cache, + "_fetch_and_synthesize", + lambda episode_index, _camera_key: {"bytes": bytes([episode_index]) * 5, "_timings": None}, + ) + return cache + + +def test_byte_cache_does_not_evict_retained_episode(monkeypatch, tmp_path): + with _fake_cache(monkeypatch, tmp_path, byte_budget=10) as cache: + cache.retain_episode(0) + cache.ensure_ready(0) + cache.ensure_ready(1) + cache.ensure_ready(2) + + assert (0, "camera") in cache._cache + assert (1, "camera") not in cache._cache + assert cache.resident_bytes <= cache.byte_budget + + +def test_byte_cache_rejects_retained_set_larger_than_budget(monkeypatch, tmp_path): + with _fake_cache(monkeypatch, tmp_path, byte_budget=4) as cache: + cache.retain_episode(0) + + with pytest.raises(MemoryError, match="byte budget"): + cache.ensure_ready(0) + + +def test_decoder_count_has_independent_limit(monkeypatch, tmp_path): + opened = [] + + class FakeDecoder: + pass + + def open_decoder(_data): + decoder = FakeDecoder() + opened.append(decoder) + return decoder + + monkeypatch.setattr("lerobot.streaming.episode_video.open_video_decoder", open_decoder) + with _fake_cache(monkeypatch, tmp_path, byte_budget=20, max_open_decoders=1) as cache: + first = cache.get_decoder(0, "camera") + second = cache.get_decoder(1, "camera") + + assert first is not second + assert cache.open_decoder_count == 1 + + +def test_releasing_episode_allows_immediate_eviction(monkeypatch, tmp_path): + with _fake_cache(monkeypatch, tmp_path, byte_budget=5) as cache: + cache.retain_episode(0) + cache.ensure_ready(0) + cache.release_episode(0) + cache.ensure_ready(1) + + assert (0, "camera") not in cache._cache + assert (1, "camera") in cache._cache + + +def test_range_fetcher_closes_handles_from_all_worker_threads(tmp_path): + (tmp_path / "video.mp4").write_bytes(b"0123456789") + fetcher = ThreadLocalRangeFetcher(tmp_path) + barrier = threading.Barrier(2) + + def read_from_worker(offset): + barrier.wait() + return fetcher.read_range("video.mp4", offset, 1) + + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [pool.submit(read_from_worker, offset) for offset in range(2)] + assert [future.result() for future in futures] == [b"0", b"1"] + + handles = list(fetcher._all_handles.values()) + assert len(handles) == 2 + fetcher.close() + + assert not fetcher._all_handles + assert all(handle.closed for handle in handles) + + +def test_http_failure_log_does_not_write_credentials(tmp_path, monkeypatch): + log_path = tmp_path / "http-failures.jsonl" + monkeypatch.setenv("LEROBOT_HTTP_FAILURE_LOG", str(log_path)) + + _log_http_failure( + backend="native-http", + method="GET", + url="https://cdn.example/private/video.mp4?token=url-secret", + headers={ + "Authorization": "Bearer header-secret", + "Range": "bytes=0-10", + "X-Request-Id": "safe-request-id", + }, + elapsed_s=0.1, + status_code=403, ) - assert_hf_hub_range_cache_branch() - - -def test_hf_hub_branch_assertion_rejects_plain_install(monkeypatch): - class FakeDist: - def read_text(self, name): - assert name == "direct_url.json" - return json.dumps({"url": "https://github.com/huggingface/huggingface_hub.git"}) - - monkeypatch.setattr( - "lerobot.datasets.episode_video_streaming.metadata.distribution", lambda _: FakeDist() - ) - - with pytest.raises(AssertionError): - assert_hf_hub_range_cache_branch() + record = json.loads(log_path.read_text()) + assert record["host"] == "cdn.example" + assert record["path"] == "/private/video.mp4" + assert record["request_id"] == "safe-request-id" + assert "secret" not in log_path.read_text() diff --git a/tests/datasets/test_exact_coverage_pool.py b/tests/datasets/test_exact_coverage_pool.py index 476c17f76..e31a7354d 100644 --- a/tests/datasets/test_exact_coverage_pool.py +++ b/tests/datasets/test_exact_coverage_pool.py @@ -16,7 +16,7 @@ from collections import Counter -from lerobot.datasets.episode_video_streaming import ExactCoveragePool +from lerobot.streaming.episode_video 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) diff --git a/tests/datasets/test_language.py b/tests/datasets/test_language.py index 52c7b3708..e45050919 100644 --- a/tests/datasets/test_language.py +++ b/tests/datasets/test_language.py @@ -25,6 +25,7 @@ from lerobot.datasets.language import ( # noqa: E402 language_persistent_arrow_type, validate_camera_field, ) +from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset # noqa: E402 from lerobot.datasets.utils import DEFAULT_DATA_PATH # noqa: E402 @@ -171,3 +172,16 @@ def test_lerobot_dataset_passes_language_columns_through(tmp_path, empty_lerobot assert first[LANGUAGE_EVENTS] == [event] assert second[LANGUAGE_PERSISTENT] == persistent assert second[LANGUAGE_EVENTS] == [] + + streamed = { + int(item["index"]): item + for item in StreamingLeRobotDataset( + repo_id=dataset.repo_id, + root=root, + shuffle=False, + ) + } + assert streamed[0][LANGUAGE_PERSISTENT] == persistent + assert streamed[0][LANGUAGE_EVENTS] == [event] + assert streamed[1][LANGUAGE_PERSISTENT] == persistent + assert streamed[1][LANGUAGE_EVENTS] == [] diff --git a/tests/datasets/test_streaming.py b/tests/datasets/test_streaming.py index db167f657..cd1e9acbe 100644 --- a/tests/datasets/test_streaming.py +++ b/tests/datasets/test_streaming.py @@ -13,64 +13,16 @@ # 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. -import numpy as np import pytest import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset -from lerobot.datasets.utils import safe_shard from lerobot.utils.constants import ACTION from tests.fixtures.constants import DUMMY_REPO_ID -def get_frames_expected_order(streaming_ds: StreamingLeRobotDataset) -> list[int]: - """Replicates the shuffling logic of StreamingLeRobotDataset to get the expected order of indices.""" - rng = np.random.default_rng(streaming_ds.seed) - buffer_size = streaming_ds.buffer_size - num_shards = streaming_ds.num_shards - - shards_indices = [] - for shard_idx in range(num_shards): - shard = streaming_ds.hf_dataset.shard(num_shards, index=shard_idx) - shard_indices = [item["index"] for item in shard] - shards_indices.append(shard_indices) - - shard_iterators = {i: iter(s) for i, s in enumerate(shards_indices)} - - buffer_indices_generator = streaming_ds._iter_random_indices(rng, buffer_size) - - frames_buffer = [] - expected_indices = [] - - while shard_iterators: # While there are still available shards - available_shard_keys = list(shard_iterators.keys()) - if not available_shard_keys: - break - - # Call _infinite_generator_over_elements with current available shards (key difference!) - shard_key = next(streaming_ds._infinite_generator_over_elements(rng, available_shard_keys)) - - try: - frame_index = next(shard_iterators[shard_key]) - - if len(frames_buffer) == buffer_size: - i = next(buffer_indices_generator) - expected_indices.append(frames_buffer[i]) - frames_buffer[i] = frame_index - else: - frames_buffer.append(frame_index) - - except StopIteration: - del shard_iterators[shard_key] # Remove exhausted shard - - rng.shuffle(frames_buffer) - expected_indices.extend(frames_buffer) - - return expected_indices - - def test_single_frame_consistency(tmp_path, lerobot_dataset_factory): """Test if are correctly accessed""" ds_num_frames = 400 @@ -92,7 +44,7 @@ def test_single_frame_consistency(tmp_path, lerobot_dataset_factory): key_checks = [] for _ in range(ds_num_frames): streaming_frame = next(streaming_ds) - frame_idx = streaming_frame["index"] + frame_idx = int(streaming_frame["index"]) target_frame = ds[frame_idx] for key in streaming_frame: @@ -141,22 +93,16 @@ def test_frames_order_over_epochs(tmp_path, lerobot_dataset_factory, shuffle): repo_id=repo_id, root=local_path, buffer_size=buffer_size, seed=seed, shuffle=shuffle ) - first_epoch_indices = [frame["index"] for frame in streaming_ds] - expected_indices = get_frames_expected_order(streaming_ds) - - assert first_epoch_indices == expected_indices, "First epoch indices do not match expected indices" - - expected_indices = get_frames_expected_order(streaming_ds) + first_epoch_indices = [int(frame["index"]) for frame in streaming_ds] + assert sorted(first_epoch_indices) == list(range(ds_num_frames)) for _ in range(n_epochs): - streaming_indices = [frame["index"] for frame in streaming_ds] - frames_match = all( - s_index == e_index for s_index, e_index in zip(streaming_indices, expected_indices, strict=True) - ) + streaming_indices = [int(frame["index"]) for frame in streaming_ds] + assert sorted(streaming_indices) == list(range(ds_num_frames)) if shuffle: - assert not frames_match + assert streaming_indices != first_epoch_indices else: - assert frames_match + assert streaming_indices == first_epoch_indices @pytest.mark.parametrize( @@ -196,22 +142,16 @@ def test_frames_order_with_shards(tmp_path, lerobot_dataset_factory, shuffle): max_num_shards=4, ) - first_epoch_indices = [frame["index"] for frame in streaming_ds] - expected_indices = get_frames_expected_order(streaming_ds) - - assert first_epoch_indices == expected_indices, "First epoch indices do not match expected indices" + first_epoch_indices = [int(frame["index"]) for frame in streaming_ds] + assert sorted(first_epoch_indices) == list(range(ds_num_frames)) for _ in range(n_epochs): - streaming_indices = [ - frame["index"] for frame in streaming_ds - ] # NOTE: this is the same as first_epoch_indices - frames_match = all( - s_index == e_index for s_index, e_index in zip(streaming_indices, expected_indices, strict=True) - ) + streaming_indices = [int(frame["index"]) for frame in streaming_ds] + assert sorted(streaming_indices) == list(range(ds_num_frames)) if shuffle: - assert not frames_match + assert streaming_indices != first_epoch_indices else: - assert frames_match + assert streaming_indices == first_epoch_indices @pytest.mark.parametrize( @@ -261,7 +201,7 @@ def test_frames_with_delta_consistency(tmp_path, lerobot_dataset_factory, state_ for i in range(ds_num_frames): streaming_frame = next(streaming_ds) - frame_idx = streaming_frame["index"] + frame_idx = int(streaming_frame["index"]) target_frame = ds[frame_idx] assert set(streaming_frame.keys()) == set(target_frame.keys()), ( @@ -344,20 +284,11 @@ def test_frames_with_delta_consistency_with_shards( max_num_shards=4, ) - iter(streaming_ds) - - num_shards = 4 - shards_indices = [] - for shard_idx in range(num_shards): - shard = safe_shard(streaming_ds.hf_dataset, shard_idx, num_shards) - shard_indices = [item["index"] for item in shard] - shards_indices.append(shard_indices) - streaming_ds = iter(streaming_ds) for i in range(ds_num_frames): streaming_frame = next(streaming_ds) - frame_idx = streaming_frame["index"] + frame_idx = int(streaming_frame["index"]) target_frame = ds[frame_idx] assert set(streaming_frame.keys()) == set(target_frame.keys()), ( diff --git a/tests/datasets/test_streaming_factory.py b/tests/datasets/test_streaming_factory.py new file mode 100644 index 000000000..3126473d2 --- /dev/null +++ b/tests/datasets/test_streaming_factory.py @@ -0,0 +1,58 @@ +#!/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 types import SimpleNamespace + +import pytest + +pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") + +from lerobot.configs.default import DatasetConfig +from lerobot.datasets import factory + + +def test_factory_wires_production_streaming_settings(monkeypatch): + captured = {} + + class DummyStreamingDataset: + def __init__(self, *args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + self.meta = SimpleNamespace(camera_keys=[], depth_keys=[], stats={}) + + monkeypatch.setattr(factory, "LeRobotDatasetMetadata", lambda *args, **kwargs: object()) + monkeypatch.setattr(factory, "resolve_delta_timestamps", lambda *args, **kwargs: {"action": [0.0]}) + monkeypatch.setattr(factory, "StreamingLeRobotDataset", DummyStreamingDataset) + dataset_config = DatasetConfig( + repo_id="owner/dataset", + streaming=True, + streaming_data_root="memory://payload", + streaming_episode_pool_size=7, + streaming_prefetch_episodes=3, + streaming_byte_budget_gb=2.5, + ) + cfg = SimpleNamespace( + dataset=dataset_config, + trainable_config=object(), + num_workers=0, + tolerance_s=1e-4, + ) + + dataset = factory.make_dataset(cfg) + + assert isinstance(dataset, DummyStreamingDataset) + assert captured["args"] == ("owner/dataset",) + assert captured["kwargs"]["data_root"] == "memory://payload" + assert captured["kwargs"]["episode_pool_size"] == 7 + assert captured["kwargs"]["prefetch_episodes"] == 3 + assert captured["kwargs"]["byte_budget_gb"] == 2.5 + assert captured["kwargs"]["max_num_shards"] == 1 + assert captured["kwargs"]["return_uint8"] is True + assert captured["kwargs"]["repeat"] is True diff --git a/tests/datasets/test_streaming_production.py b/tests/datasets/test_streaming_production.py new file mode 100644 index 000000000..1db2c6751 --- /dev/null +++ b/tests/datasets/test_streaming_production.py @@ -0,0 +1,519 @@ +#!/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 + +from itertools import islice +from pathlib import Path + +import fsspec +import pytest +import torch + +pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") + +from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset +from lerobot.utils.utils import cycle +from tests.fixtures.constants import DUMMY_REPO_ID + + +def _indices(dataset: StreamingLeRobotDataset) -> list[int]: + return [int(item["index"]) for item in dataset] + + +def _assert_item_equal(left: dict, right: dict) -> None: + assert left.keys() == right.keys() + for key in left: + if isinstance(left[key], torch.Tensor): + assert torch.equal(left[key], right[key]), key + else: + assert left[key] == right[key], key + + +def test_streaming_matches_map_style_with_exact_coverage(tmp_path: Path, lerobot_dataset_factory) -> None: + root = tmp_path / "dataset" + map_dataset = lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=4, + total_frames=40, + use_videos=False, + ) + streaming = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + shuffle=False, + buffer_size=3, + ) + + samples = list(streaming) + + assert len(samples) == len(map_dataset) + assert sorted(int(sample["index"]) for sample in samples) == list(range(len(map_dataset))) + for sample in samples: + _assert_item_equal(sample, map_dataset[int(sample["index"])]) + + +def test_streaming_rgb_video_matches_map_style(tmp_path: Path, lerobot_dataset_factory) -> None: + root = tmp_path / "dataset" + map_dataset = lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=2, + total_frames=20, + ) + streaming = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + shuffle=False, + buffer_size=2, + ) + + for sample in streaming: + reference = map_dataset[int(sample["index"])] + assert sample.keys() == reference.keys() + for camera_key in map_dataset.meta.camera_keys: + assert torch.equal(sample[camera_key], reference[camera_key]), ( + camera_key, + int(sample["index"]), + float((sample[camera_key] - reference[camera_key]).abs().max()), + ) + + +def test_streaming_applies_rgb_transforms_and_preserves_uint8( + tmp_path: Path, lerobot_dataset_factory +) -> None: + root = tmp_path / "dataset" + + def flip_width(image: torch.Tensor) -> torch.Tensor: + return image.flip(-1) + + map_dataset = lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=2, + total_frames=10, + image_transforms=flip_width, + return_uint8=True, + ) + streaming = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + shuffle=False, + buffer_size=2, + image_transforms=flip_width, + return_uint8=True, + ) + + sample = next(iter(streaming)) + reference = map_dataset[int(sample["index"])] + for camera_key in map_dataset.meta.camera_keys: + assert sample[camera_key].dtype == torch.uint8 + assert torch.equal(sample[camera_key], reference[camera_key]) + + +def test_streaming_honors_episode_subset(tmp_path: Path, lerobot_dataset_factory) -> None: + root = tmp_path / "dataset" + map_dataset = lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=5, + total_frames=50, + use_videos=False, + ) + selected = [1, 3] + streaming = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + episodes=selected, + shuffle=False, + buffer_size=2, + ) + + indices = _indices(streaming) + expected = [ + index + for episode in selected + for index in range( + map_dataset.meta.episodes[episode]["dataset_from_index"], + map_dataset.meta.episodes[episode]["dataset_to_index"], + ) + ] + + assert sorted(indices) == sorted(expected) + + +def test_streaming_reads_episode_parquet_from_configured_fsspec_root( + tmp_path: Path, lerobot_dataset_factory +) -> None: + root = tmp_path / "metadata" + map_dataset = lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=3, + total_frames=30, + use_videos=False, + ) + remote_root = "memory://streaming-production" + filesystem = fsspec.filesystem("memory") + for path in (root / "data").glob("*/*.parquet"): + relative = path.relative_to(root).as_posix() + filesystem.put(str(path), f"streaming-production/{relative}") + + streaming = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + data_root=remote_root, + shuffle=False, + buffer_size=2, + ) + + assert sorted(_indices(streaming)) == list(range(len(map_dataset))) + + +def test_streaming_reads_video_bytes_from_configured_fsspec_root( + tmp_path: Path, lerobot_dataset_factory +) -> None: + root = tmp_path / "metadata" + map_dataset = lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=2, + total_frames=10, + ) + namespace = f"streaming-video-{tmp_path.name}" + remote_root = f"memory://{namespace}" + filesystem = fsspec.filesystem("memory") + for path in [*(root / "data").glob("*/*.parquet"), *(root / "videos").glob("*/*/*.mp4")]: + relative = path.relative_to(root).as_posix() + filesystem.put(str(path), f"{namespace}/{relative}") + + streaming = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + data_root=remote_root, + shuffle=False, + buffer_size=2, + ) + + sample = next(iter(streaming)) + reference = map_dataset[int(sample["index"])] + for camera_key in map_dataset.meta.camera_keys: + assert torch.equal(sample[camera_key], reference[camera_key]) + + +def test_streaming_rank_shards_are_disjoint(tmp_path: Path, lerobot_dataset_factory, monkeypatch) -> None: + root = tmp_path / "dataset" + map_dataset = lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=8, + total_frames=80, + use_videos=False, + ) + per_rank = [] + for rank in range(2): + monkeypatch.setenv("RANK", str(rank)) + monkeypatch.setenv("WORLD_SIZE", "2") + per_rank.append( + set( + _indices( + StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + shuffle=False, + buffer_size=2, + ) + ) + ) + ) + + assert per_rank[0].isdisjoint(per_rank[1]) + 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: + root = tmp_path / "dataset" + map_dataset = lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=8, + total_frames=80, + use_videos=False, + ) + streaming = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + shuffle=False, + buffer_size=2, + ) + 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))) + + +def test_streaming_persistent_workers_advance_epochs(tmp_path: Path, lerobot_dataset_factory) -> None: + root = tmp_path / "dataset" + map_dataset = lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=8, + total_frames=80, + use_videos=False, + ) + streaming = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + seed=23, + shuffle=True, + buffer_size=2, + ) + loader = torch.utils.data.DataLoader( + streaming, + batch_size=None, + num_workers=2, + persistent_workers=True, + ) + try: + first = [int(item["index"]) for item in loader] + second = [int(item["index"]) for item in loader] + finally: + if loader._iterator is not None: + loader._iterator._shutdown_workers() + + assert sorted(first) == list(range(len(map_dataset))) + assert sorted(second) == list(range(len(map_dataset))) + assert first != second + + +def test_streaming_worker_exception_propagates_and_workers_stop( + tmp_path: Path, lerobot_dataset_factory +) -> None: + root = tmp_path / "dataset" + lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=4, + total_frames=40, + use_videos=False, + ) + streaming = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + shuffle=False, + buffer_size=2, + ) + next((root / "data").glob("*/*.parquet")).write_bytes(b"corrupt parquet") + loader = torch.utils.data.DataLoader( + streaming, + batch_size=None, + num_workers=2, + persistent_workers=True, + ) + try: + with pytest.raises(Exception, match="Parquet"): + list(loader) + finally: + if loader._iterator is not None: + loader._iterator._shutdown_workers() + assert not any(worker.is_alive() for worker in loader._iterator._workers) + + +def test_streaming_resume_reproduces_remaining_stream(tmp_path: Path, lerobot_dataset_factory) -> None: + root = tmp_path / "dataset" + lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=5, + total_frames=50, + use_videos=False, + ) + full = _indices( + StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + seed=7, + shuffle=True, + buffer_size=3, + ) + ) + resumed = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + seed=7, + shuffle=True, + buffer_size=3, + ) + resumed.load_state_dict({"epoch": 0, "offset": 11}) + + assert _indices(resumed) == full[11:] + + +@pytest.mark.parametrize(("batch_size", "offset"), [(None, 17), (4, 20)]) +def test_streaming_worker_resume_reproduces_remaining_stream( + tmp_path: Path, + lerobot_dataset_factory, + batch_size: int | None, + offset: int, +) -> None: + root = tmp_path / "dataset" + lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=8, + total_frames=80, + use_videos=False, + ) + + def load(dataset: StreamingLeRobotDataset) -> list[int]: + loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=2) + 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"]] + + full = load( + StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + seed=31, + shuffle=True, + buffer_size=2, + ) + ) + resumed = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + seed=31, + shuffle=True, + buffer_size=2, + ) + resumed.load_state_dict({"epoch": 0, "offset": offset, "batch_size": batch_size or 1}) + + assert load(resumed) == full[offset:] + + +def test_streaming_state_dict_round_trip_mid_epoch(tmp_path: Path, lerobot_dataset_factory) -> None: + root = tmp_path / "dataset" + lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=5, + total_frames=50, + use_videos=False, + ) + source = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + seed=17, + shuffle=True, + buffer_size=3, + ) + iterator = iter(source) + consumed = [int(next(iterator)["index"]) for _ in range(13)] + state = source.state_dict() + remaining = [int(item["index"]) for item in iterator] + + restored = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + seed=17, + shuffle=True, + buffer_size=3, + ) + restored.load_state_dict(state) + + assert len(consumed) == state["offset"] + assert _indices(restored) == remaining + + +def test_streaming_worker_resume_after_epoch_boundary(tmp_path: Path, lerobot_dataset_factory) -> None: + root = tmp_path / "dataset" + lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=4, + total_frames=24, + use_videos=False, + ) + + def infinite_indices(dataset: StreamingLeRobotDataset, count: int) -> list[int]: + loader = torch.utils.data.DataLoader( + dataset, + batch_size=4, + num_workers=2, + persistent_workers=True, + ) + try: + return [ + int(index) for batch in islice(cycle(loader), (count + 3) // 4) for index in batch["index"] + ][:count] + finally: + if loader._iterator is not None: + loader._iterator._shutdown_workers() + + full = infinite_indices( + StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + seed=47, + shuffle=True, + buffer_size=2, + repeat=True, + ), + 56, + ) + offset = 32 + resumed = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + seed=47, + shuffle=True, + buffer_size=2, + repeat=True, + ) + resumed.load_state_dict({"epoch": 0, "offset": offset, "batch_size": 4}) + + assert infinite_indices(resumed, 24) == full[offset : offset + 24] + + +def test_streaming_local_training_step_smoke(tmp_path: Path, lerobot_dataset_factory) -> None: + root = tmp_path / "dataset" + lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=4, + total_frames=24, + use_videos=False, + ) + dataset = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + seed=53, + buffer_size=2, + repeat=True, + ) + loader = torch.utils.data.DataLoader(dataset, batch_size=4, num_workers=2) + iterator = iter(loader) + try: + batch = next(iterator) + finally: + iterator._shutdown_workers() + model = torch.nn.Linear(batch["action"].shape[-1], batch["action"].shape[-1]) + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) + + loss = torch.nn.functional.mse_loss(model(batch["action"]), batch["action"]) + loss.backward() + optimizer.step() + + assert torch.isfinite(loss) + assert batch["index"].shape == (4,) diff --git a/tests/datasets/test_streaming_sidecar_adapter.py b/tests/datasets/test_streaming_sidecar_adapter.py new file mode 100644 index 000000000..a429f5d54 --- /dev/null +++ b/tests/datasets/test_streaming_sidecar_adapter.py @@ -0,0 +1,47 @@ +#!/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 pathlib import Path +from types import SimpleNamespace + +import pytest + +pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") + +from lerobot.datasets.streaming_sidecar import range_backend_for_root, streaming_data_root + + +def test_hub_data_root_is_revision_qualified() -> None: + meta = SimpleNamespace(repo_id="owner/dataset", revision="commit-sha") + + root = streaming_data_root(meta, requested_root=None, configured_data_root=None) + + assert root == "hf://datasets/owner/dataset@commit-sha" + assert range_backend_for_root(root) == "native-http" + + +def test_explicit_bucket_root_is_preserved() -> None: + meta = SimpleNamespace(repo_id="owner/dataset", revision="commit-sha") + bucket = "hf://buckets/owner/dataset-bucket/prefix/" + + root = streaming_data_root(meta, requested_root=None, configured_data_root=bucket) + + assert root == bucket.rstrip("/") + assert range_backend_for_root(root) == "native-http" + + +def test_local_and_generic_remote_roots_use_fsspec(tmp_path: Path) -> None: + meta = SimpleNamespace(repo_id="owner/dataset", revision="commit-sha") + + local = streaming_data_root(meta, requested_root=tmp_path, configured_data_root=None) + + assert local == str(tmp_path) + assert range_backend_for_root(local) == "fsspec" + assert range_backend_for_root("memory://dataset") == "fsspec" diff --git a/tests/test_streaming_core_imports.py b/tests/test_streaming_core_imports.py new file mode 100644 index 000000000..e12827797 --- /dev/null +++ b/tests/test_streaming_core_imports.py @@ -0,0 +1,39 @@ +#!/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 subprocess +import sys + + +def test_streaming_core_imports_without_dataset_extra() -> None: + code = """ +import importlib.abc +import sys + +class BlockDatasets(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname == "datasets" or fullname.startswith("datasets."): + raise ModuleNotFoundError("blocked optional datasets dependency") + return None + +sys.meta_path.insert(0, BlockDatasets()) +from lerobot.streaming.episode_video import EpisodeByteCache, ExactCoveragePool +from lerobot.streaming.mp4 import Mp4Index +""" + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + timeout=30, + ) + + assert result.returncode == 0, result.stderr diff --git a/tests/test_streaming_sidecar.py b/tests/test_streaming_sidecar.py new file mode 100644 index 000000000..fdc0b63ba --- /dev/null +++ b/tests/test_streaming_sidecar.py @@ -0,0 +1,188 @@ +#!/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 shutil +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import numpy as np +import pytest +from filelock import FileLock + +from lerobot.streaming.episode_video import EpisodeVideoManifest, VideoFileRecord +from lerobot.streaming.mp4 import Mp4Index +from lerobot.streaming.sidecar import ( + SidecarLockTimeoutError, + SidecarSpec, + ensure_mp4_sidecar, + sidecar_cache_path, +) + + +def _record(path: str = "videos/camera/chunk-000/file-000.mp4", size: int = 128) -> VideoFileRecord: + arrays = np.array([0], dtype=np.int64) + index = Mp4Index( + file_path=path, + file_size=size, + ftyp=b"", + moov_offset=0, + mdat_offset=0, + mdat_payload_offset=0, + mdat_payload_size=size, + faststart=True, + codec="avc1", + timescale=1, + duration=1, + track_id=1, + width=1, + height=1, + stsd_body=b"", + sample_pts=np.array([0.0]), + sample_durations=arrays, + sample_sizes=arrays, + sample_offsets=arrays, + sync_samples=arrays, + ) + return VideoFileRecord(path, size, index) + + +def _spec(revision: str = "rev-a", size: int = 128) -> SidecarSpec: + return SidecarSpec( + repo_id="owner/dataset", + revision=revision, + data_root="hf://datasets/owner/dataset", + source_files=(("videos/camera/chunk-000/file-000.mp4", size),), + ) + + +def _write_valid(path: Path, spec: SidecarSpec) -> None: + EpisodeVideoManifest.save_file_sidecar(path, [_record(size=spec.source_files[0][1])], spec=spec) + + +def test_sidecar_cache_path_is_revision_keyed(tmp_path: Path) -> None: + first = sidecar_cache_path(tmp_path, _spec("rev-a")) + second = sidecar_cache_path(tmp_path, _spec("rev-b")) + + assert first != second + assert first.parent == second.parent + + +def test_ensure_reuses_valid_local_sidecar(tmp_path: Path) -> None: + spec = _spec() + path = sidecar_cache_path(tmp_path, spec) + _write_valid(path, spec) + build_calls = 0 + + def build(_path: Path, _spec: SidecarSpec) -> None: + nonlocal build_calls + build_calls += 1 + + resolved = ensure_mp4_sidecar(spec, tmp_path, build=build) + + assert resolved == path + assert build_calls == 0 + + +def test_ensure_downloads_valid_published_sidecar(tmp_path: Path) -> None: + spec = _spec() + published = tmp_path / "published.npz" + _write_valid(published, spec) + build_calls = 0 + + def download(path: Path, _spec: SidecarSpec) -> bool: + shutil.copyfile(published, path) + return True + + def build(_path: Path, _spec: SidecarSpec) -> None: + nonlocal build_calls + build_calls += 1 + + resolved = ensure_mp4_sidecar(spec, tmp_path / "cache", build=build, download=download) + + assert EpisodeVideoManifest.validate_file_sidecar(resolved, spec) + assert build_calls == 0 + + +@pytest.mark.parametrize("invalid_kind", ["corrupt", "stale"]) +def test_ensure_rebuilds_invalid_local_sidecar(tmp_path: Path, invalid_kind: str) -> None: + spec = _spec() + path = sidecar_cache_path(tmp_path, spec) + path.parent.mkdir(parents=True, exist_ok=True) + if invalid_kind == "corrupt": + path.write_bytes(b"not-an-npz") + else: + _write_valid(path, _spec(revision="other-revision")) + build_calls = 0 + + def build(target: Path, target_spec: SidecarSpec) -> None: + nonlocal build_calls + build_calls += 1 + _write_valid(target, target_spec) + + resolved = ensure_mp4_sidecar(spec, tmp_path, build=build) + + assert EpisodeVideoManifest.validate_file_sidecar(resolved, spec) + assert build_calls == 1 + + +def test_concurrent_ensure_builds_once(tmp_path: Path) -> None: + spec = _spec() + start = threading.Barrier(2) + build_calls = 0 + build_lock = threading.Lock() + + def build(path: Path, target_spec: SidecarSpec) -> None: + nonlocal build_calls + with build_lock: + build_calls += 1 + _write_valid(path, target_spec) + + def ensure() -> Path: + start.wait() + return ensure_mp4_sidecar(spec, tmp_path, build=build) + + with ThreadPoolExecutor(max_workers=2) as pool: + paths = list(pool.map(lambda _: ensure(), range(2))) + + assert paths[0] == paths[1] + assert build_calls == 1 + + +def test_failed_build_does_not_replace_existing_file(tmp_path: Path) -> None: + spec = _spec() + path = sidecar_cache_path(tmp_path, spec) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"old-corrupt-file") + + def build(target: Path, _spec: SidecarSpec) -> None: + target.write_bytes(b"partial") + raise RuntimeError("build failed") + + with pytest.raises(RuntimeError, match="build failed"): + ensure_mp4_sidecar(spec, tmp_path, build=build) + + assert path.read_bytes() == b"old-corrupt-file" + assert not list(path.parent.glob(f".{path.name}.*.tmp.npz")) + + +def test_lock_timeout_is_actionable(tmp_path: Path) -> None: + spec = _spec() + path = sidecar_cache_path(tmp_path, spec) + lock_path = path.with_suffix(f"{path.suffix}.lock") + lock_path.parent.mkdir(parents=True, exist_ok=True) + + with ( + FileLock(lock_path), + pytest.raises(SidecarLockTimeoutError, match="Timed out waiting"), + ): + ensure_mp4_sidecar(spec, tmp_path, build=_write_valid, lock_timeout_s=0.01) diff --git a/tests/utils/test_train_utils.py b/tests/utils/test_train_utils.py index ccd769bd0..be6e7f029 100644 --- a/tests/utils/test_train_utils.py +++ b/tests/utils/test_train_utils.py @@ -24,6 +24,7 @@ from lerobot.common.train_utils import ( get_step_identifier, load_training_batch_size, load_training_num_processes, + load_training_num_workers, load_training_state, load_training_step, push_checkpoint_to_hub, @@ -90,6 +91,16 @@ def test_load_training_batch_size_absent_returns_none(tmp_path, optimizer, sched assert load_training_batch_size(tmp_path) is None +def test_save_training_state_records_num_workers(tmp_path, optimizer, scheduler): + save_training_state(tmp_path, 10, optimizer, scheduler, num_workers=6) + assert load_training_num_workers(tmp_path) == 6 + + +def test_load_training_num_workers_absent_returns_none(tmp_path, optimizer, scheduler): + save_training_state(tmp_path, 10, optimizer, scheduler) + assert load_training_num_workers(tmp_path) is None + + def test_update_last_checkpoint(tmp_path): checkpoint = tmp_path / "0005" checkpoint.mkdir() diff --git a/uv.lock b/uv.lock index 7d873d8f5..5390e7a19 100644 --- a/uv.lock +++ b/uv.lock @@ -2829,7 +2829,10 @@ dependencies = [ { name = "cmake" }, { name = "draccus" }, { name = "einops" }, + { name = "filelock" }, + { name = "fsspec" }, { name = "gymnasium" }, + { name = "httpx" }, { name = "huggingface-hub" }, { name = "numpy" }, { name = "opencv-python-headless" }, @@ -3302,7 +3305,9 @@ requires-dist = [ { name = "faker", marker = "extra == 'sarm'", specifier = ">=33.0.0,<35.0.0" }, { name = "fastapi", marker = "extra == 'phone'", specifier = "<1.0" }, { name = "feetech-servo-sdk", marker = "extra == 'feetech'", specifier = ">=1.0.0,<2.0.0" }, + { name = "filelock", specifier = ">=3.12.0,<4.0.0" }, { name = "foxglove-sdk", marker = "extra == 'viz'", specifier = ">=0.25.1,<0.26.0" }, + { name = "fsspec", specifier = ">=2023.5.0,<2027.0.0" }, { name = "grpcio", marker = "extra == 'grpcio-dep'", specifier = ">=1.73.1,<2.0.0" }, { name = "grpcio", marker = "extra == 'reachy2'", specifier = "<=1.73.1" }, { name = "grpcio-tools", marker = "extra == 'dev'", specifier = ">=1.73.1,<2.0.0" }, @@ -3313,6 +3318,7 @@ requires-dist = [ { name = "hebi-py", marker = "extra == 'phone'", specifier = ">=2.8.0,<2.12.0" }, { name = "hf-libero", marker = "sys_platform == 'linux' and extra == 'libero'", specifier = ">=0.1.4,<0.2.0" }, { name = "hidapi", marker = "extra == 'gamepad'", specifier = ">=0.14.0,<0.15.0" }, + { name = "httpx", specifier = ">=0.27.0,<1.0.0" }, { name = "huggingface-hub", specifier = ">=1.0.0,<2.0.0" }, { name = "ipykernel", marker = "extra == 'notebook'", specifier = ">=6.0.0,<7.0.0" }, { name = "jsonlines", marker = "extra == 'dataset'", specifier = ">=4.0.0,<5.0.0" },