onf/tests/test_run_sr.py

"""Tests for scripts/run_sr.py — the Python success-rate eval driver.

NO GPU, no real eval: every test either exercises the driver's pure logic (yaml parsing, the
modes.sh parser, gate math, LPT scheduling, resume/completeness decisions) against fabricated
fixtures, or monkeypatches the one external dependency that would otherwise require hardware
(torch.cuda) / a subprocess (subprocess.Popen, forced to raise so a --dry-run that
accidentally launches something fails loudly instead of silently doing the wrong thing).

scripts/run_sr.py is not a package module (it lives outside src/), so it is loaded by file
path (see _load_run_sr below).
"""

from __future__ import annotations

import dataclasses
import importlib.util
import json
import math
import os
import re
import subprocess
import sys
from pathlib import Path

import pytest
import torch

REPO = Path(__file__).resolve().parents[1]


def _load_run_sr():
    path = REPO / "scripts" / "run_sr.py"
    spec = importlib.util.spec_from_file_location("run_sr", path)
    mod = importlib.util.module_from_spec(spec)
    # register BEFORE exec: run_sr.py's dataclasses (with from __future__ import annotations)
    # resolve their type hints via sys.modules[cls.__module__] at class-body execution time.
    sys.modules[spec.name] = mod
    spec.loader.exec_module(mod)
    return mod


run_sr = _load_run_sr()

# One shared LadderConfig for every test that only READS the ladder. Tests that need a different
# results root build their own with _config(results_root=...).
CONFIG = run_sr.LadderConfig.load()
LADDER = CONFIG.doc
PATHS = CONFIG.paths
SUITES = CONFIG.suites
AXES = CONFIG.scorers.axes


def _config(results_root: Path | None = None) -> "run_sr.LadderConfig":
    return run_sr.LadderConfig.load(results_root=results_root)


def _mkcell(name: str, expected_n: int, **kw) -> "run_sr.Cell":
    """A test cell whose tag is derived the way production derives it. name only needs to make the
    cell DISTINCT from its neighbours, so it rides in as the variant suffix; read the real tag off
    cell.tag, never off name."""
    kw.setdefault("rung", "t")
    kw.setdefault("suite", "object")
    kw.setdefault("axis", "Objects_Layout")
    kw.setdefault("mode", "base")
    kw.setdefault("policy", "t")
    kw.setdefault("bench", CONFIG.scorers.bench(kw["suite"]))
    if "comparators" in kw:
        kw["comparators"] = tuple(kw["comparators"])
    return run_sr.Cell(variant=name, expected_n=expected_n, **kw)


def _score(cell, successes: int, actual_n: int) -> "run_sr.CellScore":
    return run_sr.CellScore(cell=cell, count=run_sr.DirCount(successes, actual_n))


def _verdict(score, allow_partial: bool) -> "run_sr.GateVerdict":
    """The production verdict path: completeness first, then the cell's own gate."""
    return run_sr.CellScorer(CONFIG).verdict(score, allow_partial)


def _iter_declared_cells():
    """(rung_name, cell_doc) for every hand-declared cell (skips R4's auto_generate: full_grid)."""
    for rung_name, rung_doc in LADDER["rungs"].items():
        if rung_doc.get("auto_generate"):
            continue
        for cell_doc in rung_doc["cells"]:
            yield rung_name, cell_doc


# ======================================================================================================
# ladder yaml: parses, every cell names a real suite/axis, expected counts match the task's table
# ======================================================================================================
def test_ladder_yaml_parses_and_has_all_rungs():
    # R1 (smoke) and R4 (the full 4-suite x 7-axis grid). A wider rung set here
    # (R2/R3/R5/E*/F*/G*/H* ...) was the pre-reduction ablation ladder; asserting the CURRENT set
    # pins this down instead of letting the test silently pass against a stale, larger yaml.
    assert set(LADDER["rungs"]) == {"R1", "R4"}


def test_ladder_paths_block_declares_every_artifact_the_driver_drives():
    """LadderPaths is the single place the driver learns where its external artifacts live, so a
    missing entry must fail loudly at load rather than as a KeyError deep inside a launch."""
    with pytest.raises(KeyError, match="paths block is missing"):
        run_sr.LadderPaths.from_doc({"paths": {"score_py": "x"}})


def test_every_declared_cell_names_a_real_suite_and_axis():
    for rung_name, cell_doc in _iter_declared_cells():
        assert cell_doc["suite"] in SUITES, f"{rung_name}: unknown suite {cell_doc['suite']!r}"
        assert cell_doc["axis"] in AXES, f"{rung_name}: unknown axis {cell_doc['axis']!r}"


def test_every_declared_cell_mode_is_a_known_modes_sh_recipe():
    for rung_name, cell_doc in _iter_declared_cells():
        mode = cell_doc.get("mode")
        if mode is not None:
            assert mode in CONFIG.modes.modes, f"{rung_name} {cell_doc['tag']}: unknown mode {mode!r}"


def test_expected_episode_counts_match_the_task_table():
    expected = {
        "object": {"Background_Textures": 248, "Camera_Viewpoints": 396, "Language_Instructions": 354,
                   "Light_Conditions": 296, "Objects_Layout": 403, "Robot_Initial_States": 398, "Sensor_Noise": 422},
        "spatial": {"Background_Textures": 266, "Camera_Viewpoints": 376, "Language_Instructions": 390,
                    "Light_Conditions": 291, "Objects_Layout": 385, "Robot_Initial_States": 350, "Sensor_Noise": 351},
        "goal": {"Background_Textures": 281, "Camera_Viewpoints": 408, "Language_Instructions": 410,
                 "Light_Conditions": 278, "Objects_Layout": 425, "Robot_Initial_States": 409, "Sensor_Noise": 379},
        "long": {"Background_Textures": 289, "Camera_Viewpoints": 419, "Language_Instructions": 382,
                 "Light_Conditions": 273, "Objects_Layout": 312, "Robot_Initial_States": 393, "Sensor_Noise": 449},
    }
    assert LADDER["expected_episodes"] == expected


def test_unfiltered_cell_expected_n_matches_expected_episodes_table():
    """A cell without its own instance-id filter must inherit the full suite x axis denominator."""
    for rung_name, cell_doc in _iter_declared_cells():
        if cell_doc.get("filter") not in (None,) and cell_doc["axis"] != "Light_Conditions":
            continue   # e.g. R1's 20-instance smoke filter legitimately shrinks expected_n
        suite, axis = cell_doc["suite"], cell_doc["axis"]
        assert cell_doc["expected_n"] == LADDER["expected_episodes"][suite][axis], (rung_name, cell_doc["tag"])


