mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-28 12:15:59 +00:00
Compare commits
44 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 21321d7a27 | |||
| fdf6014214 | |||
| 4db03ed535 | |||
| 79cfb52a71 | |||
| c0e9b0bbff | |||
| b5a13e43ce | |||
| 1e7e0b6de5 | |||
| a3edab661b | |||
| 564ba6395a | |||
| 85086fed7a | |||
| ee06d1005f | |||
| 1c47809bf6 | |||
| 3d70b21aac | |||
| b85620657f | |||
| fbfc861cf2 | |||
| 06aa6a0425 | |||
| be64ded80f | |||
| 88843ed675 | |||
| f2b5c4a47b | |||
| 9202fcea96 | |||
| ef47c35178 | |||
| 6d6c82eb8c | |||
| 9201be92cb | |||
| 0064a06205 | |||
| 710171ccac | |||
| 0f8257443c | |||
| 3a09d0c48a | |||
| 03fc5e3ea9 | |||
| 28c3e095bf | |||
| 5bfb749a9b | |||
| 51c023a7a1 | |||
| 51ea18cb7a | |||
| 04ab43b8d2 | |||
| cdfe192491 | |||
| 3451e53452 | |||
| 30849ce74f | |||
| 7d6907c444 | |||
| d99e1fe89d | |||
| 7fcde61b69 | |||
| bdfe8f8ce9 | |||
| 34d0495d03 | |||
| 834c282631 | |||
| f132885cbc | |||
| d0686be2f5 |
@@ -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
|
||||
|
||||
@@ -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;">
|
||||
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
@@ -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,
|
||||
)
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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."""
|
||||
@@ -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")
|
||||
@@ -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"
|
||||
)
|
||||
@@ -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
|
||||
@@ -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
|
||||
)
|
||||
@@ -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()),
|
||||
)
|
||||
@@ -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}")
|
||||
@@ -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
|
||||
@@ -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})
|
||||
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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] == []
|
||||
|
||||
@@ -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()), (
|
||||
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
|
||||
@@ -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" },
|
||||
|
||||
Reference in New Issue
Block a user