onf-v3/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 recovery 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, recovery 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 recovery-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 (``finally`` + 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
    :mod:`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.

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 sentinel] [--resume]
    python scripts/run_sr.py --cells goal:Background_Textures,long:Camera_Viewpoints --gpus 0,1 --mode sentinel
    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 atexit
import ast
import dataclasses
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 datetime import datetime, timezone
from pathlib import Path
from typing import Any

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"


# =====================================================================================================
# yaml loading
# =====================================================================================================
def load_yaml(path: Path) -> dict:
    with open(path) as f:
        return yaml.safe_load(f) or {}


def load_ladder(path: Path | None = None) -> dict:
    return load_yaml(path or LADDER_YAML)


def resolve_paths(ladder_doc: dict) -> dict[str, Path]:
    """``ladder_doc["paths"]`` (all relative to $ONF_ROOT) -> absolute Paths. 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)."""
    p = ladder_doc["paths"]
    return {k: REPO_ROOT / v for k, v in p.items()}


@dataclasses.dataclass(frozen=True)
class SuiteInfo:
    name: str
    bench: str
    stablevla_ckpt: Path
    fwm_dir: Path
    light_filter: Path


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


# =====================================================================================================
# 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.
# =====================================================================================================
_module_cache: dict[str, Any] = {}


def _load_module(name: str, path: Path):
    cached = _module_cache.get(str(path))
    if cached is not None:
        return cached
    import importlib.util

    spec = importlib.util.spec_from_file_location(name, path)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod   # register before exec so dataclasses/typing self-references resolve
    spec.loader.exec_module(mod)
    _module_cache[str(path)] = mod
    return mod


def score_mod(paths: dict[str, Path]):
    return _load_module("sr_score", paths["score_py"])


def mcnemar_mod(paths: dict[str, Path]):
    return _load_module("sr_mcnemar", paths["mcnemar_py"])


def default_results_root(paths: dict[str, Path]) -> Path:
    return Path(score_mod(paths).results_root())


# =====================================================================================================
# evals/common/modes.sh parser — the recovery-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*:-([^}]*)\}"?$')

_mode_recipes_cache: dict[str, dict[str, dict[str, str]]] = {}


def parse_mode_recipes(modes_sh_text: str) -> dict[str, dict[str, str]]:
    """Parse ``onf_set_mode()``'s ``case "${mode}" in ... esac`` block into
    ``{mode_name: {ENV_VAR: value, ...}}``. Every recipe includes the ``JOINT_RECOVERY=0 SENTINEL=0``
    baseline the function always exports before dispatch (mirrors ``export JOINT_RECOVERY=0
    SENTINEL=0`` at the top of ``onf_set_mode``), then whatever the matched arm overrides."""
    m = _CASE_BLOCK_RE.search(modes_sh_text)
    if not m:
        raise ValueError('modes.sh: could not find `case "${mode}" in ... esac` block')
    body = m.group(1)

    recipes: dict[str, dict[str, str]] = {}
    parts = _ARM_SPLIT_RE.split(body)
    # parts[0] is text before the first arm header (comments); then alternating (pattern, arm_text).
    for i in range(1, len(parts), 2):
        pattern, arm_text = parts[i], parts[i + 1]
        env = {"JOINT_RECOVERY": "0", "SENTINEL": "0"}
        env.update(_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]:
    text = arm_text.split(";;")[0]                      # stop at the arm terminator
    lines = [ln.split("#", 1)[0] for ln in text.splitlines()]   # strip comments (no '#' in any value)
    joined = ""
    for ln in lines:
        stripped = ln.rstrip()
        joined += (stripped[:-1] + " ") if stripped.endswith("\\") else (stripped + " ")
    out: dict[str, str] = {}
    for tok in joined.split():
        if tok == "export":
            continue
        m = _ASSIGN_RE.match(tok)
        if m:
            out[m.group(1)] = _expand_value(m.group(2))
    return out


def _expand_value(val: str) -> str:
    m = _PARAM_DEFAULT_RE.match(val)
    if m:
        return m.group(1)
    return val.strip('"')


def mode_env(mode: str, paths: dict[str, Path]) -> dict[str, str]:
    modes_sh = paths["modes_sh"]
    key = str(modes_sh)
    if key not in _mode_recipes_cache:
        _mode_recipes_cache[key] = parse_mode_recipes(modes_sh.read_text())
    recipes = _mode_recipes_cache[key]
    if mode not in recipes:
        raise ValueError(f"unknown MODE={mode!r}; not found in {modes_sh} (known: {sorted(recipes)})")
    return dict(recipes[mode])


# =====================================================================================================
# ladder cell model
# =====================================================================================================
@dataclasses.dataclass(frozen=True)
class Comparator:
    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")


@dataclasses.dataclass(frozen=True)
class Gate:
    kind: str                       # min_successes | max_regress_pp | exact_match
    vs: str = "base"
    min_successes: int | None = None
    max_regress_pp: float | None = None

    @classmethod
    def from_dict(cls, d: dict) -> Gate:
        return cls(kind=d["kind"], vs=d.get("vs", "base"),
                    min_successes=d.get("min_successes"), max_regress_pp=d.get("max_regress_pp"))


@dataclasses.dataclass
class Cell:
    rung: str
    suite: str
    axis: str
    mode: str
    expected_n: int
    filter_file: Path | None = None
    extra_env: dict[str, str] = dataclasses.field(default_factory=dict)
    comparators: list[Comparator] = dataclasses.field(default_factory=list)
    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)

    def bench(self, paths: dict[str, Path]) -> str:
        return score_mod(paths).SUITE_BENCH[self.suite]

    @property
    def category_value(self) -> str:
        return self.axis.replace("_", " ")

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

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


def default_mode_for_axis(axis: str, ladder_doc: dict) -> str:
    m = ladder_doc["axis_default_mode"]
    return m.get(axis, m["default"])


def resolve_filter(filt: str | None, axis: str, suite: str, suites: dict[str, SuiteInfo]) -> Path | None:
    if filt in (None, "auto"):
        if filt == "auto" or axis == "Light_Conditions":
            return suites[suite].light_filter
        return None
    return REPO_ROOT / filt


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 cell_with_mode(cell: "Cell", mode: str) -> "Cell":
    """Swap ``cell``'s mode. The tag follows automatically -- it is derived, not stored."""
    return dataclasses.replace(cell, mode=mode)


def build_cell(rung_name: str, cell_doc: dict, ladder_doc: dict, suites: dict[str, SuiteInfo]) -> Cell:
    suite, axis = cell_doc["suite"], cell_doc["axis"]
    mode = cell_doc.get("mode") or default_mode_for_axis(axis, ladder_doc)
    comparators = [Comparator(**c) for c in cell_doc.get("comparators", [])]
    gate_doc = cell_doc.get("gate") or ladder_doc.get("default_gate")
    return Cell(
        rung=rung_name, suite=suite, axis=axis, mode=mode, variant=cell_doc.get("variant", ""),
        expected_n=int(cell_doc["expected_n"]),
        filter_file=resolve_filter(cell_doc.get("filter"), axis, suite, suites),
        extra_env={k: str(v) for k, v in (cell_doc.get("env") or {}).items()},
        comparators=comparators, gate=Gate.from_dict(gate_doc) if gate_doc else None,
        num_clients=int(cell_doc.get("num_clients", ladder_doc["defaults"]["num_clients"])),
    )


def resolve_r4_comparators(suite: str, axis: str, results_root: Path, paths: dict[str, Path]) -> list[Comparator]:
    """R4's comparators aren't hand-enumerated (28 cells would go stale the moment a new baseline
    lands): probe whichever PRE-EXISTING baseline result directories are already on disk, so a run
    never re-measures a baseline that has 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 SUITE_BENCH/count_dir so this can never disagree with `score.py` about what a tag
    directory scores."""
    sm = score_mod(paths)
    bench = sm.SUITE_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}")]
    out, seen = [], set()
    for label, tag in candidates:
        if (label, tag) in seen:
            continue
        d = results_root / f"plus_{bench}" / axis / tag
        if d.is_dir():
            s, n = sm.count_dir(str(d))
            if n:
                out.append(Comparator(label=label, tag=tag, successes=s, n=n))
                seen.add((label, tag))
    return out


def generate_full_grid_cells(rung_name: str, rung_doc: dict, ladder_doc: dict, suites: dict[str, SuiteInfo],
                              results_root: Path, paths: dict[str, Path]) -> list[Cell]:
    cells = []
    sm = score_mod(paths)
    gate = Gate.from_dict(rung_doc["gate"])
    for suite in ("object", "spatial", "goal", "long"):
        for axis in sm.AXES:
            mode = default_mode_for_axis(axis, ladder_doc)
            cells.append(Cell(
                rung=rung_name, suite=suite, axis=axis, mode=mode,
                expected_n=int(ladder_doc["expected_episodes"][suite][axis]),
                filter_file=resolve_filter(None, axis, suite, suites),
                comparators=resolve_r4_comparators(suite, axis, results_root, paths),
                gate=gate, num_clients=ladder_doc["defaults"]["num_clients"],
            ))
    return cells


def rung_cells(ladder_doc: dict, rung_name: str, suites: dict[str, SuiteInfo],
                results_root: Path, paths: dict[str, Path]) -> list[Cell]:
    rung_doc = ladder_doc["rungs"][rung_name]
    if rung_doc.get("auto_generate") == "full_grid":
        return generate_full_grid_cells(rung_name, rung_doc, ladder_doc, suites, results_root, paths)
    return [build_cell(rung_name, c, ladder_doc, suites) for c in rung_doc["cells"]]


def find_rung_cell(ladder_doc: dict, suite: str, axis: str, suites: dict[str, SuiteInfo]) -> Cell | None:
    """Best-effort lookup used by ``--cells``: if an existing rung already declares this exact
    (suite, axis) pair, reuse its comparators/gate so an ad hoc cell isn't scored in a vacuum."""
    for rung_name, rung_doc in ladder_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 build_cell(rung_name, cell_doc, ladder_doc, suites)
    return None


def parse_cells_arg(cells_arg: str, mode: str | None, ladder_doc: dict, suites: dict[str, SuiteInfo],
                     results_root: Path, paths: dict[str, Path]) -> list[Cell]:
    sm = score_mod(paths)
    out = []
    for token in (t.strip() for t in cells_arg.split(",")):
        if not token:
            continue
        suite, _, axis = token.partition(":")
        if suite not in suites:
            raise ValueError(f"--cells: unknown suite {suite!r} in {token!r} (known: {sorted(suites)})")
        if axis not in sm.AXES:
            raise ValueError(f"--cells: unknown axis {axis!r} in {token!r} (known: {sm.AXES})")
        resolved_mode = mode or default_mode_for_axis(axis, ladder_doc)
        expected_n = int(ladder_doc["expected_episodes"][suite][axis])
        matched = find_rung_cell(ladder_doc, suite, axis, suites)
        comparators = matched.comparators if matched else resolve_r4_comparators(suite, axis, results_root, paths)
        gate = matched.gate if matched else Gate.from_dict(ladder_doc["default_gate"])
        out.append(Cell(
            rung="adhoc", suite=suite, axis=axis, mode=resolved_mode, expected_n=expected_n,
            filter_file=resolve_filter(None, axis, suite, suites), comparators=comparators, gate=gate,
            num_clients=ladder_doc["defaults"]["num_clients"],
        ))
    return out


# =====================================================================================================
# env composition (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 parent-of-parent, then the
    usual install prefixes."""
    candidates = []
    conda_exe = os.environ.get("CONDA_EXE")
    if conda_exe:
        candidates.append(Path(conda_exe).resolve().parent.parent / "envs" / name)
    for base in (Path.home() / "miniconda3", Path.home() / "anaconda3", Path.home() / "miniforge3",
                 Path("/opt/conda"), Path("/opt/miniconda3")):
        candidates.append(base / "envs" / name)
    for c in candidates:
        if (c / "bin" / "python").exists():
            return c
    return None


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


# 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 (compose_cell_env degrades 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"
)


def nvml_preload_path() -> str | None:
    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")))


# =====================================================================================================
# graph-artifact resolution -- ONE resolved directory per suite, shared by JQ_GRAPH_DIR (entry,
# onf.recovery.targets) and SN_GRAPH_DIR (sentinel, onf.sentinel.sentinel), instead of letting each
# fall through its OWN independent default. Before this, a cell that never set either var still got
# lucky (both defaulted through the same onf.config.Paths.graph(suite) call) -- but that agreement was
# incidental, not enforced, and evals/gr00t/run_gr00t.py's independent SN_GRAPH_DIR override (see
# run_gr00t.require_graph_artifacts) is exactly the case where it silently stopped holding.
# =====================================================================================================
def resolve_graph_dir(suite: str) -> Path:
    """The graph-artifact dir (``g_nodes.npz`` / ``g_edges.npz`` / ``g_head.npz`` / ``g_track.npz``)
    for ``suite`` -- the SAME resolution ``onf.recovery.targets`` and ``onf.sentinel.sentinel`` fall
    back to internally (``onf.config.Paths.graph(suite)``, which itself honors the ``GR_GRAPH_DIR``
    override). Computed here so both drivers can bind ``JQ_GRAPH_DIR``/``SN_GRAPH_DIR`` to the exact
    same path up front, and log/record it, rather than letting two independently-resolved env vars
    drift apart."""
    from onf.config import default_paths
    return default_paths().graph(suite)


def graph_artifact_status(graph_dir: Path) -> dict:
    """Inspect a resolved graph dir for the two artifacts the recovery ↔ sentinel share:
    ``g_head.npz`` (the trained readout head, entry's WHERE) and ``g_track.npz`` (the fitted belief
    filter, sentinel's WHEN/WHERE). Returns
    ``{"graph_dir", "ok", "missing": [...], "graph_hash", "g_head_mtime"}`` -- ``graph_hash`` is read
    straight off ``g_head.npz``'s own stamp (no need to load nodes/edges/the network just to log it),
    ``g_head_mtime`` is that file's mtime as an ISO-8601 UTC string. Never raises: a suite with no
    graph built yet gets ``ok=False`` and both fields None, not a crash -- callers (both drivers) fail
    loud on ``not ok`` for any non-``base`` mode before launching anything GPU-resident."""
    import numpy as np
    from onf.graph 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 = None
    g_head_mtime = None
    g_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 e:  # noqa: BLE001 -- report, never crash a status probe
            missing.append(f"{head} (unreadable: {e})")
        g_head_mtime = datetime.fromtimestamp(
            head.stat().st_mtime, tz=timezone.utc).isoformat(timespec="seconds")
    if track.exists():
        g_track_mtime = datetime.fromtimestamp(
            track.stat().st_mtime, tz=timezone.utc).isoformat(timespec="seconds")
    return {"graph_dir": str(graph_dir), "ok": not missing, "missing": missing,
            "graph_hash": graph_hash, "g_head_mtime": g_head_mtime, "g_track_mtime": g_track_mtime}


def compose_cell_env(cell: Cell, suites: dict[str, SuiteInfo], paths: dict[str, Path]) -> dict[str, str]:
    libero_home = libero_home_dir()
    task_classification = 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(suites[cell.suite].fwm_dir),
        "LIBERO_HOME": str(libero_home),
        "LIBERO_CONFIG_PATH": str(libero_home / "libero"),
        "LIBERO_PLUS_TASK_CLASSIFICATION": str(task_classification),
        "MUJOCO_GL": "osmesa", "PYOPENGL_PLATFORM": "osmesa",
    }
    conda_prefix = resolve_conda_env_prefix(STABLEVLA_CONDA_ENV)
    if conda_prefix is not None:
        env["MAGICK_HOME"] = str(conda_prefix)
        ld = os.environ.get("LD_LIBRARY_PATH", "")
        env["LD_LIBRARY_PATH"] = f"{conda_prefix / 'lib'}:{ld}" if ld else str(conda_prefix / "lib")
    nvml_preload = nvml_preload_path()
    if nvml_preload is not None:
        # This box's kernel nvidia module (570.172.08) and userspace NVML (580.173) are mismatched --
        # 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 -- see graph/track-mode runs, which
        # allocate more and hit this far more often than a plain `base` rollout ever does). The SAME
        # preload that makes `nvidia-smi` usable on this box fixes it for every child process too.
        # Every launched process (server AND client shards) inherits this via compose_cell_env, not an
        # ad-hoc one-off export on a single command.
        existing = os.environ.get("LD_PRELOAD", "")
        env["LD_PRELOAD"] = f"{nvml_preload}:{existing}" if existing else nvml_preload
    if cell.filter_file is not None:
        env["LIBERO_INSTANCE_ID_FILTER"] = str(cell.filter_file)
    env.update(mode_env(cell.mode, paths))
    # Bind JQ_GRAPH_DIR (entry) and SN_GRAPH_DIR (sentinel) to the SAME resolved directory -- see
    # resolve_graph_dir's docstring. Set unconditionally (cheap: no I/O beyond the eventual mtime
    # stat in graph_artifact_status); `base` mode reads neither var so this is a no-op for it.
    resolved_graph_dir = str(resolve_graph_dir(cell.suite))
    env["JQ_GRAPH_DIR"] = resolved_graph_dir
    env["SN_GRAPH_DIR"] = resolved_graph_dir
    env.update(cell.extra_env)   # explicit per-cell overrides (a yaml cell's `env:` block) win last
    return env


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


def server_pythonpath(paths: dict[str, Path]) -> str:
    return _pythonpath_for([STABLEVLA_ROOT, paths["harness_dir"], ONF_SRC])


def client_pythonpath(paths: dict[str, Path]) -> str:
    return _pythonpath_for([paths["harness_dir"], STABLEVLA_ROOT, libero_home_dir(), ONF_SRC, paths["shim_dir"]])


# =====================================================================================================
# 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.
# =====================================================================================================
@dataclasses.dataclass(frozen=True)
class GpuStatus:
    index: int
    free_gib: float
    total_gib: float
    ok: bool
    reason: str = ""


def gpu_report(indices: list[int], min_free_gib: float) -> list[GpuStatus]:
    import torch

    out = []
    n = torch.cuda.device_count()
    for i in indices:
        if i >= n or i < 0:
            out.append(GpuStatus(i, 0.0, 0.0, False, f"GPU {i} does not exist (device_count={n})"))
            continue
        free, total = torch.cuda.mem_get_info(i)
        free_gib, total_gib = free / 1024**3, total / 1024**3
        ok = free_gib >= min_free_gib
        reason = "" if ok else f"only {free_gib:.1f} GiB free (< {min_free_gib:.1f} GiB threshold)"
        out.append(GpuStatus(i, free_gib, total_gib, ok, reason))
    return out


@dataclasses.dataclass(frozen=True)
class LoadStatus:
    load1: float
    ncpu: int
    threshold: float
    ok: bool


def load_report(max_factor: float) -> LoadStatus:
    load1 = os.getloadavg()[0]
    ncpu = os.cpu_count() or 1
    threshold = ncpu * max_factor
    return LoadStatus(load1=load1, ncpu=ncpu, threshold=threshold, ok=load1 <= threshold)


def check_libero_plus(timeout_s: float = 120.0, ladder_doc: dict | None = None,
                       paths: dict[str, Path] | None = None) -> tuple[bool, str]:
    ladder_doc = ladder_doc or load_ladder()
    paths = paths or resolve_paths(ladder_doc)
    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 False, (f"external/LIBERO-plus missing/empty at {libero_home} "
                        f"(no libero/libero/benchmark/__init__.py) -- run: bash scripts/setup_external.sh")
    py = resolve_stablevla_python()
    if py is None:
        return 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["LIBERO_HOME"] = str(libero_home)
    env["LIBERO_CONFIG_PATH"] = str(libero_home / "libero")
    env["PYTHONPATH"] = _pythonpath_for([libero_home])
    env["PYTHONNOUSERSITE"] = "1"
    env["OMP_NUM_THREADS"] = env["MKL_NUM_THREADS"] = env["OPENBLAS_NUM_THREADS"] = "4"
    try:
        out = subprocess.run([str(py), "-c", snippet], cwd=REPO_ROOT, env=env,
                              capture_output=True, text=True, timeout=timeout_s)
    except Exception as e:  # report, don't crash --check-env
        return False, f"LIBERO-Plus check crashed: {e}"
    if out.returncode != 0:
        tail = "\n".join(out.stderr.strip().splitlines()[-15:])
        return False, f"LIBERO-Plus check FAILED (rc={out.returncode}):\n{tail}"
    lines = out.stdout.strip().splitlines()
    m = re.search(r"\{.*\}", lines[-1]) if lines else None
    if not m:
        return False, f"LIBERO-Plus check: could not parse output: {out.stdout!r}"
    try:
        counts = ast.literal_eval(m.group(0))
    except Exception as e:
        return False, f"LIBERO-Plus check: could not parse counts dict: {e}"
    # 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.
    sm = score_mod(paths)
    expected = {sm.SUITE_BENCH[s]: n["Robot_Initial_States"] for s, n in ladder_doc["expected_episodes"].items()}
    mismatches = {k: (counts.get(k), v) for k, v in expected.items() if counts.get(k) != v}
    if mismatches:
        return False, f"LIBERO-Plus counts mismatch: got {counts}, expected {expected}"
    return True, f"counts OK: {counts}"


# =====================================================================================================
# 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.
# =====================================================================================================
_ledger_lock = threading.Lock()


def default_ledger_path() -> Path:
    from onf.config import default_paths
    return default_paths().outputs("sr", "ledger.jsonl")


def append_ledger(record: dict, path: Path | None = None) -> None:
    path = path or default_ledger_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    line = json.dumps(record, default=str)
    with _ledger_lock, open(path, "a") as f:
        f.write(line + "\n")


def read_ledger(path: Path | None = None) -> list[dict]:
    path = path or default_ledger_path()
    if not path.exists():
        return []
    out = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                out.append(json.loads(line))
    return out


def latest_matching_record(records: list[dict], tag: str, mode: str | None = None) -> dict | None:
    """Last (most recent -- JSONL is append-ordered) record for `tag`, or None.

    ``mode``, if given, must ALSO match the record's ``mode`` field. Tags now bake the mode into their
    own string (see :func:`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 by a
    tag string that happens to collide."""
    rec = None
    for r in records:
        if r.get("tag") == tag and (mode is None or r.get("mode") == mode):
            rec = r
    return rec


def latest_record(tag: str, path: Path | None = None, mode: str | None = None) -> dict | None:
    return latest_matching_record(read_ledger(path), tag, mode)


# =====================================================================================================
# log-directory naming — shared by BOTH drivers (run_sr.py's own run_one_cell and
# evals/gr00t/run_gr00t.py, which imports this module by file path and calls this function directly)
# so a StableVLA cell and a GR00T cell can never collide, and so mode/suite/axis are always legible
# from the directory name instead of just a tag suffix a human has to decode. Before this, run_sr.py
# stamped LOCAL time while run_gr00t.py stamped UTC, and run_gr00t.py's own root additionally used the
# literal "gr00tn17" instead of its own --tag-prefix -- two runs differing only by mode/suite/axis
# truncated each other's shard logs into the same directory.
# =====================================================================================================
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 :func:`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."""
    tag = cell_tag(policy, mode, suite, axis.replace(" ", "_"), variant)
    return root / "logs" / f"{utc_stamp()}_{tag}"


# =====================================================================================================
# 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 (policy, mode, suite, axis,
# expected episode count, resolved graph dir, its graph_hash + g_head.npz mtime, the git commit, and
# the full mode env).
# =====================================================================================================
def write_run_json(result_dir: Path, **fields: object) -> 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


def _git_sha() -> str | None:
    try:
        out = subprocess.run(["git", "rev-parse", "HEAD"], cwd=REPO_ROOT,
                              capture_output=True, text=True, timeout=10)
    except (OSError, subprocess.SubprocessError):
        return None
    return out.stdout.strip() if out.returncode == 0 else None


def _git_dirty() -> bool | None:
    try:
        out = subprocess.run(["git", "status", "--porcelain"], cwd=REPO_ROOT,
                              capture_output=True, text=True, timeout=10)
    except (OSError, subprocess.SubprocessError):
        return None
    return bool(out.stdout.strip()) if out.returncode == 0 else None


# =====================================================================================================
# scoring / completeness / gates — reuse score.py's count_dir + mcnemar.py's load()/mcnemar_exact()
# =====================================================================================================
@dataclasses.dataclass
class CellScore:
    cell: Cell
    successes: int
    actual_n: int
    complete: bool
    rate: float


def score_cell(cell: Cell, results_root: Path, paths: dict[str, Path]) -> CellScore:
    d = cell.result_dir(results_root, paths)
    s, n = score_mod(paths).count_dir(str(d)) if d.exists() else (0, 0)
    rate = 100.0 * s / n if n else float("nan")
    return CellScore(cell=cell, successes=s, actual_n=n, complete=(n == cell.expected_n), rate=rate)


@dataclasses.dataclass(frozen=True)
class GateVerdict:
    passed: bool
    message: str


def evaluate_gate(score: CellScore, allow_partial: bool) -> GateVerdict:
    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)")
    gate = cell.gate
    if gate is None:
        return GateVerdict(True, "no gate configured")
    comp = cell.comparator(gate.vs)
    if comp is None:
        return GateVerdict(True, f"SKIP: no comparator labeled {gate.vs!r} to gate against")
    if gate.kind == "exact_match":
        ok = score.successes == comp.successes and score.actual_n == comp.n
        return GateVerdict(ok, f"exact_match vs {comp.label} ({comp.successes}/{comp.n}): "
                                f"got {score.successes}/{score.actual_n}")
    if gate.kind == "min_successes":
        ok = score.successes >= gate.min_successes
        return GateVerdict(ok, f"min_successes={gate.min_successes} vs {comp.label}: got {score.successes}")
    if gate.kind == "max_regress_pp":
        delta = score.rate - comp.rate
        ok = delta >= -gate.max_regress_pp
        return GateVerdict(ok, f"delta vs {comp.label} = {delta:+.1f}pp (floor -{gate.max_regress_pp:.1f}pp)")
    raise ValueError(f"unknown gate kind {gate.kind!r}")