def test_suites_yaml_axes_match_score_py_axes():
    doc = run_sr.load_yaml(PATHS.suites_yaml)
    spaced_to_underscored = [a.replace(" ", "_") for a in doc["axes"]]
    assert spaced_to_underscored == AXES


# ======================================================================================================
# evals/common/modes.sh parser — env composition must match the CONTRACT below (this pins the
# PARSER's correctness; the parser reads modes.sh directly, so it can never itself drift from the
# file -- this guards against a parsing regression instead).
# ======================================================================================================
MODES_SH_CONTRACT = {
    "base": {},
}


@pytest.mark.parametrize("mode", sorted(MODES_SH_CONTRACT))
def test_mode_env_matches_modes_sh_contract(mode):
    assert CONFIG.modes.env_for(mode) == MODES_SH_CONTRACT[mode]


def test_mode_env_unknown_mode_raises():
    with pytest.raises(ValueError, match="unknown MODE"):
        CONFIG.modes.env_for("not_a_real_mode")


def test_mode_env_returns_a_fresh_copy_each_call():
    """A caller that mutates its recipe must not corrupt the registry for the next cell."""
    first = CONFIG.modes.env_for("base")
    first["SOMETHING"] = "999"
    assert "SOMETHING" not in CONFIG.modes.env_for("base")


# ======================================================================================================
# cell resolution / env composition
# ======================================================================================================
def test_compose_cell_env_extra_env_override_wins():
    """A cell's per-cell env: (yaml) / extra_env override must win over the builder's own defaults."""
    cell = _mkcell("t", 10, extra_env={"LOAD_VISION": "1"})
    env = run_sr.EnvBuilder(CONFIG).for_cell(cell)
    assert env["LOAD_VISION"] == "1"         # extra_env override, not the builder's own default


def test_compose_cell_env_sets_qndf_dir_per_suite():
    builder = run_sr.EnvBuilder(CONFIG)
    for cell in CONFIG.cells_for_rung("R4"):
        env = builder.for_cell(cell)
        assert env["QNDF_DIR"] == str(SUITES[cell.suite].fwm_dir)


def test_resolve_graph_dir_matches_onf_config_paths_graph():
    from onf.config import default_paths
    for suite in SUITES:
        assert run_sr.resolve_graph_dir(suite) == default_paths().graph(suite)


def test_graph_artifact_status_reports_missing_files_not_ok(tmp_path):
    status = run_sr.graph_artifact_status(tmp_path / "no_such_graph_dir")
    assert not status.ok
    assert status.graph_hash is None
    assert status.g_head_mtime is None
    assert len(status.missing) == 2   # g_head.npz + g_track.npz


def test_graph_artifact_status_ok_when_both_artifacts_present(tmp_path):
    import numpy as np
    from onf.graph.core import schema as graph_schema

    d = tmp_path / "graph_dir"
    d.mkdir()
    np.savez(d / graph_schema.HEAD_NPZ, graph_hash=np.array("deadbeef"))
    (d / graph_schema.TRACK_NPZ).write_bytes(b"")   # presence is all graph_artifact_status checks
    status = run_sr.graph_artifact_status(d)
    assert status.ok
    assert status.missing == ()
    assert status.graph_hash == "deadbeef"
    assert status.g_head_mtime is not None
    assert status.as_record()["graph_hash"] == "deadbeef"


# ======================================================================================================
# 1f: log-directory naming -- shared by both drivers, UTC-stamped, encodes policy/mode/suite/axis so
# two cells that differ only by one of those can never collide/truncate each other's shard logs.
# ======================================================================================================
def test_make_log_dir_embeds_policy_mode_suite_axis_and_normalizes_axis_spaces():
    d = run_sr.make_log_dir("stablevla", "sentinel", "long", "Sensor Noise", root=Path("/x"))
    # the log dir name is exactly the cell tag, so logs and results are findable from each other
    assert d.parent == Path("/x/logs")
    assert d.name.endswith("_stablevla_sentinel_long_Sensor_Noise")
    # UTC timestamp prefix: YYYYmmdd_HHMMSS_...
    assert re.match(r"^\d{8}_\d{6}_stablevla_sentinel_long_Sensor_Noise$", d.name), d.name


# ======================================================================================================
# 1e: run.json provenance -- every result dir gets one, naming policy/mode/suite/axis/expected_n/
# resolved graph dir/graph_hash/g_head mtime/git commit/the full mode env.
# ======================================================================================================
def test_write_run_json_round_trips_expected_fields(tmp_path):
    out = run_sr.write_run_json(
        tmp_path / "result_dir", policy="stablevla", mode="sentinel", suite="long",
        axis="Sensor_Noise", expected_n=449, graph_dir="/g", graph_hash="abc123",
        g_head_mtime="2026-08-01T00:00:00+00:00", git_commit="deadbeef", mode_env={"SENTINEL": "1"},
    )
    assert out == tmp_path / "result_dir" / "run.json"
    payload = json.loads(out.read_text())
    assert payload["policy"] == "stablevla"
    assert payload["mode"] == "sentinel"
    assert payload["suite"] == "long"
    assert payload["axis"] == "Sensor_Noise"
    assert payload["expected_n"] == 449
    assert payload["graph_dir"] == "/g"
    assert payload["graph_hash"] == "abc123"
    assert payload["g_head_mtime"] == "2026-08-01T00:00:00+00:00"
    assert payload["git_commit"] == "deadbeef"
    assert payload["mode_env"] == {"SENTINEL": "1"}
    assert "written_at" in payload


def test_light_conditions_axis_auto_resolves_per_suite_filter():
    cell = CONFIG.cell_from_doc("adhoc", {"suite": "goal", "axis": "Light_Conditions", "mode": "base",
                                           "tag": "t", "expected_n": 278})
    assert cell.filter_file == SUITES["goal"].light_filter


def test_non_light_axis_has_no_filter_by_default():
    cell = CONFIG.cell_from_doc("adhoc", {"suite": "goal", "axis": "Background_Textures", "mode": "base",
                                           "tag": "t", "expected_n": 281})
    assert cell.filter_file is None


