onf-c1/scripts/run_sr.py

#!/usr/bin/env python
"""scripts/run_sr.py — Python success-rate eval driver for ONF LIBERO-Plus cells.

One Python entrypoint for launching+scoring a ladder cell (suite x axis x mode): sets up
the OSMesa render env, polls the policy server's port for readiness, sets PYTHONNOUSERSITE to dodge
a user-site torch that shadows the pinned stablevla one, writes into the
results/plus_<suite>/<Axis_With_Underscores>/<TAG>/ layout the mp4 filenames encode success into,
and passes through QNDF_DIR / LIBERO_INSTANCE_ID_FILTER. It:

  * resolves a ladder cell (suite, axis, mode) into an env + CLI recipe by reading
    configs/sr_ladder.yaml (the cell table: expected episode counts, comparator dirs, gates,
    GPU hints) and PARSING evals/common/modes.sh for the mode env recipes (rather than
    re-declaring them — a parsed value can never drift from the bash launchers it replaces);
  * launches one GPU-resident StableVLA policy server + N sharded sim clients per cell via
    subprocess, always tearing the server down (a context manager + a signal handler, so Ctrl-C
    can't orphan a GPU-resident process);
  * is polite about shared hardware: refuses a GPU with too little free memory
    (torch.cuda.mem_get_info — never nvidia-smi, which is broken on this box) and refuses a
    large launch when the box is under load (/proc/loadavg);
  * enforces completeness (mp4 count vs. the yaml's expected episode count) rather than assuming it,
    and treats a partial result dir as poison for score.py — --resume purges and reruns it;
  * appends one append-only record per cell attempt to outputs/sr/ledger.jsonl (see
    onf.graph.report's StageLogger for the run-dir/manifest convention this mirrors);
  * scores with the EXISTING evals/libero_plus/score.py / mcnemar.py (imported, not
    reimplemented) and prints a paired McNemar discordance alongside the marginal rate — on this
    deterministic bench, effects are 1-4 episodes and a marginal rate alone is not interpretable.

SHAPE
    LadderConfig is the one object that knows the ladder: resolved paths, the external
    scorers, the modes.sh recipes, the suite table, the defaults, and the results root. Everything
    else takes it in the constructor instead of threading
    (ladder_doc, suites, results_root, paths) through twenty signatures. The mutable
    process/port/ledger state that used to live in module globals now lives on
    ProcessSupervisor / PortAllocator / Ledger, one instance per run.

SECOND CONSUMER
    evals/gr00t/run_gr00t.py loads THIS module by file path and reuses its ladder/mode/naming/
    provenance layer so the two drivers can never disagree about which mode an axis defaults to,
    what a result dir is called, or what a ledger row looks like. Anything it imports is public API
    here; see that file's own RUN_SR. call sites.

Usage:
    python scripts/run_sr.py --check-env
    python scripts/run_sr.py --rung R4 --gpus 0,1 --dry-run
    python scripts/run_sr.py --rung R4 --gpus 0,1 [--mode base] [--resume]
    python scripts/run_sr.py --cells goal:Background_Textures,long:Camera_Viewpoints --gpus 0,1 --mode base
    python scripts/run_sr.py --score --rung R4
    python scripts/run_sr.py --ledger

Exit codes: 0 = ok / all gates met; 1 = operational error (preflight refusal, crash, nothing to do);
2 = ran/scored fine but a gate was NOT met (mirrors the graph CLI convention: a clean run that fails
its bar is not the same failure mode as a broken run).
"""

from __future__ import annotations

import argparse
import ast
import atexit
import dataclasses
import enum
import importlib.util
import itertools
import json
import os
import re
import shutil
import signal
import socket
import subprocess
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from contextlib import ExitStack
from dataclasses import dataclass, field, replace
from datetime import datetime, timezone
from pathlib import Path
from types import FrameType, ModuleType
from typing import Any, Mapping, Sequence, TextIO

import yaml

REPO_ROOT = Path(__file__).resolve().parents[1]
LADDER_YAML = REPO_ROOT / "configs" / "sr_ladder.yaml"
STABLEVLA_CONDA_ENV = "stablevla"          # fixed operational fact: the conda env name this box uses
STABLEVLA_ROOT = REPO_ROOT / "stablevla"
ONF_SRC = REPO_ROOT / "src"

# This box's kernel nvidia driver module and userspace NVML library are mismatched (measured:
# kernel 570.172.08, userspace 580.173) -- nvidia-smi itself is unusable without preloading a
# userspace NVML that matches the kernel module, and so is torch's CUDA caching allocator whenever it
# needs to grow its pool (nvmlInit_v2 hard-asserts on the mismatch -- see the ONF_ROOT-external
# HAMLET-Isaac-GR00T repo's own vulkan570_env.sh fix for the identical class of problem).
# ONF_NVML_PRELOAD overrides the path outright; failing that, the known-good file from that sibling
# repo's fix is used if present on this exact box, so a checkout elsewhere without that sibling repo
# just silently gets no LD_PRELOAD (a degrade to the OLD, occasionally-crashing behaviour, not a hard
# failure).
NVML_PRELOAD_CANDIDATE = Path(
    "/srv/data/HAMLET-Isaac-GR00T/.vkenv/nvidia570/lib/libnvidia-ml.so.570.172.08"
)


# =====================================================================================================
# enums -- the strings that used to be compared inline
# =====================================================================================================
class ExitCode(enum.IntEnum):
    """This driver's process exit codes. int so sys.exit / main can return one directly."""

    OK = 0
    ERROR = 1
    GATE_FAILED = 2


class GateKind(str, enum.Enum):
    """The gate kinds sr_ladder.yaml may declare. Constructing one validates the yaml value."""

    MIN_SUCCESSES = "min_successes"
    MAX_REGRESS_PP = "max_regress_pp"
    EXACT_MATCH = "exact_match"


class CellAction(str, enum.Enum):
    """What ResumePlanner decided to do with one cell's existing result dir."""

    RUN = "run"
    SKIP = "skip"
    PURGE = "purge"
    ERROR = "error"


class RunStatus(str, enum.Enum):
    """A ledger row's terminal state. RUNNING is written at launch and superseded on completion."""

    RUNNING = "running"
    COMPLETE = "complete"
    INCOMPLETE = "incomplete"
    CRASHED = "crashed"


# =====================================================================================================
# yaml loading + the resolved ladder paths/tables
# =====================================================================================================
def load_yaml(path: Path) -> dict:
    """Read a YAML file into a dict ({} when the document is empty)."""
    with open(path) as f:
        return yaml.safe_load(f) or {}


@dataclass(frozen=True)
class LadderPaths:
    """Absolute paths to every external artifact this driver drives.

    The single place the driver learns where score.py / mcnemar.py / modes.sh / the policy server /
    the harness live -- nothing about them is hardcoded outside this file's bootstrap constants
    (REPO_ROOT, LADDER_YAML).
    """

    score_py: Path
    mcnemar_py: Path
    modes_sh: Path
    policy_server: Path
    harness_dir: Path
    shim_dir: Path
    suites_yaml: Path

    @classmethod
    def from_doc(cls, doc: Mapping[str, Any], root: Path = REPO_ROOT) -> "LadderPaths":
        """Resolve sr_ladder.yaml's paths: block (all entries relative to the repo root).

        Raises:
            KeyError: The block is missing a declared path, naming every one that is absent -- a
                clear failure here beats a KeyError from deep inside a launch.
        """
        raw = doc["paths"]
        names = [f.name for f in dataclasses.fields(cls)]
        missing = [n for n in names if n not in raw]
        if missing:
            raise KeyError(f"sr_ladder.yaml: paths block is missing {missing}")
        return cls(**{n: root / raw[n] for n in names})


@dataclass(frozen=True)
class SuiteInfo:
    """One suite's row of configs/suites.yaml, with every path already made absolute."""

    name: str
    bench: str
    stablevla_ckpt: Path
    fwm_dir: Path
    light_filter: Path


@dataclass(frozen=True)
class Defaults:
    """sr_ladder.yaml's defaults: block."""

    num_clients: int
    base_port: int
    episodes_per_min: float
    server_ready_timeout_s: float

    @classmethod
    def from_doc(cls, doc: Mapping[str, Any]) -> "Defaults":
        return cls(
            num_clients=int(doc["num_clients"]),
            base_port=int(doc["base_port"]),
            episodes_per_min=float(doc["episodes_per_min"]),
            server_ready_timeout_s=float(doc["server_ready_timeout_s"]),
        )


@dataclass(frozen=True)
class PreflightLimits:
    """sr_ladder.yaml's preflight: block -- the thresholds that make this driver a polite
    neighbour on a shared box."""

    min_free_gib: float
    max_loadavg_factor: float
    large_rung_episodes: int

    @classmethod
    def from_doc(cls, doc: Mapping[str, Any]) -> "PreflightLimits":
        return cls(
            min_free_gib=float(doc["min_free_gib"]),
            max_loadavg_factor=float(doc["max_loadavg_factor"]),
            large_rung_episodes=int(doc["large_rung_episodes"]),
        )


# =====================================================================================================
# reuse evals/libero_plus/score.py + mcnemar.py rather than reimplementing the mp4-glob scorer / the
# paired McNemar test. Both are standalone scripts (no package __init__), so load them by file path.
# =====================================================================================================
@dataclass(frozen=True)
class DirCount:
    """score.py's (successes, n) for one result directory."""

    successes: int
    n: int

    @property
    def rate(self) -> float:
        """Success rate in percent; nan for an empty directory."""
        return 100.0 * self.successes / self.n if self.n else float("nan")


