feat(dataset): integrate episode streaming into training

tests/test_streaming_sidecar.py

@@ -0,0 +1,188 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+from __future__ import annotations
+
+import shutil
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import numpy as np
+import pytest
+from filelock import FileLock
+
+from lerobot.streaming.episode_video import EpisodeVideoManifest, VideoFileRecord
+from lerobot.streaming.mp4 import Mp4Index
+from lerobot.streaming.sidecar import (
+    SidecarLockTimeoutError,
+    SidecarSpec,
+    ensure_mp4_sidecar,
+    sidecar_cache_path,
+)
+
+
+def _record(path: str = "videos/camera/chunk-000/file-000.mp4", size: int = 128) -> VideoFileRecord:
+    arrays = np.array([0], dtype=np.int64)
+    index = Mp4Index(
+        file_path=path,
+        file_size=size,
+        ftyp=b"",
+        moov_offset=0,
+        mdat_offset=0,
+        mdat_payload_offset=0,
+        mdat_payload_size=size,
+        faststart=True,
+        codec="avc1",
+        timescale=1,
+        duration=1,
+        track_id=1,
+        width=1,
+        height=1,
+        stsd_body=b"",
+        sample_pts=np.array([0.0]),
+        sample_durations=arrays,
+        sample_sizes=arrays,
+        sample_offsets=arrays,
+        sync_samples=arrays,
+    )
+    return VideoFileRecord(path, size, index)
+
+
+def _spec(revision: str = "rev-a", size: int = 128) -> SidecarSpec:
+    return SidecarSpec(
+        repo_id="owner/dataset",
+        revision=revision,
+        data_root="hf://datasets/owner/dataset",
+        source_files=(("videos/camera/chunk-000/file-000.mp4", size),),
+    )
+
+
+def _write_valid(path: Path, spec: SidecarSpec) -> None:
+    EpisodeVideoManifest.save_file_sidecar(path, [_record(size=spec.source_files[0][1])], spec=spec)
+
+
+def test_sidecar_cache_path_is_revision_keyed(tmp_path: Path) -> None:
+    first = sidecar_cache_path(tmp_path, _spec("rev-a"))
+    second = sidecar_cache_path(tmp_path, _spec("rev-b"))
+
+    assert first != second
+    assert first.parent == second.parent
+
+
+def test_ensure_reuses_valid_local_sidecar(tmp_path: Path) -> None:
+    spec = _spec()
+    path = sidecar_cache_path(tmp_path, spec)
+    _write_valid(path, spec)
+    build_calls = 0
+
+    def build(_path: Path, _spec: SidecarSpec) -> None:
+        nonlocal build_calls
+        build_calls += 1
+
+    resolved = ensure_mp4_sidecar(spec, tmp_path, build=build)
+
+    assert resolved == path
+    assert build_calls == 0
+
+
+def test_ensure_downloads_valid_published_sidecar(tmp_path: Path) -> None:
+    spec = _spec()
+    published = tmp_path / "published.npz"
+    _write_valid(published, spec)
+    build_calls = 0
+
+    def download(path: Path, _spec: SidecarSpec) -> bool:
+        shutil.copyfile(published, path)
+        return True
+
+    def build(_path: Path, _spec: SidecarSpec) -> None:
+        nonlocal build_calls
+        build_calls += 1
+
+    resolved = ensure_mp4_sidecar(spec, tmp_path / "cache", build=build, download=download)
+
+    assert EpisodeVideoManifest.validate_file_sidecar(resolved, spec)
+    assert build_calls == 0
+
+
+@pytest.mark.parametrize("invalid_kind", ["corrupt", "stale"])
+def test_ensure_rebuilds_invalid_local_sidecar(tmp_path: Path, invalid_kind: str) -> None:
+    spec = _spec()
+    path = sidecar_cache_path(tmp_path, spec)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if invalid_kind == "corrupt":
+        path.write_bytes(b"not-an-npz")
+    else:
+        _write_valid(path, _spec(revision="other-revision"))
+    build_calls = 0
+
+    def build(target: Path, target_spec: SidecarSpec) -> None:
+        nonlocal build_calls
+        build_calls += 1
+        _write_valid(target, target_spec)
+
+    resolved = ensure_mp4_sidecar(spec, tmp_path, build=build)
+
+    assert EpisodeVideoManifest.validate_file_sidecar(resolved, spec)
+    assert build_calls == 1
+
+
+def test_concurrent_ensure_builds_once(tmp_path: Path) -> None:
+    spec = _spec()
+    start = threading.Barrier(2)
+    build_calls = 0
+    build_lock = threading.Lock()
+
+    def build(path: Path, target_spec: SidecarSpec) -> None:
+        nonlocal build_calls
+        with build_lock:
+            build_calls += 1
+        _write_valid(path, target_spec)
+
+    def ensure() -> Path:
+        start.wait()
+        return ensure_mp4_sidecar(spec, tmp_path, build=build)
+
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        paths = list(pool.map(lambda _: ensure(), range(2)))
+
+    assert paths[0] == paths[1]
+    assert build_calls == 1
+
+
+def test_failed_build_does_not_replace_existing_file(tmp_path: Path) -> None:
+    spec = _spec()
+    path = sidecar_cache_path(tmp_path, spec)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_bytes(b"old-corrupt-file")
+
+    def build(target: Path, _spec: SidecarSpec) -> None:
+        target.write_bytes(b"partial")
+        raise RuntimeError("build failed")
+
+    with pytest.raises(RuntimeError, match="build failed"):
+        ensure_mp4_sidecar(spec, tmp_path, build=build)
+
+    assert path.read_bytes() == b"old-corrupt-file"
+    assert not list(path.parent.glob(f".{path.name}.*.tmp.npz"))
+
+
+def test_lock_timeout_is_actionable(tmp_path: Path) -> None:
+    spec = _spec()
+    path = sidecar_cache_path(tmp_path, spec)
+    lock_path = path.with_suffix(f"{path.suffix}.lock")
+    lock_path.parent.mkdir(parents=True, exist_ok=True)
+
+    with (
+        FileLock(lock_path),
+        pytest.raises(SidecarLockTimeoutError, match="Timed out waiting"),
+    ):
+        ensure_mp4_sidecar(spec, tmp_path, build=_write_valid, lock_timeout_s=0.01)