def test_cell_category_value_is_axis_with_spaces():
    cell = _mkcell("t", 10, axis="Robot_Initial_States")
    assert cell.category_value == "Robot Initial States"


def test_cell_result_dir_layout():
    cell = _mkcell("sr_test", 10, suite="goal", axis="Background_Textures")
    d = cell.result_dir(Path("/results"))
    assert d == Path("/results/plus_libero_goal/Background_Textures") / cell.tag


# ======================================================================================================
# 1a: a rung tag must bake in its resolved MODE, and overriding --mode must re-key the tag's mode
# suffix along with it -- NOT leave the tag naming the original mode. That mismatch is the actual
# wrong-number hazard from the task: --rung R4 then --rung R4 --mode base --resume used to write
# into the SAME result directory (tag had no mode in it) and the ledger's tag-only lookup then reported
# the FIRST run's numbers under the SECOND run's mode name.
# ======================================================================================================
def test_cell_tag_embeds_policy_mode_suite_axis():
    assert run_sr.cell_tag("stablevla", "base", "object", "Robot_Initial_States") \
        == "stablevla_base_object_Robot_Initial_States"
    # the optional variant separates a filtered smoke pool from the full axis it subsets
    assert run_sr.cell_tag("gr00t", "probe", "long", "Sensor_Noise", "smoke") \
        == "gr00t_probe_long_Sensor_Noise_smoke"


def test_r1_cell_tag_bakes_in_its_declared_mode():
    cells = CONFIG.cells_for_rung("R1")
    assert cells[0].mode == "base"
    assert "_base_" in cells[0].tag


def test_r4_generated_tags_bake_in_their_axis_default_mode():
    for c in CONFIG.cells_for_rung("R4"):
        assert f"_{c.mode}_" in c.tag, c.tag


def test_r1_mode_override_re_keys_the_tag_not_just_the_mode():
    cells = CONFIG.cells_for_rung("R1")
    assert cells[0].mode == "base"
    overridden = [c.with_mode("probe") for c in cells]
    assert overridden[0].mode == "probe"
    assert overridden[0].tag != cells[0].tag                 # THE FIX: tag changes with mode
    assert "_probe_" in overridden[0].tag
    assert "_base_" not in overridden[0].tag                  # no stale mode fragment left behind
    # everything else about the cell (expected_n, comparators, gate, filter) survives the swap
    assert overridden[0].expected_n == cells[0].expected_n
    assert overridden[0].comparators == cells[0].comparators


def test_cell_with_mode_round_trips_through_two_overrides():
    """Overriding mode twice (e.g. two separate CLI invocations against the same rung) must never
    leave a stale mode fragment baked into the tag from an EARLIER override."""
    cell = run_sr.Cell(rung="t", suite="object", axis="Robot_Initial_States", bench="libero_object",
                        mode="probe", expected_n=398)
    once = cell.with_mode("other")
    assert once.tag == "stablevla_other_object_Robot_Initial_States"
    twice = once.with_mode("base")
    assert twice.tag == "stablevla_base_object_Robot_Initial_States"


def test_cell_is_frozen_so_a_tag_cannot_drift_from_its_mode():
    cell = _mkcell("t", 10)
    with pytest.raises(dataclasses.FrozenInstanceError):
        cell.mode = "sentinel"


def test_resolve_cells_with_mode_override_produces_mode_bearing_tags():
    """End-to-end through the same path main() uses for --rung ... --mode ...."""
    default_cells = _driver(["--rung", "R1", "--gpus", "0"])._resolve_cells()
    assert "_base_" in default_cells[0].tag

    overridden_cells = _driver(["--rung", "R1", "--gpus", "0", "--mode", "probe"])._resolve_cells()
    assert overridden_cells[0].mode == "probe"
    assert "_probe_" in overridden_cells[0].tag
    assert overridden_cells[0].tag != default_cells[0].tag


def _driver(argv, ledger=None) -> "run_sr.SRDriver":
    return run_sr.SRDriver(run_sr.build_argparser().parse_args(argv), ledger=ledger)


# ======================================================================================================
# comparator / gate math
# ======================================================================================================
def test_comparator_rate():
    c = run_sr.Comparator(label="base", tag="x", successes=244, n=281)
    assert c.rate == pytest.approx(100 * 244 / 281)


def test_comparator_rate_zero_n_is_nan():
    c = run_sr.Comparator(label="base", tag="x", successes=0, n=0)
    assert math.isnan(c.rate)


def test_gate_from_dict_rejects_an_unknown_kind():
    """GateKind turns a typo'd yaml gate into a load-time error instead of a silent no-op."""
    with pytest.raises(ValueError):
        run_sr.Gate.from_dict({"kind": "not_a_real_gate"})


def test_gate_max_regress_pp_pass_and_fail():
    comp = run_sr.Comparator(label="base", tag="b", successes=50, n=100)   # 50.0%
    cell = _mkcell("t", 100, gate=run_sr.Gate(kind=run_sr.GateKind.MAX_REGRESS_PP, vs="base",
                                               max_regress_pp=2.0), comparators=[comp])
    assert _verdict(_score(cell, 49, 100), allow_partial=False).passed
    assert _verdict(_score(cell, 70, 100), allow_partial=False).passed
    assert not _verdict(_score(cell, 40, 100), allow_partial=False).passed


def test_gate_exact_match():
    comp = run_sr.Comparator(label="base", tag="b", successes=302, n=403)
    cell = _mkcell("t", 403, gate=run_sr.Gate(kind=run_sr.GateKind.EXACT_MATCH, vs="base"),
                   comparators=[comp])
    assert _verdict(_score(cell, 302, 403), allow_partial=False).passed
    assert not _verdict(_score(cell, 301, 403), allow_partial=False).passed


def test_gate_min_successes():
    comp = run_sr.Comparator(label="base", tag="b", successes=9, n=11)
    cell = _mkcell("t", 11, gate=run_sr.Gate(kind=run_sr.GateKind.MIN_SUCCESSES, vs="base",
                                              min_successes=8), comparators=[comp])
    assert _verdict(_score(cell, 8, 11), allow_partial=False).passed
    assert not _verdict(_score(cell, 7, 11), allow_partial=False).passed


