admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-27 11:46:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(dataset): integrate episode streaming into training

Browse Source
This commit is contained in:
Pepijn
2026-07-23 16:10:39 +02:00
parent b85620657f
commit 3d70b21aac
34 changed files with 2771 additions and 942 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/_toctree.yml
+2
View File
@@ -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
docs/source/lerobot-dataset-v3.mdx
+56 -3
View File
@@ -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)
```
<div style="display:flex; justify-content:center; gap:12px; flex-wrap:wrap;">
docs/source/training_dataset_streaming.mdx
+127
View File
@@ -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.
examples/training/train_with_streaming.py
+18 -6
View File
@@ -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.
pyproject.toml
+3
View File
@@ -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
scripts/bench_episode_byte_cache.py
+18 -101
View File
@@ -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()
scripts/bench_streaming_dataset.py
+179
View File
@@ -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()
scripts/build_mp4_sidecar.py
+38 -36
View File
@@ -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__":
src/lerobot/common/train_utils.py
+26 -2
View File
@@ -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)
src/lerobot/configs/default.py
+14
View File
@@ -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(
src/lerobot/datasets/episode_parquet.py
+106
View File
@@ -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"
)
src/lerobot/datasets/factory.py
+35 -3
View File
@@ -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,6 +192,25 @@ def make_train_eval_datasets(
ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None
)
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,
@@ -190,6 +220,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,
)
@@ -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,
)
src/lerobot/datasets/streaming_dataset.py
+402 -560
View File
File diff suppressed because it is too large Load Diff
src/lerobot/datasets/streaming_sidecar.py
+141
View File
@@ -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,
)
src/lerobot/datasets/video_utils.py
+8 -7
View File
@@ -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 "<in-memory video>"
# 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"
)
src/lerobot/scripts/lerobot_train.py
+75 -4
View File
@@ -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 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,
)
src/lerobot/streaming/__init__.py
+9
View File
@@ -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."""
src/lerobot/datasets/episode_video_streaming.py → src/lerobot/streaming/episode_video.py
+155 -63
View File
@@ -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.streaming.mp4 import Mp4Index, Mp4SampleSlice, fetch_mp4_index, synthesize_mp4
if TYPE_CHECKING:
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets.mp4 import Mp4Index, Mp4SampleSlice, fetch_mp4_index, synthesize_mp4
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:
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"]))
entry["decoder"] = decoder
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"])
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
src/lerobot/datasets/mp4.py → src/lerobot/streaming/mp4.py
+2
View File
@@ -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
src/lerobot/streaming/sidecar.py
+176
View File
@@ -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
tests/configs/test_default.py
+13
View File
@@ -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})
tests/datasets/test_depth.py
+5 -6
View File
@@ -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
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_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)
tests/datasets/test_episode_parquet_reader.py
+110
View File
@@ -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]
tests/datasets/test_episode_video_streaming.py
+117 -28
View File
@@ -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()
tests/datasets/test_exact_coverage_pool.py
+1 -1
View File
@@ -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)
tests/datasets/test_language.py
+14
View File
@@ -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] == []
tests/datasets/test_streaming.py
+15 -84
View File
@@ -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()), (
tests/datasets/test_streaming_factory.py
+58
View File
@@ -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
tests/datasets/test_streaming_production.py
+519
View File
@@ -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,)
tests/datasets/test_streaming_sidecar_adapter.py
+47
View File
@@ -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"
tests/test_streaming_core_imports.py
+39
View File
@@ -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
tests/test_streaming_sidecar.py
+188
View File
@@ -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)
tests/utils/test_train_utils.py
+11
View File
@@ -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()
uv.lock
Generated
+6
View File
@@ -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" },