Compare commits

...

44 Commits

Author SHA1 Message Date
Pepijn 21321d7a27 Merge remote-tracking branch 'origin/main' into codex/episode-video-streaming-byte-cache
# Conflicts:
#	src/lerobot/scripts/lerobot_train.py
2026-07-28 11:18:45 +02:00
Pepijn fdf6014214 Remove streaming benchmark prototypes 2026-07-28 11:14:10 +02:00
Pepijn 4db03ed535 Derive streaming decoder cap from episode pool 2026-07-27 20:49:19 +02:00
Pepijn 79cfb52a71 Harden production dataset streaming pipeline 2026-07-27 16:07:48 +02:00
Pepijn c0e9b0bbff Preserve private Hub auth in streaming readers 2026-07-27 09:41:57 +02:00
Pepijn b5a13e43ce Merge origin/main into streaming byte-cache branch 2026-07-27 09:37:33 +02:00
Pepijn 1e7e0b6de5 Fix streaming Parquet imports and video fallback 2026-07-24 16:41:54 +02:00
Pepijn a3edab661b chore(streaming): remove obsolete benchmark flags 2026-07-24 13:36:17 +02:00
Pepijn 564ba6395a Merge remote-tracking branch 'origin/main' into codex/episode-video-streaming-byte-cache 2026-07-24 13:26:42 +02:00
Pepijn 85086fed7a refactor(streaming): split episode data plane 2026-07-24 13:26:36 +02:00
Pepijn ee06d1005f refactor(dataset): make streaming pool rank-owned 2026-07-23 17:12:04 +02:00
Pepijn 1c47809bf6 chore: Merge origin/main into streaming branch 2026-07-23 16:10:46 +02:00
Pepijn 3d70b21aac feat(dataset): integrate episode streaming into training 2026-07-23 16:10:39 +02:00
Pepijn b85620657f chore: Merge origin/main into streaming branch 2026-07-23 14:27:30 +02:00
Pepijn fbfc861cf2 refactor(streaming): exact coverage is the only pool mode
Drop the with-replacement sampled path: delete run_pool_stream_simulation
and the --coverage flag; the streaming keep-up sim always uses
run_exact_coverage_stream (ExactCoveragePool), so every frame of every
episode is decoded exactly once per epoch. --pool-samples-per-episode is
kept as a deprecated no-op so existing commands still parse (exact mode
evicts an episode only when all its frames are emitted, so a turnover
cadence no longer applies).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-03 15:07:11 +02:00
Pepijn 06aa6a0425 feat(streaming): exact-once epoch coverage for the byte-cache episode pool
The pool path sampled frames with replacement and never guaranteed a full
epoch (episodes rotated on a fixed cadence; frames drawn randomly, none
tracked). Add ExactCoveragePool: a deterministic planner that enumerates
every frame of every episode exactly once per epoch while keeping at most
pool_size episodes resident, so batch mixing stays high (uniform draw over
all remaining frames in the pool) but coverage is complete and reproducible.

Mechanics (the "evict only when all frames sampled" model): episodes are
admitted in a seeded global permutation; each resident episode carries a
seeded frame-index shuffle; each draw picks a resident episode with
probability proportional to its remaining frames and pops one; an episode
is evicted only when its last frame is emitted, then a new one is admitted;
the epoch ends when admission is exhausted and every resident episode drains.
Order is a pure function of (seed, epoch) -> resumable by deterministic
fast-forward. The planner does no I/O and exposes admission_order so callers
can prefetch episodes ahead of the sampling frontier.

Wired into the benchmark as --coverage {sampled,exact}: run_exact_coverage_stream
prefetches stream_prefetch_episodes beyond the frontier so a freshly admitted
episode's bytes are resident before it is drawn, then decodes each frame once,
paced to target.

Tests: 7 planner unit tests (exact-once coverage incl. a 45k-frame epoch,
pool-size bound, per-(seed,epoch) determinism with coverage preserved,
admission/eviction events, coupon-collector mixing, zero-length episodes) and
a mocked-cache structural test of run_exact_coverage_stream asserting
every-frame-once-per-camera plus the prefetch-before-decode invariant.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-03 15:01:35 +02:00
Pepijn be64ded80f perf(streaming): sub-range parallel fetch + non-blocking pool replacement
The 64-vs-128-worker benchmark pair proved a per-host throughput ceiling
(~270 MiB/s) on the HF bucket path: doubling connections exactly halved
per-connection speed (4.8 -> 2.2 MiB/s) and left the aggregate flat,
while per-episode latency doubled (5.7s -> 12s) and keep-up worsened.
Steady-state demand (148 MiB/s) is well below the ceiling; the keep-up
misses come entirely from consumer stalls (refill_wait 14-19s of ~84s):
the sim blocks the training hot path on ensure_ready() for the FIFO-head
replacement while episodes take 5.7-12s to arrive.

Two fixes:

- Non-blocking replacements: EpisodeByteCache.is_ready() (all cameras
  cached or futures done, no blocking) and the stream sim now swaps a
  replacement only when it is already resident, deferring otherwise;
  fetch capacity (~2x demand) repays the debt on later batches. A
  deferred_swaps metric is reported.
- Sub-range parallel fetch (native-http): --range-subranges N splits one
  camera GET into N concurrent sub-range GETs. Under a per-host ceiling
  this adds no bandwidth but divides per-episode latency by ~N. Keep
  workers x subranges near the ~64-connection saturation point (e.g.
  --workers 16 --range-subranges 4).