def test_gate_missing_comparator_label_skips_not_fails():
    cell = _mkcell("t", 10, gate=run_sr.Gate(kind=run_sr.GateKind.MAX_REGRESS_PP, vs="nonexistent",
                                              max_regress_pp=2.0), comparators=[])
    verdict = _verdict(_score(cell, 5, 10), allow_partial=False)
    assert verdict.passed and "SKIP" in verdict.message


def test_gate_missing_comparator_skips_even_a_min_successes_gate():
    """A comparator's absence means the ladder has nothing to compare against yet -- failing the cell
    for that would make a fresh checkout's first run red for a reason unrelated to the run."""
    cell = _mkcell("t", 10, gate=run_sr.Gate(kind=run_sr.GateKind.MIN_SUCCESSES, vs="base",
                                              min_successes=999), comparators=[])
    verdict = _verdict(_score(cell, 0, 10), allow_partial=False)
    assert verdict.passed and "SKIP" in verdict.message


def test_gate_none_always_passes():
    cell = _mkcell("t", 10, gate=None)
    assert _verdict(_score(cell, 1, 10), allow_partial=False).passed


# ======================================================================================================
# completeness: mp4 count vs expected_n
# ======================================================================================================
def test_completeness_rule_marks_short_cell_incomplete_and_excludes_by_default():
    comp = run_sr.Comparator(label="base", tag="b", successes=50, n=100)
    cell = _mkcell("t", 100, gate=run_sr.Gate(kind=run_sr.GateKind.MAX_REGRESS_PP, vs="base",
                                               max_regress_pp=2.0), comparators=[comp])
    verdict = _verdict(_score(cell, 90, 90), allow_partial=False)
    assert not verdict.passed
    assert "INCOMPLETE" in verdict.message


def test_completeness_rule_allow_partial_falls_through_to_the_real_gate():
    comp = run_sr.Comparator(label="base", tag="b", successes=50, n=100)
    cell = _mkcell("t", 100, gate=run_sr.Gate(kind=run_sr.GateKind.MAX_REGRESS_PP, vs="base",
                                               max_regress_pp=2.0), comparators=[comp])
    assert _verdict(_score(cell, 90, 90), allow_partial=True).passed


def test_score_cell_against_fabricated_results_tree_marks_short_dir_incomplete(tmp_path):
    results_root = tmp_path / "results"
    cell = _mkcell("sr_test", 10)
    d = results_root / "plus_libero_object" / "Objects_Layout" / cell.tag
    d.mkdir(parents=True)
    for i in range(3):
        (d / f"rollout_task_episode{i}_success.mp4").write_bytes(b"x")
    score = run_sr.CellScorer(_config(results_root)).score(cell)
    assert (score.successes, score.actual_n, score.complete) == (3, 3, False)


def test_score_cell_missing_dir_is_zero_not_a_crash(tmp_path):
    cell = _mkcell("sr_missing", 10)
    score = run_sr.CellScorer(_config(tmp_path / "results")).score(cell)
    assert (score.successes, score.actual_n, score.complete) == (0, 0, False)


# ======================================================================================================
# ledger — append/read round-trip
# ======================================================================================================
def _rec(**kw) -> "run_sr.LedgerRecord":
    kw.setdefault("suite", "object")
    kw.setdefault("bench", "libero_object")
    kw.setdefault("axis", "Objects_Layout")
    kw.setdefault("mode", "base")
    return run_sr.LedgerRecord(**kw)


def test_ledger_append_read_round_trip(tmp_path):
    ledger = run_sr.Ledger(tmp_path / "ledger.jsonl")
    ledger.append(_rec(tag="a", status=run_sr.RunStatus.RUNNING, ts="t0"))
    ledger.append(_rec(tag="a", status=run_sr.RunStatus.COMPLETE, ts="t1", actual_n=10))
    records = ledger.read()
    assert [(r.tag, r.status, r.actual_n) for r in records] == [
        ("a", run_sr.RunStatus.RUNNING, None),
        ("a", run_sr.RunStatus.COMPLETE, 10),
    ]
    assert ledger.latest("a").actual_n == 10
    assert ledger.latest("a", records=records).actual_n == 10
    assert ledger.latest("nope") is None


def test_ledger_record_stamps_ts_when_the_caller_leaves_it_blank():
    payload = _rec(tag="a", status=run_sr.RunStatus.RUNNING).to_dict()
    assert payload["ts"]
    assert payload["status"] == "running"       # the VALUE, not the enum repr


def test_ledger_record_from_dict_tolerates_an_older_row(tmp_path):
    """Rows written before a schema change must still read: unknown keys are dropped, missing
    identity fields become "", and an unrecognised status reads as crashed rather than crashing."""
    row = {"tag": "x", "status": "some_status_we_retired", "who_knows": 1, "actual_n": 3}
    rec = run_sr.LedgerRecord.from_dict(row)
    assert rec.tag == "x" and rec.suite == "" and rec.actual_n == 3
    assert rec.status is run_sr.RunStatus.CRASHED


# ======================================================================================================
# 1a (ledger end): Ledger.latest must ALSO compare mode when given one, not tag alone -- the other
# half of the --resume wrong-number fix (tags now bake mode in via cell_tag, but the ledger lookup
# itself is defense in depth against a record written before that fix, or any other tag collision).
# ======================================================================================================
def test_latest_matching_record_requires_matching_mode_when_given(tmp_path):
    ledger = run_sr.Ledger(tmp_path / "l.jsonl")
    records = [
        _rec(tag="x", mode="base", status=run_sr.RunStatus.COMPLETE, successes=1),
        _rec(tag="x", mode="sentinel", status=run_sr.RunStatus.COMPLETE, successes=2),
    ]
    assert ledger.latest("x", "base", records).successes == 1
    assert ledger.latest("x", "sentinel", records).successes == 2
    # tag-only (mode=None, the default) still returns the LAST matching record, unfiltered by mode --
    # exact prior behaviour, so a caller that doesn't pass mode is unaffected.
    assert ledger.latest("x", None, records).successes == 2


def test_latest_matching_record_mode_mismatch_is_not_a_match(tmp_path):
    ledger = run_sr.Ledger(tmp_path / "l.jsonl")
    records = [_rec(tag="x", mode="base", status=run_sr.RunStatus.COMPLETE)]
    assert ledger.latest("x", "sentinel", records) is None


def test_read_ledger_missing_file_returns_empty(tmp_path):
    assert run_sr.Ledger(tmp_path / "nope.jsonl").read() == []