def mcnemar_report(cell: Cell, results_root: Path, paths: dict[str, Path], label: str | None = None) -> str | None:
    comp_label = label or (cell.gate.vs if cell.gate else None)
    if comp_label is None:
        return None
    comp = cell.comparator(comp_label)
    if comp is None:
        return None
    mc = mcnemar_mod(paths)
    comp_dir = results_root / f"plus_{cell.bench(paths)}" / cell.axis / comp.tag
    cell_dir = cell.result_dir(results_root, paths)
    if not comp_dir.exists() or not cell_dir.exists():
        return None
    a, b = mc.load(str(comp_dir)), mc.load(str(cell_dir))
    keys = sorted(set(a) & set(b))
    if not keys:
        return f"McNemar vs {comp.label}: no shared episodes"
    import numpy as np

    av = np.array([a[k] for k in keys])
    bv = np.array([b[k] for k in keys])
    b_only = int((av & (~bv)).sum())       # comparator-only success
    c_only = int(((~av) & bv).sum())       # this cell-only success
    p = mc.mcnemar_exact(b_only, c_only)
    return f"McNemar b/c={b_only}/{c_only} p={p:.2e} n={len(keys)}"


def format_result_line(score: CellScore, verdict: GateVerdict, mcnemar_str: str | None) -> str:
    cell = score.cell
    comp = cell.comparator(cell.gate.vs) if cell.gate else (cell.comparators[0] if cell.comparators else None)
    delta_str = f", delta vs {comp.label} = {score.rate - comp.rate:+.1f}pp" if (comp and score.actual_n) else ""
    mc_str = f", {mcnemar_str}" if mcnemar_str else ""
    rate_str = f"{score.rate:.1f}%" if score.actual_n else "--"
    verdict_str = "PASS" if verdict.passed else "FAIL"
    return (f"RESULT {cell.tag}: {score.successes}/{score.actual_n} = {rate_str}"
            f"  ({delta_str.lstrip(', ')}{mc_str})  [{verdict_str}: {verdict.message}]")


