tests/test_batch_runner.py

"""Tests for fpgm.datagen.batch.runner.BatchRunner.

``fpgm.datagen.pipeline`` does not exist yet (a concurrent change lands it);
per this task's own brief, production code
(``fpgm/datagen/batch/runner.py``) imports it unconditionally and
unconditionally-deferred (see that module's docstring for why). These tests
develop against a stub injected into ``sys.modules["fpgm.datagen.pipeline"]``
-- never an import-guard in the production module itself.

Every test here uses an in-thread ``ThreadPoolExecutor`` (via
``BatchRunner``'s own ``executor_factory`` hook) instead of a real
``ProcessPoolExecutor``, specifically *because* threads share this process's
``sys.modules`` with the stub, while a real spawned child process would not
see it. The real multi-process/env-var-ordering path is exercised directly
against ``_run_shard`` and separately by the task's own manual GPU/CUDA
verification -- this file is about ``BatchRunner``'s orchestration logic
(sharding, pilot policy, resume, timeout-as-failure, pruning), which is
identical either way.
"""

from __future__ import annotations

import json
import sys
import time
import types
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path

import h5py
import pytest

from fpgm.config_datagen import DatagenProfile
from fpgm.datagen.batch.gpu_pool import GpuWorkerPool
from fpgm.datagen.batch.runner import BatchRunner

REPO_ROOT = Path(__file__).resolve().parents[1]
REAL_MESH = REPO_ROOT / "assets" / "scenes" / "8756300955" / "manipulated.glb"

RESOLVABLE_TASK = "Put brick in drawer shelf and close drawer"
UNPARSEABLE_TASK = "Do anything you like that takes multiple steps to complete."


# --------------------------------------------------------------------------- #
# Stub fpgm.datagen.pipeline -- see module docstring for why this exists
# --------------------------------------------------------------------------- #


def _install_pipeline_stub(behavior: dict[str, str]) -> types.ModuleType:
    """``behavior``: {uuid: "ok" | "fail" | "hang"}, default "ok" for unlisted uuids."""

    @dataclass(frozen=True)
    class StageResult:
        name: str
        status: str
        payload: dict
        elapsed_s: float
        error: str | None = None

    @dataclass(frozen=True)
    class EpisodeReport:
        spec: object
        stages: tuple
        gates: dict
        windows: tuple

        @property
        def status(self) -> str:
            return "ok" if all(s.status == "ok" for s in self.stages) else "failed"

        def to_json(self) -> dict:
            return {
                "stages": [
                    {"name": s.name, "status": s.status, "payload": s.payload,
                     "elapsed_s": s.elapsed_s, "error": s.error}
                    for s in self.stages
                ],
                "gates": self.gates,
                "windows": list(self.windows),
            }

    call_log: list[str] = []

    class EpisodePipeline:
        STAGES = ("s2", "s3", "s4", "s6", "s8")

        def __init__(self, profile, registry=None) -> None:
            self.profile = profile
            self.registry = registry

        def run(self, spec, stages=None, *, force: bool = False) -> EpisodeReport:
            call_log.append(spec.uuid)
            # One reusable model per registry, regardless of how many episodes
            # this registry processes -- exactly the reuse behavior the real
            # pipeline is meant to exhibit (see ModelRegistry's own docstring).
            if self.registry is not None:
                self.registry._get_or_build(("stub_model",), lambda: object())

            mode = behavior.get(spec.uuid, "ok")
            if mode == "hang":
                time.sleep(30)
                mode = "ok"
            if mode == "fail":
                failed_stage = StageResult("s2", "failed", {}, 0.01, error="stub failure")
                return EpisodeReport(spec=spec, stages=(failed_stage,), gates={}, windows=())
            window = {"window": "w0", "overlaps_pose_gap": False, "high_hole_fraction": False}
            return EpisodeReport(
                spec=spec,
                stages=(StageResult("s2", "ok", {"n": 1}, 0.01),),
                gates={"dummy_gate": {"pass": True, "value": 1.0}},
                windows=(window,),
            )

    module = types.ModuleType("fpgm.datagen.pipeline")
    module.EpisodePipeline = EpisodePipeline
    module.EpisodeReport = EpisodeReport
    module.StageResult = StageResult
    module.call_log = call_log  # test-only escape hatch
    return module