class ExternalScorers:
    """Lazy, cached accessor for score.py / mcnemar.py.

    Every scoring question in this driver goes through here, so it can never disagree with
    score.py about what a tag directory scores or which axes exist. Loaded on first use: a
    --ledger or --help invocation never pays for it.
    """

    def __init__(self, paths: LadderPaths) -> None:
        self._paths = paths
        self._score: ModuleType | None = None
        self._mcnemar: ModuleType | None = None

    @property
    def score(self) -> ModuleType:
        if self._score is None:
            self._score = self._load("sr_score", self._paths.score_py)
        return self._score

    @property
    def mcnemar(self) -> ModuleType:
        if self._mcnemar is None:
            self._mcnemar = self._load("sr_mcnemar", self._paths.mcnemar_py)
        return self._mcnemar

    @property
    def axes(self) -> list[str]:
        """The perturbation axes, in score.py's own order."""
        return list(self.score.AXES)

    def bench(self, suite: str) -> str:
        """Suite short-name -> LIBERO benchmark name (score.py's SUITE_BENCH)."""
        return self.score.SUITE_BENCH[suite]

    def default_results_root(self) -> Path:
        return Path(self.score.results_root())

    def count_dir(self, directory: Path) -> DirCount:
        """(successes, n) for a result dir; (0, 0) when it does not exist."""
        if not directory.exists():
            return DirCount(0, 0)
        successes, n = self.score.count_dir(str(directory))
        return DirCount(int(successes), int(n))

    def episode_outcomes(self, directory: Path) -> dict[str, bool]:
        """{episode_key: succeeded} for the paired McNemar test."""
        return self.mcnemar.load(str(directory))

    def exact_p(self, b_only: int, c_only: int) -> float:
        return self.mcnemar.mcnemar_exact(b_only, c_only)

    @staticmethod
    def _load(name: str, path: Path) -> ModuleType:
        spec = importlib.util.spec_from_file_location(name, path)
        if spec is None or spec.loader is None:
            raise ImportError(f"cannot load {name} from {path}")
        module = importlib.util.module_from_spec(spec)
        sys.modules[name] = module   # register before exec so dataclass/typing self-references resolve
        spec.loader.exec_module(module)
        return module


# =====================================================================================================
# evals/common/modes.sh parser — the mode env recipes are PARSED from the bash file (the
# single source of truth for onf_set_mode()), not re-declared in yaml, so this driver can never drift
# out of sync with the launchers it replaces. See tests/test_run_sr.py for a check that the parser's
# output matches the CONTRACT values documented in modes.sh's own header comment.
# =====================================================================================================
_CASE_BLOCK_RE = re.compile(r'case\s+"\$\{mode\}"\s+in\n(.*?)\n[ \t]*esac', re.DOTALL)
_ARM_SPLIT_RE = re.compile(r"^[ \t]*([A-Za-z_]\w*(?:\|[A-Za-z_]\w*)*)\)[ \t]*\n", re.MULTILINE)
_ASSIGN_RE = re.compile(r"^([A-Za-z_]\w*)=(.*)$")
_PARAM_DEFAULT_RE = re.compile(r'^"?\$\{[A-Za-z_]\w*:-([^}]*)\}"?$')


def parse_mode_recipes(modes_sh_text: str) -> dict[str, dict[str, str]]:
    """Parse onf_set_mode()'s case "${mode}" in ... esac block.

    Args:
        modes_sh_text: The contents of evals/common/modes.sh.

    Returns:
        {mode_name: {ENV_VAR: value}} -- whatever the matched arm exports.

    Raises:
        ValueError: The case block could not be found.
    """
    match = _CASE_BLOCK_RE.search(modes_sh_text)
    if not match:
        raise ValueError('modes.sh: could not find `case "${mode}" in ... esac` block')

    recipes: dict[str, dict[str, str]] = {}
    # parts[0] is text before the first arm header (comments); then alternating (pattern, arm_text).
    parts = _ARM_SPLIT_RE.split(match.group(1))
    for i in range(1, len(parts), 2):
        pattern, arm_text = parts[i], parts[i + 1]
        env = _parse_arm_exports(arm_text)
        for name in pattern.split("|"):
            recipes[name] = env
    return recipes


def _parse_arm_exports(arm_text: str) -> dict[str, str]:
    """One case arm's body -> its NAME=value exports, honouring backslash continuations."""
    text = arm_text.split(";;")[0]                                   # stop at the arm terminator
    lines = [line.split("#", 1)[0].rstrip() for line in text.splitlines()]   # no '#' in any value
    joined = " ".join(line[:-1] if line.endswith("\\") else line for line in lines)
    out: dict[str, str] = {}
    for token in joined.split():
        if token == "export":
            continue
        match = _ASSIGN_RE.match(token)
        if match:
            out[match.group(1)] = _expand_value(match.group(2))
    return out


def _expand_value(value: str) -> str:
    """"${VAR:-default}" -> default; anything else just loses its surrounding quotes."""
    match = _PARAM_DEFAULT_RE.match(value)
    return match.group(1) if match else value.strip('"')


class ModeRegistry:
    """The mode env recipes, parsed once from modes.sh."""

    def __init__(self, modes_sh: Path) -> None:
        self._path = modes_sh
        self._recipes = parse_mode_recipes(modes_sh.read_text())

    @property
    def modes(self) -> list[str]:
        return sorted(self._recipes)

    def env_for(self, mode: str) -> dict[str, str]:
        """A fresh copy of mode's env recipe.

        Raises:
            ValueError: mode has no arm in modes.sh.
        """
        if mode not in self._recipes:
            raise ValueError(f"unknown MODE={mode!r}; not found in {self._path} (known: {self.modes})")
        return dict(self._recipes[mode])


# =====================================================================================================
# ladder cell model
# =====================================================================================================
@dataclass(frozen=True)
class Comparator:
    """A previously-measured result directory this cell is scored against."""

    label: str
    tag: str
    successes: int
    n: int

    @property
    def rate(self) -> float:
        return 100.0 * self.successes / self.n if self.n else float("nan")


@dataclass(frozen=True)
class GateVerdict:
    """Whether a cell cleared its bar, and the one-line reason."""

    passed: bool
    message: str


@dataclass(frozen=True)
class Gate:
    """One cell's pass/fail bar, as declared in sr_ladder.yaml."""

    kind: GateKind
    vs: str = "base"
    min_successes: int | None = None
    max_regress_pp: float | None = None

    @classmethod
    def from_dict(cls, doc: Mapping[str, Any]) -> "Gate":
        """Build from a yaml gate block. Raises ValueError on an unknown kind."""
        return cls(
            kind=GateKind(doc["kind"]),
            vs=doc.get("vs", "base"),
            min_successes=doc.get("min_successes"),
            max_regress_pp=doc.get("max_regress_pp"),
        )

    def evaluate(self, score: "CellScore") -> GateVerdict:
        """Apply this gate to a scored cell.

        A gate with no matching comparator on disk PASSES with a SKIP note -- including a
        min_successes gate. That is deliberate and load-bearing: the comparator's absence means
        the ladder has nothing to compare against yet, and failing a cell for that would make a
        fresh checkout's first run red for a reason that has nothing to do with the run.
        """
        comparator = score.cell.comparator(self.vs)
        if comparator is None:
            return GateVerdict(True, f"SKIP: no comparator labeled {self.vs!r} to gate against")
        if self.kind is GateKind.EXACT_MATCH:
            passed = score.successes == comparator.successes and score.actual_n == comparator.n
            return GateVerdict(
                passed,
                f"exact_match vs {comparator.label} ({comparator.successes}/{comparator.n}): "
                f"got {score.successes}/{score.actual_n}",
            )
        if self.kind is GateKind.MIN_SUCCESSES:
            return GateVerdict(
                score.successes >= self.min_successes,
                f"min_successes={self.min_successes} vs {comparator.label}: got {score.successes}",
            )
        delta = score.rate - comparator.rate
        return GateVerdict(
            delta >= -self.max_regress_pp,
            f"delta vs {comparator.label} = {delta:+.1f}pp (floor -{self.max_regress_pp:.1f}pp)",
        )


def cell_tag(policy: str, mode: str, suite: str, axis: str, variant: str = "") -> str:
    """The one on-disk name for a cell: {policy}_{mode}_{suite}_{axis}[_{variant}].

    Both drivers (this one and evals/gr00t/run_gr00t.py) route every result dir, log dir and
    ledger row through this function, so a name always says which policy and which mode produced it.
    Baking the mode in is what stops --resume's ledger lookup from reporting one mode's numbers
    under another's; variant separates a filtered smoke pool from the full axis it subsets.
    """
    tag = f"{policy}_{mode}_{suite}_{axis}"
    return f"{tag}_{variant}" if variant else tag


def utc_stamp() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")


def make_log_dir(policy: str, mode: str, suite: str, axis: str, variant: str = "",
                  root: Path = REPO_ROOT) -> Path:
    """logs/{UTC ts}_{tag}/ for this cell's cell_tag -- the log dir and the result dir
    therefore carry the SAME name, so one is findable from the other. Axis is normalized to
    underscores so a space-separated category name ("Sensor Noise") stays one path component.

    Shared by BOTH drivers. Before this, run_sr.py stamped LOCAL time while run_gr00t.py stamped UTC,
    and run_gr00t.py additionally used the literal "gr00tn17" instead of its own --policy --
    two runs differing only by mode/suite/axis truncated each other's shard logs into one directory.
    """
    tag = cell_tag(policy, mode, suite, axis.replace(" ", "_"), variant)
    return root / "logs" / f"{utc_stamp()}_{tag}"


@dataclass(frozen=True)
class Cell:
    """One ladder cell: a (suite, axis, mode) triple plus everything needed to run and score
    it. Frozen -- with_mode is the only way to vary one, so a cell's tag can never fall
    out of step with the mode that produced it."""

    rung: str
    suite: str
    axis: str
    bench: str
    mode: str
    expected_n: int
    filter_file: Path | None = None
    extra_env: Mapping[str, str] = field(default_factory=dict)
    comparators: tuple[Comparator, ...] = ()
    gate: Gate | None = None
    num_clients: int = 8
    policy: str = "stablevla"
    variant: str = ""

    @property
    def tag(self) -> str:
        return cell_tag(self.policy, self.mode, self.suite, self.axis, self.variant)

    @property
    def key(self) -> str:
        """suite:axis -- what --cells matches on."""
        return f"{self.suite}:{self.axis}"

    @property
    def category_value(self) -> str:
        """The axis as the harness's --args.category_value wants it (spaces, not underscores)."""
        return self.axis.replace("_", " ")

    def result_dir(self, results_root: Path) -> Path:
        return results_root / f"plus_{self.bench}" / self.axis / self.tag

    def comparator_dir(self, results_root: Path, comparator: Comparator) -> Path:
        return results_root / f"plus_{self.bench}" / self.axis / comparator.tag

    def comparator(self, label: str) -> Comparator | None:
        return next((c for c in self.comparators if c.label == label), None)

    def with_mode(self, mode: str) -> "Cell":
        """Swap this cell's mode. The tag follows automatically -- it is derived, not stored."""
        return replace(self, mode=mode)


@dataclass(frozen=True)
class CellScore:
    """What one cell's result directory currently contains."""

    cell: Cell
    count: DirCount

    @property
    def successes(self) -> int:
        return self.count.successes

    @property
    def actual_n(self) -> int:
        return self.count.n

    @property
    def rate(self) -> float:
        return self.count.rate

    @property
    def complete(self) -> bool:
        return self.count.n == self.cell.expected_n