# =====================================================================================================
# 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.
# =====================================================================================================
def lpt_schedule(cells: list[Cell], gpus: list[int]) -> dict[int, list[Cell]]:
    order = sorted(cells, key=lambda c: c.expected_n, reverse=True)
    load = dict.fromkeys(gpus, 0)
    assign: dict[int, list[Cell]] = {g: [] for g in gpus}
    for c in order:
        g = min(load, key=lambda k: (load[k], k))
        assign[g].append(c)
        load[g] += c.expected_n
    return assign


def rolling_episodes_per_min(cell: Cell, ladder_doc: dict, ledger_records: list[dict]) -> float:
    matches = [r for r in ledger_records
               if r.get("suite") == cell.suite and r.get("axis") == cell.axis and r.get("mode") == cell.mode
               and r.get("status") == "complete" and r.get("episodes_per_min")]
    if matches:
        vals = [float(r["episodes_per_min"]) for r in matches[-5:]]
        return sum(vals) / len(vals)
    return float(ladder_doc["defaults"]["episodes_per_min"])


def project_makespan(assign: dict[int, list[Cell]], ladder_doc: dict, ledger_records: list[dict]) -> dict:
    per_gpu_minutes: dict[int, float] = {}
    cell_lines = []
    for gpu, cells in assign.items():
        total = 0.0
        for c in cells:
            epm = rolling_episodes_per_min(c, ladder_doc, ledger_records)
            mins = c.expected_n / epm if epm else float("inf")
            total += mins
            cell_lines.append((gpu, c, mins))
        per_gpu_minutes[gpu] = total
    makespan = max(per_gpu_minutes.values()) if per_gpu_minutes else 0.0
    return {"per_gpu_minutes": per_gpu_minutes, "cell_lines": cell_lines, "makespan_min": makespan}


