test_subset.py

"""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