def test_print_ledger_empty(tmp_path, capsys):
    rc = _driver(["--ledger"], ledger=run_sr.Ledger(tmp_path / "empty.jsonl")).run()
    assert rc == run_sr.ExitCode.OK
    assert "empty" in capsys.readouterr().out


def test_print_ledger_with_records(tmp_path, capsys):
    ledger = run_sr.Ledger(tmp_path / "ledger.jsonl")
    ledger.append(_rec(ts="t0", tag="x", status=run_sr.RunStatus.COMPLETE, actual_n=5,
                       expected_n=5, successes=4, rate=80.0, wall_s=12.3, episodes_per_min=5.0,
                       gpu=0))
    rc = _driver(["--ledger"], ledger=ledger).run()
    assert rc == run_sr.ExitCode.OK
    out = capsys.readouterr().out
    assert "x" in out and "complete" in out


# ======================================================================================================
# resume: ResumePlanner.plan is a PURE decision; .apply does the actual purge.
# ======================================================================================================
def _planner(results_root: Path, ledger_path: Path) -> "run_sr.ResumePlanner":
    return run_sr.ResumePlanner(_config(results_root), run_sr.Ledger(ledger_path))


def test_plan_cell_action_no_existing_dir_runs(tmp_path):
    cell = _mkcell("x", 5)
    plan = _planner(tmp_path, tmp_path / "l.jsonl").plan(cell, resume=False, records=[])
    assert plan.action is run_sr.CellAction.RUN


def test_plan_cell_action_existing_dir_without_resume_errors(tmp_path):
    cell = _mkcell("x", 5)
    d = tmp_path / "plus_libero_object" / "Objects_Layout" / cell.tag
    d.mkdir(parents=True)
    (d / "rollout_t_episode0_success.mp4").write_bytes(b"x")
    plan = _planner(tmp_path, tmp_path / "l.jsonl").plan(cell, resume=False, records=[])
    assert plan.action is run_sr.CellAction.ERROR


def test_plan_cell_action_resume_skips_ledger_confirmed_complete(tmp_path):
    cell = _mkcell("x", 5)   # _mkcell defaults mode="base" -- the ledger record's mode must match it.
    d = tmp_path / "plus_libero_object" / "Objects_Layout" / cell.tag
    d.mkdir(parents=True)
    for i in range(5):
        (d / f"rollout_t_episode{i}_success.mp4").write_bytes(b"x")
    records = [_rec(tag=cell.tag, mode="base", status=run_sr.RunStatus.COMPLETE)]
    plan = _planner(tmp_path, tmp_path / "l.jsonl").plan(cell, resume=True, records=records)
    assert plan.action is run_sr.CellAction.SKIP
    assert d.exists()   # untouched


def test_plan_cell_action_resume_does_not_skip_on_a_different_modes_ledger_record(tmp_path):
    """The 1a regression this whole fix targets: a complete ledger record for the SAME tag but a
    DIFFERENT mode (e.g. left over from a --rung R4 run, now re-invoked as --rung R4 --mode base)
    must NOT be read as confirming completeness for this cell's mode -- that is exactly how the first
    run's numbers used to get silently reported under the second run's mode name."""
    cell = _mkcell("x", 5, mode="base")
    d = tmp_path / "plus_libero_object" / "Objects_Layout" / cell.tag
    d.mkdir(parents=True)
    for i in range(5):
        (d / f"rollout_t_episode{i}_success.mp4").write_bytes(b"x")
    records = [_rec(tag=cell.tag, mode="sentinel", status=run_sr.RunStatus.COMPLETE)]  # different mode
    plan = _planner(tmp_path, tmp_path / "l.jsonl").plan(cell, resume=True, records=records)
    assert plan.action is run_sr.CellAction.PURGE   # NOT skip


def test_plan_cell_action_resume_purges_partial_dir_despite_stale_complete_ledger(tmp_path):
    """The hazard from the task: a ledger record claiming 'complete' must NOT be trusted blindly --
    the mp4 count on disk is re-verified against expected_n every time."""
    cell = _mkcell("x", 5)
    d = tmp_path / "plus_libero_object" / "Objects_Layout" / cell.tag
    d.mkdir(parents=True)
    (d / "rollout_t_episode0_success.mp4").write_bytes(b"x")   # only 1 of expected 5
    records = [_rec(tag=cell.tag, status=run_sr.RunStatus.COMPLETE)]   # stale/wrong ledger claim
    planner = _planner(tmp_path, tmp_path / "l.jsonl")
    plan = planner.plan(cell, resume=True, records=records)
    assert plan.action is run_sr.CellAction.PURGE
    assert d.exists()   # plan is pure -- nothing deleted yet
    planner.apply(plan)
    assert not d.exists()


def test_plan_cell_action_resume_without_ledger_record_purges(tmp_path):
    cell = _mkcell("x", 5)
    d = tmp_path / "plus_libero_object" / "Objects_Layout" / cell.tag
    d.mkdir(parents=True)
    (d / "rollout_t_episode0_success.mp4").write_bytes(b"x")
    plan = _planner(tmp_path, tmp_path / "l.jsonl").plan(cell, resume=True, records=[])
    assert plan.action is run_sr.CellAction.PURGE


def test_apply_cell_action_is_a_noop_for_run_and_skip(tmp_path):
    cell = _mkcell("x", 5)
    d = tmp_path / "plus_libero_object" / "Objects_Layout" / cell.tag
    d.mkdir(parents=True)
    (d / "rollout_t_episode0_success.mp4").write_bytes(b"x")
    planner = _planner(tmp_path, tmp_path / "l.jsonl")
    planner.apply(run_sr.CellPlan(cell, run_sr.CellAction.RUN, ""))
    planner.apply(run_sr.CellPlan(cell, run_sr.CellAction.SKIP, ""))
    assert d.exists()


# ======================================================================================================
# LPT scheduling
# ======================================================================================================
def _assign(cells, gpus):
    return run_sr.Schedule.build(cells, gpus, CONFIG, []).assignments