Verified: sub-range span math + order-preserving concat and is_ready
semantics (unit-level, network stubbed).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-03 14:05:46 +02:00
Pepijn 88843ed675 Fix episode pool stream benchmark pacing 2026-06-22 20:40:33 +02:00
Pepijn f2b5c4a47b Add distributed episode pool benchmark summaries 2026-06-22 17:08:02 +02:00
Pepijn 9202fcea96 Fix pool sampling camera timestamps 2026-06-22 16:44:37 +02:00
Pepijn ef47c35178 Benchmark random sampling from episode pool 2026-06-22 16:26:26 +02:00
Pepijn 6d6c82eb8c Add GOP window range benchmark 2026-06-22 15:10:21 +02:00
Pepijn 9201be92cb Log failed HTTP range attempts 2026-06-22 12:42:30 +02:00
Pepijn 0064a06205 Instrument HfFileSystem range requests 2026-06-22 12:07:59 +02:00
Pepijn 710171ccac Clarify native HTTP exception timing 2026-06-22 12:02:12 +02:00
Pepijn 0f8257443c Retry native HTTP timeout statuses 2026-06-22 11:35:28 +02:00
Pepijn 3a09d0c48a Track native HTTP failed attempt timing 2026-06-22 11:27:08 +02:00
Pepijn 03fc5e3ea9 Allow dynamic range timing counters 2026-06-19 09:43:57 +02:00
Pepijn 28c3e095bf Report native HTTP chunk timing 2026-06-19 09:34:22 +02:00
Pepijn 5bfb749a9b Log episode cache fill progress 2026-06-18 18:43:43 +02:00
Pepijn 51c023a7a1 Tune native HTTP range diagnostics 2026-06-17 21:50:05 +02:00
Pepijn 51ea18cb7a Allow native HTTP sidecar range diagnostics 2026-06-17 21:36:57 +02:00
Pepijn 04ab43b8d2 Report range read timing breakdown 2026-06-17 21:20:08 +02:00
Pepijn cdfe192491 Remove random frame benchmark path 2026-06-17 21:14:42 +02:00
Pepijn 3451e53452 Use HfFileSystem for sidecar episode benchmark 2026-06-17 21:01:43 +02:00
Pepijn 30849ce74f Report memory usage in cache benchmarks 2026-06-17 20:54:12 +02:00
Pepijn 7d6907c444 Add random frame range fetch benchmark 2026-06-17 20:48:46 +02:00
Pepijn d99e1fe89d Report episode cache fill stage timings 2026-06-17 20:29:57 +02:00
Pepijn 7fcde61b69 Report full dataset estimate in episode cache benchmark 2026-06-17 20:25:21 +02:00
Pepijn bdfe8f8ce9 Use full MP4 sidecar for episode cache benchmark 2026-06-17 20:22:04 +02:00
Pepijn 34d0495d03 Retry transient native HTTP range failures 2026-06-17 20:19:54 +02:00
Pepijn 834c282631 Make episode cache benchmark fetch-only by default 2026-06-17 20:16:30 +02:00
Pepijn f132885cbc Pin Hub range cache and datasets main sources 2026-06-17 19:46:41 +02:00
Pepijn d0686be2f5 Add episode video streaming byte cache 2026-06-17 19:31:02 +02:00
35 changed files with 5666 additions and 721 deletions
+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
+42 -3
View File
@@ -131,15 +131,54 @@ 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.
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;">
+121
View File
@@ -0,0 +1,121 @@
# 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 owns a deterministic, frame-balanced set of complete episodes. One logical
exact-coverage pool per rank mixes those episodes while visiting every selected frame once per
rank-local coverage epoch. Parquet columns and compressed MP4 byte ranges are prefetched from the
same byte-aware admission frontier. Temporal history and future windows are resolved inside the
complete episode, including the same boundary padding masks as map-style loading.
When `--num_workers` is nonzero, training uses one dedicated DataLoader process per rank. Its
bounded result queue holds decoded batches while the policy trains. The configured worker count is
instead used as internal Parquet and byte-range fetch concurrency, so increasing it does not create
independent samplers or multiply the cache budget.
The default map-style `LeRobotDataset` behavior is unchanged when `--dataset.streaming=false`.
## 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 | Maximum complete episodes mixed by each rank |
| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool |
| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per rank |
| `streaming_decode_threads` | 2 | Parallel sample assembly and video decode workers |
| `streaming_decoded_queue_size` | 8 | Decoded samples buffered ahead, in planner order |
| `streaming_max_open_decoders` | pool × cameras | Independent open-decoder LRU cap per rank |
| `streaming_native_http_connections` | unset | Native HTTP connection cap per rank |
| `streaming_native_http_subranges` | 1 | Concurrent subranges per native HTTP range read |
| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root |
The active episode set is capped by both episode count and the exact synthesized mini-MP4 sizes
computed from the sidecar. An episode larger than the complete rank budget fails before training
fetches its payload. The cache, decoder LRU, and decoded-batch queue remain independently bounded.
Start with a smaller pool or budget on memory-constrained hosts:
```bash
lerobot-train \
--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 \
--dataset.streaming_decode_threads=2 \
--dataset.streaming_decoded_queue_size=8 \
--num_workers=4 \
--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
batch size changes ownership or batch boundaries. For sample-exact comparisons, resume with the
same world size and batch size. Keep the same internal fetch concurrency when comparing performance.
Streaming checkpoints created by the earlier multi-worker sampler record a different ownership
topology and are rejected for sample-exact resume. Start a new run from the saved policy weights
rather than claiming that its dataset stream resumes exactly.
## 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 so every rank owns at least one selected
episode.
- **The byte budget is exceeded:** lower the episode pool, increase the per-rank byte budget, or
use an explicitly prepared payload layout with smaller episode ranges.
- **Remote reads cannot authenticate:** verify the normal Hugging Face token or fsspec credentials
are available in every worker environment. Credentials are never embedded in the sidecar.
- **Refill stalls are high:** compare p95/p99 batch wait, reduce network contention, raise prefetch
gradually, and verify that the dataset sidecar matches the exact revision.
+22 -8
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,18 +68,28 @@ 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(
dataset,
num_workers=4,
# One worker owns the rank-level pool. Internal fetch concurrency is configured through
# StreamingLeRobotDataset.max_num_shards (the lerobot-train CLI derives it from num_workers).
num_workers=1,
batch_size=16,
pin_memory=device.type != "cpu",
drop_last=True,
prefetch_factor=2, # loads batches with multiprocessing while policy trains
prefetch_factor=2, # bounded decoded-batch queue while the policy trains
persistent_workers=True,
)
# Run training loop.
+4
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
@@ -436,6 +439,7 @@ exclude_dirs = [
skips = ["B101", "B311", "B404", "B603", "B615"]
[tool.typos]
default.extend-words = { trak = "trak" }
default.extend-ignore-re = [
"(?Rm)^.*(#|//)\\s*spellchecker:disable-line$", # spellchecker:disable-line
"(?s)(#|//)\\s*spellchecker:off.*?\\n\\s*(#|//)\\s*spellchecker:on", # spellchecker:<on|off>
+95
View File
@@ -0,0 +1,95 @@
#!/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 argparse
import time
from pathlib import Path
import fsspec
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
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", 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=None)
parser.add_argument("--max-probe-mb", type=int, default=64)
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, spec: SidecarSpec) -> list[str]:
if not spec.data_root.startswith("hf://"):
raise ValueError("--push currently supports only hf:// data roots")
fs = fsspec.filesystem("hf")
remote = published_sidecar_url(spec)
fs.put(str(Path(local_path)), remote)
return [remote]
def main() -> None:
args = parse_args()
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))
)
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()
build_mp4_sidecar(
args.output,
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(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__":
main()
+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)
+32
View File
@@ -44,6 +44,22 @@ 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 rank-level exact-coverage sampler.
streaming_episode_pool_size: int = 32
# Complete episodes fetched ahead of the current admission frontier.
streaming_prefetch_episodes: int = 8
# Hard per-rank cap for synthesized episode-video bytes.
streaming_byte_budget_gb: float = 8.0
# Parallel sample assembly/decode workers and their bounded in-order result queue.
streaming_decode_threads: int = 2
streaming_decoded_queue_size: int = 8
# Independent decoder-state cap. None covers every camera in the configured episode pool.
streaming_max_open_decoders: int | None = None
# Per-rank native HTTP limits. None preserves the fetcher's worker-derived default.
streaming_native_http_connections: int | None = None
streaming_native_http_subranges: int = 1
# Fraction of episodes held out per task for offline evaluation (0.0 = disabled).
eval_split: float = 0.0
@@ -54,6 +70,22 @@ 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.streaming_decode_threads <= 0:
raise ValueError("streaming_decode_threads must be positive")
if self.streaming_decoded_queue_size <= 0:
raise ValueError("streaming_decoded_queue_size must be positive")
if self.streaming_max_open_decoders is not None and self.streaming_max_open_decoders <= 0:
raise ValueError("streaming_max_open_decoders must be positive")
if self.streaming_native_http_connections is not None and self.streaming_native_http_connections <= 0:
raise ValueError("streaming_native_http_connections must be positive")
if self.streaming_native_http_subranges <= 0:
raise ValueError("streaming_native_http_subranges must be positive")
if self.episodes is not None:
if any(ep < 0 for ep in self.episodes):
raise ValueError(
+58 -14
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,21 @@ 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,
video_backend=cfg.dataset.video_backend,
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,
decode_threads=cfg.dataset.streaming_decode_threads,
decoded_queue_size=cfg.dataset.streaming_decoded_queue_size,
max_open_decoders=cfg.dataset.streaming_max_open_decoders,
native_http_connections=cfg.dataset.streaming_native_http_connections,
native_http_subranges=cfg.dataset.streaming_native_http_subranges,
repeat=True,
)
else:
raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.")
@@ -138,7 +152,10 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
def make_train_eval_datasets(
cfg: TrainPipelineConfig,
) -> tuple[LeRobotDataset | MultiLeRobotDataset, LeRobotDataset | None]:
) -> tuple[
LeRobotDataset | StreamingLeRobotDataset | MultiLeRobotDataset,
LeRobotDataset | None,
]:
"""Create train and optional eval datasets by splitting episodes based on eval_split.
The last ceil(n_episodes * eval_split) episodes per task are held out for evaluation.
@@ -181,17 +198,43 @@ def make_train_eval_datasets(
ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None
)
train_dataset = LeRobotDataset(
cfg.dataset.repo_id,
root=cfg.dataset.root,
episodes=train_episodes,
delta_timestamps=delta_timestamps,
image_transforms=train_image_transforms,
revision=cfg.dataset.revision,
video_backend=cfg.dataset.video_backend,
return_uint8=True,
tolerance_s=cfg.tolerance_s,
)
if cfg.dataset.streaming:
train_dataset = StreamingLeRobotDataset(
cfg.dataset.repo_id,
root=cfg.dataset.root,
episodes=train_episodes,
delta_timestamps=delta_timestamps,
image_transforms=train_image_transforms,
revision=cfg.dataset.revision,
max_num_shards=max(1, cfg.num_workers),
tolerance_s=cfg.tolerance_s,
return_uint8=True,
depth_output_unit=cfg.dataset.depth_output_unit,
video_backend=cfg.dataset.video_backend,
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,
decode_threads=cfg.dataset.streaming_decode_threads,
decoded_queue_size=cfg.dataset.streaming_decoded_queue_size,
max_open_decoders=cfg.dataset.streaming_max_open_decoders,
native_http_connections=cfg.dataset.streaming_native_http_connections,
native_http_subranges=cfg.dataset.streaming_native_http_subranges,
repeat=True,
)
else:
train_dataset = LeRobotDataset(
cfg.dataset.repo_id,
root=cfg.dataset.root,
episodes=train_episodes,
delta_timestamps=delta_timestamps,
image_transforms=train_image_transforms,
revision=cfg.dataset.revision,
video_backend=cfg.dataset.video_backend,
return_uint8=True,
depth_output_unit=cfg.dataset.depth_output_unit,
tolerance_s=cfg.tolerance_s,
)
eval_dataset = LeRobotDataset(
cfg.dataset.repo_id,
@@ -202,6 +245,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,
)
File diff suppressed because it is too large Load Diff
+149
View File
@@ -0,0 +1,149 @@
# 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.manifest 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,
token: str | bool | None = None,
) -> 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,
token=token,
)
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,
token: str | bool | None = None,
) -> bool:
if Path(spec.data_root).expanduser().is_dir():
return False
source_url = published_sidecar_url(spec, cache_root)
storage_options = {"token": token} if token is not None and source_url.startswith("hf://") else {}
filesystem, source = fsspec.core.url_to_fs(source_url, **storage_options)
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,
token: str | bool | None = None,
) -> 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,
token=token,
),
download=lambda path, target_spec: download_published_sidecar(
path,
target_spec,
cache_root=cache_root,
token=token,
),
lock_timeout_s=lock_timeout_s,
)
+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"
)
+73 -6
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,
@@ -71,9 +72,13 @@ from lerobot.utils.utils import (
from .lerobot_eval import eval_policy_all
def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
def _dataloader_worker_kwargs(
cfg: TrainPipelineConfig,
*,
num_workers: int | None = None,
) -> dict[str, Any]:
"""Return worker-only DataLoader options, disabling them for single-process loading."""
workers_enabled = cfg.num_workers > 0
workers_enabled = (cfg.num_workers if num_workers is None else num_workers) > 0
return {
"prefetch_factor": cfg.prefetch_factor if workers_enabled else None,
"persistent_workers": cfg.persistent_workers and workers_enabled,
@@ -211,7 +216,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()
@@ -469,6 +474,53 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
else:
shuffle = True
sampler = None
# One dedicated DataLoader process owns the rank-level planner/cache. Its bounded result
# queue provides decoded-batch prefetch; max_num_shards controls the internal Parquet/range
# fetch executors independently.
train_num_workers = min(cfg.num_workers, 1)
if cfg.num_workers > 1 and is_main_process:
logging.info(
"Using one streaming DataLoader worker per rank; %d configured workers remain "
"available as the dataset's internal fetch concurrency.",
cfg.num_workers,
)
rank_frames = dataset.num_frames_for_rank(
accelerator.process_index,
accelerator.num_processes,
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
@@ -476,14 +528,17 @@ 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,
**_dataloader_worker_kwargs(cfg),
**_dataloader_worker_kwargs(
cfg,
num_workers=train_num_workers if cfg.dataset.streaming else cfg.num_workers,
),
)
# Build eval dataloader if a held-out split exists
@@ -514,7 +569,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
# Prepare everything with accelerator
accelerator.wait_for_everyone()
if eval_dataloader is not None:
if cfg.dataset.streaming:
# StreamingLeRobotDataset already assigns complete episodes disjointly to ranks and workers.
# Preparing this loader would make Accelerate shard its batches a second time.
if eval_dataloader is not None:
policy, optimizer, lr_scheduler, eval_dataloader = accelerator.prepare(
policy, optimizer, lr_scheduler, eval_dataloader
)
else:
policy, optimizer, lr_scheduler = accelerator.prepare(policy, optimizer, lr_scheduler)
elif eval_dataloader is not None:
policy, optimizer, dataloader, lr_scheduler, eval_dataloader = accelerator.prepare(
policy, optimizer, dataloader, lr_scheduler, eval_dataloader
)
@@ -577,6 +641,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
@@ -673,6 +739,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,
)
+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."""
+519
View File
@@ -0,0 +1,519 @@
# 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 io
import logging
import threading
import time
from collections import OrderedDict
from collections.abc import Callable
from concurrent.futures import Future, ThreadPoolExecutor
from pathlib import Path
from types import SimpleNamespace
from typing import Any
from lerobot.streaming.manifest import EpisodeVideoManifest
from lerobot.streaming.mp4 import Mp4SampleSlice, synthesize_mp4
from lerobot.streaming.range_fetch import make_range_fetcher
logger = logging.getLogger(__name__)
class EpisodeByteCache:
def __init__(
self,
manifest: EpisodeVideoManifest,
data_root: str | Path,
*,
byte_budget: int = 80 * 1024**3,
workers: int = 8,
range_backend: str = "fsspec",
native_http_connections: int | None = None,
native_http_timeout: float = 60.0,
native_http_retries: int = 4,
native_http_subranges: int = 1,
open_decoders: bool = True,
max_open_decoders: int = 64,
video_backend: str = "torchcodec",
tolerance_s: float = 1e-4,
token: str | bool | None = None,
):
if byte_budget <= 0:
raise ValueError("byte_budget must be positive")
if max_open_decoders <= 0:
raise ValueError("max_open_decoders must be positive")
if video_backend == "video_reader":
video_backend = "pyav"
if video_backend not in {"torchcodec", "pyav"}:
raise ValueError(f"Unsupported video backend: {video_backend}")
if tolerance_s <= 0:
raise ValueError("tolerance_s must be positive")
self.manifest = manifest
self.fetcher = make_range_fetcher(
data_root,
range_backend=range_backend,
workers=workers,
native_http_connections=native_http_connections,
native_http_timeout=native_http_timeout,
native_http_retries=native_http_retries,
native_http_subranges=native_http_subranges,
token=token,
)
self.byte_budget = byte_budget
self.open_decoders = open_decoders
self.max_open_decoders = max_open_decoders
self.video_backend = video_backend
self.tolerance_s = tolerance_s
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._decoder_locks: dict[tuple[int, str], threading.Lock] = {}
self._futures: dict[tuple[int, str], Future[dict[str, Any]]] = {}
self._retained_episodes: dict[int, int] = {}
self._decoder_fallback_count = 0
self._fallback_decoders: set[tuple[int, str]] = set()
self._fallback_warning_emitted = False
self._bytes = 0
self._lock = threading.Lock()
self._timing_totals = {
"lookup_s": 0.0,
"fetch_s": 0.0,
"synthesize_s": 0.0,
"store_s": 0.0,
"jobs": 0.0,
}
def close(self) -> None:
self._pool.shutdown(wait=True, cancel_futures=True)
with self._lock:
decoders = list(self._decoders.values())
self._cache.clear()
self._decoders.clear()
self._decoder_locks.clear()
self._futures.clear()
self._retained_episodes.clear()
self._fallback_decoders.clear()
self._bytes = 0
for decoder in decoders:
_close_decoder(decoder)
self.fetcher.close()
def __enter__(self) -> EpisodeByteCache:
return self
def __exit__(self, *_exc) -> None:
self.close()
def submit_prefetch(self, episode_index: int) -> None:
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)
@property
def decoder_fallback_count(self) -> int:
with self._lock:
return self._decoder_fallback_count
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).
Lets a consumer swap in replacements only when they are already resident, instead of
blocking the training hot path on a remote fetch (head-of-line stall).
"""
for camera_key in self.manifest.video_keys:
key = (episode_index, camera_key)
with self._lock:
if key in self._cache:
continue
future = self._futures.get(key)
if future is None or not future.done():
return False
return True
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) -> Any:
key = (episode_index, camera_key)
entry = self._get_entry(episode_index, camera_key)
with self._lock:
decoder = self._decoders.get(key)
if decoder is not None:
self._decoders.move_to_end(key)
return decoder
decoder = self._open_decoder(key, entry["bytes"])
with self._lock:
existing = self._decoders.get(key)
if existing is not None:
self._decoders.move_to_end(key)
_close_decoder(decoder)
return existing
self._decoders[key] = decoder
while len(self._decoders) > self.max_open_decoders:
evicted_key, evicted_decoder = self._decoders.popitem(last=False)
self._decoder_locks.pop(evicted_key, None)
self._fallback_decoders.discard(evicted_key)
_close_decoder(evicted_decoder)
return decoder
def _open_decoder(self, key: tuple[int, str], data: bytes) -> Any:
try:
if self.video_backend == "torchcodec":
return open_video_decoder(io.BytesIO(data))
return open_video_decoder(io.BytesIO(data), backend=self.video_backend)
except Exception as primary_error:
if self.video_backend != "torchcodec":
raise
try:
decoder = open_video_decoder(io.BytesIO(data), backend="pyav")
except Exception as fallback_error:
raise RuntimeError(
"Both TorchCodec and PyAV rejected synthesized episode video "
f"{key}: TorchCodec error: {primary_error}"
) from fallback_error
with self._lock:
self._decoder_fallback_count += 1
self._fallback_decoders.add(key)
should_warn = not self._fallback_warning_emitted
self._fallback_warning_emitted = True
if should_warn:
logger.warning(
"TorchCodec rejected a synthesized episode MP4; using the bounded PyAV "
"decoder fallback for affected videos. First error: %s",
primary_error,
)
else:
logger.debug("Using PyAV decoder fallback for synthesized episode video %s", key)
return decoder
def get_frames(self, episode_index: int, camera_key: str, timestamps: list[float]):
key = (episode_index, camera_key)
span = self.manifest.lookup(episode_index, camera_key)
local_ts = [ts - span.source_start_pts for ts in timestamps]
decoder, release = self._decoder_for_frames(episode_index, camera_key)
with self._lock:
uses_pyav_timestamps = self.video_backend == "pyav" and key not in self._fallback_decoders
decode_lock = self._decoder_locks.setdefault(key, threading.Lock())
try:
with decode_lock:
if uses_pyav_timestamps:
return decoder.get_frames_played_at(local_ts, tolerance_s=self.tolerance_s).data
metadata = decoder.metadata
fps = getattr(metadata, "average_fps", None)
if fps is None:
duration = max(getattr(metadata, "end_stream_seconds", 0.0), 1e-9)
fps = metadata.num_frames / duration
num_frames = getattr(metadata, "num_frames", None)
if num_frames is None:
duration = max(getattr(metadata, "end_stream_seconds", 0.0), 1e-9)
num_frames = round(duration * fps)
last_index = max(0, int(num_frames) - 1)
indices = [min(max(round(ts * fps), 0), last_index) for ts in local_ts]
return decoder.get_frames_at(indices=indices).data
finally:
if release is not None:
release()
def _decoder_for_frames(
self, episode_index: int, camera_key: str
) -> tuple[Any, Callable[[], None] | None]:
key = (episode_index, camera_key)
while True:
decoder = self.get_decoder(episode_index, camera_key)
acquire = getattr(decoder, "acquire", None)
if acquire is None:
return decoder, None
try:
acquire()
except RuntimeError:
# The decoder was evicted between lookup and lease acquisition. Remove a stale
# cached reference if it raced with close, then retry with a fresh decoder.
with self._lock:
if self._decoders.get(key) is decoder:
self._decoders.pop(key)
self._decoder_locks.pop(key, None)
self._fallback_decoders.discard(key)
continue
return decoder, decoder.release
def timing_summary(self) -> dict[str, float]:
with self._lock:
summary = dict(self._timing_totals)
summary["decoder_fallbacks"] = float(self._decoder_fallback_count)
fetcher_summary = getattr(self.fetcher, "timing_summary", None)
if fetcher_summary is not None:
summary.update(fetcher_summary())
return summary
def _submit(self, episode_index: int, camera_key: str) -> Future[dict[str, Any]]:
key = (episode_index, camera_key)
with self._lock:
if key in self._cache:
future: Future[dict[str, Any]] = Future()
future.set_result(self._cache[key])
return future
future = self._futures.get(key)
if future is None:
future = self._pool.submit(self._fetch_and_synthesize, episode_index, camera_key)
self._futures[key] = future
return future
def _get_entry(self, episode_index: int, camera_key: str) -> dict[str, Any]:
key = (episode_index, camera_key)
with self._lock:
entry = self._cache.get(key)
if entry is not None:
self._cache.move_to_end(key)
return entry
future = self._submit(episode_index, camera_key)
entry = future.result()
store_start = time.perf_counter()
with self._lock:
self._futures.pop(key, None)
existing = self._cache.get(key)
if existing is not None:
self._cache.move_to_end(key)
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"]
self._timing_totals["fetch_s"] += timings["fetch_s"]
self._timing_totals["synthesize_s"] += timings["synthesize_s"]
self._timing_totals["store_s"] += time.perf_counter() - store_start
self._timing_totals["jobs"] += 1
return entry
def _evict_locked(self) -> None:
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"])
decoder = self._decoders.pop(key, None)
self._decoder_locks.pop(key, None)
self._fallback_decoders.discard(key)
if decoder is not None:
_close_decoder(decoder)
def _fetch_and_synthesize(self, episode_index: int, camera_key: str) -> dict[str, Any]:
lookup_start = time.perf_counter()
span = self.manifest.lookup(episode_index, camera_key)
file_record = self.manifest.file_lookup(span.file_id)
sample_slice = Mp4SampleSlice(
sample_lo=span.sample_lo,
sample_hi=span.sample_hi,
byte_offset=span.mdat_offset,
byte_length=span.mdat_length,
source_start_pts=span.source_start_pts,
)
lookup_s = time.perf_counter() - lookup_start
fetch_start = time.perf_counter()
payload = self.fetcher.read_range(file_record.file_path, span.mdat_offset, span.mdat_length)
fetch_s = time.perf_counter() - fetch_start
if len(payload) != span.mdat_length:
raise OSError(
f"Short read for {file_record.file_path}: expected {span.mdat_length}, got {len(payload)}"
)
synthesize_start = time.perf_counter()
mp4_bytes = synthesize_mp4(file_record.mp4, sample_slice, payload)
synthesize_s = time.perf_counter() - synthesize_start
entry: dict[str, Any] = {
"bytes": mp4_bytes,
"_timings": {
"lookup_s": lookup_s,
"fetch_s": fetch_s,
"synthesize_s": synthesize_s,
},
}
return entry
class _PyAVVideoDecoder:
"""Small seekable PyAV adapter matching the TorchCodec calls used by the byte cache."""
def __init__(self, file_like_or_bytesio: Any):
import av
self._source = file_like_or_bytesio
self._container = av.open(file_like_or_bytesio)
self._stream = self._container.streams.video[0]
average_rate = self._stream.average_rate
if average_rate is None:
raise ValueError("PyAV video stream does not expose an average frame rate")
self._fps = float(average_rate)
duration = (
float(self._stream.duration * self._stream.time_base)
if self._stream.duration is not None
else 0.0
)
self.metadata = SimpleNamespace(
average_fps=self._fps,
num_frames=int(self._stream.frames or round(duration * self._fps)),
begin_stream_seconds=0.0,
end_stream_seconds=duration,
)
self._decode_lock = threading.Lock()
self._state_lock = threading.Lock()
self._users = 0
self._close_requested = False
self._closed = False
def acquire(self) -> None:
with self._state_lock:
if self._close_requested or self._closed:
raise RuntimeError("PyAV decoder is closing")
self._users += 1
def release(self) -> None:
with self._state_lock:
self._users -= 1
if self._users < 0:
raise RuntimeError("Unbalanced PyAV decoder release")
if self._users == 0 and self._close_requested:
self._close_resources()
def get_frames_at(self, *, indices: list[int]) -> SimpleNamespace:
if not indices:
import torch
return SimpleNamespace(data=torch.empty((0, 3, 0, 0), dtype=torch.uint8))
timestamps = [index / self._fps for index in indices]
return self._get_frames_played_at(timestamps, tolerance_s=0.5 / self._fps + 1e-6)
def get_frames_played_at(
self,
timestamps: list[float],
*,
tolerance_s: float,
) -> SimpleNamespace:
return self._get_frames_played_at(timestamps, tolerance_s=tolerance_s)
def _get_frames_played_at(
self,
timestamps: list[float],
*,
tolerance_s: float,
) -> SimpleNamespace:
import torch
first_ts = min(timestamps)
last_ts = max(timestamps)
loaded_frames: list[torch.Tensor] = []
loaded_ts: list[float] = []
with self._decode_lock:
self._container.seek(
round(first_ts / self._stream.time_base) - 1,
backward=True,
any_frame=False,
stream=self._stream,
)
for frame in self._container.decode(self._stream):
if frame.pts is None:
continue
current_ts = float(frame.pts * self._stream.time_base)
array = frame.to_ndarray(format="rgb24")
loaded_frames.append(torch.from_numpy(array).permute(2, 0, 1).contiguous())
loaded_ts.append(current_ts)
if current_ts >= last_ts:
break
if not loaded_frames:
raise ValueError(f"PyAV decoded no frames for timestamps {timestamps}")
query_ts = torch.tensor(timestamps)
loaded_ts_tensor = torch.tensor(loaded_ts)
distances = torch.cdist(query_ts[:, None], loaded_ts_tensor[:, None], p=1)
minimum, closest = distances.min(1)
if not (minimum <= tolerance_s).all():
raise ValueError(
f"PyAV frame timestamps exceed tolerance {tolerance_s}: "
f"queries={query_ts}, decoded={loaded_ts_tensor}"
)
return SimpleNamespace(data=torch.stack([loaded_frames[index] for index in closest]))
def close(self) -> None:
with self._state_lock:
self._close_requested = True
if self._users == 0:
self._close_resources()
def _close_resources(self) -> None:
if self._closed:
return
self._container.close()
close = getattr(self._source, "close", None)
if close is not None:
close()
self._closed = True
def _close_decoder(decoder: Any) -> None:
close = getattr(decoder, "close", None)
if close is not None:
try:
close()
except Exception:
logger.debug("Failed to close video decoder", exc_info=True)
def open_video_decoder(file_like_or_bytesio, frame_mappings=None, *, backend: str = "torchcodec"):
if frame_mappings is not None:
raise ValueError("Synthesized episode videos use a local timeline; pass frame_mappings=None.")
if backend == "pyav":
return _PyAVVideoDecoder(file_like_or_bytesio)
if backend != "torchcodec":
raise ValueError(f"Unsupported video backend: {backend}")
from torchcodec.decoders import VideoDecoder
return VideoDecoder(file_like_or_bytesio, seek_mode="approximate")
+140
View File
@@ -0,0 +1,140 @@
# 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
"""Pure episode-scoped Parquet reads for training-time dataset streaming."""
from __future__ import annotations
import posixpath
import threading
import time
from collections.abc import Sequence
from pathlib import Path
from typing import Any
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],
token: str | bool | None = None,
max_retries: int = 4,
retry_backoff_s: float = 0.05,
):
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")
)
data_root_str = str(data_root)
if max_retries < 0:
raise ValueError("max_retries must be non-negative")
if retry_backoff_s < 0:
raise ValueError("retry_backoff_s must be non-negative")
self._max_retries = max_retries
self._retry_backoff_s = retry_backoff_s
self._open_lock = threading.Lock() if data_root_str.startswith("hf://") else None
storage_options = {"token": token} if token is not None and data_root_str.startswith("hf://") else {}
self._filesystem, self._root_path = fsspec.core.url_to_fs(data_root_str, **storage_options)
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._open_with_retry(path) 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
def _open_with_retry(self, path: str) -> Any:
for attempt in range(self._max_retries + 1):
try:
if self._open_lock is None:
return self._filesystem.open(path, "rb")
with self._open_lock:
return self._filesystem.open(path, "rb")
except FileNotFoundError:
if attempt == self._max_retries:
raise
self._filesystem.invalidate_cache(posixpath.dirname(path))
time.sleep(self._retry_backoff_s * 2**attempt)
raise RuntimeError("unreachable")
@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"
)
+168
View File
@@ -0,0 +1,168 @@
# 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 collections.abc import Mapping, Sequence
import numpy as np
class ExactCoveragePool:
"""Deterministic, exactly-once frame coverage over a byte-cache episode pool.
A with-replacement pool never guarantees a full
epoch: frames are drawn randomly and episodes rotate on a fixed cadence. This planner instead
enumerates *every frame of every episode exactly once per epoch* while keeping at most
``pool_size`` episodes resident, so batch mixing stays high but coverage is complete and
reproducible.
Mechanics (this is the "evict only when all frames sampled" model):
- Episodes are admitted in a seeded global permutation until either ``pool_size`` or the
optional indexed-byte budget is reached.
- Each resident episode carries a seeded shuffle of its own frame indices.
- Each draw picks a resident episode with probability proportional to its *remaining* frames
(i.e. a uniform draw over all remaining frames in the pool, the map-style ideal) and pops
one frame.
- An episode is evicted only when its last frame is emitted; a new episode is then admitted.
- The epoch ends when the admission order is exhausted and every resident episode is drained.
Newly admitted episodes are surfaced via :attr:`newly_admitted` (drain it to drive prefetch)
and evictions via :attr:`evicted` (drain to release cache bytes). The planner does no I/O and
is fully unit-testable. It yields ``(episode_index, frame_index)``; map to a decode timestamp
with ``frame_index / max(frame_count - 1, 1)``.
Determinism: the order is a pure function of ``(seed, epoch)``, the episode frame counts, and
optional byte sizes/budget. Resume is a deterministic fast-forward: re-instantiate with the
same inputs and skip ``n`` samples (tabular only, no decode).
"""
def __init__(
self,
episode_frame_counts: Sequence[tuple[int, int]],
pool_size: int,
*,
seed: int,
epoch: int = 0,
episode_byte_sizes: Mapping[int, int] | None = None,
byte_budget: int | None = None,
):
self._counts = {int(ep): int(n) for ep, n in episode_frame_counts if int(n) > 0}
self._rng = np.random.default_rng([seed, epoch])
order = np.array(sorted(self._counts), dtype=np.int64)
self._rng.shuffle(order)
self.pool_size = max(1, pool_size)
self._byte_budget = byte_budget
if byte_budget is not None and byte_budget <= 0:
raise ValueError("byte_budget must be positive")
if byte_budget is not None and episode_byte_sizes is None:
raise ValueError("episode_byte_sizes are required when byte_budget is set")
self._byte_sizes = {
episode: int(episode_byte_sizes[episode]) if episode_byte_sizes is not None else 0
for episode in self._counts
}
if any(size < 0 for size in self._byte_sizes.values()):
raise ValueError("episode byte sizes must be non-negative")
if byte_budget is not None:
oversized = next(
((episode, size) for episode, size in self._byte_sizes.items() if size > byte_budget),
None,
)
if oversized is not None:
episode, size = oversized
raise ValueError(
f"Episode {episode} requires {size} bytes, exceeding the byte budget {byte_budget}"
)
# Preserve the full seeded order for benchmark/tooling compatibility. Byte-aware admission
# may temporarily skip an entry, but every episode remains in this deterministic frontier.
self.admission_order: list[int] = order.tolist()
self._pending: list[int] = list(self.admission_order)
self._admitted_count = 0
self._remaining: dict[int, tuple[np.ndarray, int]] = {}
self._remaining_total = 0
self._resident_bytes = 0
self.newly_admitted: list[int] = []
self.evicted: list[int] = []
self._admit_available()
def _admit_available(self) -> None:
while len(self._remaining) < self.pool_size and self._pending:
available_bytes = None if self._byte_budget is None else self._byte_budget - self._resident_bytes
pending_index = next(
(
index
for index, episode in enumerate(self._pending)
if available_bytes is None or self._byte_sizes[episode] <= available_bytes
),
None,
)
if pending_index is None:
return
episode = self._pending.pop(pending_index)
frame_count = self._counts[episode]
frames = np.arange(frame_count, dtype=np.int64)
self._rng.shuffle(frames)
self._remaining[episode] = (frames, frame_count)
self._remaining_total += frame_count
self._resident_bytes += self._byte_sizes[episode]
self._admitted_count += 1
self.newly_admitted.append(episode)
@property
def remaining_total(self) -> int:
return self._remaining_total
@property
def admitted_count(self) -> int:
"""Number of episodes pulled from the admission order so far (pool fills + rotations)."""
return self._admitted_count
@property
def resident(self) -> list[int]:
return list(self._remaining)
@property
def resident_bytes(self) -> int:
return self._resident_bytes
def prefetch_candidates(self, count: int) -> list[int]:
"""Return the next deterministic pending frontier without admitting it."""
if count <= 0:
return []
return self._pending[:count]
def __iter__(self) -> ExactCoveragePool:
return self
def __next__(self) -> tuple[int, int]:
if self._remaining_total == 0:
raise StopIteration
# Uniform draw over all remaining frames in the pool: walk the residents by cumulative
# remaining count. O(pool_size) per draw (~1024) -> negligible next to decode.
target = int(self._rng.integers(self._remaining_total))
chosen = None
for ep, (_frames, remaining) in self._remaining.items():
if target < remaining:
chosen = ep
break
target -= remaining
frames, remaining = self._remaining[chosen]
remaining -= 1
frame_index = int(frames[remaining])
self._remaining_total -= 1
if remaining == 0:
del self._remaining[chosen]
self.evicted.append(chosen)
self._resident_bytes -= self._byte_sizes[chosen]
self._admit_available()
else:
self._remaining[chosen] = (frames, remaining)
return chosen, frame_index
+351
View File
@@ -0,0 +1,351 @@
# 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 json
import logging
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any
import numpy as np
from lerobot.streaming.mp4 import (
Mp4Index,
Mp4SampleSlice,
fetch_mp4_index,
synthesized_mp4_size,
)
from lerobot.streaming.range_fetch import make_range_fetcher
if TYPE_CHECKING:
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.streaming.sidecar import SidecarSpec
@dataclass(frozen=True)
class EpisodeVideoSpan:
file_id: int
mdat_offset: int
mdat_length: int
first_pts: float
last_pts: float
frame_count: int
sample_lo: int
sample_hi: int
source_start_pts: float
@dataclass(frozen=True)
class VideoFileRecord:
file_path: str
file_size: int
mp4: Mp4Index
class EpisodeVideoManifest:
_FILE_SIDECAR_CACHE: dict[str, tuple[tuple[int, int], dict[str, VideoFileRecord]]] = {}
_FILE_SIDECAR_CACHE_LOCK = threading.Lock()
def __init__(
self,
*,
video_keys: list[str],
files: list[VideoFileRecord],
spans: dict[str, np.ndarray],
):
self.video_keys = list(video_keys)
self._camera_to_id = {key: idx for idx, key in enumerate(self.video_keys)}
self.files = files
self.spans = spans
@classmethod
def build(
cls,
meta: LeRobotDatasetMetadata,
data_root: str | Path,
*,
episode_indices: list[int] | range | None = None,
range_backend: str = "fsspec",
workers: int = 8,
header_probe_bytes: int = 4 * 1024 * 1024,
max_probe_bytes: int = 64 * 1024 * 1024,
keyframe_pad_s: float = 0.1,
keyframe_pad_fraction: float = 0.05,
sidecar_path: str | Path | None = None,
token: str | bool | None = None,
) -> EpisodeVideoManifest:
meta.ensure_readable()
video_keys = list(meta.video_keys)
if episode_indices is None:
episode_indices = range(int(meta.total_episodes))
rel_paths = sorted(
{str(meta.get_video_file_path(ep_idx, key)) for ep_idx in episode_indices for key in video_keys}
)
path_to_id = {path: idx for idx, path in enumerate(rel_paths)}
if sidecar_path is None:
files = cls._build_file_records(
rel_paths,
data_root,
range_backend=range_backend,
workers=workers,
header_probe_bytes=header_probe_bytes,
max_probe_bytes=max_probe_bytes,
token=token,
)
else:
records = cls.load_file_sidecar(sidecar_path)
missing = [path for path in rel_paths if path not in records]
if missing:
raise ValueError(
f"Sidecar {sidecar_path} is missing {len(missing)} files, first: {missing[0]}"
)
files = [records[path] for path in rel_paths]
total = int(meta.total_episodes)
num_cameras = len(video_keys)
spans: dict[str, np.ndarray] = {
"file_id": np.zeros((total, num_cameras), dtype=np.int32),
"mdat_offset": np.zeros((total, num_cameras), dtype=np.int64),
"mdat_length": np.zeros((total, num_cameras), dtype=np.int64),
"first_pts": np.zeros((total, num_cameras), dtype=np.float64),
"last_pts": np.zeros((total, num_cameras), dtype=np.float64),
"frame_count": np.zeros((total, num_cameras), dtype=np.int32),
"sample_lo": np.zeros((total, num_cameras), dtype=np.int32),
"sample_hi": np.zeros((total, num_cameras), dtype=np.int32),
"source_start_pts": np.zeros((total, num_cameras), dtype=np.float64),
}
for ep_idx in episode_indices:
ep = meta.episodes[ep_idx]
for cam_idx, key in enumerate(video_keys):
rel_path = str(meta.get_video_file_path(ep_idx, key))
file_id = path_to_id[rel_path]
mp4 = files[file_id].mp4
from_ts = float(ep[f"videos/{key}/from_timestamp"])
to_ts = float(ep[f"videos/{key}/to_timestamp"])
sample_slice = mp4.sample_slice(
from_ts,
to_ts,
keyframe_pad_s=keyframe_pad_s,
keyframe_pad_fraction=keyframe_pad_fraction,
file_size=files[file_id].file_size,
)
spans["file_id"][ep_idx, cam_idx] = file_id
spans["mdat_offset"][ep_idx, cam_idx] = sample_slice.byte_offset
spans["mdat_length"][ep_idx, cam_idx] = sample_slice.byte_length
spans["first_pts"][ep_idx, cam_idx] = from_ts
spans["last_pts"][ep_idx, cam_idx] = to_ts
spans["frame_count"][ep_idx, cam_idx] = sample_slice.sample_hi - sample_slice.sample_lo + 1
spans["sample_lo"][ep_idx, cam_idx] = sample_slice.sample_lo
spans["sample_hi"][ep_idx, cam_idx] = sample_slice.sample_hi
spans["source_start_pts"][ep_idx, cam_idx] = sample_slice.source_start_pts
return cls(video_keys=video_keys, files=files, spans=spans)
@staticmethod
def _build_file_records(
rel_paths: list[str],
data_root: str | Path,
*,
range_backend: str,
workers: int,
header_probe_bytes: int,
max_probe_bytes: int,
token: str | bool | None,
) -> list[VideoFileRecord]:
fetcher = make_range_fetcher(
data_root,
range_backend=range_backend,
workers=workers,
token=token,
)
def build_file(path: str) -> VideoFileRecord:
file_size = fetcher.info_size(path)
mp4 = fetch_mp4_index(
path,
fetcher.read_range,
file_size=file_size,
header_probe_bytes=header_probe_bytes,
max_probe_bytes=max_probe_bytes,
)
return VideoFileRecord(path, file_size, mp4)
try:
with ThreadPoolExecutor(max_workers=workers) as pool:
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()
@classmethod
def write_file_sidecar(
cls,
sidecar_path: str | Path,
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,
max_probe_bytes: int = 64 * 1024 * 1024,
token: str | bool | None = None,
) -> None:
records = cls._build_file_records(
sorted(set(rel_paths)),
data_root,
range_backend=range_backend,
workers=workers,
header_probe_bytes=header_probe_bytes,
max_probe_bytes=max_probe_bytes,
token=token,
)
cls.save_file_sidecar(sidecar_path, records, spec=spec)
@staticmethod
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": 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
],
}
arrays = {}
for file_idx, record in enumerate(records):
arrays[f"{file_idx}/sample_pts"] = record.mp4.sample_pts
arrays[f"{file_idx}/sample_durations"] = record.mp4.sample_durations
arrays[f"{file_idx}/sample_sizes"] = record.mp4.sample_sizes
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]:
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 and cached[0] == signature:
return cached[1]
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 = {
name: data[f"{file_idx}/{name}"]
for name in [
"sample_pts",
"sample_durations",
"sample_sizes",
"sample_offsets",
"sync_samples",
]
}
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] = (signature, records)
return records
def camera_id(self, camera_key: str) -> int:
return self._camera_to_id[camera_key]
def lookup(self, episode_index: int, camera_key: str) -> EpisodeVideoSpan:
cam = self.camera_id(camera_key)
return EpisodeVideoSpan(
file_id=int(self.spans["file_id"][episode_index, cam]),
mdat_offset=int(self.spans["mdat_offset"][episode_index, cam]),
mdat_length=int(self.spans["mdat_length"][episode_index, cam]),
first_pts=float(self.spans["first_pts"][episode_index, cam]),
last_pts=float(self.spans["last_pts"][episode_index, cam]),
frame_count=int(self.spans["frame_count"][episode_index, cam]),
sample_lo=int(self.spans["sample_lo"][episode_index, cam]),
sample_hi=int(self.spans["sample_hi"][episode_index, cam]),
source_start_pts=float(self.spans["source_start_pts"][episode_index, cam]),
)
def file_lookup(self, file_id: int) -> VideoFileRecord:
return self.files[file_id]
def mp4_index(self, episode_index: int, camera_key: str) -> Mp4Index:
return self.files[self.lookup(episode_index, camera_key).file_id].mp4
def sample_slice(self, episode_index: int, camera_key: str) -> Mp4SampleSlice:
span = self.lookup(episode_index, camera_key)
return Mp4SampleSlice(
sample_lo=span.sample_lo,
sample_hi=span.sample_hi,
byte_offset=span.mdat_offset,
byte_length=span.mdat_length,
source_start_pts=span.source_start_pts,
)
def episode_byte_size(self, episode_index: int) -> int:
"""Exact synthesized video bytes retained while an episode is active."""
return sum(
synthesized_mp4_size(
self.mp4_index(episode_index, camera_key),
self.sample_slice(episode_index, camera_key),
)
for camera_key in self.video_keys
)
+707
View File
@@ -0,0 +1,707 @@
# 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
"""MP4 indexing and in-memory episode synthesis primitives."""
from __future__ import annotations
import struct
from collections.abc import Callable, Iterable
from dataclasses import dataclass
import numpy as np
@dataclass(frozen=True)
class Box:
type: bytes
start: int
header_size: int
end: int
@property
def payload_start(self) -> int:
return self.start + self.header_size
@property
def size(self) -> int:
return self.end - self.start
@dataclass(frozen=True)
class Mp4SampleSlice:
sample_lo: int
sample_hi: int
byte_offset: int
byte_length: int
source_start_pts: float
@dataclass(frozen=True)
class Mp4Index:
file_path: str
file_size: int
ftyp: bytes
moov_offset: int
mdat_offset: int
mdat_payload_offset: int
mdat_payload_size: int
faststart: bool
codec: str
timescale: int
duration: int
track_id: int
width: int
height: int
stsd_body: bytes
sample_pts: np.ndarray
sample_durations: np.ndarray
sample_sizes: np.ndarray
sample_offsets: np.ndarray
sync_samples: np.ndarray
def sample_slice(
self,
from_ts: float,
to_ts: float,
*,
keyframe_pad_s: float = 0.1,
keyframe_pad_fraction: float = 0.05,
file_size: int | None = None,
) -> Mp4SampleSlice:
if to_ts < from_ts:
raise ValueError(f"Invalid timestamp span: {from_ts=} {to_ts=}")
if len(self.sample_pts) == 0:
raise ValueError(f"{self.file_path} contains no indexed samples")
pad = max(keyframe_pad_s, (to_ts - from_ts) * keyframe_pad_fraction)
lo_ts = max(0.0, from_ts - pad)
hi_ts = to_ts + pad
lo = int(np.searchsorted(self.sample_pts, lo_ts, side="left"))
hi = int(np.searchsorted(self.sample_pts, hi_ts, side="right")) - 1
lo = min(max(lo, 0), len(self.sample_pts) - 1)
hi = min(max(hi, lo), len(self.sample_pts) - 1)
if len(self.sync_samples):
prev_sync = self.sync_samples[self.sync_samples <= lo]
if len(prev_sync):
lo = int(prev_sync[-1])
else:
lo = int(self.sync_samples[0])
if lo > hi:
hi = lo
offsets = self.sample_offsets[lo : hi + 1]
sizes = self.sample_sizes[lo : hi + 1]
slice_lo = int(offsets.min())
slice_hi = int((offsets + sizes).max())
if file_size is not None:
slice_hi = min(slice_hi, int(file_size))
return Mp4SampleSlice(
sample_lo=lo,
sample_hi=hi,
byte_offset=slice_lo,
byte_length=slice_hi - slice_lo,
source_start_pts=float(self.sample_pts[lo]),
)
def to_dict(self) -> dict:
return {
"file_path": self.file_path,
"file_size": self.file_size,
"ftyp": self.ftyp.hex(),
"moov_offset": self.moov_offset,
"mdat_offset": self.mdat_offset,
"mdat_payload_offset": self.mdat_payload_offset,
"mdat_payload_size": self.mdat_payload_size,
"faststart": self.faststart,
"codec": self.codec,
"timescale": self.timescale,
"duration": self.duration,
"track_id": self.track_id,
"width": self.width,
"height": self.height,
"stsd_body": self.stsd_body.hex(),
}
@classmethod
def from_dict(cls, data: dict, arrays: dict[str, np.ndarray]) -> Mp4Index:
return cls(
file_path=data["file_path"],
file_size=int(data["file_size"]),
ftyp=bytes.fromhex(data["ftyp"]),
moov_offset=int(data["moov_offset"]),
mdat_offset=int(data["mdat_offset"]),
mdat_payload_offset=int(data["mdat_payload_offset"]),
mdat_payload_size=int(data["mdat_payload_size"]),
faststart=bool(data["faststart"]),
codec=data["codec"],
timescale=int(data["timescale"]),
duration=int(data["duration"]),
track_id=int(data["track_id"]),
width=int(data["width"]),
height=int(data["height"]),
stsd_body=bytes.fromhex(data["stsd_body"]),
sample_pts=arrays["sample_pts"],
sample_durations=arrays["sample_durations"],
sample_sizes=arrays["sample_sizes"],
sample_offsets=arrays["sample_offsets"],
sync_samples=arrays["sync_samples"],
)
def fetch_mp4_index(
path: str,
read_range: Callable[[str, int, int], bytes],
*,
file_size: int,
header_probe_bytes: int = 4 * 1024 * 1024,
max_probe_bytes: int = 64 * 1024 * 1024,
) -> Mp4Index:
probe_size = min(header_probe_bytes, file_size)
while True:
data = read_range(path, 0, probe_size)
top = list(iter_boxes(data, 0, len(data), absolute_base=0, allow_truncated=True))
has_mdat = any(box.type == b"mdat" for box in top)
has_moov = any(box.type == b"moov" and box.end <= len(data) for box in top)
if has_mdat and has_moov:
return parse_mp4_index(path, data, file_size=file_size)
if probe_size >= min(max_probe_bytes, file_size):
if has_mdat and not has_moov:
tail_index = _fetch_tail_moov_index(path, read_range, data, top, file_size, max_probe_bytes)
if tail_index is not None:
return tail_index
missing = []
if not has_mdat:
missing.append("mdat")
if not has_moov:
missing.append("moov")
raise ValueError(
f"Could not find complete {'/'.join(missing)} in first {probe_size} bytes of {path}"
)
probe_size = min(probe_size * 2, max_probe_bytes, file_size)
def _fetch_tail_moov_index(
path: str,
read_range: Callable[[str, int, int], bytes],
prefix: bytes,
top_boxes: list[Box],
file_size: int,
max_probe_bytes: int,
) -> Mp4Index | None:
mdat_box = _one(top_boxes, b"mdat")
if mdat_box is None or mdat_box.end >= file_size:
return None
tail_offset = mdat_box.end
tail_length = min(max_probe_bytes, file_size - tail_offset)
tail = read_range(path, tail_offset, tail_length)
tail_boxes = list(iter_boxes(tail, 0, len(tail), absolute_base=tail_offset, allow_truncated=True))
moov_box = next(
(box for box in tail_boxes if box.type == b"moov" and box.end <= tail_offset + len(tail)), None
)
if moov_box is None:
return None
ftyp_box = _one(top_boxes, b"ftyp", required=False)
ftyp = (
prefix[ftyp_box.start : ftyp_box.end]
if ftyp_box is not None
else _box(b"ftyp", b"isom\0\0\2\0isomiso2mp41")
)
moov_start = moov_box.payload_start - tail_offset
moov_end = moov_box.end - tail_offset
return _parse_mp4_index_from_layout(
path,
file_size=file_size,
ftyp=ftyp,
moov_offset=moov_box.start,
moov=tail[moov_start:moov_end],
mdat_box=mdat_box,
)
def parse_mp4_index(path: str, data: bytes, *, file_size: int | None = None) -> Mp4Index:
if file_size is None:
file_size = len(data)
top = list(iter_boxes(data, 0, len(data), absolute_base=0, allow_truncated=True))
ftyp_box = _one(top, b"ftyp", required=False)
moov_box = _one(top, b"moov")
mdat_box = _one(top, b"mdat")
if moov_box.end > len(data):
raise ValueError(f"{path}: moov box is truncated")
moov = data[moov_box.payload_start : moov_box.end]
ftyp = (
data[ftyp_box.start : ftyp_box.end]
if ftyp_box is not None
else _box(b"ftyp", b"isom\0\0\2\0isomiso2mp41")
)
return _parse_mp4_index_from_layout(
path,
file_size=file_size,
ftyp=ftyp,
moov_offset=moov_box.start,
moov=moov,
mdat_box=mdat_box,
)
def _parse_mp4_index_from_layout(
path: str,
*,
file_size: int,
ftyp: bytes,
moov_offset: int,
moov: bytes,
mdat_box: Box,
) -> Mp4Index:
mvhd_timescale, mvhd_duration = _parse_mvhd(_find_descendant(moov, [b"mvhd"]))
trak_box, trak_payload = _find_video_trak(moov)
_ = trak_box
tkhd = _parse_tkhd(_find_descendant(trak_payload, [b"tkhd"]))
mdhd_timescale, mdhd_duration = _parse_mdhd(_find_descendant(trak_payload, [b"mdia", b"mdhd"]))
stbl = _find_descendant(trak_payload, [b"mdia", b"minf", b"stbl"])
stsd = _find_child(stbl, b"stsd")
stsd_body = stbl[stsd.payload_start : stsd.end]
codec = _parse_stsd_codec(stsd_body)
stts = _parse_stts(_payload(stbl, b"stts"))
sample_sizes = _parse_stsz(_payload(stbl, b"stsz"))
stsc = _parse_stsc(_payload(stbl, b"stsc"))
chunk_offsets = _parse_chunk_offsets(stbl)
sync_samples = _parse_stss(stbl, len(sample_sizes))
sample_durations = _expand_stts(stts, len(sample_sizes))
sample_pts_units = np.empty(len(sample_durations), dtype=np.int64)
if len(sample_durations):
sample_pts_units[0] = 0
if len(sample_durations) > 1:
sample_pts_units[1:] = np.cumsum(sample_durations[:-1], dtype=np.int64)
sample_pts = sample_pts_units.astype(np.float64) / float(mdhd_timescale)
sample_offsets = _sample_offsets(stsc, chunk_offsets, sample_sizes)
return Mp4Index(
file_path=path,
file_size=file_size,
ftyp=ftyp,
moov_offset=moov_offset,
mdat_offset=mdat_box.start,
mdat_payload_offset=mdat_box.payload_start,
mdat_payload_size=mdat_box.end - mdat_box.payload_start
if mdat_box.end <= file_size
else file_size - mdat_box.payload_start,
faststart=moov_offset < mdat_box.start,
codec=codec,
timescale=mdhd_timescale,
duration=mdhd_duration or mvhd_duration,
track_id=tkhd["track_id"],
width=tkhd["width"],
height=tkhd["height"],
stsd_body=stsd_body,
sample_pts=sample_pts,
sample_durations=sample_durations,
sample_sizes=sample_sizes,
sample_offsets=sample_offsets,
sync_samples=sync_samples,
)
def synthesize_mp4(index: Mp4Index, sample_slice: Mp4SampleSlice, mdat_payload: bytes) -> bytes:
lo = sample_slice.sample_lo
hi = sample_slice.sample_hi + 1
if lo < 0 or hi > len(index.sample_sizes) or lo >= hi:
raise ValueError(f"Invalid sample range [{lo}, {hi}) for {index.file_path}")
offsets = index.sample_offsets[lo:hi]
sizes = index.sample_sizes[lo:hi]
rel_offsets = offsets - sample_slice.byte_offset
if int(rel_offsets.min()) != 0:
raise ValueError("Sample slice must start at the minimum referenced sample offset")
if int((rel_offsets + sizes).max()) > len(mdat_payload):
raise ValueError("Sample slice does not cover all referenced samples")
durations = index.sample_durations[lo:hi]
sync = index.sync_samples[(index.sync_samples >= lo) & (index.sync_samples < hi)] - lo + 1
moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=0)
header_size = len(index.ftyp) + len(moov)
mdat_header_size = 8 if len(mdat_payload) + 8 <= 0xFFFFFFFF else 16
moov = _make_moov(
index,
durations,
sizes,
rel_offsets,
sync,
mdat_data_offset=header_size + mdat_header_size,
)
return index.ftyp + moov + _box(b"mdat", mdat_payload)
def synthesized_mp4_size(index: Mp4Index, sample_slice: Mp4SampleSlice) -> int:
"""Return the exact synthesized mini-MP4 size without fetching its media payload."""
lo = sample_slice.sample_lo
hi = sample_slice.sample_hi + 1
if lo < 0 or hi > len(index.sample_sizes) or lo >= hi:
raise ValueError(f"Invalid sample range [{lo}, {hi}) for {index.file_path}")
offsets = index.sample_offsets[lo:hi]
sizes = index.sample_sizes[lo:hi]
rel_offsets = offsets - sample_slice.byte_offset
if int(rel_offsets.min()) != 0:
raise ValueError("Sample slice must start at the minimum referenced sample offset")
if int((rel_offsets + sizes).max()) > sample_slice.byte_length:
raise ValueError("Sample slice does not cover all referenced samples")
durations = index.sample_durations[lo:hi]
sync = index.sync_samples[(index.sync_samples >= lo) & (index.sync_samples < hi)] - lo + 1
moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=0)
header_size = len(index.ftyp) + len(moov)
mdat_header_size = 8 if sample_slice.byte_length + 8 <= 0xFFFFFFFF else 16
moov = _make_moov(
index,
durations,
sizes,
rel_offsets,
sync,
mdat_data_offset=header_size + mdat_header_size,
)
return len(index.ftyp) + len(moov) + mdat_header_size + sample_slice.byte_length
def iter_boxes(
data: bytes,
start: int,
end: int,
*,
absolute_base: int = 0,
allow_truncated: bool = False,
) -> Iterable[Box]:
pos = start
while pos + 8 <= end:
size = struct.unpack_from(">I", data, pos)[0]
typ = data[pos + 4 : pos + 8]
header_size = 8
if size == 1:
if pos + 16 > end:
break
size = struct.unpack_from(">Q", data, pos + 8)[0]
header_size = 16
elif size == 0:
size = end - pos
if size < header_size:
break
box_end = pos + size
if box_end > end and not allow_truncated:
break
yield Box(typ, absolute_base + pos, header_size, absolute_base + box_end)
pos = box_end
def _find_video_trak(moov: bytes) -> tuple[Box, bytes]:
for trak in _children(moov, 0, len(moov)):
if trak.type != b"trak":
continue
payload = moov[trak.payload_start : trak.end]
hdlr = _find_descendant(payload, [b"mdia", b"hdlr"])
if hdlr[8:12] == b"vide":
return trak, payload
raise ValueError("No video track found")
def _find_descendant(data: bytes, path: list[bytes]) -> bytes:
current = data
for typ in path:
box = _find_child(current, typ)
current = current[box.payload_start : box.end]
return current
def _find_child(data: bytes, typ: bytes) -> Box:
for box in _children(data, 0, len(data)):
if box.type == typ:
return box
raise ValueError(f"Missing MP4 box {typ.decode('latin1')}")
def _children(data: bytes, start: int, end: int) -> Iterable[Box]:
return iter_boxes(data, start, end, absolute_base=0)
def _one(boxes: list[Box], typ: bytes, *, required: bool = True) -> Box | None:
matches = [box for box in boxes if box.type == typ]
if not matches and required:
raise ValueError(f"Missing MP4 box {typ.decode('latin1')}")
return matches[0] if matches else None
def _payload(parent: bytes, typ: bytes) -> bytes:
box = _find_child(parent, typ)
return parent[box.payload_start : box.end]
def _parse_mvhd(payload: bytes) -> tuple[int, int]:
version = payload[0]
if version == 1:
return struct.unpack_from(">IQ", payload, 20)
return struct.unpack_from(">II", payload, 12)
def _parse_mdhd(payload: bytes) -> tuple[int, int]:
version = payload[0]
if version == 1:
return struct.unpack_from(">IQ", payload, 20)
return struct.unpack_from(">II", payload, 12)
def _parse_tkhd(payload: bytes) -> dict[str, int]:
version = payload[0]
if version == 1:
track_id = struct.unpack_from(">I", payload, 20)[0]
duration = struct.unpack_from(">Q", payload, 28)[0]
width, height = struct.unpack_from(">II", payload, 88)
else:
track_id = struct.unpack_from(">I", payload, 12)[0]
duration = struct.unpack_from(">I", payload, 20)[0]
width, height = struct.unpack_from(">II", payload, 76)
return {"track_id": track_id, "duration": duration, "width": width >> 16, "height": height >> 16}
def _parse_stsd_codec(stsd_body: bytes) -> str:
if len(stsd_body) < 16:
return "unknown"
return stsd_body[12:16].decode("latin1")
def _parse_stts(payload: bytes) -> list[tuple[int, int]]:
count = struct.unpack_from(">I", payload, 4)[0]
out = []
offset = 8
for _ in range(count):
out.append(struct.unpack_from(">II", payload, offset))
offset += 8
return out
def _expand_stts(entries: list[tuple[int, int]], sample_count: int) -> np.ndarray:
values = np.empty(sample_count, dtype=np.int64)
pos = 0
for count, delta in entries:
values[pos : pos + count] = delta
pos += count
if pos != sample_count:
raise ValueError(f"stts describes {pos} samples, stsz describes {sample_count}")
return values
def _parse_stsz(payload: bytes) -> np.ndarray:
sample_size, sample_count = struct.unpack_from(">II", payload, 4)
if sample_size:
return np.full(sample_count, sample_size, dtype=np.int64)
offset = 12
values = np.empty(sample_count, dtype=np.int64)
for idx in range(sample_count):
values[idx] = struct.unpack_from(">I", payload, offset)[0]
offset += 4
return values
def _parse_stsc(payload: bytes) -> list[tuple[int, int, int]]:
count = struct.unpack_from(">I", payload, 4)[0]
out = []
offset = 8
for _ in range(count):
out.append(struct.unpack_from(">III", payload, offset))
offset += 12
return out
def _parse_chunk_offsets(stbl: bytes) -> np.ndarray:
with_stco = None
with_co64 = None
for box in _children(stbl, 0, len(stbl)):
if box.type == b"stco":
with_stco = stbl[box.payload_start : box.end]
elif box.type == b"co64":
with_co64 = stbl[box.payload_start : box.end]
if with_co64 is not None:
count = struct.unpack_from(">I", with_co64, 4)[0]
return np.array(
[struct.unpack_from(">Q", with_co64, 8 + idx * 8)[0] for idx in range(count)], dtype=np.int64
)
if with_stco is None:
raise ValueError("Missing stco/co64 chunk offsets")
count = struct.unpack_from(">I", with_stco, 4)[0]
return np.array(
[struct.unpack_from(">I", with_stco, 8 + idx * 4)[0] for idx in range(count)], dtype=np.int64
)
def _parse_stss(stbl: bytes, sample_count: int) -> np.ndarray:
for box in _children(stbl, 0, len(stbl)):
if box.type == b"stss":
payload = stbl[box.payload_start : box.end]
count = struct.unpack_from(">I", payload, 4)[0]
return np.array(
[struct.unpack_from(">I", payload, 8 + idx * 4)[0] - 1 for idx in range(count)],
dtype=np.int64,
)
return np.arange(sample_count, dtype=np.int64)
def _sample_offsets(
stsc: list[tuple[int, int, int]], chunk_offsets: np.ndarray, sample_sizes: np.ndarray
) -> np.ndarray:
if not stsc:
raise ValueError("stsc is empty")
offsets = np.empty(len(sample_sizes), dtype=np.int64)
sample_idx = 0
for entry_idx, (first_chunk, samples_per_chunk, _desc_idx) in enumerate(stsc):
next_first = stsc[entry_idx + 1][0] if entry_idx + 1 < len(stsc) else len(chunk_offsets) + 1
for chunk_number in range(first_chunk, next_first):
if chunk_number < 1 or chunk_number > len(chunk_offsets):
raise ValueError("stsc references a chunk outside stco/co64")
chunk_pos = int(chunk_offsets[chunk_number - 1])
for _ in range(samples_per_chunk):
if sample_idx >= len(sample_sizes):
return offsets
offsets[sample_idx] = chunk_pos
chunk_pos += int(sample_sizes[sample_idx])
sample_idx += 1
if sample_idx != len(sample_sizes):
raise ValueError(f"stsc describes {sample_idx} samples, stsz describes {len(sample_sizes)}")
return offsets
def _make_moov(
index: Mp4Index,
durations: np.ndarray,
sizes: np.ndarray,
rel_offsets: np.ndarray,
sync_samples: np.ndarray,
*,
mdat_data_offset: int,
) -> bytes:
duration = int(durations.sum())
stco_values = [int(mdat_data_offset + value) for value in rel_offsets]
if any(value > 0xFFFFFFFF for value in stco_values):
offset_box = _co64(stco_values)
else:
offset_box = _stco(stco_values)
stbl = _box(
b"stbl",
_box(b"stsd", index.stsd_body)
+ _stts(durations)
+ _stsc_one_sample_per_chunk(len(sizes))
+ _stsz(sizes)
+ offset_box
+ (_stss(sync_samples) if len(sync_samples) else b""),
)
minf = _box(b"minf", _vmhd() + _dinf() + stbl)
mdia = _box(b"mdia", _mdhd(index.timescale, duration) + _hdlr() + minf)
trak = _box(b"trak", _tkhd(index.track_id, duration, index.width, index.height) + mdia)
return _box(b"moov", _mvhd(index.timescale, duration, index.track_id + 1) + trak)
def _full_box(typ: bytes, version: int, flags: int, payload: bytes = b"") -> bytes:
return _box(typ, bytes([version]) + flags.to_bytes(3, "big") + payload)
def _box(typ: bytes, payload: bytes) -> bytes:
size = len(payload) + 8
if size <= 0xFFFFFFFF:
return struct.pack(">I4s", size, typ) + payload
return struct.pack(">I4sQ", 1, typ, size + 8) + payload
def _mvhd(timescale: int, duration: int, next_track_id: int) -> bytes:
matrix = struct.pack(">9I", 0x00010000, 0, 0, 0, 0x00010000, 0, 0, 0, 0x40000000)
payload = (
struct.pack(">IIII", 0, 0, timescale, duration)
+ struct.pack(">IHH", 0x00010000, 0x0100, 0)
+ b"\0" * 8
+ matrix
+ b"\0" * 24
+ struct.pack(">I", next_track_id)
)
return _full_box(b"mvhd", 0, 0, payload)
def _tkhd(track_id: int, duration: int, width: int, height: int) -> bytes:
matrix = struct.pack(">9I", 0x00010000, 0, 0, 0, 0x00010000, 0, 0, 0, 0x40000000)
payload = (
struct.pack(">IIIII", 0, 0, track_id, 0, duration)
+ b"\0" * 8
+ struct.pack(">hhhh", 0, 0, 0, 0)
+ matrix
+ struct.pack(">II", width << 16, height << 16)
)
return _full_box(b"tkhd", 0, 7, payload)
def _mdhd(timescale: int, duration: int) -> bytes:
return _full_box(b"mdhd", 0, 0, struct.pack(">IIIIH", 0, 0, timescale, duration, 0x55C4) + b"\0\0")
def _hdlr() -> bytes:
return _full_box(b"hdlr", 0, 0, b"\0" * 4 + b"vide" + b"\0" * 12 + b"VideoHandler\0")
def _vmhd() -> bytes:
return _full_box(b"vmhd", 0, 1, struct.pack(">HHHH", 0, 0, 0, 0))
def _dinf() -> bytes:
url = _full_box(b"url ", 0, 1)
dref = _full_box(b"dref", 0, 0, struct.pack(">I", 1) + url)
return _box(b"dinf", dref)
def _stts(durations: np.ndarray) -> bytes:
runs = []
for duration in durations.tolist():
if runs and runs[-1][1] == int(duration):
runs[-1][0] += 1
else:
runs.append([1, int(duration)])
payload = struct.pack(">I", len(runs)) + b"".join(
struct.pack(">II", count, delta) for count, delta in runs
)
return _full_box(b"stts", 0, 0, payload)
def _stsc_one_sample_per_chunk(sample_count: int) -> bytes:
return _full_box(b"stsc", 0, 0, struct.pack(">IIII", 1, 1, 1, 1))
def _stsz(sizes: np.ndarray) -> bytes:
return _full_box(
b"stsz",
0,
0,
struct.pack(">II", 0, len(sizes)) + b"".join(struct.pack(">I", int(size)) for size in sizes.tolist()),
)
def _stco(values: list[int]) -> bytes:
return _full_box(
b"stco", 0, 0, struct.pack(">I", len(values)) + b"".join(struct.pack(">I", v) for v in values)
)
def _co64(values: list[int]) -> bytes:
return _full_box(
b"co64", 0, 0, struct.pack(">I", len(values)) + b"".join(struct.pack(">Q", v) for v in values)
)
def _stss(values: np.ndarray) -> bytes:
return _full_box(
b"stss",
0,
0,
struct.pack(">I", len(values)) + b"".join(struct.pack(">I", int(value)) for value in values.tolist()),
)
+743
View File
@@ -0,0 +1,743 @@
# 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 contextlib
import json
import os
import posixpath
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import UTC, datetime
from pathlib import Path
from types import MethodType
from typing import Any
from urllib.parse import quote, urljoin, urlparse
from uuid import uuid4
import fsspec
import httpx
from huggingface_hub import HfApi, HfFileSystem, constants
from huggingface_hub.utils import get_session, hf_raise_for_status
_HTTP_FAILURE_LOG_LOCK = threading.Lock()
def _get_header(headers: Any, name: str) -> str | None:
if hasattr(headers, "get"):
return headers.get(name)
lower_name = name.lower()
for key, value in headers.items():
if key.lower() == lower_name:
return value
return None
def _ensure_request_id(headers: dict[str, str]) -> str:
request_id = _get_header(headers, "X-Amzn-Trace-Id") or _get_header(headers, "X-Request-Id")
if request_id is None:
request_id = str(uuid4())
headers["X-Amzn-Trace-Id"] = request_id
return request_id
def _log_http_failure(
*,
backend: str,
method: str,
url: str,
headers: dict[str, str],
elapsed_s: float,
status_code: int | None = None,
exception: Exception | None = None,
attempt: int | None = None,
response_headers: Any | None = None,
) -> None:
log_path = os.environ.get("LEROBOT_HTTP_FAILURE_LOG")
if not log_path:
return
parsed = urlparse(url)
record = {
"ts": datetime.now(UTC).isoformat(),
"backend": backend,
"method": method,
"host": parsed.netloc,
"path": parsed.path,
"range": _get_header(headers, "Range") or _get_header(headers, "range"),
"request_id": _get_header(headers, "X-Amzn-Trace-Id") or _get_header(headers, "X-Request-Id"),
"elapsed_s": round(elapsed_s, 6),
}
if attempt is not None:
record["attempt"] = attempt
if status_code is not None:
record["status_code"] = status_code
if exception is not None:
record["exception_type"] = type(exception).__name__
record["exception"] = str(exception)
if response_headers is not None:
record["response_request_id"] = (
_get_header(response_headers, "x-request-id")
or _get_header(response_headers, "x-amz-cf-id")
or _get_header(response_headers, "x-amz-request-id")
)
record["cache_status"] = (
_get_header(response_headers, "x-cache")
or _get_header(response_headers, "cf-cache-status")
or _get_header(response_headers, "x-hf-cache")
)
record["content_range"] = _get_header(response_headers, "content-range")
record["content_length"] = _get_header(response_headers, "content-length")
path = Path(log_path).expanduser()
with _HTTP_FAILURE_LOG_LOCK:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("a") as out:
out.write(json.dumps(record, sort_keys=True) + "\n")
class ThreadLocalRangeFetcher:
"""Range reader that gives each worker thread independent file handles."""
def __init__(
self,
data_root: str | Path,
*,
block_size: int = 2**20,
cache_type: str = "none",
token: str | bool | None = None,
):
self.data_root = str(data_root).rstrip("/")
storage_options = {"token": token} if token is not None and self.data_root.startswith("hf://") else {}
self.fs, self._root_path = fsspec.core.url_to_fs(self.data_root, **storage_options)
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,
"range_bytes": 0.0,
"range_open_s": 0.0,
"range_seek_s": 0.0,
"range_read_s": 0.0,
}
def _url(self, relative_path: str) -> str:
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)
if handles is None:
handles = {}
self._local.handles = handles
handle = handles.get(relative_path)
if handle is None or getattr(handle, "closed", False):
handle = self.fs.open(
self._url(relative_path), "rb", block_size=self.block_size, cache_type=self.cache_type
)
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:
return int(self.fs.info(self._url(relative_path))["size"])
def read_range(self, relative_path: str, offset: int, length: int) -> bytes:
open_start = time.perf_counter()
handle = self._handle(relative_path)
open_s = time.perf_counter() - open_start
seek_start = time.perf_counter()
handle.seek(offset)
seek_s = time.perf_counter() - seek_start
read_start = time.perf_counter()
data = handle.read(length)
read_s = time.perf_counter() - read_start
self._record_timing(
range_jobs=1.0,
range_bytes=float(len(data)),
range_open_s=open_s,
range_seek_s=seek_s,
range_read_s=read_s,
)
return data
def _record_timing(self, **kwargs: float) -> None:
with self._timing_lock:
for key, value in kwargs.items():
self._timing_totals[key] = self._timing_totals.get(key, 0.0) + value
def _instrument_hf_handle(self, handle: Any) -> None:
if getattr(handle, "_lerobot_range_timing", False):
return
if not hasattr(handle, "_request_with_retry"):
return
def request_with_retry(
handle_self,
method: str,
url: str,
*,
headers: dict[str, str],
follow_redirects: bool | None = None,
max_retries: int = 5,
) -> httpx.Response:
from huggingface_hub.hf_file_system import _RANGE_RETRY_EXCEPTIONS, _RANGE_RETRY_STATUS_CODES
method_key = method.lower()
sleep_time = 1.0
retry_attempts = 0.0
retry_sleep_s = 0.0
failed_attempt_s = 0.0
exception_attempts = 0.0
extra_counts: dict[str, float] = {}
call_start = time.perf_counter()
request_kwargs: dict[str, Any] = {
"headers": headers,
"timeout": constants.HF_HUB_DOWNLOAD_TIMEOUT,
}
_ensure_request_id(headers)
if follow_redirects is not None:
request_kwargs["follow_redirects"] = follow_redirects
for attempt in range(max_retries + 1):
attempt_start = time.perf_counter()
try:
response = get_session().request(method, url, **request_kwargs)
except _RANGE_RETRY_EXCEPTIONS as exc:
attempt_s = time.perf_counter() - attempt_start
failed_attempt_s += attempt_s
exception_attempts += 1.0
key = f"range_hffs_{method_key}_exception_{type(exc).__name__}"
extra_counts[key] = extra_counts.get(key, 0.0) + 1.0
_log_http_failure(
backend="hffs",
method=method,
url=url,
headers=headers,
elapsed_s=attempt_s,
exception=exc,
attempt=attempt,
)
if attempt == max_retries:
self._record_hffs_request_timing(
method_key,
time.perf_counter() - call_start,
retry_attempts,
retry_sleep_s,
failed_attempt_s,
exception_attempts,
None,
0,
extra_counts,
)
raise
else:
elapsed = time.perf_counter() - attempt_start
if response.status_code not in _RANGE_RETRY_STATUS_CODES or attempt == max_retries:
self._record_hffs_request_timing(
method_key,
time.perf_counter() - call_start,
retry_attempts,
retry_sleep_s,
failed_attempt_s,
exception_attempts,
response.status_code,
len(response.content),
extra_counts,
)
return response
failed_attempt_s += elapsed
key = f"range_hffs_{method_key}_failed_status_{response.status_code}"
extra_counts[key] = extra_counts.get(key, 0.0) + 1.0
response.close()
_log_http_failure(
backend="hffs",
method=method,
url=url,
headers=headers,
elapsed_s=elapsed,
status_code=response.status_code,
attempt=attempt,
response_headers=response.headers,
)
time.sleep(sleep_time)
retry_attempts += 1.0
retry_sleep_s += sleep_time
sleep_time = min(8.0, sleep_time * 2)
raise RuntimeError("unreachable")
handle._request_with_retry = MethodType(request_with_retry, handle)
handle._lerobot_range_timing = True
def _record_hffs_request_timing(
self,
method_key: str,
total_s: float,
retry_attempts: float,
retry_sleep_s: float,
failed_attempt_s: float,
exception_attempts: float,
status_code: int | None,
byte_count: int,
extra_counts: dict[str, float],
) -> None:
timings = {
f"range_hffs_{method_key}_requests": 1.0,
f"range_hffs_{method_key}_s": total_s,
f"range_hffs_{method_key}_retries": retry_attempts,
f"range_hffs_{method_key}_retry_sleep_s": retry_sleep_s,
f"range_hffs_{method_key}_failed_attempt_s": failed_attempt_s,
f"range_hffs_{method_key}_exception_attempts": exception_attempts,
f"range_hffs_{method_key}_bytes": float(byte_count),
}
if status_code is not None:
timings[f"range_hffs_{method_key}_status_{status_code}"] = 1.0
timings.update(extra_counts)
self._record_timing(**timings)
def timing_summary(self) -> dict[str, float]:
with self._timing_lock:
return dict(self._timing_totals)
def close(self) -> None:
with self._handles_lock:
handles = list(self._all_handles.values())
self._all_handles.clear()
for handle in handles:
with contextlib.suppress(Exception):
handle.close()
local_handles = getattr(self._local, "handles", None)
if local_handles is not None:
local_handles.clear()
class NativeHTTPRangeFetcher:
"""Direct pooled HTTP range reader for hf:// paths."""
_GLOBAL_SOURCE_URLS: dict[tuple[str, str], str] = {}
_GLOBAL_RESOLVED_URLS: dict[tuple[str, str], str] = {}
_GLOBAL_SIZES: dict[tuple[str, str], int] = {}
_GLOBAL_LOCK = threading.Lock()
_RETRYABLE_EXCEPTIONS = (
httpx.ConnectError,
httpx.ConnectTimeout,
httpx.ReadError,
httpx.ReadTimeout,
httpx.RemoteProtocolError,
httpx.PoolTimeout,
)
_RETRYABLE_STATUS_CODES = {408, 425, 429, 500, 502, 503, 504}
def __init__(
self,
data_root: str | Path,
*,
max_connections: int = 32,
timeout: float = 60.0,
max_retries: int = 4,
subrange_parts: int = 1,
subrange_min_bytes: int = 8 * 1024 * 1024,
token: str | bool | None = None,
):
self.data_root = str(data_root).rstrip("/")
if not self.data_root.startswith("hf://"):
raise ValueError("NativeHTTPRangeFetcher only supports hf:// roots")
self.max_retries = max_retries
# Sub-range parallelism: split one large GET into `subrange_parts` concurrent GETs.
# Under a per-host throughput ceiling this adds no aggregate bandwidth, but divides
# per-request latency by ~parts - keep (in-flight jobs x parts) near the ceiling's
# connection sweet spot (~64 on the observed HF bucket path) rather than raising both.
self.subrange_parts = max(1, subrange_parts)
self.subrange_min_bytes = max(1, subrange_min_bytes)
self._subrange_pool = (
ThreadPoolExecutor(max_workers=max_connections, thread_name_prefix="subrange")
if self.subrange_parts > 1
else None
)
self.api = HfApi(token=token)
self.fs: HfFileSystem | None = None
self._bucket_id: str | None = None
self._bucket_prefix = ""
if self.data_root.startswith("hf://buckets/"):
bucket_root = self.data_root.removeprefix("hf://buckets/")
parts = bucket_root.split("/", 2)
if len(parts) < 2:
raise ValueError(f"Invalid bucket root: {self.data_root}")
self._bucket_id = f"{parts[0]}/{parts[1]}"
self._bucket_prefix = parts[2].strip("/") if len(parts) == 3 else ""
else:
self.fs = HfFileSystem(token=token)
self.client = httpx.Client(
timeout=timeout,
limits=httpx.Limits(max_connections=max_connections, max_keepalive_connections=max_connections),
follow_redirects=False,
)
self._resolved_urls: dict[str, str] = {}
self._source_urls: dict[str, str] = {}
self._sizes: dict[str, int] = {}
self._lock = threading.Lock()
self._timing_lock = threading.Lock()
self._timing_totals = {
"range_jobs": 0.0,
"range_bytes": 0.0,
"range_resolve_s": 0.0,
"range_header_s": 0.0,
"range_first_byte_s": 0.0,
"range_body_s": 0.0,
"range_retry_attempts": 0.0,
"range_retry_sleep_s": 0.0,
"range_failed_requests": 0.0,
}
def _request(self, method: str, url: str, **kwargs) -> httpx.Response:
last_exc: Exception | None = None
for attempt in range(self.max_retries + 1):
try:
return self.client.request(method, url, **kwargs)
except self._RETRYABLE_EXCEPTIONS as exc:
last_exc = exc
if attempt >= self.max_retries:
break
time.sleep(min(0.5 * 2**attempt, 5.0))
if last_exc is None:
raise RuntimeError("HTTP request failed without an exception")
raise last_exc
def _cache_key(self, relative_path: str) -> tuple[str, str]:
return self.data_root, relative_path
def _path(self, relative_path: str) -> str:
return f"{self.data_root}/{relative_path}"
def _bucket_path(self, relative_path: str) -> str:
if self._bucket_prefix:
return f"{self._bucket_prefix}/{relative_path}"
return relative_path
def _headers_for(self, request_url: str, source_url: str) -> dict[str, str]:
headers = self.api._build_hf_headers()
if urlparse(request_url).netloc != urlparse(source_url).netloc:
headers.pop("authorization", None)
headers.pop("Authorization", None)
return headers
def _source_url(self, relative_path: str) -> str:
with self._lock:
source = self._source_urls.get(relative_path)
if source is not None:
return source
key = self._cache_key(relative_path)
with self._GLOBAL_LOCK:
source = self._GLOBAL_SOURCE_URLS.get(key)
if source is None:
if self._bucket_id is not None:
source = (
f"{constants.ENDPOINT}/buckets/{self._bucket_id}/resolve/"
f"{quote(self._bucket_path(relative_path))}"
)
else:
if self.fs is None:
raise RuntimeError("HfFileSystem fallback was not initialized")
source = self.fs.url(self._path(relative_path))
with self._GLOBAL_LOCK:
self._GLOBAL_SOURCE_URLS[key] = source
with self._lock:
self._source_urls[relative_path] = source
return source
def _resolve_url(self, relative_path: str, *, refresh: bool = False) -> str:
with self._lock:
if not refresh and relative_path in self._resolved_urls:
return self._resolved_urls[relative_path]
key = self._cache_key(relative_path)
if not refresh:
with self._GLOBAL_LOCK:
resolved = self._GLOBAL_RESOLVED_URLS.get(key)
size = self._GLOBAL_SIZES.get(key)
if resolved is not None:
with self._lock:
self._resolved_urls[relative_path] = resolved
if size is not None:
self._sizes[relative_path] = size
return resolved
source = self._source_url(relative_path)
response = self._request("HEAD", source, headers=self.api._build_hf_headers(), follow_redirects=False)
try:
hf_raise_for_status(response)
location = response.headers.get("Location")
resolved = urljoin(source, location) if location else source
with self._lock:
self._resolved_urls[relative_path] = resolved
if "Content-Length" in response.headers:
self._sizes[relative_path] = int(response.headers["Content-Length"])
with self._GLOBAL_LOCK:
self._GLOBAL_RESOLVED_URLS[key] = resolved
if "Content-Length" in response.headers:
self._GLOBAL_SIZES[key] = int(response.headers["Content-Length"])
return resolved
finally:
response.close()
def info_size(self, relative_path: str) -> int:
with self._lock:
size = self._sizes.get(relative_path)
if size is not None:
return size
key = self._cache_key(relative_path)
with self._GLOBAL_LOCK:
size = self._GLOBAL_SIZES.get(key)
if size is not None:
with self._lock:
self._sizes[relative_path] = size
return size
resolved = self._resolve_url(relative_path)
source = self._source_url(relative_path)
response = self._request(
"HEAD", resolved, headers=self._headers_for(resolved, source), follow_redirects=True
)
try:
hf_raise_for_status(response)
size = int(response.headers["Content-Length"])
with self._lock:
self._sizes[relative_path] = size
with self._GLOBAL_LOCK:
self._GLOBAL_SIZES[key] = size
return size
finally:
response.close()
def read_range(self, relative_path: str, offset: int, length: int) -> bytes:
parts = self.subrange_parts
if self._subrange_pool is None or parts <= 1 or length < 2 * self.subrange_min_bytes:
return self._read_range_single(relative_path, offset, length)
parts = min(parts, max(1, length // self.subrange_min_bytes))
if parts <= 1:
return self._read_range_single(relative_path, offset, length)
step = (length + parts - 1) // parts
spans = [(offset + i * step, min(step, length - i * step)) for i in range(parts)]
futures = [
self._subrange_pool.submit(self._read_range_single, relative_path, span_off, span_len)
for span_off, span_len in spans
]
return b"".join(future.result() for future in futures)
def _read_range_single(self, relative_path: str, offset: int, length: int) -> bytes:
resolve_start = time.perf_counter()
resolved = self._resolve_url(relative_path)
source = self._source_url(relative_path)
resolve_s = time.perf_counter() - resolve_start
headers = self._headers_for(resolved, source)
headers["Range"] = f"bytes={offset}-{offset + length - 1}"
payload, status_code, timings = self._read_range_response(resolved, headers)
if status_code == 403:
refresh_start = time.perf_counter()
resolved = self._resolve_url(relative_path, refresh=True)
resolve_s += time.perf_counter() - refresh_start
headers = self._headers_for(resolved, source)
headers["Range"] = f"bytes={offset}-{offset + length - 1}"
payload, status_code, retry_timings = self._read_range_response(resolved, headers)
for key, value in retry_timings.items():
timings[key] += value
if status_code == 403:
raise PermissionError(f"HTTP range request returned 403 after URL refresh: {relative_path}")
if status_code != 206:
raise RuntimeError(f"HTTP range request returned {status_code} after retries: {relative_path}")
self._record_timing(
range_jobs=1.0,
range_bytes=float(len(payload)),
range_resolve_s=resolve_s,
**{f"range_status_{status_code}": 1.0},
**timings,
)
return payload
def _read_range_response(self, url: str, headers: dict[str, str]) -> tuple[bytes, int, dict[str, float]]:
last_exc: Exception | None = None
retry_attempts = 0.0
retry_sleep_s = 0.0
failed_attempt_s = 0.0
exception_attempts = 0.0
exception_counts: dict[str, float] = {}
_ensure_request_id(headers)
for attempt in range(self.max_retries + 1):
attempt_start = time.perf_counter()
try:
payload, status_code, timings = self._read_range_response_once(url, headers)
if status_code in self._RETRYABLE_STATUS_CODES:
attempt_s = time.perf_counter() - attempt_start
failed_attempt_s += attempt_s
exception_attempts += 1.0
status_key = f"range_failed_status_{status_code}"
exception_counts[status_key] = exception_counts.get(status_key, 0.0) + 1.0
_log_http_failure(
backend="native-http",
method="GET",
url=url,
headers=headers,
elapsed_s=attempt_s,
status_code=status_code,
attempt=attempt,
)
if attempt >= self.max_retries:
timings["range_retry_attempts"] = retry_attempts
timings["range_retry_sleep_s"] = retry_sleep_s
timings["range_failed_attempt_s"] = failed_attempt_s
timings["range_exception_attempts"] = exception_attempts
timings.update(exception_counts)
return payload, status_code, timings
retry_attempts += 1.0
sleep_s = min(0.5 * 2**attempt, 5.0)
retry_sleep_s += sleep_s
time.sleep(sleep_s)
continue
timings["range_retry_attempts"] = retry_attempts
timings["range_retry_sleep_s"] = retry_sleep_s
timings["range_failed_attempt_s"] = failed_attempt_s
timings["range_exception_attempts"] = exception_attempts
timings.update(exception_counts)
return payload, status_code, timings
except self._RETRYABLE_EXCEPTIONS as exc:
last_exc = exc
attempt_s = time.perf_counter() - attempt_start
failed_attempt_s += attempt_s
exception_attempts += 1.0
exception_key = f"range_exception_{type(exc).__name__}"
exception_counts[exception_key] = exception_counts.get(exception_key, 0.0) + 1.0
_log_http_failure(
backend="native-http",
method="GET",
url=url,
headers=headers,
elapsed_s=attempt_s,
exception=exc,
attempt=attempt,
)
if attempt >= self.max_retries:
break
retry_attempts += 1.0
sleep_s = min(0.5 * 2**attempt, 5.0)
retry_sleep_s += sleep_s
time.sleep(sleep_s)
self._record_timing(
range_failed_requests=1.0,
range_retry_attempts=retry_attempts,
range_retry_sleep_s=retry_sleep_s,
range_failed_attempt_s=failed_attempt_s,
range_exception_attempts=exception_attempts,
**exception_counts,
)
if last_exc is None:
raise RuntimeError("HTTP range request failed without an exception")
raise last_exc
def _read_range_response_once(
self, url: str, headers: dict[str, str]
) -> tuple[bytes, int, dict[str, float]]:
header_start = time.perf_counter()
with self.client.stream("GET", url, headers=headers) as response:
header_s = time.perf_counter() - header_start
if response.status_code == 403 or response.status_code in self._RETRYABLE_STATUS_CODES:
return (
b"",
response.status_code,
{
"range_header_s": header_s,
"range_first_byte_s": 0.0,
"range_body_s": 0.0,
},
)
hf_raise_for_status(response)
chunks = []
first_byte_s = 0.0
first_chunk = True
chunk_gap_s = 0.0
chunk_count = 0.0
previous_chunk_at = body_start = time.perf_counter()
for chunk in response.iter_bytes():
now = time.perf_counter()
if first_chunk:
first_byte_s = now - body_start
first_chunk = False
chunk_gap_s += now - previous_chunk_at
previous_chunk_at = now
chunk_count += 1.0
chunks.append(chunk)
body_s = time.perf_counter() - body_start
join_start = time.perf_counter()
payload = b"".join(chunks)
join_s = time.perf_counter() - join_start
return (
payload,
response.status_code,
{
"range_header_s": header_s,
"range_first_byte_s": first_byte_s,
"range_body_s": body_s,
"range_join_s": join_s,
"range_chunks": chunk_count,
"range_chunk_gap_s": chunk_gap_s,
},
)
def _record_timing(self, **kwargs: float) -> None:
with self._timing_lock:
for key, value in kwargs.items():
self._timing_totals[key] = self._timing_totals.get(key, 0.0) + value
def timing_summary(self) -> dict[str, float]:
with self._timing_lock:
return dict(self._timing_totals)
def close(self) -> None:
if self._subrange_pool is not None:
self._subrange_pool.shutdown(wait=False, cancel_futures=True)
self.client.close()
def make_range_fetcher(
data_root: str | Path,
*,
range_backend: str,
workers: int,
native_http_connections: int | None = None,
native_http_timeout: float = 60.0,
native_http_retries: int = 4,
native_http_subranges: int = 1,
token: str | bool | None = None,
):
if range_backend == "fsspec":
return ThreadLocalRangeFetcher(data_root, token=token)
if range_backend == "native-http":
max_connections = native_http_connections or max(8, workers)
return NativeHTTPRangeFetcher(
data_root,
max_connections=max_connections,
timeout=native_http_timeout,
max_retries=native_http_retries,
subrange_parts=native_http_subranges,
token=token,
)
raise ValueError(f"Unknown range backend: {range_backend}")
+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.manifest 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
+22
View File
@@ -36,3 +36,25 @@ def test_dataset_config_none_episodes_ok():
def test_dataset_config_empty_episodes_ok():
DatasetConfig(repo_id="user/repo", episodes=[])
def test_dataset_config_derives_streaming_decoder_limit_by_default():
assert DatasetConfig(repo_id="user/repo").streaming_max_open_decoders is None
@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"),
("streaming_decode_threads", 0, "decode_threads"),
("streaming_decoded_queue_size", 0, "decoded_queue_size"),
("streaming_max_open_decoders", 0, "max_open_decoders"),
("streaming_native_http_connections", 0, "native_http_connections"),
("streaming_native_http_subranges", 0, "native_http_subranges"),
],
)
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})
+9 -10
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
from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
stream_dataset = StreamingLeRobotDataset(
repo_id=DUMMY_REPO_ID, root=tmp_path / "ds", depth_output_unit=output_unit
)
stream_depth = next(iter(stream_dataset))[DEPTH_KEY]
assert torch.allclose(stream_depth, torch.full_like(stream_depth, expected))
stream_dataset = StreamingLeRobotDataset(
repo_id=DUMMY_REPO_ID, root=tmp_path / "ds", depth_output_unit=output_unit
)
stream_item = next(iter(stream_dataset))
stream_depth = stream_item[DEPTH_KEY]
reference_depth = read_dataset[int(stream_item["index"])][DEPTH_KEY]
assert torch.allclose(stream_depth, reference_depth)
assert torch.allclose(stream_depth, torch.full_like(stream_depth, expected), rtol=0.05)
@@ -0,0 +1,154 @@
#!/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
from unittest.mock import Mock
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.streaming.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]
def test_reader_forwards_explicit_token_to_hf_filesystem(monkeypatch) -> None:
url_to_fs = Mock(return_value=(fsspec.filesystem("memory"), "datasets/private@revision"))
monkeypatch.setattr(fsspec.core, "url_to_fs", url_to_fs)
EpisodeParquetReader(
"hf://datasets/private@revision",
columns=("episode_index",),
token="hf_test_token",
)
url_to_fs.assert_called_once_with("hf://datasets/private@revision", token="hf_test_token")
def test_reader_retries_transient_remote_file_not_found(tmp_path: Path, monkeypatch) -> None:
path = tmp_path / "episode.parquet"
pq.write_table(_table([0, 0]), path)
filesystem = Mock()
attempts = 0
def open_remote(*_args, **_kwargs):
nonlocal attempts
attempts += 1
if attempts == 1:
raise FileNotFoundError("partial HfFileSystem dircache")
return path.open("rb")
filesystem.open.side_effect = open_remote
filesystem.invalidate_cache = Mock()
monkeypatch.setattr(fsspec.core, "url_to_fs", Mock(return_value=(filesystem, "root")))
reader = EpisodeParquetReader(
"hf://datasets/private@revision",
columns=("episode_index", "frame_index"),
max_retries=1,
retry_backoff_s=0,
)
table = reader.read_episode("data/file.parquet", episode_index=0, expected_rows=2)
assert len(table) == 2
assert attempts == 2
filesystem.invalidate_cache.assert_called_once()
@@ -0,0 +1,334 @@
#!/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
import json
import struct
import threading
import time
from concurrent.futures import ThreadPoolExecutor
import numpy as np
import pytest
from lerobot.streaming.episode_cache import EpisodeByteCache
from lerobot.streaming.manifest import EpisodeVideoManifest
from lerobot.streaming.mp4 import (
_box,
_co64,
_dinf,
_hdlr,
_mdhd,
_mvhd,
_stco,
_stsc_one_sample_per_chunk,
_stss,
_stsz,
_stts,
_tkhd,
_vmhd,
parse_mp4_index,
synthesize_mp4,
synthesized_mp4_size,
)
from lerobot.streaming.range_fetch import ThreadLocalRangeFetcher, _log_http_failure
def _minimal_mp4(sample_offsets: list[int], *, use_co64: bool = False) -> bytes:
ftyp = _box(b"ftyp", b"isom\0\0\2\0isomiso2mp41")
sizes = np.array([10, 10, 10], dtype=np.int64)
durations = np.array([1000, 1000, 1000], dtype=np.int64)
stsd_body = struct.pack(">II", 0, 1) + struct.pack(">I4s", 16, b"avc1") + b"\0" * 8
offsets = _co64(sample_offsets) if use_co64 else _stco(sample_offsets)
stbl = _box(
b"stbl",
_box(b"stsd", stsd_body)
+ _stts(durations)
+ _stsc_one_sample_per_chunk(len(sizes))
+ _stsz(sizes)
+ offsets
+ _stss(np.array([1], dtype=np.int64)),
)
minf = _box(b"minf", _vmhd() + _dinf() + stbl)
mdia = _box(b"mdia", _mdhd(1000, 3000) + _hdlr() + minf)
trak = _box(b"trak", _tkhd(1, 3000, 64, 48) + mdia)
moov = _box(b"moov", _mvhd(1000, 3000, 2) + trak)
mdat_payload_start = 10_000
free_size = mdat_payload_start - 8 - len(ftyp) - len(moov)
assert free_size >= 8
free = _box(b"free", b"\0" * (free_size - 8))
return ftyp + moov + free + _box(b"mdat", b"x" * 128)
def test_episode_slice_uses_min_max_sample_offsets_for_reordered_chunks():
mp4 = parse_mp4_index("test.mp4", _minimal_mp4([10_000, 10_050, 10_025]))
sample_slice = mp4.sample_slice(0.0, 2.0, keyframe_pad_s=0, keyframe_pad_fraction=0)
assert sample_slice.byte_offset == 10_000
assert sample_slice.byte_length == 60
assert sample_slice.sample_lo == 0
assert sample_slice.sample_hi == 2
def test_synthesized_mp4_rebases_one_chunk_per_sample_offsets():
mp4 = parse_mp4_index("test.mp4", _minimal_mp4([10_000, 10_050, 10_025]))
sample_slice = mp4.sample_slice(0.0, 2.0, keyframe_pad_s=0, keyframe_pad_fraction=0)
mini = synthesize_mp4(mp4, sample_slice, b"x" * sample_slice.byte_length)
mini_index = parse_mp4_index("mini.mp4", mini)
expected = np.array([0, 50, 25], dtype=np.int64) + mini_index.mdat_payload_offset
np.testing.assert_array_equal(mini_index.sample_offsets, expected)
np.testing.assert_array_equal(mini_index.sample_sizes, np.array([10, 10, 10]))
def test_synthesized_mp4_size_matches_materialized_bytes():
mp4 = parse_mp4_index("test.mp4", _minimal_mp4([10_000, 10_050, 10_025]))
sample_slice = mp4.sample_slice(0.0, 2.0, keyframe_pad_s=0, keyframe_pad_fraction=0)
mini = synthesize_mp4(mp4, sample_slice, b"x" * sample_slice.byte_length)
assert synthesized_mp4_size(mp4, sample_slice) == len(mini)
def test_parser_accepts_co64_chunk_offsets():
mp4 = parse_mp4_index("test.mp4", _minimal_mp4([10_000, 10_050, 10_025], use_co64=True))
np.testing.assert_array_equal(mp4.sample_offsets, np.array([10_000, 10_050, 10_025]))
def _fake_cache(
monkeypatch,
tmp_path,
*,
byte_budget=8,
max_open_decoders=1,
video_backend="torchcodec",
):
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,
video_backend=video_backend,
)
monkeypatch.setattr(
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_cache.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_decoder_eviction_and_cache_shutdown_close_backend_resources(monkeypatch, tmp_path):
opened = []
class FakeDecoder:
def __init__(self):
self.closed = False
def close(self):
self.closed = True
def open_decoder(_data):
decoder = FakeDecoder()
opened.append(decoder)
return decoder
monkeypatch.setattr("lerobot.streaming.episode_cache.open_video_decoder", open_decoder)
with _fake_cache(monkeypatch, tmp_path, byte_budget=20, max_open_decoders=1) as cache:
cache.get_decoder(0, "camera")
cache.get_decoder(1, "camera")
assert opened[0].closed
assert not opened[1].closed
assert opened[1].closed
def test_decoder_falls_back_to_pyav_when_torchcodec_rejects_mini_mp4(monkeypatch, tmp_path):
opened_backends = []
class FakeDecoder:
pass
def open_decoder(_data, frame_mappings=None, *, backend="torchcodec"):
assert frame_mappings is None
opened_backends.append(backend)
if backend == "torchcodec":
raise ValueError("No valid stream found")
return FakeDecoder()
monkeypatch.setattr("lerobot.streaming.episode_cache.open_video_decoder", open_decoder)
with _fake_cache(monkeypatch, tmp_path, video_backend="torchcodec") as cache:
decoder = cache.get_decoder(0, "camera")
assert isinstance(decoder, FakeDecoder)
assert opened_backends == ["torchcodec", "pyav"]
assert cache.decoder_fallback_count == 1
def test_torchcodec_frame_indices_are_clamped_to_decoder_bounds(monkeypatch, tmp_path):
requested_indices = []
class FakeDecoder:
metadata = type("Metadata", (), {"average_fps": 30.0, "num_frames": 10})()
def get_frames_at(self, *, indices):
requested_indices.extend(indices)
return type("Frames", (), {"data": indices})()
with _fake_cache(monkeypatch, tmp_path) as cache:
monkeypatch.setattr(
cache.manifest,
"lookup",
lambda *_args: type("Span", (), {"source_start_pts": 0.0})(),
)
monkeypatch.setattr(cache, "_decoder_for_frames", lambda *_args: (FakeDecoder(), None))
cache.get_frames(0, "camera", [-0.1, 1.0])
assert requested_indices == [0, 9]
def test_frame_reads_serialize_access_to_each_decoder(monkeypatch, tmp_path):
state_lock = threading.Lock()
active = 0
max_active = 0
class FakeDecoder:
metadata = type("Metadata", (), {"average_fps": 30.0, "num_frames": 10})()
def get_frames_at(self, *, indices):
nonlocal active, max_active
with state_lock:
active += 1
max_active = max(max_active, active)
try:
time.sleep(0.01)
return type("Frames", (), {"data": indices})()
finally:
with state_lock:
active -= 1
with _fake_cache(monkeypatch, tmp_path) as cache:
monkeypatch.setattr(
cache.manifest,
"lookup",
lambda *_args: type("Span", (), {"source_start_pts": 0.0})(),
)
decoder = FakeDecoder()
monkeypatch.setattr(cache, "_decoder_for_frames", lambda *_args: (decoder, None))
with ThreadPoolExecutor(max_workers=2) as executor:
futures = [executor.submit(cache.get_frames, 0, "camera", [0.1]) for _ in range(2)]
for future in futures:
future.result()
assert max_active == 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,
)
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()
+146
View File
@@ -0,0 +1,146 @@
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""ExactCoveragePool: exactly-once frame coverage over a bounded episode pool."""
from collections import Counter
import pytest
from lerobot.streaming.episode_pool import ExactCoveragePool
EPISODES = [(0, 5), (1, 3), (2, 8), (3, 1), (4, 6), (5, 4), (6, 7), (7, 2)]
TOTAL = sum(n for _, n in EPISODES)
EXPECTED = Counter((ep, i) for ep, n in EPISODES for i in range(n))
def _drain(pool):
out, max_resident = [], 0
while True:
try:
out.append(next(pool))
except StopIteration:
break
max_resident = max(max_resident, len(pool.resident))
return out, max_resident
def test_exact_once_coverage():
out, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=42))
assert len(out) == TOTAL
assert Counter(out) == EXPECTED # every (episode, frame) exactly once, no dups/misses
def test_pool_never_exceeds_size():
_, max_resident = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=42))
assert max_resident <= 3
def test_deterministic_per_seed_and_epoch():
a, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=7))
b, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=7))
c, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=8))
d, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=7, epoch=1))
assert a == b
assert a != c and a != d # seed and epoch both change the order
assert Counter(c) == EXPECTED and Counter(d) == EXPECTED # ... but coverage is preserved
def test_admission_and_eviction_events():
pool = ExactCoveragePool(EPISODES, pool_size=3, seed=0)
admitted_ever, evicted_ever = set(), set()
# first three episodes admitted at construction
admitted_ever.update(pool.newly_admitted)
assert len(admitted_ever) == 3
while True:
pool.newly_admitted.clear()
pool.evicted.clear()
try:
next(pool)
except StopIteration:
break
admitted_ever.update(pool.newly_admitted)
evicted_ever.update(pool.evicted)
assert admitted_ever == {ep for ep, _ in EPISODES} # every episode admitted exactly once
# every episode except the pool_size still resident at the end is evicted on exhaustion
assert len(evicted_ever) >= len(EPISODES) - 3
def test_uniform_mixing_matches_coupon_collector():
# 64 equal episodes, pool 64, first 64 draws -> ~64*(1-(1-1/64)^64) ~= 41 distinct
big = [(e, 100) for e in range(64)]
pool = ExactCoveragePool(big, pool_size=64, seed=0)
head = [next(pool)[0] for _ in range(64)]
assert len(set(head)) >= 30 # far above sequential (=1); ~41 expected
def test_large_epoch_bounded_and_complete():
big = [(e, 90) for e in range(500)]
out, max_resident = _drain(ExactCoveragePool(big, pool_size=64, seed=3))
assert len(out) == 500 * 90
assert len(set(out)) == 500 * 90 # exactly once
assert max_resident <= 64
def test_zero_length_episodes_skipped():
pool = ExactCoveragePool([(0, 3), (1, 0), (2, 2)], pool_size=8, seed=0)
out, _ = _drain(pool)
assert Counter(out) == Counter({(0, 0): 1, (0, 1): 1, (0, 2): 1, (2, 0): 1, (2, 1): 1})
def test_byte_aware_admission_never_exceeds_budget():
sizes = {0: 7, 1: 6, 2: 4, 3: 3, 4: 2}
pool = ExactCoveragePool(
EPISODES[:5],
pool_size=4,
seed=9,
episode_byte_sizes=sizes,
byte_budget=10,
)
out = []
max_resident_bytes = pool.resident_bytes
while pool.remaining_total:
out.append(next(pool))
max_resident_bytes = max(max_resident_bytes, pool.resident_bytes)
assert Counter(out) == Counter((ep, frame) for ep, count in EPISODES[:5] for frame in range(count))
assert max_resident_bytes <= 10
assert len(pool.admission_order) == len(EPISODES[:5])
def test_byte_aware_admission_rejects_one_oversized_episode():
with pytest.raises(ValueError, match="Episode 1.*byte budget"):
ExactCoveragePool(
[(0, 2), (1, 3)],
pool_size=2,
seed=0,
episode_byte_sizes={0: 4, 1: 11},
byte_budget=10,
)
def test_prefetch_candidates_follow_deterministic_pending_frontier():
pool = ExactCoveragePool(
EPISODES,
pool_size=2,
seed=17,
episode_byte_sizes={episode: 1 for episode, _ in EPISODES},
byte_budget=2,
)
candidates = pool.prefetch_candidates(3)
assert len(candidates) == 3
assert not set(candidates) & set(pool.resident)
assert candidates == pool.prefetch_candidates(3)
+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] == []
+24 -91
View File
@@ -16,7 +16,6 @@
from types import SimpleNamespace
from unittest.mock import Mock
import numpy as np
import pytest
import torch
@@ -24,74 +23,32 @@ pytest.importorskip("datasets", reason="datasets is required (install lerobot[da
import lerobot.datasets.streaming_dataset as streaming_dataset_module
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
@pytest.mark.parametrize("token", ["hf_test_token", True, False])
@pytest.mark.parametrize("from_local", [False, True])
def test_streaming_dataset_forwards_hub_token_only_for_remote_data(tmp_path, monkeypatch, token, from_local):
def test_streaming_dataset_forwards_token_to_metadata_without_retaining_it(
tmp_path, monkeypatch, token, from_local
):
requested_root = tmp_path / "local" if from_local else None
metadata = SimpleNamespace(
repo_id=DUMMY_REPO_ID,
root=requested_root or tmp_path / "snapshot",
revision=streaming_dataset_module.CODEBASE_VERSION,
_version=streaming_dataset_module.CODEBASE_VERSION,
features={},
total_episodes=0,
video_keys=[],
depth_keys=[],
image_keys=[],
rescale_depth_stats=Mock(),
)
metadata_cls = Mock(return_value=metadata)
load_dataset = Mock(return_value=SimpleNamespace(num_shards=1))
ensure_sidecar = Mock(return_value=None)
monkeypatch.setattr(streaming_dataset_module, "LeRobotDatasetMetadata", metadata_cls)
monkeypatch.setattr(streaming_dataset_module, "load_dataset", load_dataset)
monkeypatch.setattr(streaming_dataset_module, "ensure_dataset_mp4_sidecar", ensure_sidecar)
dataset = StreamingLeRobotDataset(DUMMY_REPO_ID, root=requested_root, token=token)
@@ -102,10 +59,7 @@ def test_streaming_dataset_forwards_hub_token_only_for_remote_data(tmp_path, mon
force_cache_sync=False,
token=token,
)
if from_local:
assert "token" not in load_dataset.call_args.kwargs
else:
assert load_dataset.call_args.kwargs["token"] is token
assert ensure_sidecar.call_args.kwargs["token"] is (None if from_local else token)
assert not hasattr(dataset, "_token")
@@ -130,7 +84,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:
@@ -179,22 +133,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(
@@ -234,22 +182,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(
@@ -299,7 +241,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()), (
@@ -382,20 +324,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()), (
+70
View File
@@ -0,0 +1,70 @@
#!/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,
video_backend="pyav",
streaming_data_root="memory://payload",
streaming_episode_pool_size=7,
streaming_prefetch_episodes=3,
streaming_byte_budget_gb=2.5,
streaming_decode_threads=2,
streaming_decoded_queue_size=5,
streaming_max_open_decoders=17,
streaming_native_http_connections=9,
streaming_native_http_subranges=3,
)
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"]["decode_threads"] == 2
assert captured["kwargs"]["decoded_queue_size"] == 5
assert captured["kwargs"]["max_open_decoders"] == 17
assert captured["kwargs"]["native_http_connections"] == 9
assert captured["kwargs"]["native_http_subranges"] == 3
assert captured["kwargs"]["max_num_shards"] == 1
assert captured["kwargs"]["video_backend"] == "pyav"
assert captured["kwargs"]["return_uint8"] is True
assert captured["kwargs"]["repeat"] is True
+645
View File
@@ -0,0 +1,645 @@
#!/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 threading
import time
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, _balanced_episode_shards
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_parallel_decode_queue_preserves_planner_order(
tmp_path: Path,
lerobot_dataset_factory,
monkeypatch,
) -> None:
root = tmp_path / "dataset"
lerobot_dataset_factory(
root=root,
repo_id=DUMMY_REPO_ID,
total_episodes=4,
total_frames=40,
use_videos=False,
)
expected = _indices(
StreamingLeRobotDataset(
DUMMY_REPO_ID,
root=root,
seed=17,
buffer_size=3,
decode_threads=1,
decoded_queue_size=1,
)
)
parallel = StreamingLeRobotDataset(
DUMMY_REPO_ID,
root=root,
seed=17,
buffer_size=3,
decode_threads=3,
decoded_queue_size=5,
)
original_make_item = parallel._make_episode_item
state_lock = threading.Lock()
active = 0
max_active = 0
def delayed_make_item(*args, **kwargs):
nonlocal active, max_active
with state_lock:
active += 1
max_active = max(max_active, active)
try:
frame_index = int(args[2])
time.sleep(0.005 if frame_index % 3 == 0 else 0.001)
return original_make_item(*args, **kwargs)
finally:
with state_lock:
active -= 1
monkeypatch.setattr(parallel, "_make_episode_item", delayed_make_item)
assert _indices(parallel) == expected
assert 1 < max_active <= parallel.decode_threads
def test_default_decoder_limit_covers_the_configured_episode_pool(
tmp_path: Path,
lerobot_dataset_factory,
) -> None:
root = tmp_path / "dataset"
lerobot_dataset_factory(
root=root,
repo_id=DUMMY_REPO_ID,
total_episodes=2,
total_frames=20,
)
streaming = StreamingLeRobotDataset(
DUMMY_REPO_ID,
root=root,
episode_pool_size=7,
)
assert streaming.max_open_decoders == 7 * len(streaming.meta.video_keys)
overridden = StreamingLeRobotDataset(
DUMMY_REPO_ID,
root=root,
episode_pool_size=7,
max_open_decoders=5,
)
assert overridden.max_open_decoders == 5
@pytest.mark.parametrize("video_backend", ["torchcodec", "pyav"])
def test_streaming_rgb_video_matches_map_style(
tmp_path: Path,
lerobot_dataset_factory,
video_backend: str,
) -> None:
root = tmp_path / "dataset"
map_dataset = lerobot_dataset_factory(
root=root,
repo_id=DUMMY_REPO_ID,
total_episodes=2,
total_frames=20,
video_backend=video_backend,
)
streaming = StreamingLeRobotDataset(
DUMMY_REPO_ID,
root=root,
shuffle=False,
buffer_size=2,
video_backend=video_backend,
)
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_rejects_episode_larger_than_rank_byte_budget(
tmp_path: Path, lerobot_dataset_factory
) -> None:
root = tmp_path / "dataset"
lerobot_dataset_factory(
root=root,
repo_id=DUMMY_REPO_ID,
total_episodes=2,
total_frames=10,
)
streaming = StreamingLeRobotDataset(
DUMMY_REPO_ID,
root=root,
shuffle=False,
buffer_size=2,
byte_budget_gb=1 / 1024**3,
)
with pytest.raises(ValueError, match="Episode .*byte budget"):
next(iter(streaming))
def test_streaming_rank_shards_are_disjoint(tmp_path: Path, lerobot_dataset_factory, monkeypatch) -> None:
root = tmp_path / "dataset"
map_dataset = lerobot_dataset_factory(
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_rank_shards_are_greedily_balanced_by_frame_count() -> None:
shards = _balanced_episode_shards(
[0, 1, 2, 3, 4],
{0: 100, 1: 90, 2: 20, 3: 10, 4: 5},
world_size=2,
)
assert {episode for shard in shards for episode in shard} == {0, 1, 2, 3, 4}
assert set(shards[0]).isdisjoint(shards[1])
totals = [sum({0: 100, 1: 90, 2: 20, 3: 10, 4: 5}[episode] for episode in shard) for shard in shards]
assert max(totals) - min(totals) <= 15
def test_streaming_rejects_multiple_sampling_workers(tmp_path: Path, lerobot_dataset_factory) -> None:
root = tmp_path / "dataset"
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)
with pytest.raises(RuntimeError, match="one DataLoader worker per rank"):
list(loader)
def test_streaming_persistent_workers_advance_epochs(tmp_path: Path, lerobot_dataset_factory) -> None:
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=1,
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=1,
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=1)
if batch_size is None:
return [int(item["index"]) for item in loader]
return [int(index) for batch in loader for index in batch["index"]]
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=1,
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=1)
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,)
@@ -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"
+41
View File
@@ -0,0 +1,41 @@
#!/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_cache import EpisodeByteCache
from lerobot.streaming.episode_pool import ExactCoveragePool
from lerobot.streaming.manifest import EpisodeVideoManifest
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
+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.manifest 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)
+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()
Generated
+6
View File
@@ -2822,7 +2822,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" },
@@ -3295,7 +3298,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" },
@@ -3306,6 +3311,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" },