@dataclass(frozen=True)
class CellReport:
    """A scored cell plus its verdict and (optionally) its paired McNemar line."""

    score: CellScore
    verdict: GateVerdict
    mcnemar: str | None = None

    @property
    def line(self) -> str:
        """The one RESULT ... line printed per cell."""
        cell = self.score.cell
        comparator = (
            cell.comparator(cell.gate.vs) if cell.gate
            else (cell.comparators[0] if cell.comparators else None)
        )
        parts: list[str] = []
        if comparator is not None and self.score.actual_n:
            parts.append(f"delta vs {comparator.label} = {self.score.rate - comparator.rate:+.1f}pp")
        if self.mcnemar:
            parts.append(self.mcnemar)
        rate = f"{self.score.rate:.1f}%" if self.score.actual_n else "--"
        status = "PASS" if self.verdict.passed else "FAIL"
        return (
            f"RESULT {cell.tag}: {self.score.successes}/{self.score.actual_n} = {rate}"
            f"  ({', '.join(parts)})  [{status}: {self.verdict.message}]"
        )


# =====================================================================================================
# the ladder itself
# =====================================================================================================
class LadderConfig:
    """Everything sr_ladder.yaml + suites.yaml + modes.sh say, resolved once.

    Constructed once per driver invocation and handed to every collaborator, replacing the
    (ladder_doc, suites, results_root, paths) four-tuple that used to be threaded through most of
    this module's signatures.
    """

    def __init__(self, doc: Mapping[str, Any], results_root: Path | None = None,
                 root: Path = REPO_ROOT) -> None:
        self.doc = doc
        self.root = root
        self.paths = LadderPaths.from_doc(doc, root)
        self.scorers = ExternalScorers(self.paths)
        self.modes = ModeRegistry(self.paths.modes_sh)
        self.defaults = Defaults.from_doc(doc["defaults"])
        self.limits = PreflightLimits.from_doc(doc["preflight"])
        self.suites = self._load_suites()
        self.results_root = Path(results_root) if results_root else self.scorers.default_results_root()

    @classmethod
    def load(cls, ladder_yaml: Path | None = None, results_root: Path | None = None) -> "LadderConfig":
        return cls(load_yaml(ladder_yaml or LADDER_YAML), results_root)

    def suite(self, name: str) -> SuiteInfo:
        if name not in self.suites:
            raise ValueError(f"unknown suite {name!r} (known: {sorted(self.suites)})")
        return self.suites[name]

    def default_mode(self, axis: str) -> str:
        """The per-axis default recipe from axis_default_mode. Shared with the GR00T driver, so
        the two can never disagree about which mode an axis defaults to."""
        table = self.doc["axis_default_mode"]
        return table.get(axis, table["default"])

    def expected_episodes(self, suite: str, axis: str) -> int:
        return int(self.doc["expected_episodes"][suite][axis])

    def cells_for_rung(self, rung_name: str) -> list[Cell]:
        """Every cell of a named rung, either declared or auto-generated as the full grid.

        Raises:
            ValueError: No such rung.
        """
        if rung_name not in self.doc["rungs"]:
            raise ValueError(f"unknown rung {rung_name!r}; known: {sorted(self.doc['rungs'])}")
        rung_doc = self.doc["rungs"][rung_name]
        if rung_doc.get("auto_generate") == "full_grid":
            return self._full_grid_cells(rung_name, Gate.from_dict(rung_doc["gate"]))
        return [self.cell_from_doc(rung_name, c) for c in rung_doc["cells"]]

    def cells_from_tokens(self, spec: str, mode: str | None = None) -> list[Cell]:
        """--cells suite:axis,... -> ad hoc cells.

        If an existing rung already declares this exact (suite, axis) pair, its comparators/gate are
        reused so an ad hoc cell isn't scored in a vacuum.

        Raises:
            ValueError: An unknown suite or axis token.
        """
        axes = self.scorers.axes
        cells: list[Cell] = []
        for token in (t.strip() for t in spec.split(",")):
            if not token:
                continue
            suite, _, axis = token.partition(":")
            if suite not in self.suites:
                raise ValueError(f"--cells: unknown suite {suite!r} in {token!r} (known: {sorted(self.suites)})")
            if axis not in axes:
                raise ValueError(f"--cells: unknown axis {axis!r} in {token!r} (known: {axes})")
            declared = self._find_declared_cell(suite, axis)
            cells.append(Cell(
                rung="adhoc", suite=suite, axis=axis, bench=self.scorers.bench(suite),
                mode=mode or self.default_mode(axis),
                expected_n=self.expected_episodes(suite, axis),
                filter_file=self._resolve_filter(None, axis, suite),
                comparators=(declared.comparators if declared
                             else self.probe_comparators(suite, axis)),
                gate=declared.gate if declared else Gate.from_dict(self.doc["default_gate"]),
                num_clients=self.defaults.num_clients,
            ))
        return cells

    def probe_comparators(self, suite: str, axis: str) -> tuple[Comparator, ...]:
        """The PRE-EXISTING baseline result dirs already on disk for this (suite, axis).

        The full grid's comparators aren't hand-enumerated (28 cells would go stale the moment a new
        baseline lands): probe what is there, so a run never re-measures a baseline that has already
        been measured. Those directory names are historical, from whichever run produced them -- they
        are disk facts, not modes this repo can produce. Reuses score.py's count_dir so this
        can never disagree with it about what a tag directory scores.
        """
        bench = self.scorers.bench(suite)
        if axis == "Robot_Initial_States":
            candidates = [("base", f"onfbase_{bench}_{axis}"), ("onf", f"onfv1_{bench}_{axis}")]
        else:
            candidates = [
                ("base", f"cat_{axis}_base"),
                ("sentinel", f"cat_{axis}_sentinel"),
                ("base", f"lp_{bench}_{axis}"),
                ("prior", f"cndf_{bench}_{axis}"),
            ]
        found: list[Comparator] = []
        seen: set[tuple[str, str]] = set()
        for label, tag in candidates:
            if (label, tag) in seen:
                continue
            directory = self.results_root / f"plus_{bench}" / axis / tag
            if not directory.is_dir():
                continue
            count = self.scorers.count_dir(directory)
            if count.n:
                found.append(Comparator(label=label, tag=tag, successes=count.successes, n=count.n))
                seen.add((label, tag))
        return tuple(found)

    def cell_from_doc(self, rung_name: str, cell_doc: Mapping[str, Any]) -> Cell:
        """One cells: entry -> a Cell, resolving its filter, gate and client count."""
        suite, axis = cell_doc["suite"], cell_doc["axis"]
        gate_doc = cell_doc.get("gate") or self.doc.get("default_gate")
        return Cell(
            rung=rung_name, suite=suite, axis=axis, bench=self.scorers.bench(suite),
            mode=cell_doc.get("mode") or self.default_mode(axis),
            expected_n=int(cell_doc["expected_n"]),
            variant=cell_doc.get("variant", ""),
            filter_file=self._resolve_filter(cell_doc.get("filter"), axis, suite),
            extra_env={k: str(v) for k, v in (cell_doc.get("env") or {}).items()},
            comparators=tuple(Comparator(**c) for c in cell_doc.get("comparators", [])),
            gate=Gate.from_dict(gate_doc) if gate_doc else None,
            num_clients=int(cell_doc.get("num_clients", self.defaults.num_clients)),
        )

    def _full_grid_cells(self, rung_name: str, gate: Gate) -> list[Cell]:
        return [
            Cell(
                rung=rung_name, suite=suite, axis=axis, bench=self.scorers.bench(suite),
                mode=self.default_mode(axis),
                expected_n=self.expected_episodes(suite, axis),
                filter_file=self._resolve_filter(None, axis, suite),
                comparators=self.probe_comparators(suite, axis),
                gate=gate, num_clients=self.defaults.num_clients,
            )
            for suite in ("object", "spatial", "goal", "long")
            for axis in self.scorers.axes
        ]

    def _find_declared_cell(self, suite: str, axis: str) -> Cell | None:
        for rung_name, rung_doc in self.doc["rungs"].items():
            if rung_doc.get("auto_generate"):
                continue
            for cell_doc in rung_doc["cells"]:
                if cell_doc["suite"] == suite and cell_doc["axis"] == axis:
                    return self.cell_from_doc(rung_name, cell_doc)
        return None

    def _resolve_filter(self, filt: str | None, axis: str, suite: str) -> Path | None:
        if filt in (None, "auto"):
            if filt == "auto" or axis == "Light_Conditions":
                return self.suite(suite).light_filter
            return None
        return self.root / filt

    def _load_suites(self) -> dict[str, SuiteInfo]:
        doc = load_yaml(self.paths.suites_yaml)
        return {
            name: SuiteInfo(
                name=name,
                bench=info["task_suite"],
                stablevla_ckpt=self.root / info["stablevla_ckpt"],
                fwm_dir=self.root / "data" / info["fwm_dir"],
                light_filter=self.root / info["light_filter"],
            )
            for name, info in doc["suites"].items()
        }


# =====================================================================================================
# child-process environment (see the module docstring). This dict (minus the per-process
# CUDA_VISIBLE_DEVICES/PYTHONPATH layered on at launch time) is what gets recorded verbatim in the
# ledger for reproducibility.
# =====================================================================================================
def resolve_conda_env_prefix(name: str) -> Path | None:
    """Resolve a conda env's prefix without sourcing a shell: $CONDA_EXE's parent-of-parent, then
    the usual install prefixes."""
    candidates: list[Path] = []
    conda_exe = os.environ.get("CONDA_EXE")
    if conda_exe:
        candidates.append(Path(conda_exe).resolve().parent.parent / "envs" / name)
    candidates += [
        base / "envs" / name
        for base in (Path.home() / "miniconda3", Path.home() / "anaconda3",
                     Path.home() / "miniforge3", Path("/opt/conda"), Path("/opt/miniconda3"))
    ]
    return next((c for c in candidates if (c / "bin" / "python").exists()), None)


def resolve_stablevla_python() -> Path | None:
    prefix = resolve_conda_env_prefix(STABLEVLA_CONDA_ENV)
    return (prefix / "bin" / "python") if prefix else None


def nvml_preload_path() -> str | None:
    """The NVML shim to LD_PRELOAD into every child, or None on a box that doesn't need it.
    See NVML_PRELOAD_CANDIDATE for why."""
    override = os.environ.get("ONF_NVML_PRELOAD")
    if override:
        return override if Path(override).exists() else None
    return str(NVML_PRELOAD_CANDIDATE) if NVML_PRELOAD_CANDIDATE.exists() else None