def test_lpt_schedule_exact_greedy_assignment():
    cells = [_mkcell("a", 100), _mkcell("b", 90), _mkcell("c", 50), _mkcell("d", 10), _mkcell("e", 10)]
    assign = _assign(cells, [0, 1])
    assert [c.variant for c in assign[0]] == ["a", "d", "e"]     # 100 + 10 + 10 = 120
    assert [c.variant for c in assign[1]] == ["b", "c"]          # 90 + 50 = 140
    totals = {g: sum(c.expected_n for c in cs) for g, cs in assign.items()}
    assert totals == {0: 120, 1: 140}


def test_lpt_schedule_processes_largest_first_on_a_single_gpu():
    assign = _assign([_mkcell("small", 1), _mkcell("big", 1000)], [0])
    assert [c.variant for c in assign[0]] == ["big", "small"]


def test_lpt_schedule_covers_every_cell_exactly_once():
    cells = [_mkcell(str(i), i + 1) for i in range(9)]
    assign = _assign(cells, [0, 1, 2])
    scheduled = [c.variant for cs in assign.values() for c in cs]
    assert sorted(scheduled, key=int) == sorted((c.variant for c in cells), key=int)


def test_schedule_makespan_is_the_busiest_track():
    cells = [_mkcell("a", 600), _mkcell("b", 60)]
    sched = run_sr.Schedule.build(cells, [0, 1], CONFIG, [])
    epm = CONFIG.defaults.episodes_per_min
    assert sched.makespan_min == pytest.approx(600 / epm)
    assert {round(v, 6) for v in sched.per_gpu_minutes.values()} == {
        round(600 / epm, 6), round(60 / epm, 6)}


# ======================================================================================================
# rolling episodes/min
# ======================================================================================================
def test_rolling_episodes_per_min_uses_ledger_history():
    cell = _mkcell("t", 100, suite="goal", axis="Background_Textures", mode="sentinel")
    records = [
        _rec(tag="a", suite="goal", axis="Background_Textures", mode="sentinel",
             status=run_sr.RunStatus.COMPLETE, episodes_per_min=10.0),
        _rec(tag="a", suite="goal", axis="Background_Textures", mode="sentinel",
             status=run_sr.RunStatus.COMPLETE, episodes_per_min=20.0),
    ]
    assert run_sr.Schedule.rolling_rate(cell, CONFIG, records) == 15.0


def test_rolling_episodes_per_min_falls_back_to_yaml_default_with_no_history():
    cell = _mkcell("t", 100)
    assert run_sr.Schedule.rolling_rate(cell, CONFIG, []) == CONFIG.defaults.episodes_per_min


def test_rolling_episodes_per_min_ignores_other_cells():
    cell = _mkcell("t", 100, suite="goal", axis="Background_Textures", mode="sentinel")
    records = [_rec(tag="a", suite="long", axis="Camera_Viewpoints", mode="sentinel",
                    status=run_sr.RunStatus.COMPLETE, episodes_per_min=999.0)]
    assert run_sr.Schedule.rolling_rate(cell, CONFIG, records) == CONFIG.defaults.episodes_per_min


def test_rolling_episodes_per_min_ignores_incomplete_runs():
    """An incomplete run's throughput is not a throughput -- it is a partial run divided by a full
    wall clock, and using it would under-project every future makespan for that cell."""
    cell = _mkcell("t", 100, suite="goal", axis="Background_Textures", mode="sentinel")
    records = [_rec(tag="a", suite="goal", axis="Background_Textures", mode="sentinel",
                    status=run_sr.RunStatus.INCOMPLETE, episodes_per_min=999.0)]
    assert run_sr.Schedule.rolling_rate(cell, CONFIG, records) == CONFIG.defaults.episodes_per_min


# ======================================================================================================
# --cells parsing
# ======================================================================================================
def test_parse_cells_arg_builds_cells_in_order_with_explicit_mode():
    cells = CONFIG.cells_from_tokens("goal:Background_Textures,long:Camera_Viewpoints", "sentinel")
    assert [(c.suite, c.axis, c.mode) for c in cells] == [
        ("goal", "Background_Textures", "sentinel"),
        ("long", "Camera_Viewpoints", "sentinel"),
    ]
    # this exact (suite, axis) pair isn't hand-declared in either remaining rung (R1/R4 -- R4 is
    # auto_generate, skipped by _find_declared_cell), so this falls through to the dynamic
    # probe_comparators disk lookup -- exercised here, and asserted non-empty because the
    # lp_libero_goal_Background_Textures family is a real, present result dir (see
    # test_r4_comparators_resolved_dynamically_from_disk).
    assert cells[0].comparator("base") is not None
    assert cells[0].tag == "stablevla_sentinel_goal_Background_Textures"


def test_parse_cells_arg_defaults_mode_from_axis_when_omitted():
    cells = CONFIG.cells_from_tokens("goal:Robot_Initial_States", None)
    # Read the expectation from the ladder rather than restating it: the whole point of this test is
    # that an omitted --mode falls through to axis_default_mode, not that the default is any one value.
    assert cells[0].mode == CONFIG.default_mode("Robot_Initial_States")


def test_parse_cells_arg_unknown_axis_raises():
    with pytest.raises(ValueError):
        CONFIG.cells_from_tokens("goal:NotAnAxis", "base")


def test_parse_cells_arg_unknown_suite_raises():
    with pytest.raises(ValueError):
        CONFIG.cells_from_tokens("notasuite:Background_Textures", "base")


# ======================================================================================================
# R4 (full grid) auto-generation
# ======================================================================================================
def test_r4_generates_the_full_28_cell_grid():
    cells = CONFIG.cells_for_rung("R4")
    assert len(cells) == 28
    assert {c.suite for c in cells} == {"object", "spatial", "goal", "long"}
    assert {c.axis for c in cells} == set(AXES)
    for c in cells:
        assert c.mode == CONFIG.default_mode(c.axis)


def test_r4_comparators_resolved_dynamically_from_disk():
    comps = CONFIG.probe_comparators("goal", "Background_Textures")
    labels = {c.label for c in comps}
    assert "base" in labels
    base = next(c for c in comps if c.label == "base")
    assert base.n == 281   # lp_libero_goal_Background_Textures on disk


def test_r4_comparators_empty_when_results_root_absent(tmp_path):
    assert _config(tmp_path / "nowhere").probe_comparators("goal", "Background_Textures") == ()


def test_cells_for_rung_unknown_name_raises():
    with pytest.raises(ValueError, match="unknown rung"):
        CONFIG.cells_for_rung("NOPE")