@pytest.fixture()
def pipeline_stub(monkeypatch):
    installed: dict[str, types.ModuleType] = {}

    def _install(behavior: dict[str, str] | None = None):
        module = _install_pipeline_stub(behavior or {})
        monkeypatch.setitem(sys.modules, "fpgm.datagen.pipeline", module)
        installed["module"] = module
        return module

    yield _install
    installed.clear()


# --------------------------------------------------------------------------- #
# Profile / episode fixtures
# --------------------------------------------------------------------------- #


def _write_flows_h5(path: Path, uuid: str, task: str, camera_serial: str = "cam1") -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with h5py.File(path, "w") as fh:
        fh.attrs["uuid"] = uuid
        fh.attrs["scene_id"] = "8756300955"
        fh.attrs["current_task"] = task
        fh.attrs["ext1_cam_serial"] = camera_serial


def _write_profile_yaml(tmp_path: Path, **batch_overrides) -> Path:
    batch = {
        "camera_role": "ext1",
        "pilot_fraction": 0.5,
        "pilot_min_episodes": 1,
        "pilot_pass_threshold": 0.6,
        "episode_timeout_s": 2.0,
        "high_hole_fraction_gate": 0.10,
        "prune_intermediates": True,
    }
    batch.update(batch_overrides)
    import yaml

    raw = {
        "paths": {
            "droid_raw_root": str(tmp_path / "droid_raw"),
            "flows_root": str(tmp_path / "flows"),
            "cameras_root": str(tmp_path / "cameras"),
            "output_root": str(tmp_path / "outputs"),
            "checkpoints_root": str(tmp_path / "checkpoints"),
            "urdf": str(tmp_path / "fake.urdf"),
            "python_executable": sys.executable,
        },
        "scene_assets": {"meshes": {"8756300955": {"manipulated": str(REAL_MESH)}}},
        "gpu": {
            "device_ids": None, "min_free_mb": 0, "workers_per_device": 1,
            "max_workers": None, "threads_per_worker": 2,
            "acquire_timeout_s": 5.0, "poll_interval_s": 0.05,
        },
        "batch": batch,
        "publish": {"bucket_url": "hf://buckets/test/test"},
    }
    yaml_path = tmp_path / "profile.yaml"
    yaml_path.write_text(yaml.safe_dump(raw))
    return yaml_path


def _make_episodes(
    tmp_path: Path, n_ok: int, *, n_unparseable: int = 0
) -> tuple[DatagenProfile, Path, list[str]]:
    yaml_path = _write_profile_yaml(tmp_path)
    profile = DatagenProfile.from_yaml(yaml_path)
    uuids = []
    for i in range(n_ok):
        uuid = f"AUTOLab+ok{i:03d}+2024-01-01-00h-00m-00s"
        (profile.paths.droid_raw_root / uuid).mkdir(parents=True, exist_ok=True)
        _write_flows_h5(profile.paths.flows_h5(uuid), uuid, RESOLVABLE_TASK)
        uuids.append(uuid)
    for i in range(n_unparseable):
        uuid = f"AUTOLab+skip{i:03d}+2024-01-01-00h-00m-00s"
        (profile.paths.droid_raw_root / uuid).mkdir(parents=True, exist_ok=True)
        _write_flows_h5(profile.paths.flows_h5(uuid), uuid, UNPARSEABLE_TASK)
        uuids.append(uuid)
    return profile, yaml_path, uuids


def _thread_executor_factory(n: int) -> ThreadPoolExecutor:
    return ThreadPoolExecutor(max_workers=n)