def libero_home_dir() -> Path:
    return Path(os.environ.get("LIBERO_HOME", str(REPO_ROOT / "external" / "LIBERO-plus")))


def resolve_graph_dir(suite: str, env: Mapping[str, str] | None = None) -> Path:
    """The graph-artifact dir (g_nodes.npz / g_edges.npz / g_head.npz / g_track.npz) for suite.

    Args:
        suite: Suite name, for the default location.
        env: The CHILD's environment. Its GR_GRAPH_DIR is what the child will actually load, and
            passing it is not optional in practice: onf.config.Paths.graph honors GR_GRAPH_DIR from
            os.environ, and the driver never sets that variable on ITSELF -- it builds a per-cell
            env dict and hands it to the subprocess. So resolving without env silently reported and
            hard-checked outputs/<suite>/latest/artifacts for every arm that overrides the dir,
            which is every belief arm on this branch. The recorded graph_hash and g_head_mtime for
            those runs describe a directory the run never opened, and an arm pointed at an
            incomplete dir would have passed the pre-launch check and died hours in.

    Returns:
        The directory the child will load.
    """
    from onf.config import default_paths
    override = (env or {}).get("GR_GRAPH_DIR", "")
    return Path(override) if override else default_paths().graph(suite)


@dataclass(frozen=True)
class GraphArtifacts:
    """What a resolved graph dir actually contains, for the pre-launch check and the run record."""

    graph_dir: Path
    ok: bool
    missing: tuple[str, ...]
    graph_hash: str | None = None
    g_head_mtime: str | None = None
    g_track_mtime: str | None = None

    def as_record(self) -> dict[str, Any]:
        """The subset that goes into run.json / the ledger row."""
        return {
            "graph_dir": str(self.graph_dir),
            "graph_hash": self.graph_hash,
            "g_head_mtime": self.g_head_mtime,
            "g_track_mtime": self.g_track_mtime,
        }


def graph_artifact_status(graph_dir: Path) -> GraphArtifacts:
    """Inspect a resolved graph dir for the two artifacts the sentinel needs.

    Args:
        graph_dir: A resolved graph-artifact directory.

    Returns:
        The status. graph_hash is read straight off g_head.npz's own stamp (no need to load
        nodes/edges/the network just to log it); the mtimes are ISO-8601 UTC strings.

    Never raises: a suite with no graph built yet gets ok=False and None fields, not a crash.
    Both drivers fail loud on not ok for any non-base mode BEFORE launching anything
    GPU-resident -- a missing g_track.npz otherwise means an unfit belief-filter kernel and
    onf.graph.run.readout.BasinReadout raising at the first check, hours into a run.
    """
    import numpy as np
    from onf.graph.core import schema as graph_schema

    graph_dir = Path(graph_dir)
    head = graph_dir / graph_schema.HEAD_NPZ
    track = graph_dir / graph_schema.TRACK_NPZ
    missing = [str(p) for p in (head, track) if not p.exists()]
    graph_hash = head_mtime = track_mtime = None

    if head.exists():
        try:
            with np.load(head, allow_pickle=False) as raw:
                if "graph_hash" in raw.files:
                    graph_hash = str(raw["graph_hash"])
        except Exception as exc:   # noqa: BLE001 -- report, never crash a status probe
            missing.append(f"{head} (unreadable: {exc})")
        head_mtime = _iso_mtime(head)
    if track.exists():
        track_mtime = _iso_mtime(track)

    return GraphArtifacts(
        graph_dir=graph_dir, ok=not missing, missing=tuple(missing),
        graph_hash=graph_hash, g_head_mtime=head_mtime, g_track_mtime=track_mtime,
    )


def _iso_mtime(path: Path) -> str:
    return datetime.fromtimestamp(path.stat().st_mtime, tz=timezone.utc).isoformat(timespec="seconds")


class EnvBuilder:
    """Builds the environment every child process of a cell inherits.

    The conda prefix, NVML preload and PYTHONPATHs are resolved ONCE per driver rather than per cell:
    they cannot change mid-run, and each involves filesystem probing.
    """

    def __init__(self, config: LadderConfig) -> None:
        self._config = config
        self._libero_home = libero_home_dir()
        self._conda_prefix = resolve_conda_env_prefix(STABLEVLA_CONDA_ENV)
        self._nvml_preload = nvml_preload_path()
        self._server_pythonpath = self._pythonpath([STABLEVLA_ROOT, config.paths.harness_dir, ONF_SRC])
        self._client_pythonpath = self._pythonpath(
            [config.paths.harness_dir, STABLEVLA_ROOT, self._libero_home, ONF_SRC, config.paths.shim_dir]
        )

    @property
    def libero_home(self) -> Path:
        return self._libero_home

    @property
    def server_pythonpath(self) -> str:
        return self._server_pythonpath

    @property
    def client_pythonpath(self) -> str:
        return self._client_pythonpath

    def for_cell(self, cell: Cell) -> dict[str, str]:
        """The full env recipe for one cell, recorded verbatim in the ledger and run.json."""
        task_classification = (
            self._libero_home / "libero" / "libero" / "benchmark" / "task_classification.json"
        )
        env: dict[str, str] = {
            "PYTHONNOUSERSITE": "1",
            "OMP_NUM_THREADS": "4", "MKL_NUM_THREADS": "4", "OPENBLAS_NUM_THREADS": "4",
            "LOAD_VISION": "0",
            "QNDF_DIR": str(self._config.suite(cell.suite).fwm_dir),
            "LIBERO_HOME": str(self._libero_home),
            "LIBERO_CONFIG_PATH": str(self._libero_home / "libero"),
            "LIBERO_PLUS_TASK_CLASSIFICATION": str(task_classification),
            "MUJOCO_GL": "osmesa", "PYOPENGL_PLATFORM": "osmesa",
        }
        if self._conda_prefix is not None:
            env["MAGICK_HOME"] = str(self._conda_prefix)
            env["LD_LIBRARY_PATH"] = self._prepend(
                str(self._conda_prefix / "lib"), os.environ.get("LD_LIBRARY_PATH", "")
            )
        if self._nvml_preload is not None:
            # torch's CUDA caching allocator calls nvmlInit_v2 whenever it needs to GROW its pool and
            # hard-asserts on failure (a nondeterministic, memory-pressure-dependent crash: only the
            # shard that happens to need more pool at the wrong moment dies -- graph/track-mode runs
            # allocate more and hit this far more often than a plain base rollout). Every launched
            # process (server AND client shards) inherits the fix from here, not an ad-hoc export.
            env["LD_PRELOAD"] = self._prepend(self._nvml_preload, os.environ.get("LD_PRELOAD", ""))
        if cell.filter_file is not None:
            env["LIBERO_INSTANCE_ID_FILTER"] = str(cell.filter_file)

        env.update(self._config.modes.env_for(cell.mode))
        env.update(cell.extra_env)   # explicit per-cell overrides (a yaml cell's env: block) win last
        return env

    @staticmethod
    def _prepend(value: str, existing: str) -> str:
        return f"{value}:{existing}" if existing else value

    @staticmethod
    def _pythonpath(dirs: Sequence[Path]) -> str:
        parts = [str(d) for d in dirs]
        existing = os.environ.get("PYTHONPATH")
        if existing:
            parts.append(existing)
        return os.pathsep.join(parts)


# =====================================================================================================
# GPU / load preflight — torch.cuda ONLY. nvidia-smi is broken on this box (NVML library/kernel
# version mismatch); torch.cuda.device_count()/mem_get_info() work fine regardless.
# =====================================================================================================
@dataclass(frozen=True)
class GpuStatus:
    index: int
    free_gib: float
    total_gib: float
    ok: bool
    reason: str = ""


@dataclass(frozen=True)
class LoadStatus:
    load1: float
    ncpu: int
    threshold: float

    @property
    def ok(self) -> bool:
        return self.load1 <= self.threshold


@dataclass(frozen=True)
class CheckResult:
    ok: bool
    message: str


class Preflight:
    """The "is this box ready, and are we being polite" checks."""

    def __init__(self, config: LadderConfig) -> None:
        self._config = config
        self._limits = config.limits

    def gpus(self, indices: Sequence[int] | None = None) -> list[GpuStatus]:
        """Free-memory status per GPU. indices=None reports every visible device."""
        import torch

        device_count = torch.cuda.device_count()
        wanted = list(indices) if indices is not None else list(range(device_count))
        statuses = []
        for index in wanted:
            if not 0 <= index < device_count:
                statuses.append(GpuStatus(
                    index, 0.0, 0.0, False,
                    f"GPU {index} does not exist (device_count={device_count})",
                ))
                continue
            free, total = torch.cuda.mem_get_info(index)
            free_gib, total_gib = free / 1024**3, total / 1024**3
            ok = free_gib >= self._limits.min_free_gib
            reason = "" if ok else (
                f"only {free_gib:.1f} GiB free (< {self._limits.min_free_gib:.1f} GiB threshold)"
            )
            statuses.append(GpuStatus(index, free_gib, total_gib, ok, reason))
        return statuses

    def load(self) -> LoadStatus:
        ncpu = os.cpu_count() or 1
        return LoadStatus(load1=os.getloadavg()[0], ncpu=ncpu,
                          threshold=ncpu * self._limits.max_loadavg_factor)

    def libero_plus(self, timeout_s: float = 120.0) -> CheckResult:
        """Verify the LIBERO-Plus checkout resolves and reports the episode counts the ladder expects."""
        libero_home = libero_home_dir()
        benchmark_init = libero_home / "libero" / "libero" / "benchmark" / "__init__.py"
        if not libero_home.is_dir() or not benchmark_init.exists():
            return CheckResult(False, (
                f"external/LIBERO-plus missing/empty at {libero_home} "
                f"(no libero/libero/benchmark/__init__.py) -- clone it, see docs/SETUP.md step 1"
            ))
        python = resolve_stablevla_python()
        if python is None:
            return CheckResult(False, f"cannot locate the `{STABLEVLA_CONDA_ENV}` conda env's python")

        snippet = ("from libero.libero.benchmark import get_ids_by_category as g; "
                   "d = g('Robot Initial States'); print({k: len(v[0]) for k, v in d.items()})")
        env = dict(os.environ)
        env.update({
            "LIBERO_HOME": str(libero_home),
            "LIBERO_CONFIG_PATH": str(libero_home / "libero"),
            "PYTHONPATH": EnvBuilder._pythonpath([libero_home]),
            "PYTHONNOUSERSITE": "1",
            "OMP_NUM_THREADS": "4", "MKL_NUM_THREADS": "4", "OPENBLAS_NUM_THREADS": "4",
        })
        try:
            proc = subprocess.run([str(python), "-c", snippet], cwd=self._config.root, env=env,
                                   capture_output=True, text=True, timeout=timeout_s)
        except Exception as exc:   # noqa: BLE001 -- report, don't crash --check-env
            return CheckResult(False, f"LIBERO-Plus check crashed: {exc}")
        if proc.returncode != 0:
            tail = "\n".join(proc.stderr.strip().splitlines()[-15:])
            return CheckResult(False, f"LIBERO-Plus check FAILED (rc={proc.returncode}):\n{tail}")

        counts = self._parse_counts(proc.stdout)
        if counts is None:
            return CheckResult(False, f"LIBERO-Plus check: could not parse output: {proc.stdout!r}")
        # Expected "Robot Initial States" counts per bench, read from the ladder yaml's
        # expected_episodes table (via SUITE_BENCH) rather than a second hardcoded copy.
        expected = {
            self._config.scorers.bench(suite): table["Robot_Initial_States"]
            for suite, table in self._config.doc["expected_episodes"].items()
        }
        if any(counts.get(bench) != n for bench, n in expected.items()):
            return CheckResult(False, f"LIBERO-Plus counts mismatch: got {counts}, expected {expected}")
        return CheckResult(True, f"counts OK: {counts}")

    @staticmethod
    def _parse_counts(stdout: str) -> dict[str, int] | None:
        lines = stdout.strip().splitlines()
        match = re.search(r"\{.*\}", lines[-1]) if lines else None
        if not match:
            return None
        try:
            return ast.literal_eval(match.group(0))
        except (ValueError, SyntaxError):
            return None