# =====================================================================================================
# 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.
#
# `plan_cell_action` is a PURE decision (no filesystem writes) so `--dry-run` can preview it safely;
# `apply_cell_action` performs the purge and is only ever called from the real (non-dry-run) path.
# =====================================================================================================
def plan_cell_action(cell: Cell, results_root: Path, paths: dict[str, Path],
                      resume: bool, ledger_records: list[dict] | None = None) -> tuple[str, str]:
    """Returns (action, message) where action is one of:
    ``run`` (nothing there yet), ``skip`` (ledger-confirmed complete, matches expected_n),
    ``purge`` (partial/stale content -- apply_cell_action must remove it before re-running),
    ``error`` (content exists and --resume was not passed)."""
    d = cell.result_dir(results_root, paths)
    has_content = d.exists() and any(d.rglob("*.mp4"))
    if not has_content:
        return "run", "no existing result dir"
    if not resume:
        return "error", f"result dir already has content at {d} (pass --resume)"
    records = ledger_records if ledger_records is not None else read_ledger()
    rec = latest_matching_record(records, cell.tag, cell.mode)
    if rec and rec.get("status") == "complete":
        _s, n = score_mod(paths).count_dir(str(d))
        if n == cell.expected_n:
            return "skip", f"ledger says complete and {n} mp4s match expected_n={cell.expected_n}"
    return "purge", "partial/stale result dir -- would be purged and re-run"