def _make_runner(profile: DatagenProfile, yaml_path: Path, *, n_devices: int = 2) -> BatchRunner:
    from fpgm.config_datagen import GpuPoolConfig

    measured = {i: 50000 for i in range(n_devices)}
    gpu_pool = GpuWorkerPool(
        GpuPoolConfig(
            min_free_mb=0, workers_per_device=1, acquire_timeout_s=5.0, poll_interval_s=0.02
        ),
        query_fn=lambda: measured,
    )
    return BatchRunner(
        profile, yaml_path, gpu_pool=gpu_pool, executor_factory=_thread_executor_factory
    )


# --------------------------------------------------------------------------- #
# Tests
# --------------------------------------------------------------------------- #


class TestDiscoveryAndJobs:
    def test_discover_episode_uuids(self, tmp_path, pipeline_stub):
        pipeline_stub()
        profile, yaml_path, uuids = _make_episodes(tmp_path, n_ok=3)
        runner = _make_runner(profile, yaml_path)
        assert runner.discover_episode_uuids() == sorted(uuids)

    def test_build_jobs_defaults_camera_role_from_profile(self, tmp_path, pipeline_stub):
        pipeline_stub()
        profile, yaml_path, uuids = _make_episodes(tmp_path, n_ok=2)
        runner = _make_runner(profile, yaml_path)
        jobs = runner.build_jobs()
        assert all(j.camera_role == "ext1" for j in jobs)


class TestFullRun:
    def test_all_episodes_recorded_ok_and_sharded_with_per_episode_registry_release(
        self, tmp_path, pipeline_stub
    ):
        module = pipeline_stub()
        profile, yaml_path, uuids = _make_episodes(tmp_path, n_ok=6)
        # workers_per_device=1 -> concurrency=2 with 2 devices
        runner = _make_runner(profile, yaml_path, n_devices=2)

        summary = runner.run(mode="full")

        assert {e["uuid"] for e in summary.episodes} == set(uuids)
        assert all(e["status"] == "ok" for e in summary.episodes)
        assert sorted(module.call_log) == sorted(uuids)

        # This assertion was inverted deliberately, and the inversion is the
        # point of the test.
        #
        # It used to require ONE model construction per shard -- cross-episode
        # registry reuse. That contract was given up on measured evidence: in
        # a real worker process, episode 1 ok, episode 2 ok, episode 3 died at
        # "30.51 GiB in use ... 29.55 GiB allocated by PyTorch". Almost all of
        # it *allocated* rather than *reserved*, i.e. live references retained
        # across episodes, which no allocator flag or empty_cache() can
        # reclaim. ``_run_shard`` now releases the registry after every
        # episode, so constructions scale with episodes, not with shards.
        #
        # The count is asserted to *rise* rather than merely "not be 1", so
        # that silently restoring reuse (and the leak with it) fails here
        # rather than passing a weaker check.
        shard_ids = {e["shard_id"] for e in summary.episodes}
        assert len(shard_ids) == 2
        for shard_id in shard_ids:
            counts = sorted(
                e["model_construction_count"]
                for e in summary.episodes if e["shard_id"] == shard_id
            )
            assert counts == list(range(1, len(counts) + 1)), (
                "each episode must rebuild after the previous one's release, so a shard's "
                f"construction counts should be 1,2,3,... -- got {counts}"
            )

        # checkpoint actually written to disk
        assert summary.json_path.exists()
        on_disk = json.loads(summary.json_path.read_text())
        assert on_disk["n_ok"] == 6

    def test_failed_episode_status_and_gate_extraction(self, tmp_path, pipeline_stub):
        profile, yaml_path, uuids = _make_episodes(tmp_path, n_ok=2)
        pipeline_stub({uuids[0]: "fail"})
        runner = _make_runner(profile, yaml_path, n_devices=1)

        summary = runner.run(mode="full")
        by_uuid = {e["uuid"]: e for e in summary.episodes}
        assert by_uuid[uuids[0]]["status"] == "failed"
        assert by_uuid[uuids[1]]["status"] == "ok"