# ======================================================================================================
# GPU / load preflight
# ======================================================================================================
def _preflight(min_free_gib=12.0, max_loadavg_factor=2.5) -> "run_sr.Preflight":
    config = _config()
    config.limits = dataclasses.replace(
        config.limits, min_free_gib=min_free_gib, max_loadavg_factor=max_loadavg_factor)
    return run_sr.Preflight(config)


def test_gpu_report_refuses_busy_gpu(monkeypatch):
    monkeypatch.setattr(torch.cuda, "device_count", lambda: 1)
    monkeypatch.setattr(torch.cuda, "mem_get_info", lambda i: (5 * 1024**3, 140 * 1024**3))
    statuses = _preflight(min_free_gib=12.0).gpus([0])
    assert len(statuses) == 1
    assert not statuses[0].ok
    assert "GiB free" in statuses[0].reason


def test_gpu_report_allows_free_gpu(monkeypatch):
    monkeypatch.setattr(torch.cuda, "device_count", lambda: 1)
    monkeypatch.setattr(torch.cuda, "mem_get_info", lambda i: (50 * 1024**3, 140 * 1024**3))
    assert _preflight(min_free_gib=12.0).gpus([0])[0].ok


def test_gpu_report_nonexistent_index_is_not_ok(monkeypatch):
    monkeypatch.setattr(torch.cuda, "device_count", lambda: 1)
    monkeypatch.setattr(torch.cuda, "mem_get_info", lambda i: (50 * 1024**3, 140 * 1024**3))
    assert not _preflight(min_free_gib=12.0).gpus([5])[0].ok


def test_gpu_report_defaults_to_every_visible_device(monkeypatch):
    monkeypatch.setattr(torch.cuda, "device_count", lambda: 3)
    monkeypatch.setattr(torch.cuda, "mem_get_info", lambda i: (50 * 1024**3, 140 * 1024**3))
    assert [s.index for s in _preflight().gpus()] == [0, 1, 2]


def test_load_report_ok_and_busy(monkeypatch):
    monkeypatch.setattr(run_sr.os, "getloadavg", lambda: (1.0, 1.0, 1.0))
    monkeypatch.setattr(run_sr.os, "cpu_count", lambda: 8)
    ok = _preflight(max_loadavg_factor=1.5).load()
    assert ok.ok and ok.threshold == 12.0

    monkeypatch.setattr(run_sr.os, "getloadavg", lambda: (20.0, 20.0, 20.0))
    assert not _preflight(max_loadavg_factor=1.5).load().ok


def test_check_libero_plus_reports_missing_tree(monkeypatch, tmp_path):
    monkeypatch.setenv("LIBERO_HOME", str(tmp_path / "does_not_exist"))
    result = _preflight().libero_plus()
    assert not result.ok
    assert "missing/empty" in result.message


# ======================================================================================================
# --dry-run: no subprocess, no filesystem mutation, even with --resume against a partial dir
# ======================================================================================================
def _no_launch(monkeypatch, tmp_path, n_gpus=2):
    def _boom(*_a, **_k):
        raise AssertionError("subprocess.Popen must not be called during --dry-run")
    monkeypatch.setattr(subprocess, "Popen", _boom)
    monkeypatch.setattr(torch.cuda, "device_count", lambda: n_gpus)
    monkeypatch.setattr(torch.cuda, "mem_get_info", lambda i: (100 * 1024**3, 140 * 1024**3))
    monkeypatch.setattr(run_sr.os, "getloadavg", lambda: (0.1, 0.1, 0.1))
    return run_sr.Ledger(tmp_path / "ledger.jsonl")


def test_dry_run_never_calls_subprocess_popen(monkeypatch, tmp_path):
    ledger = _no_launch(monkeypatch, tmp_path)
    results_root = tmp_path / "results"
    rc = _driver(["--rung", "R1", "--gpus", "0,1", "--dry-run",
                  "--results-root", str(results_root)], ledger=ledger).run()
    assert rc == run_sr.ExitCode.OK


def test_dry_run_with_resume_does_not_delete_a_partial_dir(monkeypatch, tmp_path):
    """Regression guard: --dry-run --resume must PREVIEW the purge, never perform it."""
    ledger = _no_launch(monkeypatch, tmp_path, n_gpus=1)
    results_root = tmp_path / "results"
    cells = _config(results_root).cells_for_rung("R1")
    partial_dir = cells[0].result_dir(results_root)
    partial_dir.mkdir(parents=True)
    (partial_dir / "rollout_t_episode0_success.mp4").write_bytes(b"x")   # 1 of expected 11

    rc = _driver(["--rung", "R1", "--gpus", "0", "--dry-run", "--resume",
                  "--results-root", str(results_root)], ledger=ledger).run()
    assert rc == run_sr.ExitCode.OK
    assert partial_dir.exists()   # NOT purged -- dry-run is read-only


def test_dry_run_refuses_when_every_gpu_is_busy(monkeypatch, tmp_path):
    _no_launch(monkeypatch, tmp_path)
    monkeypatch.setattr(torch.cuda, "mem_get_info", lambda i: (1 * 1024**3, 140 * 1024**3))
    rc = _driver(["--rung", "R1", "--gpus", "0,1", "--dry-run",
                  "--results-root", str(tmp_path / "r")],
                 ledger=run_sr.Ledger(tmp_path / "l.jsonl")).run()
    assert rc == run_sr.ExitCode.ERROR


# ======================================================================================================
# exit codes (main() wiring)
# ======================================================================================================
R1_TAG = "stablevla_base_object_Robot_Initial_States_smoke"   # rung R1's first cell, tag derived by
                                                                # cell_tag(policy, mode, suite, axis, variant)


def test_main_exit_code_2_when_gate_not_met(tmp_path):
    results_root = tmp_path / "results"
    cell_dir = results_root / "plus_libero_object" / "Robot_Initial_States" / R1_TAG
    cell_dir.mkdir(parents=True)
    # 1/11 successes vs the min_successes=8 gate (comparator base is 9/11) -- must fail
    (cell_dir / "rollout_t_episode0_success.mp4").write_bytes(b"x")
    for i in range(1, 11):
        (cell_dir / f"rollout_t_episode{i}_failure.mp4").write_bytes(b"x")
    rc = run_sr.main(["--score", "--rung", "R1", "--results-root", str(results_root)])
    assert rc == 2