def apply_cell_action(action: str, cell: Cell, results_root: Path, paths: dict[str, Path]) -> None:
    """Perform the filesystem side effect for a ``purge`` verdict. Never called on the --dry-run path."""
    if action == "purge":
        shutil.rmtree(cell.result_dir(results_root, paths))


# =====================================================================================================
# launch orchestration — one GPU-resident policy server + N sharded sim clients per cell.
# =====================================================================================================
_active_procs_lock = threading.Lock()
_active_procs: set[subprocess.Popen] = set()


def _register_proc(p: subprocess.Popen) -> None:
    with _active_procs_lock:
        _active_procs.add(p)


def _unregister_proc(p: subprocess.Popen) -> None:
    with _active_procs_lock:
        _active_procs.discard(p)


def _kill_all_active(signum, _frame) -> None:  # signal handler signature (frame unused)
    with _active_procs_lock:
        procs = list(_active_procs)
    for p in procs:
        try:
            p.terminate()
        except Exception:  # best-effort teardown, never let this handler raise
            pass
    time.sleep(2)
    for p in procs:
        try:
            if p.poll() is None:
                p.kill()
        except Exception:
            pass
    sys.exit(128 + signum)


def install_signal_handlers() -> None:
    signal.signal(signal.SIGINT, _kill_all_active)
    signal.signal(signal.SIGTERM, _kill_all_active)


