"""Real shared-file video/data fixtures: subset correctness and ordinary LeRobot loading.""" import json import shutil import sys from pathlib import Path import av import numpy as np import pandas as pd import pyarrow as pa import pyarrow.parquet as pq import pytest sys.path.insert(0, str(Path(__file__).parent)) from subset import EpisodeSubset, trim_video # noqa: E402 from worker import Executor, PreservingWriter, video_path # noqa: E402 from lerobot.annotations.steerable_pipeline.reader import iter_episodes from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging from lerobot.datasets.feature_utils import create_empty_dataset_info from lerobot.datasets.io_utils import write_info, write_tasks from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.utils.constants import DEFAULT_FEATURES @pytest.fixture def shared_source(tmp_path, request): n = getattr(request, "param", 10) root = tmp_path / "source" cameras = ["observation.images.top", "observation.images.wrist"] features = { **DEFAULT_FEATURES, "action": {"dtype": "float32", "shape": [2], "names": ["joint", "joint2"]}, **{ c: { "dtype": "video", "shape": [32, 32, 3], "names": ["height", "width", "channels"], "info": { "video.height": 32, "video.width": 32, "video.codec": "h264", "video.fps": 10, "video.pix_fmt": "yuv420p", "video.is_depth_map": False, "has_audio": False, }, } for c in cameras }, } info = create_empty_dataset_info("v3.0", 10, features, True) info.total_episodes, info.total_frames, info.total_tasks = 3, 3 * n, 3 info.splits = {"train": "0:3"} write_info(info, root) write_tasks( pd.DataFrame( {"task_index": [0, 1, 2]}, index=pd.Index(["task zero", "excluded task", "task two"], name="task") ), root, ) for ci, camera in enumerate(cameras): video = root / info.video_path.format(video_key=camera, chunk_index=0, file_index=0) video.parent.mkdir(parents=True, exist_ok=True) with av.open(str(video), "w") as container: stream = container.add_stream("libx264", rate=10, options={"crf": "0", "g": "30"}) stream.width = stream.height = 32 stream.pix_fmt = "yuv420p" for idx in range(3 * n): frame = av.VideoFrame.from_ndarray( np.full((32, 32, 3), 30 + (idx // n) * 80 + ci * 10, dtype=np.uint8), format="rgb24" ) for packet in stream.encode(frame): container.mux(packet) for packet in stream.encode(): container.mux(packet) table = pa.table( { "episode_index": pa.array([i // n for i in range(3 * n)], type=pa.int64()), "frame_index": pa.array([i % n for i in range(3 * n)], type=pa.int64()), "timestamp": pa.array([(i % n) / 10 for i in range(3 * n)], type=pa.float32()), "index": pa.array(range(3 * n), type=pa.int64()), "task_index": pa.array([i // n for i in range(3 * n)], type=pa.int64()), "action": pa.array([[float(i), float(i)] for i in range(3 * n)], type=pa.list_(pa.float32(), 2)), } ) path = root / info.data_path.format(chunk_index=0, file_index=0) path.parent.mkdir(parents=True, exist_ok=True) pq.write_table(table, path, row_group_size=n) rows = [] for ep in range(3): row = { "episode_index": ep, "length": n, "tasks": [["task zero", "excluded task", "task two"][ep]], "data/chunk_index": 0, "data/file_index": 0, "dataset_from_index": ep * n, "dataset_to_index": (ep + 1) * n, "meta/episodes/chunk_index": 0, "meta/episodes/file_index": 0, "stats/action/min": [float(ep * n)], "stats/action/max": [float(ep * n + n - 1)], "stats/action/mean": [float(ep * n + (n - 1) / 2)], "stats/action/std": [float(np.std(np.arange(n)))], "stats/action/count": [n], } for camera in cameras: row.update( { f"videos/{camera}/chunk_index": 0, f"videos/{camera}/file_index": 0, f"videos/{camera}/from_timestamp": float(ep * n / 10), f"videos/{camera}/to_timestamp": float((ep + 1) * n / 10), } ) rows.append(row) meta_path = root / "meta/episodes/chunk-000/file-000.parquet" meta_path.parent.mkdir(parents=True, exist_ok=True) pq.write_table(pa.Table.from_pylist(rows), meta_path) return root, cameras, path @pytest.mark.parametrize("per_episode", [False, True]) def test_video_resolver_keeps_episode_and_file_indices_distinct(per_episode): camera = "observation.images.ego" row = {"episode_index": 123, f"videos/{camera}/chunk_index": 0, f"videos/{camera}/file_index": 7} suffix = "episode_{episode_index:06d}" if per_episode else "file-{file_index:03d}" info = {"video_path": "videos/{video_key}/chunk-{chunk_index:03d}/" + suffix + ".mp4"} expected = "episode_000123.mp4" if per_episode else "file-007.mp4" assert video_path(row, info, camera) == f"videos/{camera}/chunk-000/{expected}" def test_episode_named_videos_export_with_offsets_and_load_normally(shared_source, tmp_path): root, cameras, data_path = shared_source path = root / "meta/info.json" info = json.loads(path.read_text()) for camera in cameras: original = root / info["video_path"].format(video_key=camera, chunk_index=0, file_index=0) shutil.copyfile(original, original.with_name("episode_000002.mp4")) info["video_path"] = "videos/{video_key}/chunk-{chunk_index:03d}/episode_{episode_index:06d}.mp4" path.write_text(json.dumps(info)) output = tmp_path / "episode-layout" subset = EpisodeSubset(root, output, [2]) subset.add(2, pq.read_table(data_path)) loaded = LeRobotDataset("test/episode-layout", root=output, video_backend="pyav") assert len(loaded) == 10 and loaded[0]["task"] == "task two" assert float(loaded[0][cameras[0]].mean()) == pytest.approx(190 / 255, abs=0.02) assert subset.mapping[0]["source_episode_index"] == 2 @pytest.mark.parametrize("camera_workers", [1, 3]) def test_successful_episodes_only_all_cameras_and_real_lerobot_load(shared_source, tmp_path, camera_workers): root, cameras, data_path = shared_source records = [r for r in iter_episodes(root) if r.episode_index != 1] staging = tmp_path / "staging" for r in records: EpisodeStaging(staging, r.episode_index).write( "plan", [ { "role": "assistant", "content": f"annotated source {r.episode_index}", "style": "subtask", "timestamp": 0.0, "camera": None, "tool_calls": None, } ], ) PreservingWriter([0, 2], {"subtask"}).write_all(records, staging, root) Executor._ensure_annotation_metadata_in_info(root) output = tmp_path / "subset" builder = EpisodeSubset(root, output, [0, 1, 2], camera_workers=camera_workers) table = pq.read_table(data_path) builder.add(0, table) # Loading a checkpoint works before all completed episodes have been copied. first = LeRobotDataset("test/subset", root=output, video_backend="pyav", token=False) assert first.num_episodes == 1 and len(first) == 10 assert "annotated source 0" in str(first[0]["language_persistent"]) builder.add(2, table) result = LeRobotDataset("test/subset", root=output, video_backend="pyav", token=False) assert result.num_episodes == 2 and len(result) == 20 assert result.meta.info.splits == {"train": "0:2"} assert result.meta.total_tasks == 2 assert list(result.meta.tasks.index) == ["task zero", "task two"] assert builder.mapping == [ {"episode_index": 0, "source_episode_index": 0}, {"episode_index": 1, "source_episode_index": 2}, ] data = pq.read_table(output / "data").to_pylist() assert [r["index"] for r in data] == list(range(20)) assert {r["episode_index"] for r in data} == {0, 1} assert {r["task_index"] for r in data} == {0, 1} assert [r["action"][0] for r in data] == list(range(10)) + list(range(20, 30)) for episode in [0, 1]: for ci, camera in enumerate(cameras): path = output / builder.info["video_path"].format( video_key=camera, chunk_index=0, file_index=episode ) with av.open(str(path)) as video: frames = list(video.decode(video=0)) assert len(frames) == 10 assert float(frames[0].pts * frames[0].time_base) == 0 assert abs(frames[0].to_ndarray(format="rgb24").mean() - (30 + episode * 160 + ci * 10)) < 5 loaded = result[episode * 10] assert f"annotated source {episode * 2}" in str(loaded["language_persistent"]) assert tuple(loaded[cameras[0]].shape) == (3, 32, 32) stats = json.loads((output / "meta/stats.json").read_text()) assert stats["action"]["mean"] == [14.5] assert stats["index"]["mean"] == [9.5] assert stats["episode_index"]["max"] == [1] assert len(list((output / "videos").rglob("*.mp4"))) == 4 with pytest.raises(ValueError, match="already added"): builder.add(2, table) def test_truncated_video_fails_instead_of_copying_wrong_frames(shared_source, tmp_path): root, cameras, _ = shared_source source = root / f"videos/{cameras[0]}/chunk-000/file-000.mp4" with pytest.raises(ValueError, match="ended early"): trim_video(source, tmp_path / "too-long.mp4", 2.0, 20, 10) def test_subset_visualizer_link_uses_new_index_not_source_id(): from reporting import annotation_url from studio import COPY report = { "manifest": {"mode": COPY, "copy_scope": "annotated_episodes", "output": "test/subset"}, "status": "completed", "episodes": [{"episode_index": 900, "status": "completed"}], "written_episode_ids": [], } assert annotation_url(report) is None report["written_episode_ids"] = [0] assert annotation_url(report).endswith("/test/subset/0?tab=annotations") @pytest.mark.parametrize("fail_publish", [False, True]) @pytest.mark.parametrize("prefetch", [False, True]) def test_worker_publishes_only_successful_subset_without_hub_duplication( shared_source, tmp_path, monkeypatch, fail_publish, prefetch ): import hashlib import shutil from threading import Event from types import SimpleNamespace from unittest.mock import Mock import worker from reporting import annotation_url from studio import COPY source, cameras, _ = shared_source (source / "README.md").write_text("Original dataset attribution") (source / "LICENSE").write_text("Original license") published = tmp_path / "published" published.mkdir() api = Mock() api.list_repo_refs.return_value = SimpleNamespace(tags=[]) commits = [] export_started = Event() metadata_ready, model_loading = Event(), Event() metadata_downloads = [] def write_file(path, data): target = published / path target.parent.mkdir(parents=True, exist_ok=True) target.write_bytes(data if isinstance(data, bytes) else Path(data).read_bytes()) def commit(*args, operations, **kwargs): publishing = any(op.path_in_repo == "meta/info.json" for op in operations) if publishing: export_started.set() if fail_publish: raise RuntimeError("Simulated atomic upload failure") for op in operations: write_file(op.path_in_repo, op.path_or_fileobj) if any(op.path_in_repo == "meta/info.json" for op in operations): snapshot = LeRobotDataset("test/output", root=published, video_backend="pyav", token=False) commits.append(snapshot.num_episodes) return SimpleNamespace(oid="output-commit") api.create_commit.side_effect = commit api.upload_file.side_effect = lambda **kw: write_file(kw["path_in_repo"], kw["path_or_fileobj"]) api.get_paths_info.side_effect = lambda repo, files, **kw: [ SimpleNamespace(size=(source / f).stat().st_size) for f in files ] monkeypatch.setattr(worker, "HfApi", lambda **kw: api) monkeypatch.setattr(worker.shutil, "disk_usage", lambda *a: SimpleNamespace(free=10_000_000_000)) monkeypatch.setenv("HF_TOKEN", "local-test-token") monkeypatch.delenv("ANNOTATION_DEADLINE_EPOCH", raising=False) rows = pq.read_table(source / "meta/episodes").to_pylist() info = json.loads((source / "meta/info.json").read_text()) monkeypatch.setattr(worker, "read_selection", lambda *args, **kw: {"info": info, "episodes": rows}) def snapshot(*args, local_dir, **kw): metadata_downloads.append(local_dir) if prefetch: assert model_loading.wait(5), "Source setup did not overlap model loading" shutil.copytree(source / "meta", local_dir / "meta") for filename in ["README.md", "LICENSE"]: shutil.copyfile(source / filename, local_dir / filename) metadata_ready.set() downloads = [] def download(repo, filename, *, local_dir, **kw): downloads.append(filename) target = local_dir / filename target.parent.mkdir(parents=True, exist_ok=True) shutil.copyfile(source / filename, target) return str(target) monkeypatch.setattr(worker, "snapshot_download", snapshot) monkeypatch.setattr(worker, "hf_hub_download", download) monkeypatch.setattr(worker, "Provider", lambda *args, **kw: SimpleNamespace(concurrency=2)) def annotate(record, row, info, root, manifest, provider, staging, media=None): if record.episode_index == 1: raise ValueError("Simulated annotation validation failure") if record.episode_index == (0 if prefetch else 2): # Accelerated Flora must publish source 2 before slow source 0. # Other workflows keep their existing prefix order. assert export_started.wait(5), "Episode export did not overlap annotation" EpisodeStaging(staging, record.episode_index).write( "plan", [ { "role": "assistant", "content": f"annotated source {record.episode_index}", "style": "subtask", "timestamp": 0.0, "camera": None, "tool_calls": None, } ], ) return {"episode_index": record.episode_index, "status": "completed", "rows": 1} monkeypatch.setattr(worker, "annotate_episode", annotate) manifest = { "dataset": "test/source", "dataset_revision": "pinned", "output": "test/output", "mode": COPY, "visibility": "public", "copy_scope": "annotated_episodes", "source_license": "apache-2.0", "start": 0, "count": 3, "episode_ids_sha256": hashlib.sha256(json.dumps([0, 1, 2]).encode()).hexdigest(), "inference_budget": 1, "model": "mock", "camera": cameras[0], "features": ["Subtasks"], "costs": {"input_rate": 0.1, "output_rate": 0.2}, } if prefetch: import flora_batcher import flora_captioner class Native: execution, batch_size, kernel_versions = "accelerated", 8, {"test": "native"} def __init__(self, **kwargs): pass class Broker(Native): parallelism, gpu_batch_size, vision_frame_batching = 8, 64, True def __init__(self, *args, **kwargs): model_loading.set() assert metadata_ready.wait(5), "Metadata prefetch did not finish during setup" def close(self): pass monkeypatch.setattr(flora_captioner, "FloraCaptioner", Native) monkeypatch.setattr(flora_batcher, "CrossEpisodeFlora", Broker) monkeypatch.setenv("FLORA_EXECUTION", "accelerated") monkeypatch.setenv("FLORA_EPISODE_PARALLELISM", "8") manifest.update(captioner="flora", captioner_profile=flora_captioner.FLORA) if fail_publish: with pytest.raises(RuntimeError, match="Annotation stopped"): worker.run(manifest, tmp_path / "work") report = json.loads((published / "annotation_studio/run.json").read_text()) assert report["status"] == "partial" assert report["written_episode_ids"] == report["source_episode_map"] == [] assert report["export_error"] == {"source_episode": 2 if prefetch else 0, "type": "RuntimeError"} assert not commits assert len(metadata_downloads) == 1 api.create_tag.assert_not_called() return worker.run(manifest, tmp_path / "work") api.duplicate_repo.assert_not_called() api.create_repo.assert_called_once_with("test/output", repo_type="dataset", private=False, exist_ok=False) assert commits == [1, 2] final = LeRobotDataset("test/output", root=published, video_backend="pyav", token=False) assert final.num_episodes == 2 and len(final) == 20 source_order = [2, 0] if prefetch else [0, 2] assert f"annotated source {source_order[0]}" in str(final[0]["language_persistent"]) assert f"annotated source {source_order[1]}" in str(final[10]["language_persistent"]) assert not list(published.rglob("*.jsonl")) assert not (published / "annotation_studio/episodes").exists() report = json.loads((published / "annotation_studio/run.json").read_text()) assert report["status"] == "partial" assert report["written_episode_ids"] == [0, 1] assert [m["source_episode_index"] for m in report["source_episode_map"]] == source_order assert annotation_url(report).endswith("/test/output/0?tab=annotations") assert (published / "annotation_studio/source_README.md").read_text() == "Original dataset attribution" assert (published / "LICENSE").read_text() == "Original license" assert 'license: "apache-2.0"' in (published / "README.md").read_text() assert not list((tmp_path / "work/subset/videos").rglob("*.mp4")) assert not list((tmp_path / "work/subset/data").rglob("*.parquet")) # Both successful episodes share their source camera shards. Download once. assert len(downloads) == len(set(downloads)) == 3 assert len(metadata_downloads) == 1 assert not list((tmp_path / "work/dataset/videos").rglob("*.mp4")) api.create_tag.assert_called_once_with( "test/output", tag="v3.0", repo_type="dataset", revision="output-commit" ) def test_dense_frame_sampling_matches_source_pixels_with_one_seek(tmp_path, monkeypatch): from types import SimpleNamespace import worker camera = "observation.images.top" path = tmp_path / f"videos/{camera}/chunk-000/file-000.mp4" path.parent.mkdir(parents=True) with av.open(str(path), "w") as container: stream = container.add_stream("libx264", rate=30, options={"crf": "18", "g": "60", "bf": "2"}) stream.width = stream.height = 32 stream.pix_fmt = "yuv420p" for index in range(180): pixels = np.full((32, 32, 3), (index * 13) % 256, dtype=np.uint8) pixels[:, :16, 1] = (index * 7) % 256 frame = av.VideoFrame.from_ndarray(pixels, format="rgb24") for packet in stream.encode(frame): container.mux(packet) for packet in stream.encode(): container.mux(packet) with av.open(str(path)) as container: expected = [frame.to_image() for frame in container.decode(video=0)] original_open = av.open seeks = [] class CountSeeks: def __init__(self, *args, **kwargs): self.container = original_open(*args, **kwargs) def __enter__(self): return self def __exit__(self, *args): self.container.close() def __getattr__(self, name): return getattr(self.container, name) def seek(self, *args, **kwargs): seeks.append(args) return self.container.seek(*args, **kwargs) monkeypatch.setattr(worker.av, "open", CountSeeks) row = { "length": 120, f"videos/{camera}/from_timestamp": 1, f"videos/{camera}/chunk_index": 0, f"videos/{camera}/file_index": 0, } source = worker.SourceFrames( row, {"fps": 30, "video_path": "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4"}, tmp_path, camera, max_size=None, ) indices = list(range(0, 120, 3)) # Unsorted inputs and duplicates must preserve the requested order. requested = [indices[-1], *indices, indices[0]] images = source.frames_at(SimpleNamespace(), [i / 30 for i in requested]) assert [image.tobytes() for image in images] == [expected[30 + i].tobytes() for i in requested] assert len(seeks) == 1 # Previously 40 GOP seeks for the same pixels. source.frames_at(SimpleNamespace(), [i / 30 for i in reversed(requested)]) assert len(seeks) == 1 @pytest.mark.parametrize("shared_source", [40], indirect=True) def test_hand_actions_load_in_same_dataset_and_keep_unobserved_episode(shared_source, tmp_path): from lerobot.annotations.hands.config import ACTION_SCHEMA from lerobot.annotations.hands.writer import FEATURES as HAND_FEATURES, PARAMETERS, attach_actions root, cameras, data_path = shared_source builder = EpisodeSubset(root, tmp_path / "annotated", [0, 2]) source = pq.read_table(data_path) for episode in (0, 2): table = source.filter(pa.compute.equal(source["episode_index"], episode)) report = dict( action_schema=ACTION_SCHEMA, motion_infilling=False, gt_scale_alignment=False, rectified_frames=40, rectified_fps=10, accepted_frames=0, segments=[], ) if episode == 2: path = tmp_path / "observed.npz" fields = {name: np.zeros((30, *shape)) for name, shape in PARAMETERS.items()} np.savez( path, action=np.full((30, 156), 3.0), source_frame_index=np.arange(30), timestamp_s=np.arange(30) / 10, **fields, ) report.update(segments=[str(path)], accepted_frames=30) table, summary = attach_actions(table, report, directory=tmp_path, fps=10) builder.add(episode, table, hand_features=HAND_FEATURES) assert summary["unmeasured_frames"] == (40 if episode == 0 else 10) if episode == 0: assert "action" not in json.loads((builder.root / "meta/stats.json").read_text()) result = LeRobotDataset("test/annotated", root=builder.root, video_backend="pyav", token=False) assert len(result) == 80 and result.num_episodes == 2 assert tuple(result[0]["action"].shape) == (156,) assert tuple(result[0]["annotation.original_action"].shape) == (2,) assert not result[0]["annotation.hand_action_valid"].any() assert result[40]["annotation.hand_action_valid"].all() assert result[40]["annotation.hand_action_observed_valid"].all() assert tuple(result[40]["annotation.soma_hand_pose"].shape) == (432,) assert not result[79]["annotation.hand_action_valid"].any() table = pq.read_table(builder.root / "data") assert ( table["annotation.original_action"].to_pylist() == source.filter(pa.compute.is_in(source["episode_index"], value_set=pa.array([0, 2])))[ "action" ].to_pylist() ) stats = json.loads((builder.root / "meta/stats.json").read_text()) assert stats["action"]["mean"] == [3.0] * 156 assert stats["action"]["count"] == [30] assert len(stats["annotation.original_action"]["mean"]) == 2 assert builder.mapping == [ {"episode_index": 0, "source_episode_index": 0}, {"episode_index": 1, "source_episode_index": 2}, ] assert len(list((builder.root / "videos").rglob("*.mp4"))) == 4