# =====================================================================================================
# ledger — outputs/sr/ledger.jsonl, one append-only record per cell attempt (a "running" record at
# launch, a final complete/incomplete/crashed record after scoring). Uses onf.config.Paths.outputs()
# so it honors ONF_OUTPUTS like every other StageLogger-based run dir in this repo.
# =====================================================================================================
def default_ledger_path() -> Path:
    from onf.config import default_paths
    return default_paths().outputs("sr", "ledger.jsonl")


@dataclass
class LedgerRecord:
    """One row of ledger.jsonl. Written twice per attempt: running, then the outcome."""

    tag: str
    suite: str
    bench: str
    axis: str
    mode: str
    status: RunStatus
    ts: str = ""
    gpu: int | None = None
    port: int | None = None
    expected_n: int | None = None
    actual_n: int | None = None
    successes: int | None = None
    rate: float | None = None
    wall_s: float | None = None
    episodes_per_min: float | None = None
    git_sha: str | None = None
    git_dirty: bool | None = None
    graph_dir: str | None = None
    graph_hash: str | None = None
    g_head_mtime: str | None = None
    result_dir: str | None = None
    log_dir: str | None = None
    env: dict[str, str] = field(default_factory=dict)
    error: str | None = None

    def to_dict(self) -> dict[str, Any]:
        """JSON-ready payload; stamps ts at write time when the caller left it blank."""
        payload = dataclasses.asdict(self)
        payload["status"] = self.status.value
        payload["ts"] = self.ts or datetime.now(timezone.utc).isoformat(timespec="seconds")
        return payload

    @classmethod
    def from_dict(cls, payload: Mapping[str, Any]) -> "LedgerRecord":
        """Parse one JSONL row, tolerating rows written by an older schema: unknown keys are dropped,
        absent identity fields become "", and an unrecognised status reads as CRASHED."""
        known = {f.name for f in dataclasses.fields(cls)}
        data = {k: v for k, v in payload.items() if k in known}
        for name in ("tag", "suite", "bench", "axis", "mode"):
            data.setdefault(name, "")
        try:
            data["status"] = RunStatus(data.get("status", "crashed"))
        except ValueError:
            data["status"] = RunStatus.CRASHED
        return cls(**data)


class Ledger:
    """Append-only JSONL run record. Thread-safe: one worker thread per GPU writes to it."""

    COLUMNS = (
        "ts", "tag", "suite", "axis", "mode", "status", "actual_n", "expected_n",
        "successes", "rate", "wall_s", "episodes_per_min", "gpu",
    )

    def __init__(self, path: Path | None = None) -> None:
        self._path = path or default_ledger_path()
        self._lock = threading.Lock()

    @property
    def path(self) -> Path:
        return self._path

    def append(self, record: LedgerRecord) -> None:
        self._path.parent.mkdir(parents=True, exist_ok=True)
        line = json.dumps(record.to_dict(), default=str)
        with self._lock, open(self._path, "a") as f:
            f.write(line + "\n")

    def read(self) -> list[LedgerRecord]:
        if not self._path.exists():
            return []
        with open(self._path) as f:
            return [LedgerRecord.from_dict(json.loads(line)) for line in f if line.strip()]

    def latest(self, tag: str, mode: str | None = None,
               records: Sequence[LedgerRecord] | None = None) -> LedgerRecord | None:
        """Last (most recent -- JSONL is append-ordered) record for tag, or None.

        mode, if given, must ALSO match. Tags bake the mode into their own string (see
        cell_tag), so this is defense in depth -- it guards a ledger record written before
        that fix landed (an old-style tag with no mode suffix) from being mistaken for a match.
        """
        candidates = records if records is not None else self.read()
        match: LedgerRecord | None = None
        for record in candidates:
            if record.tag == tag and (mode is None or record.mode == mode):
                match = record
        return match

    def table(self, records: Sequence[LedgerRecord] | None = None) -> str:
        """The ledger rendered as an aligned text table."""
        source = records if records is not None else self.read()
        rows = [
            {c: ("" if getattr(r, c, None) is None else getattr(r, c)) for c in self.COLUMNS}
            for r in source
        ]
        for row in rows:
            if isinstance(row["status"], RunStatus):
                row["status"] = row["status"].value
        widths = {c: max(len(c), *(len(str(row[c])) for row in rows)) for c in self.COLUMNS}
        header = "  ".join(c.ljust(widths[c]) for c in self.COLUMNS)
        lines = [header, "-" * len(header)]
        lines += ["  ".join(str(row[c]).ljust(widths[c]) for c in self.COLUMNS) for row in rows]
        return "\n".join(lines)