class TestSkipping:
    def test_unparseable_instruction_falls_back_to_motion_not_skipped(
        self, tmp_path, pipeline_stub
    ):
        """No episode is refused for an unparseable DROID instruction any more --
        see fpgm.datagen.episode_spec's module docstring. An episode whose
        instruction fails to parse now runs the (stubbed) pipeline with
        objects_from="motion" instead of being recorded "skipped" and never
        reaching the pipeline at all -- the inverse of this test's own old name.
        """
        module = pipeline_stub()
        profile, yaml_path, uuids = _make_episodes(tmp_path, n_ok=2, n_unparseable=2)
        runner = _make_runner(profile, yaml_path, n_devices=2)

        summary = runner.run(mode="full")
        by_status = {}
        for e in summary.episodes:
            by_status.setdefault(e["status"], []).append(e["uuid"])

        assert not by_status.get("skipped")
        assert len(by_status.get("ok", [])) == 4

        objects_from_by_uuid = {e["uuid"]: e.get("objects_from") for e in summary.episodes}
        unparseable_uuids = [u for u in uuids if "skip" in u]
        parseable_uuids = [u for u in uuids if "skip" not in u]
        assert unparseable_uuids and parseable_uuids  # sanity: fixture built both kinds
        assert all(objects_from_by_uuid[u] == "motion" for u in unparseable_uuids)
        assert all(objects_from_by_uuid[u] == "text" for u in parseable_uuids)

        # The pipeline IS invoked for every episode now, unparseable ones included.
        assert set(module.call_log) == set(uuids)


class TestTimeout:
    def test_hung_episode_ends_its_shard_and_defers_the_rest(self, tmp_path, pipeline_stub):
        """A timeout must end the shard, not merely skip one episode.

        This assertion was inverted deliberately. It used to require the shard
        to carry on to the next episode, which is what the code did -- and that
        behaviour was measured to be actively harmful: ``_run_with_timeout``
        abandons its worker thread (there is no portable way to kill a thread
        stuck in native CUDA code), so the thread keeps ~20 GB of GPU memory
        alive. On the real pilot, one timeout took down the three episodes
        after it: the next reported "reserved 31.78 GB" before doing any work,
        died inside SAM3 with a bare ``KeyError('previous_stages_out')``, and
        the one after that OOM'd in 26 s. Four episodes lost, looking like
        three unrelated bugs.

        Letting the worker *process* exit is the only reliable way to reclaim
        an abandoned thread's memory. The deferred episodes are not recorded as
        failures -- they were never attempted -- and the multi-pass campaign
        re-runs them in a clean process.
        """
        profile, yaml_path, uuids = _make_episodes(tmp_path, n_ok=2)
        pipeline_stub({uuids[0]: "hang"})
        yaml_path = _write_profile_yaml(tmp_path, episode_timeout_s=0.3)
        profile = DatagenProfile.from_yaml(yaml_path)
        runner = _make_runner(profile, yaml_path, n_devices=1)

        t0 = time.time()
        summary = runner.run(mode="full")
        elapsed = time.time() - t0

        by_uuid = {e["uuid"]: e for e in summary.episodes}
        assert by_uuid[uuids[0]]["status"] == "failed"
        error_msg = by_uuid[uuids[0]]["error"].lower()
        assert "timeout" in error_msg or "exceeded" in error_msg

        # The episode behind the hung one is DEFERRED, not run in the poisoned
        # process. Either it is absent from the summary entirely, or the
        # crashed-shard path recorded it as failed -- what must never happen is
        # it being reported "ok" from a process holding a zombie allocation.
        assert by_uuid.get(uuids[1], {}).get("status") != "ok", (
            "an episode queued behind a timeout must not be executed in the "
            "same, memory-poisoned worker process"
        )
        assert elapsed < 10, "the hung episode's 30s sleep must not have stalled the batch"