def _wait_for_server(proc: subprocess.Popen, host: str, port: int, log_path: Path,
                      timeout_s: float, poll_every: float = 2.0) -> None:
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        if proc.poll() is not None:
            tail = ""
            if log_path.exists():
                tail = "\n".join(log_path.read_text(errors="replace").splitlines()[-40:])
            raise RuntimeError(f"policy server exited before ready (rc={proc.returncode}); log tail:\n{tail}")
        try:
            with socket.create_connection((host, port), timeout=1):
                return
        except OSError:
            time.sleep(poll_every)
    raise RuntimeError(f"policy server not ready after {timeout_s:.0f}s on {host}:{port}")


@dataclasses.dataclass
class LaunchResult:
    wall_s: float
    client_returncodes: list[int]
    crashed: bool


def launch_cell(cell: Cell, gpu: int, port: int, suites: dict[str, SuiteInfo], log_dir: Path,
                 results_root: Path, ladder_doc: dict, paths: dict[str, Path]) -> LaunchResult:
    t0 = time.time()
    result_dir = cell.result_dir(results_root, paths)
    result_dir.mkdir(parents=True, exist_ok=True)
    log_dir.mkdir(parents=True, exist_ok=True)

    stablevla_python = resolve_stablevla_python()
    if stablevla_python is None:
        raise RuntimeError(f"cannot locate the `{STABLEVLA_CONDA_ENV}` conda env's python")

    shared_env = dict(os.environ)
    shared_env.update(compose_cell_env(cell, suites, paths))

    server_env = dict(shared_env)
    server_env["CUDA_VISIBLE_DEVICES"] = str(gpu)
    server_env["PYTHONPATH"] = server_pythonpath(paths)
    ckpt = str(suites[cell.suite].stablevla_ckpt)

    server_log = open(log_dir / "server.log", "w")
    server_proc = subprocess.Popen(
        [str(stablevla_python), str(paths["policy_server"]), "--ckpt_path", ckpt,
         "--port", str(port), "--cuda", "0", "--use_bf16"],
        cwd=REPO_ROOT, env=server_env, stdout=server_log, stderr=subprocess.STDOUT,
    )
    _register_proc(server_proc)
    try:
        timeout_s = float(ladder_doc["defaults"]["server_ready_timeout_s"])
        _wait_for_server(server_proc, "127.0.0.1", port, log_dir / "server.log", timeout_s)

        client_env = dict(shared_env)
        client_env["CUDA_VISIBLE_DEVICES"] = str(gpu)
        client_env["PYTHONPATH"] = client_pythonpath(paths)

        client_procs = []
        for k in range(cell.num_clients):
            log_f = open(log_dir / f"client_shard{k}.log", "w")
            p = subprocess.Popen(
                [str(stablevla_python), str(paths["harness_dir"] / "eval_libero.py"),
                 "--args.pretrained-path", ckpt, "--args.host", "127.0.0.1", "--args.port", str(port),
                 "--args.task-suite-name", cell.bench(paths), "--args.category_value", cell.category_value,
                 "--args.num-trials-per-task", "1", "--args.with_state", "true",
                 "--args.shard_index", str(k), "--args.num_shards", str(cell.num_clients),
                 "--args.video-out-path", str(result_dir)],
                cwd=REPO_ROOT, env=client_env, stdout=log_f, stderr=subprocess.STDOUT,
            )
            _register_proc(p)
            client_procs.append(p)

        rcs = []
        for p in client_procs:
            rcs.append(p.wait())
            _unregister_proc(p)
        return LaunchResult(wall_s=time.time() - t0, client_returncodes=rcs, crashed=any(rc != 0 for rc in rcs))
    finally:
        _unregister_proc(server_proc)
        server_proc.terminate()
        try:
            server_proc.wait(timeout=15)
        except subprocess.TimeoutExpired:
            server_proc.kill()
            server_proc.wait(timeout=15)
        server_log.close()


# =====================================================================================================
# execution — one worker thread per GPU, each running its LPT-assigned cells sequentially. A thread
# (not process) pool is fine here: each cell's work is dominated by subprocess I/O wait.
# =====================================================================================================
_port_lock = threading.Lock()
_port_counter = itertools.count()


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


def _port_reservation_dir() -> Path:
    d = REPO_ROOT / "logs" / ".ports"
    d.mkdir(parents=True, exist_ok=True)
    return d