# =====================================================================================================
# run.json provenance — written into EVERY result directory by both drivers so any success-rate number
# can later be traced back to the exact code + artifacts that produced it.
# =====================================================================================================
def write_run_json(result_dir: Path, **fields: Any) -> Path:
    """Write <result_dir>/run.json and return its path."""
    result_dir = Path(result_dir)
    result_dir.mkdir(parents=True, exist_ok=True)
    payload = {"written_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), **fields}
    out = result_dir / "run.json"
    with open(out, "w") as f:
        json.dump(payload, f, indent=2, default=str)
    return out


@dataclass(frozen=True)
class GitState:
    """The repo's commit + dirtiness at launch. Both fields are None when git is unavailable."""

    sha: str | None
    dirty: bool | None

    @classmethod
    def probe(cls, root: Path = REPO_ROOT) -> "GitState":
        return cls(sha=cls._run(["rev-parse", "HEAD"], root), dirty=cls._dirty(root))

    @staticmethod
    def _run(args: Sequence[str], root: Path) -> str | None:
        try:
            proc = subprocess.run(["git", *args], cwd=root, capture_output=True, text=True, timeout=10)
        except (OSError, subprocess.SubprocessError):
            return None
        return proc.stdout.strip() if proc.returncode == 0 else None

    @classmethod
    def _dirty(cls, root: Path) -> bool | None:
        out = cls._run(["status", "--porcelain"], root)
        return None if out is None else bool(out)


# =====================================================================================================
# scoring / completeness / gates
# =====================================================================================================
class CellScorer:
    """Scores a cell's result dir, applies its gate, and pairs it against a comparator."""

    def __init__(self, config: LadderConfig) -> None:
        self._config = config
        self._scorers = config.scorers
        self._results_root = config.results_root

    def score(self, cell: Cell) -> CellScore:
        return CellScore(cell=cell, count=self._scorers.count_dir(cell.result_dir(self._results_root)))

    def report(self, cell: Cell, allow_partial: bool) -> CellReport:
        score = self.score(cell)
        return CellReport(
            score=score,
            verdict=self.verdict(score, allow_partial),
            mcnemar=self.mcnemar(cell) if cell.gate else None,
        )

    def verdict(self, score: CellScore, allow_partial: bool) -> GateVerdict:
        """Completeness first, then the cell's own gate. An incomplete cell fails regardless of its
        gate: a partial mp4 glob silently under-counts n and makes every rate meaningless."""
        cell = score.cell
        if not score.complete and not allow_partial:
            return GateVerdict(False, (
                f"INCOMPLETE: {score.actual_n}/{cell.expected_n} episodes "
                f"(excluded from totals; pass --allow-partial to score anyway)"
            ))
        if cell.gate is None:
            return GateVerdict(True, "no gate configured")
        return cell.gate.evaluate(score)

    def mcnemar(self, cell: Cell, label: str | None = None) -> str | None:
        """The paired McNemar discordance vs a comparator, or None when it can't be computed.

        On this deterministic bench an effect is 1-4 episodes, so the marginal rate alone is not
        interpretable -- the paired test is what says whether the same episodes flipped.
        """
        comparator_label = label or (cell.gate.vs if cell.gate else None)
        if comparator_label is None:
            return None
        comparator = cell.comparator(comparator_label)
        if comparator is None:
            return None
        comparator_dir = cell.comparator_dir(self._results_root, comparator)
        cell_dir = cell.result_dir(self._results_root)
        if not comparator_dir.exists() or not cell_dir.exists():
            return None

        import numpy as np

        baseline = self._scorers.episode_outcomes(comparator_dir)
        measured = self._scorers.episode_outcomes(cell_dir)
        keys = sorted(set(baseline) & set(measured))
        if not keys:
            return f"McNemar vs {comparator.label}: no shared episodes"
        base_vec = np.array([baseline[k] for k in keys])
        cell_vec = np.array([measured[k] for k in keys])
        b_only = int((base_vec & (~cell_vec)).sum())        # comparator-only success
        c_only = int(((~base_vec) & cell_vec).sum())        # this cell-only success
        p = self._scorers.exact_p(b_only, c_only)
        return f"McNemar b/c={b_only}/{c_only} p={p:.2e} n={len(keys)}"


# =====================================================================================================
# scheduling — longest-processing-time (order by expected_n desc, greedily assign to the
# currently-lightest GPU track) + a rolling episodes/min estimate for the --dry-run makespan.
# =====================================================================================================
@dataclass(frozen=True)
class ScheduledCell:
    gpu: int
    cell: Cell
    projected_min: float


@dataclass(frozen=True)
class Schedule:
    """An LPT assignment of cells to GPU tracks, with a projected makespan."""

    assignments: dict[int, list[Cell]]
    lines: tuple[ScheduledCell, ...]
    per_gpu_minutes: dict[int, float]

    @property
    def makespan_min(self) -> float:
        return max(self.per_gpu_minutes.values()) if self.per_gpu_minutes else 0.0

    @classmethod
    def build(cls, cells: Sequence[Cell], gpus: Sequence[int], config: LadderConfig,
              ledger_records: Sequence[LedgerRecord]) -> "Schedule":
        assignments: dict[int, list[Cell]] = {g: [] for g in gpus}
        load = dict.fromkeys(gpus, 0)
        for cell in sorted(cells, key=lambda c: c.expected_n, reverse=True):
            gpu = min(load, key=lambda k: (load[k], k))
            assignments[gpu].append(cell)
            load[gpu] += cell.expected_n

        lines: list[ScheduledCell] = []
        per_gpu_minutes: dict[int, float] = {}
        for gpu, assigned in assignments.items():
            total = 0.0
            for cell in assigned:
                rate = cls.rolling_rate(cell, config, ledger_records)
                minutes = cell.expected_n / rate if rate else float("inf")
                total += minutes
                lines.append(ScheduledCell(gpu=gpu, cell=cell, projected_min=minutes))
            per_gpu_minutes[gpu] = total
        return cls(assignments=assignments, lines=tuple(lines), per_gpu_minutes=per_gpu_minutes)

    @staticmethod
    def rolling_rate(cell: Cell, config: LadderConfig,
                     records: Sequence[LedgerRecord]) -> float:
        """Mean episodes/min over the last 5 COMPLETE runs of this exact cell, else the yaml default."""
        matches = [
            float(r.episodes_per_min)
            for r in records
            if r.suite == cell.suite and r.axis == cell.axis and r.mode == cell.mode
            and r.status is RunStatus.COMPLETE and r.episodes_per_min
        ]
        if matches:
            recent = matches[-5:]
            return sum(recent) / len(recent)
        return config.defaults.episodes_per_min


# =====================================================================================================
# resume — skip a cell only when the ledger says complete AND the mp4 count still matches; otherwise
# purge the (partial/stale) result dir, because a partial dir silently poisons score.py's glob count.
# =====================================================================================================
@dataclass(frozen=True)
class CellPlan:
    cell: Cell
    action: CellAction
    reason: str


class ResumePlanner:
    """Decides what to do with a cell's existing result dir.

    plan is a PURE decision (no filesystem writes) so --dry-run can preview it safely;
    apply performs the purge and is only ever called from the real (non-dry-run) path.
    """

    def __init__(self, config: LadderConfig, ledger: Ledger) -> None:
        self._config = config
        self._ledger = ledger
        self._results_root = config.results_root

    def plan(self, cell: Cell, resume: bool,
             records: Sequence[LedgerRecord] | None = None) -> CellPlan:
        directory = cell.result_dir(self._results_root)
        if not (directory.exists() and any(directory.rglob("*.mp4"))):
            return CellPlan(cell, CellAction.RUN, "no existing result dir")
        if not resume:
            return CellPlan(cell, CellAction.ERROR,
                            f"result dir already has content at {directory} (pass --resume)")
        record = self._ledger.latest(cell.tag, cell.mode, records)
        if record and record.status is RunStatus.COMPLETE:
            count = self._config.scorers.count_dir(directory)
            if count.n == cell.expected_n:
                return CellPlan(cell, CellAction.SKIP, (
                    f"ledger says complete and {count.n} mp4s match expected_n={cell.expected_n}"
                ))
        return CellPlan(cell, CellAction.PURGE, "partial/stale result dir -- would be purged and re-run")

    def apply(self, plan: CellPlan) -> None:
        """Perform the filesystem side effect for a purge verdict."""
        if plan.action is CellAction.PURGE:
            shutil.rmtree(plan.cell.result_dir(self._results_root))


# =====================================================================================================
# process supervision — one GPU-resident policy server + N sharded sim clients per cell.
# =====================================================================================================
class ProcessSupervisor:
    """Tracks every live child so Ctrl-C can't orphan a GPU-resident process."""

    def __init__(self) -> None:
        self._procs: set[subprocess.Popen] = set()
        self._lock = threading.Lock()

    def spawn(self, cmd: Sequence[str], env: Mapping[str, str], stdout: Any,
              cwd: Path = REPO_ROOT) -> subprocess.Popen:
        proc = subprocess.Popen(list(cmd), cwd=cwd, env=dict(env), stdout=stdout,
                                stderr=subprocess.STDOUT)
        with self._lock:
            self._procs.add(proc)
        return proc

    def forget(self, proc: subprocess.Popen) -> None:
        """Stop tracking a process the caller is now reaping itself."""
        with self._lock:
            self._procs.discard(proc)

    def terminate_all(self, grace_s: float = 2.0) -> None:
        with self._lock:
            procs = list(self._procs)
        for proc in procs:
            self._quietly(proc.terminate)
        time.sleep(grace_s)
        for proc in procs:
            if proc.poll() is None:
                self._quietly(proc.kill)

    def install_signal_handlers(self) -> None:
        """Tear everything down on SIGINT/SIGTERM. Installed only on the real launch path, so a
        --dry-run never replaces the caller's handlers."""
        def handler(signum: int, _frame: FrameType | None) -> None:
            self.terminate_all()
            sys.exit(128 + signum)

        signal.signal(signal.SIGINT, handler)
        signal.signal(signal.SIGTERM, handler)

    @staticmethod
    def _quietly(action: Any) -> None:
        try:
            action()
        except OSError:      # best-effort teardown, never let this raise
            pass


class PortAllocator:
    """Hands out policy-server ports that are free both in this process and across the box.

    A process-local counter alone is not enough: two concurrent run_sr.py invocations (a rung on
    one GPU plus an ad-hoc --cells run on another) both start at base_port, and the second
    one's clients then connect to the FIRST one's policy server. That is silent and it mis-scores --
    caught in practice with a goal Robot-Init cell whose clients reached a long-checkpoint
    server on the shared port and produced 121 episodes under the wrong policy, with nothing in any
    log saying so.

    A bind-probe alone is not enough either: a policy server takes minutes to load its checkpoint
    before it binds, so a second process probing in that window still sees the port as free. Hence
    the filesystem reservation, taken the instant a port is chosen and outliving the gap until the
    server actually listens. A reservation lasts as long as the owning process (dropped by the atexit
    hook); one left behind by a killed run is reclaimed by the next caller's pid-liveness check.
    """

    MAX_PROBES = 256

    def __init__(self, base_port: int, reservation_dir: Path | None = None) -> None:
        self._base_port = base_port
        self._dir = reservation_dir or (REPO_ROOT / "logs" / ".ports")
        self._lock = threading.Lock()
        self._counter = itertools.count()
        self._held: set[int] = set()
        atexit.register(self.release_all)

    def claim(self) -> int:
        """The next port at or above base_port that is both unbound AND unreserved.

        Raises:
            RuntimeError: No free port within MAX_PROBES of the base.
        """
        with self._lock:
            self._dir.mkdir(parents=True, exist_ok=True)
            for _ in range(self.MAX_PROBES):
                port = self._base_port + next(self._counter)
                if self._bindable(port) and self._reserve(port):
                    self._held.add(port)
                    return port
        raise RuntimeError(
            f"no free port in [{self._base_port}, {self._base_port + self.MAX_PROBES}) "
            f"-- is something leaking servers?"
        )

    def release_all(self) -> None:
        for port in list(self._held):
            (self._dir / str(port)).unlink(missing_ok=True)
            self._held.discard(port)

    @staticmethod
    def _bindable(port: int) -> bool:
        """True iff nothing is listening -- tested by actually binding, not by connecting (a
        connect-probe races with a server still loading its checkpoint and not yet accepting)."""
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            try:
                sock.bind(("0.0.0.0", port))
                return True
            except OSError:
                return False

    def _reserve(self, port: int) -> bool:
        """Cross-process claim: create <dir>/<port> exclusively, holding this pid."""
        marker = self._dir / str(port)
        try:
            fd = os.open(marker, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
        except FileExistsError:
            return self._reclaim_if_stale(marker)
        with os.fdopen(fd, "w") as f:
            f.write(f"{os.getpid()}\n")
        return True

    @staticmethod
    def _reclaim_if_stale(marker: Path) -> bool:
        """Take over a reservation whose owning pid is gone; refuse while it is alive."""
        try:
            content = marker.read_text().strip()
            owner: int | None = int(content) if content else None
        except (ValueError, OSError):
            owner = None
        if owner is not None:
            try:
                os.kill(owner, 0)      # signal 0 = liveness probe, sends nothing
                return False           # a live run_sr owns this port
            except (ProcessLookupError, PermissionError, OSError):
                pass                   # stale reservation from a dead run -- take it over
        try:
            marker.write_text(f"{os.getpid()}\n")
            return True
        except OSError:
            return False


class PolicyServer:
    """The GPU-resident StableVLA policy server for one cell, as a context manager.

    __enter__ starts it and blocks until its port accepts; __exit__ always tears it down,
    which is what stops a client-side crash from leaking a process holding GPU memory.
    """

    def __init__(self, config: LadderConfig, supervisor: ProcessSupervisor, python: Path,
                 checkpoint: Path, port: int, log_dir: Path, env: Mapping[str, str]) -> None:
        self._config = config
        self._supervisor = supervisor
        self._python = python
        self._checkpoint = checkpoint
        self._port = port
        self._log_path = log_dir / "server.log"
        self._env = dict(env)
        self._proc: subprocess.Popen | None = None
        self._log: TextIO | None = None

    def __enter__(self) -> "PolicyServer":
        self._log = open(self._log_path, "w")
        self._proc = self._supervisor.spawn(
            [str(self._python), str(self._config.paths.policy_server),
             "--ckpt_path", str(self._checkpoint), "--port", str(self._port),
             "--cuda", "0", "--use_bf16"],
            env=self._env, stdout=self._log,
        )
        try:
            self._wait_ready(self._config.defaults.server_ready_timeout_s)
        except BaseException:
            self.__exit__(None, None, None)
            raise
        return self

    def __exit__(self, *_exc: Any) -> None:
        if self._proc is not None:
            self._supervisor.forget(self._proc)
            self._proc.terminate()
            try:
                self._proc.wait(timeout=15)
            except subprocess.TimeoutExpired:
                self._proc.kill()
                self._proc.wait(timeout=15)
            self._proc = None
        if self._log is not None:
            self._log.close()
            self._log = None

    def _wait_ready(self, timeout_s: float, poll_every: float = 2.0) -> None:
        assert self._proc is not None
        deadline = time.time() + timeout_s
        while time.time() < deadline:
            if self._proc.poll() is not None:
                raise RuntimeError(
                    f"policy server exited before ready (rc={self._proc.returncode}); "
                    f"log tail:\n{self._log_tail()}"
                )
            try:
                with socket.create_connection(("127.0.0.1", self._port), timeout=1):
                    return
            except OSError:
                time.sleep(poll_every)
        raise RuntimeError(f"policy server not ready after {timeout_s:.0f}s on port {self._port}")

    def _log_tail(self, lines: int = 40) -> str:
        if not self._log_path.exists():
            return ""
        return "\n".join(self._log_path.read_text(errors="replace").splitlines()[-lines:])


@dataclass(frozen=True)
class LaunchResult:
    wall_s: float
    client_returncodes: tuple[int, ...]

    @property
    def crashed(self) -> bool:
        return any(rc != 0 for rc in self.client_returncodes)

    @property
    def crash_note(self) -> str:
        """How the RESULT line reports a crash, or "" when every client exited 0.

        A crashed cell used to print as merely INCOMPLETE, because the verdict is computed from the
        episode count alone and never sees the return codes -- only the ledger recorded the truth.
        One Light_Conditions cell died 156 s into a 45-minute run, wrote 4 of 273 episodes, and
        printed "4/4 = 100.0% [FAIL: INCOMPLETE]", which reads like a short run rather than a dead
        one.
        """
        bad = [rc for rc in self.client_returncodes if rc != 0]
        if not bad:
            return ""
        return (f"  [CRASHED: {len(bad)}/{len(self.client_returncodes)} clients exited "
                f"{sorted(set(bad))} -- this cell did not finish; the episode count below is "
                "whatever it had written by then]")


class CellRunner:
    """Runs one cell end to end: provenance, launch, score, ledger, verdict."""

    def __init__(self, config: LadderConfig, env_builder: EnvBuilder, scorer: CellScorer,
                 ledger: Ledger, ports: PortAllocator, supervisor: ProcessSupervisor,
                 allow_partial: bool = False) -> None:
        self._config = config
        self._env_builder = env_builder
        self._scorer = scorer
        self._ledger = ledger
        self._ports = ports
        self._supervisor = supervisor
        self._allow_partial = allow_partial
        self._git = GitState.probe(config.root)
        self._python = resolve_stablevla_python()

    def run(self, cell: Cell, gpu: int) -> ExitCode:
        """Run and score one cell on one GPU.

        Returns:
            OK when it completed and cleared its gate, GATE_FAILED when it completed but did
            not, ERROR when it crashed or came back incomplete. A crash is recorded in the ledger
            and returned, never raised -- one bad cell must not take the whole schedule down.
        """
        if self._python is None:
            raise RuntimeError(f"cannot locate the `{STABLEVLA_CONDA_ENV}` conda env's python")

        port = self._ports.claim()
        log_dir = make_log_dir(cell.policy, cell.mode, cell.suite, cell.axis, cell.variant,
                               self._config.root)
        env = self._env_builder.for_cell(cell)
        artifacts = self._require_graph_artifacts(cell, env)
        result_dir = cell.result_dir(self._config.results_root)

        write_run_json(
            result_dir, policy=cell.policy, mode=cell.mode, suite=cell.suite, axis=cell.axis,
            expected_n=cell.expected_n, git_commit=self._git.sha, git_dirty=self._git.dirty,
            mode_env=env, **artifacts.as_record(),
        )

        base = dict(
            tag=cell.tag, suite=cell.suite, bench=cell.bench, axis=cell.axis, mode=cell.mode,
            gpu=gpu, port=port, expected_n=cell.expected_n, git_sha=self._git.sha,
            git_dirty=self._git.dirty, env=env, result_dir=str(result_dir), log_dir=str(log_dir),
            graph_dir=str(artifacts.graph_dir), graph_hash=artifacts.graph_hash,
            g_head_mtime=artifacts.g_head_mtime,
        )
        self._ledger.append(LedgerRecord(status=RunStatus.RUNNING, **base))

        try:
            launch = self._launch(cell, gpu, port, log_dir, env)
        except Exception as exc:   # noqa: BLE001 -- record the crash, don't kill the schedule
            self._ledger.append(LedgerRecord(status=RunStatus.CRASHED, error=str(exc), **base))
            print(f"RESULT {cell.tag}: CRASHED -- {exc}", file=sys.stderr)
            return ExitCode.ERROR

        report = self._scorer.report(cell, self._allow_partial)
        score = report.score
        rate = (score.actual_n / (launch.wall_s / 60.0)) if launch.wall_s > 0 and score.actual_n else 0.0
        status = (RunStatus.CRASHED if launch.crashed
                  else (RunStatus.COMPLETE if score.complete else RunStatus.INCOMPLETE))
        self._ledger.append(LedgerRecord(
            status=status, actual_n=score.actual_n, successes=score.successes,
            rate=score.rate, wall_s=launch.wall_s, episodes_per_min=rate, **base,
        ))
        print(report.line + launch.crash_note)

        if status is not RunStatus.COMPLETE and not self._allow_partial:
            return ExitCode.ERROR
        return ExitCode.OK if report.verdict.passed else ExitCode.GATE_FAILED

    def _require_graph_artifacts(self, cell: Cell, env: Mapping[str, str]) -> GraphArtifacts:
        """One startup line naming the resolved graph dir, then FAIL HARD for any non-base mode
        missing an artifact -- same policy as evals/gr00t/run_gr00t.py's own check."""
        artifacts = graph_artifact_status(resolve_graph_dir(cell.suite, env))
        print(f"[run_sr] {cell.tag}: mode={cell.mode} graph_dir={artifacts.graph_dir} "
              f"graph_hash={artifacts.graph_hash} g_head_mtime={artifacts.g_head_mtime} "
              f"g_track_mtime={artifacts.g_track_mtime}")
        if not artifacts.ok and cell.mode != "base":
            raise RuntimeError(
                f"{cell.tag}: resolved graph dir {artifacts.graph_dir} is missing required artifacts: "
                f"{list(artifacts.missing)} -- train the graph for this suite first, or point "
                f"GR_GRAPH_DIR at a suite that already has g_head.npz + g_track.npz."
            )
        return artifacts

    def _launch(self, cell: Cell, gpu: int, port: int, log_dir: Path,
                cell_env: Mapping[str, str]) -> LaunchResult:
        assert self._python is not None
        started = time.time()
        result_dir = cell.result_dir(self._config.results_root)
        result_dir.mkdir(parents=True, exist_ok=True)
        log_dir.mkdir(parents=True, exist_ok=True)

        shared = {**os.environ, **cell_env, "CUDA_VISIBLE_DEVICES": str(gpu)}
        checkpoint = self._config.suite(cell.suite).stablevla_ckpt
        server_env = {**shared, "PYTHONPATH": self._env_builder.server_pythonpath}
        client_env = {**shared, "PYTHONPATH": self._env_builder.client_pythonpath}

        with PolicyServer(self._config, self._supervisor, self._python, checkpoint, port,
                          log_dir, server_env):
            with ExitStack() as stack:
                procs = [
                    self._supervisor.spawn(
                        self._client_cmd(cell, checkpoint, port, shard, result_dir),
                        env=client_env,
                        stdout=stack.enter_context(open(log_dir / f"client_shard{shard}.log", "w")),
                    )
                    for shard in range(cell.num_clients)
                ]
                returncodes = []
                for proc in procs:
                    returncodes.append(proc.wait())
                    self._supervisor.forget(proc)

        return LaunchResult(wall_s=time.time() - started, client_returncodes=tuple(returncodes))

    def _client_cmd(self, cell: Cell, checkpoint: Path, port: int, shard: int,
                    result_dir: Path) -> list[str]:
        assert self._python is not None
        return [
            str(self._python), str(self._config.paths.harness_dir / "eval_libero.py"),
            "--args.pretrained-path", str(checkpoint),
            "--args.host", "127.0.0.1",
            "--args.port", str(port),
            "--args.task-suite-name", cell.bench,
            "--args.category_value", cell.category_value,
            "--args.num-trials-per-task", "1",
            "--args.with_state", "true",
            "--args.shard_index", str(shard),
            "--args.num_shards", str(cell.num_clients),
            "--args.video-out-path", str(result_dir),
        ]


# =====================================================================================================
# CLI
# =====================================================================================================
class SRDriver:
    """Dispatches one CLI invocation. Owns the lazily-built LadderConfig and the ledger."""

    def __init__(self, args: argparse.Namespace, ledger: Ledger | None = None) -> None:
        self._args = args
        self._config: LadderConfig | None = None
        self._ledger = ledger or Ledger()

    @property
    def config(self) -> LadderConfig:
        """Built on first use: --ledger never needs to read suites.yaml or modes.sh."""
        if self._config is None:
            self._config = LadderConfig.load(
                Path(self._args.ladder),
                Path(self._args.results_root) if self._args.results_root else None,
            )
        return self._config

    def run(self) -> ExitCode:
        args = self._args
        if args.check_env:
            return self._check_env()
        if args.ledger:
            return self._print_ledger()

        cells = self._resolve_cells()
        if not cells:
            print("[run_sr] nothing to do -- pass --rung or --cells", file=sys.stderr)
            return ExitCode.ERROR
        if args.score:
            return self._score(cells)

        gpus = self._parse_gpus(args.gpus)
        if not gpus:
            print("[run_sr] --gpus is required to launch or dry-run a schedule", file=sys.stderr)
            return ExitCode.ERROR
        return self._launch(cells, gpus)

    # ---- cell resolution -------------------------------------------------------------------------
    def _resolve_cells(self) -> list[Cell]:
        args = self._args
        if args.rung:
            if args.rung not in self.config.doc["rungs"]:
                raise SystemExit(
                    f"unknown rung {args.rung!r}; known: {sorted(self.config.doc['rungs'])}"
                )
            cells = self.config.cells_for_rung(args.rung)
            if args.cells:
                cells = self._subset(cells, args.cells, args.rung)
            if args.mode:
                # with_mode (not a bare dataclasses.replace) so the tag's mode suffix is re-keyed too
                # -- see cell_tag's docstring for the wrong-number hazard a stale tag creates.
                cells = [c.with_mode(args.mode) for c in cells]
            return cells
        if args.cells:
            return self.config.cells_from_tokens(args.cells, args.mode)
        return []

    @staticmethod
    def _subset(cells: Sequence[Cell], spec: str, rung: str) -> list[Cell]:
        """--rung R --cells suite:axis,... runs a SUBSET of the rung, keeping each cell's own
        expected_n / filter / comparators / gate. Without this, re-running one cell means --cells
        alone, which rebuilds it as an ad hoc full-axis cell and silently drops the rung's episode
        filter -- a 12-episode smoke would become a 393-episode run."""
        wanted = {t.strip() for t in spec.split(",") if t.strip()}
        unknown = wanted - {c.key for c in cells}
        if unknown:
            raise SystemExit(
                f"--cells: {sorted(unknown)} not in rung {rung} (has: {sorted(c.key for c in cells)})"
            )
        return [c for c in cells if c.key in wanted]

    @staticmethod
    def _parse_gpus(spec: str | None) -> list[int]:
        """"0,1" -> [0, 1], preserving order and dropping repeats (a duplicated index would
        otherwise open two schedule tracks against the same device)."""
        if not spec:
            return []
        seen: dict[int, None] = {}
        for token in spec.split(","):
            if token.strip():
                seen[int(token)] = None
        return list(seen)

    # ---- subcommands -----------------------------------------------------------------------------
    def _check_env(self) -> ExitCode:
        preflight = Preflight(self.config)
        ok = True

        print("=== GPU ===")
        try:
            statuses = preflight.gpus()
            if not statuses:
                print("  no CUDA devices visible")
                ok = False
            for status in statuses:
                label = "OK" if status.ok else "BUSY"
                print(f"  gpu{status.index}: {status.free_gib:6.1f} / {status.total_gib:6.1f} GiB free"
                      f"  [{label}]")
                ok = ok and status.ok
        except Exception as exc:   # noqa: BLE001 -- a broken torch must not crash the report
            print(f"  GPU check FAILED: {exc}")
            ok = False

        print("=== load ===")
        load = preflight.load()
        print(f"  loadavg(1m)={load.load1:.2f}  ncpu={load.ncpu}  threshold={load.threshold:.1f}  "
              f"[{'OK' if load.ok else 'BUSY'}]")
        ok = ok and load.ok

        print("=== LIBERO-Plus ===")
        libero = preflight.libero_plus()
        print(f"  {libero.message}")
        ok = ok and libero.ok

        print(f"\n[check-env] overall: {'OK' if ok else 'NOT READY'}")
        return ExitCode.OK if ok else ExitCode.ERROR

    def _print_ledger(self) -> ExitCode:
        records = self._ledger.read()
        if not records:
            print(f"[run_sr] ledger is empty (no runs recorded yet): {self._ledger.path}")
            return ExitCode.OK
        print(self._ledger.table(records))
        return ExitCode.OK

    def _score(self, cells: Sequence[Cell]) -> ExitCode:
        scorer = CellScorer(self.config)
        any_scored = False
        any_failed = False
        for cell in cells:
            report = scorer.report(cell, self._args.allow_partial)
            if report.score.actual_n == 0:
                print(f"RESULT {cell.tag}: NO DATA (expected {cell.expected_n} episodes at "
                      f"{cell.result_dir(self.config.results_root)})")
                continue
            any_scored = True
            any_failed = any_failed or not report.verdict.passed
            print(report.line)
        if not any_scored:
            return ExitCode.ERROR
        return ExitCode.GATE_FAILED if any_failed else ExitCode.OK

    # ---- launch ----------------------------------------------------------------------------------
    def _launch(self, cells: Sequence[Cell], gpus: Sequence[int]) -> ExitCode:
        usable_gpus = self._validate_system(cells, gpus)
        if not usable_gpus:
            return ExitCode.ERROR

        plans = self._build_execution_plan(cells)
        if plans is None:
            return ExitCode.ERROR
        if not plans:
            print("[run_sr] nothing left to run (everything already complete)")
            return ExitCode.OK

        schedule = Schedule.build([p.cell for p in plans], usable_gpus, self.config,
                                  self._ledger.read())
        self._print_schedule(schedule, usable_gpus)

        if self._args.dry_run:
            print("\n[dry-run] no subprocess launched, no filesystem changes made.")
            return ExitCode.OK
        return self._execute_schedule(plans, schedule)

    def _validate_system(self, cells: Sequence[Cell], gpus: Sequence[int]) -> list[int]:
        """GPU + loadavg preflight. Returns the usable GPUs, or [] to refuse the launch."""
        preflight = Preflight(self.config)
        statuses = preflight.gpus(gpus)
        for status in statuses:
            if not status.ok:
                print(f"[preflight] SKIP gpu{status.index}: {status.reason}")
        sys.stdout.flush()   # keep stdout ahead of the stderr refusal below when interleaved (2>&1)

        usable = [s.index for s in statuses if s.ok]
        if not usable:
            print("[preflight] no usable GPUs -- refusing to start anything", file=sys.stderr)
            return []

        total_expected = sum(c.expected_n for c in cells)
        is_large = total_expected > self.config.limits.large_rung_episodes
        load = preflight.load()
        print(f"[preflight] load1={load.load1:.2f} ncpu={load.ncpu} threshold={load.threshold:.1f} "
              f"large_launch={is_large} (total_expected={total_expected})")
        sys.stdout.flush()
        if is_large and not load.ok:
            print(f"[preflight] REFUSING: a large launch ({total_expected} episodes) and loadavg(1m)="
                  f"{load.load1:.2f} exceeds {load.threshold:.1f} -- box looks busy with a "
                  f"neighbouring job", file=sys.stderr)
            return []

        for cell in cells:
            if cell.filter_file is not None and not cell.filter_file.exists():
                print(f"[preflight] WARNING: {cell.tag}: filter file missing at {cell.filter_file}")
        return usable

    def _build_execution_plan(self, cells: Sequence[Cell]) -> list[CellPlan] | None:
        """The per-cell resume decisions. None means refuse the whole launch."""
        planner = ResumePlanner(self.config, self._ledger)
        records = self._ledger.read()
        runnable: list[CellPlan] = []
        for cell in cells:
            plan = planner.plan(cell, self._args.resume, records)
            print(f"[plan] {cell.tag}: {plan.action.value} ({plan.reason})")
            if plan.action is CellAction.ERROR:
                print(f"[run_sr] {cell.tag}: {plan.reason}", file=sys.stderr)
                return None
            if plan.action in (CellAction.RUN, CellAction.PURGE):
                runnable.append(plan)
        return runnable

    def _execute_schedule(self, plans: Sequence[CellPlan], schedule: Schedule) -> ExitCode:
        # Apply any purges now -- never before this point, so --dry-run (even with --resume) is a
        # pure read-only preview and can never delete a partial result dir.
        planner = ResumePlanner(self.config, self._ledger)
        for plan in plans:
            planner.apply(plan)

        supervisor = ProcessSupervisor()
        supervisor.install_signal_handlers()
        runner = CellRunner(
            config=self.config,
            env_builder=EnvBuilder(self.config),
            scorer=CellScorer(self.config),
            ledger=self._ledger,
            ports=PortAllocator(self.config.defaults.base_port,
                                self.config.root / "logs" / ".ports"),
            supervisor=supervisor,
            allow_partial=self._args.allow_partial,
        )

        def run_track(gpu: int, track_cells: Sequence[Cell]) -> list[ExitCode]:
            return [runner.run(cell, gpu) for cell in track_cells]

        # One worker thread per GPU, each running its LPT-assigned cells sequentially. Threads (not
        # processes) are fine: each cell's work is dominated by subprocess I/O wait.
        codes: list[ExitCode] = []
        assignments = {g: c for g, c in schedule.assignments.items() if c}
        with ThreadPoolExecutor(max_workers=max(1, len(assignments))) as pool:
            futures = [pool.submit(run_track, gpu, cells) for gpu, cells in assignments.items()]
            for future in as_completed(futures):
                codes.extend(future.result())

        if ExitCode.GATE_FAILED in codes:
            return ExitCode.GATE_FAILED
        if any(code is ExitCode.ERROR for code in codes):
            return ExitCode.ERROR
        return ExitCode.OK

    @staticmethod
    def _print_schedule(schedule: Schedule, gpus: Sequence[int]) -> None:
        print("\n=== schedule ===")
        for line in schedule.lines:
            cell = line.cell
            print(f"  gpu{line.gpu}: {cell.tag:<45} n={cell.expected_n:<5} "
                  f"mode={cell.mode:<16} ~{line.projected_min:6.1f} min")
        print(f"projected makespan: {schedule.makespan_min:.1f} min "
              f"({schedule.makespan_min / 60.0:.1f} h) across {len(gpus)} GPU(s)")


def build_argparser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="run_sr.py", description="ONF LIBERO-Plus success-rate eval driver."
    )
    parser.add_argument("--rung", help="ladder rung name (R1, R4)")
    parser.add_argument("--cells", help="comma-separated suite:axis pairs, "
                                         "e.g. goal:Background_Textures,long:Camera_Viewpoints")
    parser.add_argument("--gpus", help="comma-separated physical GPU indices, e.g. 0,1")
    parser.add_argument("--mode", help="mode override (onf_set_mode name)")
    parser.add_argument("--dry-run", action="store_true", help="print the schedule/projection; launch nothing")
    parser.add_argument("--resume", action="store_true", help="skip complete cells, purge+rerun partial ones")
    parser.add_argument("--score", action="store_true", help="score existing results only; launch nothing")
    parser.add_argument("--ledger", action="store_true", help="print the ledger as a table and exit")
    parser.add_argument("--check-env", action="store_true", help="GPU/load/LIBERO-Plus preflight report")
    parser.add_argument("--allow-partial", action="store_true", help="include incomplete cells in totals/gates")
    parser.add_argument("--ladder", default=str(LADDER_YAML), help="path to the ladder yaml")
    parser.add_argument("--results-root", default=None, help="override the results/ root (else score.py's default)")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = build_argparser().parse_args(argv)
    return int(SRDriver(args).run())


if __name__ == "__main__":
    sys.exit(main())