class TestResume:
    def test_already_ok_episode_is_not_reprocessed(self, tmp_path, pipeline_stub):
        module = pipeline_stub()
        profile, yaml_path, uuids = _make_episodes(tmp_path, n_ok=2)
        runner = _make_runner(profile, yaml_path, n_devices=1)
        runner.run(mode="full")
        assert sorted(module.call_log) == sorted(uuids)

        module.call_log.clear()
        runner2 = _make_runner(profile, yaml_path, n_devices=1)
        # reuse the same summary file on disk (BatchSummary.load_or_new)
        summary2 = runner2.run(mode="full")
        assert module.call_log == []  # nothing reprocessed
        assert {e["uuid"] for e in summary2.episodes} == set(uuids)


class TestPilotPolicy:
    def test_pilot_mode_stops_without_running_remainder(self, tmp_path, pipeline_stub):
        module = pipeline_stub()
        profile, yaml_path, uuids = _make_episodes(tmp_path, n_ok=4)
        yaml_path = _write_profile_yaml(tmp_path, pilot_fraction=0.5, pilot_min_episodes=1)
        profile = DatagenProfile.from_yaml(yaml_path)
        runner = _make_runner(profile, yaml_path, n_devices=1)

        summary = runner.run(mode="pilot")
        assert len(module.call_log) == 2  # ceil(0.5*4) == 2
        assert len(summary.episodes) == 2
        assert summary.meta["pilot_decision"]["pilot_verdict"].startswith("stopped")

    def test_auto_mode_continues_when_pilot_passes(self, tmp_path, pipeline_stub):
        module = pipeline_stub()
        profile, yaml_path, uuids = _make_episodes(tmp_path, n_ok=4)
        yaml_path = _write_profile_yaml(tmp_path, pilot_fraction=0.5, pilot_min_episodes=1,
                                          pilot_pass_threshold=0.5)
        profile = DatagenProfile.from_yaml(yaml_path)
        runner = _make_runner(profile, yaml_path, n_devices=1)

        summary = runner.run(mode="auto")
        assert len(module.call_log) == 4  # pilot (2) + remainder (2)
        assert all(e["status"] == "ok" for e in summary.episodes)
        assert summary.meta["pilot_decision"]["pilot_verdict"].startswith("passed")

    def test_auto_mode_stops_when_pilot_fails_threshold(self, tmp_path, pipeline_stub):
        profile, yaml_path, uuids = _make_episodes(tmp_path, n_ok=4)
        # first 2 (the pilot, at pilot_fraction=0.5) fail -> 0% pass rate
        module = pipeline_stub({uuids[0]: "fail", uuids[1]: "fail"})
        yaml_path = _write_profile_yaml(tmp_path, pilot_fraction=0.5, pilot_min_episodes=1,
                                          pilot_pass_threshold=0.6)
        profile = DatagenProfile.from_yaml(yaml_path)
        runner = _make_runner(profile, yaml_path, n_devices=1)

        summary = runner.run(mode="auto")
        assert len(module.call_log) == 2  # remainder never ran
        assert summary.meta["pilot_decision"]["pilot_verdict"].startswith("FAILED")
        assert summary.meta["pilot_decision"]["pilot_pass_rate"] == 0.0


class TestPruneIntermediates:
    def test_deletes_large_buffers_only_after_ok_and_reports_bytes(self, tmp_path, pipeline_stub):
        pipeline_stub()
        profile, yaml_path, uuids = _make_episodes(tmp_path, n_ok=1)
        runner = _make_runner(profile, yaml_path, n_devices=1)

        master_dir = profile.paths.master_dir(uuids[0], "cam1")
        buf_dir = master_dir / "robot_buffers"
        buf_dir.mkdir(parents=True)
        (buf_dir / "depth_mm.npy").write_bytes(b"x" * 1000)
        (buf_dir / "seg.npy").write_bytes(b"y" * 500)
        (buf_dir / "meta.json").write_text("{}")  # must survive pruning

        summary = runner.run(mode="full")
        entry = summary.episodes[0]
        assert entry["status"] == "ok"
        assert entry["bytes_reclaimed"] == 1500
        assert not (buf_dir / "depth_mm.npy").exists()
        assert not (buf_dir / "seg.npy").exists()
        assert (buf_dir / "meta.json").exists()