def test_main_exit_code_0_when_gate_met(tmp_path):
    results_root = tmp_path / "results"
    cell_dir = results_root / "plus_libero_object" / "Robot_Initial_States" / R1_TAG
    cell_dir.mkdir(parents=True)
    # 9/11 successes -- meets the min_successes=8 gate
    for i in range(9):
        (cell_dir / f"rollout_t_episode{i}_success.mp4").write_bytes(b"x")
    for i in range(9, 11):
        (cell_dir / f"rollout_t_episode{i}_failure.mp4").write_bytes(b"x")
    rc = run_sr.main(["--score", "--rung", "R1", "--results-root", str(results_root)])
    assert rc == 0


def test_main_exit_code_1_when_nothing_scored(tmp_path):
    rc = run_sr.main(["--score", "--rung", "R1", "--results-root", str(tmp_path / "results")])
    assert rc == 1


def test_main_requires_rung_or_cells(tmp_path):
    rc = run_sr.main(["--results-root", str(tmp_path / "results")])
    assert rc == 1


def test_main_rejects_rung_and_cells_together():
    with pytest.raises(SystemExit):
        run_sr.main(["--rung", "R1", "--cells", "object:Objects_Layout", "--gpus", "0"])


def test_main_unknown_rung_raises():
    with pytest.raises(SystemExit):
        run_sr.main(["--rung", "NOPE", "--gpus", "0"])


def test_main_requires_gpus_to_launch(tmp_path):
    rc = run_sr.main(["--rung", "R1", "--results-root", str(tmp_path / "results")])
    assert rc == 1


def test_main_ledger_flag_short_circuits_before_touching_ladder_cells(tmp_path, monkeypatch):
    ledger_path = tmp_path / "ledger.jsonl"
    monkeypatch.setattr(run_sr, "default_ledger_path", lambda: ledger_path)
    rc = run_sr.main(["--ledger"])
    assert rc == 0


def test_parse_gpus_dedups_and_keeps_order():
    parse = run_sr.SRDriver._parse_gpus
    assert parse(None) == [] and parse("") == []
    assert parse("0,1") == [0, 1]
    assert parse("3,2") == [3, 2]
    assert parse(" 1 , 2 ") == [1, 2]
    # a repeated index would otherwise open two schedule tracks against the same device
    assert parse("0,0,1") == [0, 1]


# ======================================================================================================
# Port claiming — concurrent run_sr invocations must never share a policy-server port.
#
# Why this has its own test: the failure is SILENT and it mis-scores. Two run_sr processes both start
# their port counter at defaults.base_port, so the second one's sim clients connect to the FIRST one's
# already-running policy server -- a different CHECKPOINT -- and write ordinary-looking mp4s. It was
# caught in practice with a goal Robot-Init cell whose clients reached a long-checkpoint server and
# produced 121 episodes under the wrong policy, with nothing in any log to say so.
# ======================================================================================================
def test_claim_port_skips_a_bound_port(tmp_path):
    import socket as _s
    with _s.socket(_s.AF_INET, _s.SOCK_STREAM) as held:
        held.setsockopt(_s.SOL_SOCKET, _s.SO_REUSEADDR, 1)
        held.bind(("0.0.0.0", 0))
        held.listen(1)
        busy = held.getsockname()[1]
        got = run_sr.PortAllocator(busy, tmp_path / ".ports").claim()
        assert got != busy, "claim handed out a port something is already listening on"


def test_claim_port_reservation_blocks_a_concurrent_process(tmp_path):
    """A reservation held by a LIVE pid must not be handed out again. This is the half a bind-probe
    cannot cover: a policy server spends minutes loading its checkpoint before it binds, so a second
    process probing during that window still sees the port as free and claims it too."""
    base = 15900
    reservations = tmp_path / ".ports"
    reservations.mkdir(parents=True)
    (reservations / str(base)).write_text(f"{os.getpid()}\n")     # us: definitely alive
    got = run_sr.PortAllocator(base, reservations).claim()
    assert got != base, "a reservation held by a live process was handed out anyway"


def test_claim_port_reclaims_a_reservation_whose_owner_died(tmp_path):
    """A run killed mid-flight leaves its reservation files behind; they must not wedge the port
    range forever, so an unreachable owner pid is treated as stale and taken over."""
    base = 15910
    reservations = tmp_path / ".ports"
    reservations.mkdir(parents=True)
    dead_pid = 999_999                              # far above /proc/sys/kernel/pid_max on this box
    (reservations / str(base)).write_text(f"{dead_pid}\n")
    assert run_sr.PortAllocator(base, reservations).claim() == base, \
        "a reservation from a dead process was not reclaimed"


def test_claim_port_release_all_drops_every_marker(tmp_path):
    reservations = tmp_path / ".ports"
    allocator = run_sr.PortAllocator(16100, reservations)
    ports = [allocator.claim() for _ in range(3)]
    assert all((reservations / str(p)).exists() for p in ports)
    allocator.release_all()
    assert not any((reservations / str(p)).exists() for p in ports)


def test_claim_port_gives_out_a_distinct_port_every_call(tmp_path):
    allocator = run_sr.PortAllocator(16200, tmp_path / ".ports")
    ports = [allocator.claim() for _ in range(5)]
    assert len(set(ports)) == 5


# ======================================================================================================
# --rung + --cells subsetting
# ======================================================================================================
def test_rung_plus_cells_selects_a_subset_and_keeps_the_rungs_own_filter_and_n():
    """--rung R --cells suite:axis must SUBSET the rung, not rebuild the cell ad hoc: an ad hoc cell
    carries the full-axis expected_n and no instance filter, so a 12-episode smoke would silently
    become a full run."""
    cells = _driver(["--rung", "R1", "--cells", "long:Sensor_Noise", "--gpus", "0"])._resolve_cells()
    assert [(c.suite, c.axis) for c in cells] == [("long", "Sensor_Noise")]
    full = [c for c in CONFIG.cells_for_rung("R1") if (c.suite, c.axis) == ("long", "Sensor_Noise")][0]
    assert cells[0].expected_n == full.expected_n
    assert cells[0].filter_file == full.filter_file is not None


def test_rung_plus_cells_rejects_a_cell_the_rung_does_not_declare():
    driver = _driver(["--rung", "R1", "--cells", "goal:Camera_Viewpoints", "--gpus", "0"])
    with pytest.raises(SystemExit):
        driver._resolve_cells()