def _reserve_port(port: int) -> bool:
    """Cross-process claim on ``port``: create ``logs/.ports/<port>`` exclusively, holding this pid.

    A stale file (owner pid gone) is reclaimed. Returns False if a LIVE process already holds it.
    """
    f = _port_reservation_dir() / str(port)
    try:
        fd = os.open(f, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
    except FileExistsError:
        try:
            owner = int(f.read_text().strip() or -1)
        except (ValueError, OSError):
            owner = -1
        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:
            f.write_text(f"{os.getpid()}\n")
            return True
        except OSError:
            return False
    with os.fdopen(fd, "w") as fh:
        fh.write(f"{os.getpid()}\n")
    return True


def claim_port(ladder_doc: dict) -> int:
    """Next port at or above ``defaults.base_port`` that is both unbound AND unreserved.

    The counter alone is not enough: it is process-local, so 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, which is taken the instant the port is chosen and outlives the gap until
    the server actually listens.

    A reservation lasts as long as the owning process: it is dropped by the atexit hook below, and one
    left behind by a killed run is reclaimed by the next caller's pid-liveness check. Nothing needs to
    release a port mid-run -- ``_port_counter`` never hands the same number out twice anyway.
    """
    with _port_lock:
        base = int(ladder_doc["defaults"]["base_port"])
        for _ in range(256):
            port = base + next(_port_counter)
            if _port_is_free(port) and _reserve_port(port):
                _reserved_ports.add(port)
                return port
        raise RuntimeError(f"no free port in [{base}, {base + 256}) -- is something leaking servers?")


_reserved_ports: set[int] = set()
atexit.register(lambda: [(_port_reservation_dir() / str(p)).unlink(missing_ok=True)
                         for p in list(_reserved_ports)])


def run_one_cell(cell: Cell, gpu: int, ladder_doc: dict, suites: dict[str, SuiteInfo],
                  results_root: Path, paths: dict[str, Path], allow_partial: bool) -> int:
    port = claim_port(ladder_doc)
    ts = datetime.now(timezone.utc).isoformat(timespec="seconds")
    log_dir = make_log_dir(cell.policy, cell.mode, cell.suite, cell.axis, cell.variant)
    env = compose_cell_env(cell, suites, paths)
    sha, dirty = _git_sha(), _git_dirty()

    # One startup line naming the resolved graph dir (bound into env["JQ_GRAPH_DIR"] ==
    # env["SN_GRAPH_DIR"] by compose_cell_env), its mode, its graph_hash, and g_head.npz/g_track.npz
    # mtimes. FAIL HARD (same policy as evals/gr00t/run_gr00t.py's require_graph_artifacts) for any
    # non-`base` mode missing an artifact -- a missing g_track.npz means an unfit belief-filter kernel
    # and onf.sentinel.track.BasinReadout raising ValueError at the first fire, hours into a
    # GPU-resident run, rather than a fast, readable failure here.
    gstatus = graph_artifact_status(Path(env["SN_GRAPH_DIR"]))
    print(f"[run_sr] {cell.tag}: mode={cell.mode} graph_dir={gstatus['graph_dir']} "
          f"graph_hash={gstatus['graph_hash']} g_head_mtime={gstatus['g_head_mtime']} "
          f"g_track_mtime={gstatus['g_track_mtime']}")
    if not gstatus["ok"] and cell.mode != "base":
        raise RuntimeError(
            f"{cell.tag}: resolved graph dir {gstatus['graph_dir']} is missing required artifacts: "
            f"{gstatus['missing']} -- train the graph for this suite first, or point GR_GRAPH_DIR at "
            f"a suite that already has g_head.npz + g_track.npz."
        )

    result_dir = cell.result_dir(results_root, paths)
    write_run_json(
        result_dir, policy="stablevla", mode=cell.mode, suite=cell.suite, axis=cell.axis,
        expected_n=cell.expected_n, graph_dir=gstatus["graph_dir"], graph_hash=gstatus["graph_hash"],
        g_head_mtime=gstatus["g_head_mtime"], git_commit=sha, git_dirty=dirty, mode_env=env,
    )

    base_record = {
        "tag": cell.tag, "suite": cell.suite, "bench": cell.bench(paths), "axis": cell.axis,
        "mode": cell.mode, "gpu": gpu, "port": port, "expected_n": cell.expected_n,
        "git_sha": sha, "git_dirty": dirty, "env": env, "graph_dir": gstatus["graph_dir"],
        "graph_hash": gstatus["graph_hash"], "g_head_mtime": gstatus["g_head_mtime"],
        "result_dir": str(result_dir), "log_dir": str(log_dir),
    }
    append_ledger({**base_record, "ts": ts, "actual_n": None, "successes": None, "rate": None,
                   "wall_s": None, "episodes_per_min": None, "status": "running"})

    try:
        launch = launch_cell(cell, gpu, port, suites, log_dir, results_root, ladder_doc, paths)
    except Exception as e:  # record the crash, don't take the whole schedule down
        append_ledger({**base_record, "ts": datetime.now(timezone.utc).isoformat(timespec="seconds"),
                       "actual_n": None, "successes": None, "rate": None, "wall_s": None,
                       "episodes_per_min": None, "status": "crashed", "error": str(e)})
        print(f"RESULT {cell.tag}: CRASHED -- {e}", file=sys.stderr)
        return 1

    score = score_cell(cell, results_root, paths)
    epm = (score.actual_n / (launch.wall_s / 60.0)) if launch.wall_s > 0 and score.actual_n else 0.0
    status = "crashed" if launch.crashed else ("complete" if score.complete else "incomplete")
    append_ledger({**base_record, "ts": datetime.now(timezone.utc).isoformat(timespec="seconds"),
                   "actual_n": score.actual_n, "successes": score.successes, "rate": score.rate,
                   "wall_s": launch.wall_s, "episodes_per_min": epm, "status": status})

    verdict = evaluate_gate(score, allow_partial)
    mc = mcnemar_report(cell, results_root, paths) if cell.gate else None
    print(format_result_line(score, verdict, mc))
    if status != "complete" and not allow_partial:
        return 1
    return 0 if verdict.passed else 2


def execute_schedule(assign: dict[int, list[Cell]], ladder_doc: dict, suites: dict[str, SuiteInfo],
                      results_root: Path, paths: dict[str, Path], allow_partial: bool) -> int:
    def run_track(gpu: int, cells: list[Cell]) -> list[int]:
        return [run_one_cell(c, gpu, ladder_doc, suites, results_root, paths, allow_partial) for c in cells]

    codes: list[int] = []
    with ThreadPoolExecutor(max_workers=max(1, len(assign))) as ex:
        futures = {ex.submit(run_track, gpu, cells): gpu for gpu, cells in assign.items()}
        for fut in as_completed(futures):
            codes.extend(fut.result())
    if any(c == 2 for c in codes):
        return 2
    if any(c not in (0, 2) for c in codes):
        return 1
    return 0


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


def parse_gpu_list(s: str | None) -> list[int]:
    if not s:
        return []
    return [int(x) for x in s.split(",") if x.strip() != ""]


def resolve_cells(args: argparse.Namespace, ladder_doc: dict, suites: dict[str, SuiteInfo],
                   results_root: Path, paths: dict[str, Path]) -> list[Cell]:
    if args.rung:
        if args.rung not in ladder_doc["rungs"]:
            raise SystemExit(f"unknown rung {args.rung!r}; known: {sorted(ladder_doc['rungs'])}")
        cells = rung_cells(ladder_doc, args.rung, suites, results_root, paths)
        if args.cells:
            # `--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 args.cells.split(",") if t.strip()}
            unknown = wanted - {f"{c.suite}:{c.axis}" for c in cells}
            if unknown:
                raise SystemExit(f"--cells: {sorted(unknown)} not in rung {args.rung} "
                                 f"(has: {sorted(f'{c.suite}:{c.axis}' for c in cells)})")
            cells = [c for c in cells if f"{c.suite}:{c.axis}" in wanted]
        if args.mode:
            # cell_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 = [cell_with_mode(c, args.mode) for c in cells]
        return cells
    if args.cells:
        return parse_cells_arg(args.cells, args.mode, ladder_doc, suites, results_root, paths)
    return []


def cmd_check_env(ladder_doc: dict) -> int:
    pf = ladder_doc["preflight"]
    ok = True

    print("=== GPU ===")
    try:
        import torch

        n = torch.cuda.device_count()
        if n == 0:
            print("  no CUDA devices visible")
            ok = False
        for i in range(n):
            free, total = torch.cuda.mem_get_info(i)
            free_g, total_g = free / 1024**3, total / 1024**3
            busy = free_g < pf["min_free_gib"]
            ok = ok and not busy
            print(f"  gpu{i}: {free_g:6.1f} / {total_g:6.1f} GiB free  [{'BUSY' if busy else 'OK'}]")
    except Exception as e:
        print(f"  GPU check FAILED: {e}")
        ok = False

    print("=== load ===")
    load1, _load5, _load15 = os.getloadavg()
    ncpu = os.cpu_count() or 1
    threshold = ncpu * pf["max_loadavg_factor"]
    busy = load1 > threshold
    print(f"  loadavg(1m)={load1:.2f}  ncpu={ncpu}  threshold={threshold:.1f}  [{'BUSY' if busy else 'OK'}]")
    ok = ok and not busy

    print("=== LIBERO-Plus ===")
    libero_ok, libero_msg = check_libero_plus(ladder_doc=ladder_doc, paths=resolve_paths(ladder_doc))
    print(f"  {libero_msg}")
    ok = ok and libero_ok

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


def cmd_print_ledger(path: Path | None = None) -> int:
    records = read_ledger(path)
    if not records:
        print(f"[run_sr] ledger is empty (no runs recorded yet): {path or default_ledger_path()}")
        return 0
    cols = ["ts", "tag", "suite", "axis", "mode", "status", "actual_n", "expected_n", "successes",
            "rate", "wall_s", "episodes_per_min", "gpu"]
    rows = [{c: ("" if r.get(c) is None else r.get(c)) for c in cols} for r in records]
    widths = {c: max(len(c), *(len(str(row[c])) for row in rows)) for c in cols}
    header = "  ".join(c.ljust(widths[c]) for c in cols)
    print(header)
    print("-" * len(header))
    for row in rows:
        print("  ".join(str(row[c]).ljust(widths[c]) for c in cols))
    return 0


def cmd_score(cells: list[Cell], results_root: Path, paths: dict[str, Path], allow_partial: bool) -> int:
    any_gate_failed = False
    any_scored = False
    for cell in cells:
        score = score_cell(cell, results_root, paths)
        if score.actual_n == 0:
            print(f"RESULT {cell.tag}: NO DATA (expected {cell.expected_n} episodes at "
                  f"{cell.result_dir(results_root, paths)})")
            continue
        any_scored = True
        verdict = evaluate_gate(score, allow_partial)
        mc = mcnemar_report(cell, results_root, paths) if cell.gate else None
        print(format_result_line(score, verdict, mc))
        any_gate_failed = any_gate_failed or not verdict.passed
    if not any_scored:
        return 1
    return 2 if any_gate_failed else 0


def cmd_run(cells: list[Cell], gpus: list[int], ladder_doc: dict, suites: dict[str, SuiteInfo],
            results_root: Path, paths: dict[str, Path], args: argparse.Namespace) -> int:
    pf = ladder_doc["preflight"]
    gstatus = gpu_report(gpus, pf["min_free_gib"])
    for g in gstatus:
        if not g.ok:
            print(f"[preflight] SKIP gpu{g.index}: {g.reason}")
    sys.stdout.flush()   # keep stdout ahead of the stderr refusal below when interleaved (e.g. `2>&1`)
    usable_gpus = [g.index for g in gstatus if g.ok]
    if not usable_gpus:
        print("[preflight] no usable GPUs -- refusing to start anything", file=sys.stderr)
        return 1

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

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

    ledger_records = read_ledger()
    runnable = []
    for c in cells:
        action, msg = plan_cell_action(c, results_root, paths, args.resume, ledger_records)
        print(f"[plan] {c.tag}: {action} ({msg})")
        if action == "error":
            print(f"[run_sr] {c.tag}: {msg}", file=sys.stderr)
            return 1
        if action in ("run", "purge"):
            runnable.append((c, action))

    if not runnable:
        print("[run_sr] nothing left to run (everything already complete)")
        return 0

    assign = lpt_schedule([c for c, _ in runnable], usable_gpus)
    proj = project_makespan(assign, ladder_doc, ledger_records)

    print("\n=== schedule ===")
    for gpu, cell, mins in proj["cell_lines"]:
        print(f"  gpu{gpu}: {cell.tag:<45} n={cell.expected_n:<5} mode={cell.mode:<16} ~{mins:6.1f} min")
    print(f"projected makespan: {proj['makespan_min']:.1f} min ({proj['makespan_min'] / 60.0:.1f} h) "
          f"across {len(usable_gpus)} GPU(s)")

    if args.dry_run:
        print("\n[dry-run] no subprocess launched, no filesystem changes made.")
        return 0

    # 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.
    for c, action in runnable:
        apply_cell_action(action, c, results_root, paths)

    return execute_schedule(assign, ladder_doc, suites, results_root, paths, args.allow_partial)


def main(argv: list[str] | None = None) -> int:
    args = build_argparser().parse_args(argv)
    ladder_doc = load_ladder(Path(args.ladder))

    if args.check_env:
        return cmd_check_env(ladder_doc)
    if args.ledger:
        return cmd_print_ledger()

    paths = resolve_paths(ladder_doc)
    suites = load_suites(paths)
    results_root = Path(args.results_root) if args.results_root else default_results_root(paths)

    cells = resolve_cells(args, ladder_doc, suites, results_root, paths)
    if not cells:
        print("[run_sr] nothing to do -- pass --rung or --cells", file=sys.stderr)
        return 1

    if args.score:
        return cmd_score(cells, results_root, paths, args.allow_partial)

    gpus = parse_gpu_list(args.gpus)
    if not gpus:
        print("[run_sr] --gpus is required to launch or dry-run a schedule", file=sys.stderr)
        return 1

    install_signal_handlers()
    return cmd_run(cells, gpus, ladder_doc, suites, results_root, paths, args)


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