Blame for code/scripts/p2/expert_iteration.py

#!/usr/bin/env python
"""Resumable multi-round expert iteration (self-distillation) for the P2 prover.

    .venv/bin/python scripts/p2/expert_iteration.py --run-dir RUN --pool train_pool.jsonl \\
        --eval-problems minif2f_valid.jsonl --protected-problems minif2f.jsonl mobench.jsonl \\
        --base-model SNAPSHOT_DIR --rounds 3 --dry-run

The FIRST --eval-problems file is scored and drives the stop rule, so it must be a development split
(miniF2F-valid); a test split is refused unless --allow-test-selection (one pre-registered run only).

One round = eleven stages (docs/SELF_DISTILLATION.md has the recipe and the reasoning):

    select          pass-rate band selection from the previous round's measurements     (in-process)
    gen_verify      scripts/p2/pipeline_gen_verify.py: sample n proofs per problem, Lean-verify (GPU+CPU)
    gen_verify_corr [--corrections] one Lean-feedback correction round for unsolved problems (GPU+CPU)
    measure         per-problem pass rates after this round                                  (in-process)
    sft_data        scripts/p2/build_sft_data.py with the keep rules                         (CPU)
    train           scripts/sft.py + configs/p2_sft.toml (LoRA), 4 ranks                     (GPU)
    export          scripts/qwen3_export_check.py: merge LoRA into the round's input model    (GPU)
    average         [--avg-weight] align.model.weight_average: (1-w)*anchor + w*merged        (CPU)
    eval            pipeline_gen_verify on miniF2F at --eval-n / --eval-cap                  (GPU+CPU)
    report          scripts/p2/eval_report.py against the baseline run                        (CPU)
    manifest        round_manifest.json, paired McNemar vs previous round, stop decision     (in-process)

THIS DRIVER NEVER TAKES THE HEAVY LOCK AND NEVER RUNS ANYTHING ITSELF: every heavy child
(pipeline_gen_verify, sft.py, qwen3_export_check.py) takes the repo lock itself, one at a time, in
series. `--dry-run` prints the exact command lines and touches nothing.

IDEMPOTENT AND RESUMABLE PER STAGE. `RUN/round_NN/stages/<stage>.done.json` is written only after a stage
finished and its declared outputs exist. A marker stores a CHAIN HASH = H(previous stage's chain hash,
this stage's command line / parameters); a re-run skips a stage whose marker matches. If the
configuration changed so that a finished stage's hash no longer matches, the driver REFUSES (exit 3)
rather than silently mixing two configurations; `--redo ROUND:STAGE` deletes that stage's markers AND
artefacts (checkpoints, samples, exports, eval files) and every later one's, so nothing produced under the old
configuration is resumed. A crash leaves no marker, so the stage re-runs (the children are themselves resumable:
generation and verification are keyed append-only, training resumes from checkpoints, and a stage's
partial outputs are removed first where a child refuses to overwrite them).

Exit codes: 0 finished (or stopped by the pre-registered rule), 3 a stage failed / config mismatch,
4 stopped by a data condition (no problems in band, no training rows), 143 stopped by a signal.
"""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import os
import shlex
import shutil
import signal
import subprocess
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Optional

REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO))


def _load_common():
    spec = importlib.util.spec_from_file_location("p2_common", Path(__file__).with_name("p2_common.py"))
    mod = importlib.util.module_from_spec(spec)
    sys.modules["p2_common"] = mod
    spec.loader.exec_module(mod)
    return mod


C = sys.modules.get("p2_common") or _load_common()

PY = str(REPO / ".venv" / "bin" / "python")
SCRIPTS = REPO / "scripts"
STAGE_ORDER = ["select", "gen_verify", "gen_verify_corr", "measure", "sft_data", "train", "export", "average", "eval", "report", "manifest"]
EXIT_OK, EXIT_FAIL, EXIT_DATA_STOP, EXIT_SIGNAL = 0, 3, 4, 143


class DriverError(RuntimeError):
    pass


class DataStop(Exception):
    """A data condition that makes continuing meaningless (not a failure)."""


# =============================================================================== pure logic


def pass_rate_table(ver_paths, cap_tokens: int) -> dict[str, dict]:
    """Per-problem pass rate from verification records: round-0 samples at `cap_tokens` only
    (correction rounds are a different task; other caps are a different experiment)."""
    acc: dict[str, list[int]] = {}
    seen = set()
    for p in ver_paths:
        for rec in C.iter_jsonl(p):
            if int(rec["round"]) != 0 or int(rec["cap_tokens"]) != cap_tokens:
                continue
            key = C.gen_key(rec)
            if key in seen:
                raise C.ContractError(f"{p}: duplicate verification key {key}")
            seen.add(key)
            a = acc.setdefault(rec["problem_id"], [0, 0])
            a[0] += int(bool(rec["accepted"]))
            a[1] += 1
    return {pid: {"p": c / n, "c": c, "n": n} for pid, (c, n) in acc.items()}


def _h(seed: int, pid: str) -> str:
    return hashlib.sha256(f"{seed}|{pid}".encode()).hexdigest()


def select_by_band(table: dict[str, dict], pool_ids: list[str], lo: float, hi: float, *, include_zero: bool = False,
                   max_zero: Optional[int] = None, max_problems: Optional[int] = None, seed: int = 0):
    """Problems whose measured pass rate p satisfies lo < p <= hi (the band where a sample both
    succeeds sometimes and fails sometimes; Goedel-V2 uses (0, 0.75], GAR drops p = 0 and p > 0.5).

      * p > hi            -> "too_easy"   (nothing to learn, and it crowds the mix)
      * 0 < p <= lo       -> "too_hard"   (only when lo > 0)
      * p == 0            -> "unsolved"   (kept only with include_zero, at most max_zero, hash-ordered)
      * never measured    -> "unmeasured" (kept: it must be measured before it can be excluded)
    Priority when --max-problems truncates: in-band, then unmeasured, then unsolved; inside a group
    by a seeded hash of the problem id (deterministic, independent of file order).
    Returns (selected ids in priority order, {pid: reason}, counts)."""
    if not (0.0 <= lo < hi <= 1.0):
        raise ValueError(f"need 0 <= lo < hi <= 1, got lo={lo} hi={hi}")
    groups = {"in_band": [], "unmeasured": [], "unsolved": []}
    why: dict[str, str] = {}
    for pid in pool_ids:
        m = table.get(pid)
        if m is None:
            groups["unmeasured"].append(pid)
            why[pid] = "unmeasured"
        elif m["p"] > hi:
            why[pid] = "too_easy"
        elif m["p"] == 0.0:
            groups["unsolved"].append(pid)
            why[pid] = "unsolved"
        elif m["p"] <= lo:
            why[pid] = "too_hard"
        else:
            groups["in_band"].append(pid)
            why[pid] = "in_band"
    for g in groups.values():
        g.sort(key=lambda p: _h(seed, p))
    if not include_zero:
        for pid in groups["unsolved"]:
            why[pid] = "unsolved_excluded"
        groups["unsolved"] = []
    elif max_zero is not None:
        for pid in groups["unsolved"][max_zero:]:
            why[pid] = "unsolved_over_cap"
        groups["unsolved"] = groups["unsolved"][:max_zero]
    order = groups["in_band"] + groups["unmeasured"] + groups["unsolved"]
    if max_problems is not None:
        for pid in order[max_problems:]:
            why[pid] = "over_max_problems"
        order = order[:max_problems]
    counts: dict[str, int] = {}
    for r in why.values():
        counts[r] = counts.get(r, 0) + 1
    return order, why, counts


def keep_rule_args(cfg) -> list[str]:
    """The recipe's keep rules, as build_sft_data.py flags:
      * at most `max_per_problem` proofs per problem, DISTINCT (whitespace/comment-insensitive hash);
      * `rank`: shortest completion first (default) or sampling order;
      * problems solved n/n at round 0 contribute nothing (unless keep_easy);
      * duplicates fill the cap only with fill_duplicates;
      * at most `max_correction_per_problem` self-correction rows (only if corrections are on);
      * rows longer than max_seq_len are dropped, never truncated."""
    a = ["--max-per-problem", str(cfg.max_per_problem), "--max-correction-per-problem", str(cfg.max_correction_per_problem),
         "--rank", cfg.rank, "--max-seq-len", str(cfg.max_seq_len)]
    if cfg.keep_easy:
        a.append("--keep-easy")
    if cfg.fill_duplicates:
        a.append("--fill-duplicates")
    return a


def solved_map(ver_path, problem_ids, n_expected: int, cap_tokens: int) -> dict[str, bool]:
    """problem -> solved@n from a COMPLETE miniF2F-style verification file (round 0, one cap)."""
    got: dict[str, list[bool]] = {}
    for rec in C.iter_jsonl(ver_path):
        if int(rec["round"]) == 0 and int(rec["cap_tokens"]) == cap_tokens:
            got.setdefault(rec["problem_id"], []).append(bool(rec["accepted"]))
    missing = [p for p in problem_ids if len(got.get(p, [])) < n_expected]
    if missing:
        raise DriverError(f"{ver_path}: {len(missing)} problems have fewer than {n_expected} verified samples at cap {cap_tokens} "
                          f"(first: {missing[0]}); a missing sample is not a failure of the model")
    return {p: any(got[p][:n_expected]) for p in problem_ids}


def eval_summary(ver_path, problem_ids, n: int, cap_tokens: int) -> dict:
    sm = solved_map(ver_path, problem_ids, n, cap_tokens)
    c = tot = 0
    for rec in C.iter_jsonl(ver_path):
        if int(rec["round"]) == 0 and int(rec["cap_tokens"]) == cap_tokens and rec["problem_id"] in sm:
            c += int(bool(rec["accepted"]))
            tot += 1
    return {"n_problems": len(sm), "n_samples_per_problem": n, "cap_tokens": cap_tokens, "solved_at_n": sum(sm.values()),
            "solved_fraction": sum(sm.values()) / len(sm), "mean_pass_at_1": c / tot if tot else None, "samples": tot}


def stop_decision(history: list[dict], *, rounds_done: int, max_rounds: int, min_gain: int = 2, patience: int = 2,
                  max_regress: int = 3, target_solved: Optional[int] = None) -> dict:
    """The pre-registered stopping rule (docs/SELF_DISTILLATION.md section 5). `history` is one entry per
    finished round, oldest first: {"round", "n_eval", "cap", "solved"}; only entries with the same
    (n_eval, cap) as the newest are comparable (the final round may use a bigger n).

      1. target_reached : newest solved >= target_solved
      2. regression     : newest solved <= best-so-far - max_regress  (keep the best round's model)
      3. plateau        : each of the last `patience` rounds failed to beat the best BEFORE it by >= min_gain problems
      4. max_rounds     : rounds_done >= max_rounds
    otherwise continue. Gains are in problems on the SAME problem set, n and cap (1 problem = 0.41 pp on miniF2F-test)."""
    if not history:
        return {"stop": False, "reason": "no_history", "best_round": None}
    last = history[-1]
    comp = [h for h in history if (h["n_eval"], h["cap"]) == (last["n_eval"], last["cap"])]
    best = max(comp, key=lambda h: (h["solved"], -h["round"]))
    if target_solved is not None and last["solved"] >= target_solved:
        return {"stop": True, "reason": "target_reached", "best_round": best["round"]}
    prior_best = max((h["solved"] for h in comp[:-1]), default=None)
    if prior_best is not None and last["solved"] <= prior_best - max_regress:
        return {"stop": True, "reason": "regression", "best_round": best["round"]}
    if len(comp) >= patience + 1:
        flat = 0
        for i in range(len(comp) - patience, len(comp)):
            before = max(h["solved"] for h in comp[:i])
            if comp[i]["solved"] - before < min_gain:
                flat += 1
        if flat == patience:
            return {"stop": True, "reason": "plateau", "best_round": best["round"]}
    if rounds_done >= max_rounds:
        return {"stop": True, "reason": "max_rounds", "best_round": best["round"]}
    return {"stop": False, "reason": "continue", "best_round": best["round"]}


# =============================================================================== config and layout


@dataclass
class Cfg:
    run_dir: str
    pool: str
    eval_problems: list
    base_model: str
    rounds: int = 3
    start_round: int = 1
    name: str = "ei"
    # selection
    band_lo: float = 0.0
    band_hi: float = 0.75
    include_zero: bool = False
    max_zero: Optional[int] = None
    max_problems: Optional[int] = None
    seed_ver: list = field(default_factory=list)
    seed: int = 0
    # generation
    n: int = 8
    gen_cap: int = 16384
    tp: int = 4
    workers: int = 8
    corrections: bool = False
    gen_temperature: Optional[float] = None  # None: lean_gen's own default (1.0); not passed, so old stage markers stay valid
    # keep rules
    max_per_problem: int = 4
    max_correction_per_problem: int = 1
    rank: str = "shortest"
    keep_easy: bool = False
    fill_duplicates: bool = False
    max_seq_len: int = 18432
    min_rows: int = 50
    # training
    train_config: str = "configs/p2_sft.toml"
    nproc: int = 4
    lr: Optional[float] = None
    epochs: Optional[int] = None
    lora_r: int = 64
    lora_alpha: float = 128.0
    # averaging
    avg_weight: Optional[float] = None
    avg_anchor: str = "original"
    prune_export: bool = False
    # evaluation
    eval_n: int = 8
    final_eval_n: Optional[int] = None
    eval_cap: int = 30000
    eval_seed: int = 0
    baseline_run: Optional[str] = None
    # stopping rule
    min_gain: int = 2
    patience: int = 2
    max_regress: int = 3
    target_solved: Optional[int] = None
    allow_test_selection: bool = False
    protected_problems: list = field(default_factory=list)
    redo: list = field(default_factory=list)
    dry: bool = False

    def eval_n_for(self, r: int) -> int:
        return self.final_eval_n if (self.final_eval_n and r == self.rounds) else self.eval_n


def resolve_snapshot(path: str) -> str:
    """A HF `snapshots/` parent resolves to its (last) hash dir; a dir with config.json is itself."""
    p = Path(path)
    if (p / "config.json").is_file():
        return str(p)
    kids = sorted(d for d in p.iterdir() if d.is_dir()) if p.is_dir() else []
    return str(kids[-1]) if kids else str(p)


def minif2f_file(cfg: Cfg) -> str:
    """The eval problem file that holds the benchmark we score: the first --eval-problems entry."""
    return str(cfg.eval_problems[0])


class Layout:
    def __init__(self, cfg: Cfg, r: int):
        self.cfg, self.r = cfg, r
        self.root = Path(cfg.run_dir) / f"round_{r:02d}"
        self.stages = self.root / "stages"
        self.logs = self.root / "logs"
        self.sel = self.root / "select"
        self.gen = self.root / "gen"
        self.data = self.root / "sft_data"
        self.sft = self.root / "sft"
        self.export = self.root / "export"
        self.export_work = self.root / "export_check"
        self.avg = self.root / "avg"
        self.eval = self.root / "eval"
        self.report = self.root / "report"
        self.pass_rates = self.root / "pass_rates.json"
        self.manifest = self.root / "round_manifest.json"

    @property
    def model_out(self) -> Path:
        return self.avg if self.cfg.avg_weight is not None else self.export

    def eval_files(self):
        n = self.cfg.eval_n_for(self.r)
        tag = f"n{n}_c{self.cfg.eval_cap}"
        return self.eval / f"gen_{tag}.jsonl", self.eval / f"ver_{tag}.jsonl"


def model_in(cfg: Cfg, r: int) -> str:
    """The model round r samples from and trains on: the base for the first round, else round r-1's output."""
    if r <= 1:
        return resolve_snapshot(cfg.base_model)
    return str(Layout(cfg, r - 1).model_out)


def model_rev_args(cfg: Cfg, model: str, tag: str) -> list[str]:
    """lean_gen infers a revision from a HF `snapshots/<hash>` path; anything else needs --model-rev."""
    if os.path.basename(os.path.dirname(os.path.realpath(model))) == "snapshots":
        return []
    return ["--model-rev", f"{cfg.name}-{tag}"]


# =============================================================================== stages


@dataclass
class Cmd:
    argv: list
    env: dict = field(default_factory=dict)


@dataclass
class Stage:
    name: str
    r: int
    cmds: list = field(default_factory=list)  # list[Cmd]
    fn: Optional[Callable] = None  # in-process
    desc: str = ""
    params: dict = field(default_factory=dict)  # hashed with the commands (in-process stages: their only identity)
    outputs: list = field(default_factory=list)
    pre_clean: list = field(default_factory=list)
    chain: str = ""

    def identity(self) -> str:
        return json.dumps({"name": self.name, "r": self.r, "cmds": [[c.argv, sorted(c.env.items())] for c in self.cmds],
                           "params": self.params}, sort_keys=True, default=str)


LATEST = "@latest_ckpt:"


def latest_checkpoint(sft_dir: Path) -> Optional[Path]:
    """Highest micro_N checkpoint under `sft_dir/*/checkpoints` (the layout scripts/sft.py writes)."""
    cands = sorted(Path(sft_dir).glob("*/checkpoints/micro_*"), key=lambda p: int(p.name.split("_")[-1]) if p.name.split("_")[-1].isdigit() else -1)
    cands = [c for c in cands if (c / "optimizer_rank0.pt").exists()]
    return cands[-1] if cands else None


def _sha_or_none(p) -> Optional[str]:
    return C.sha256_file(p) if os.path.isfile(p) else None


def build_stages(cfg: Cfg, r: int) -> list:
    L = Layout(cfg, r)
    gen_model = model_in(cfg, r)
    mrev = model_rev_args(cfg, gen_model, f"r{r - 1}")
    pool_ids = None  # read lazily inside the stage
    prev_table = Layout(cfg, r - 1).pass_rates if r > cfg.start_round else None
    ev_gen, ev_ver = L.eval_files()
    en = cfg.eval_n_for(r)
    gen_r0, ver_r0 = L.gen / "gen_r0.jsonl", L.gen / "ver_r0.jsonl"
    gen_r1, ver_r1 = L.gen / "gen_r1.jsonl", L.gen / "ver_r1.jsonl"
    sel_problems = L.sel / "problems.jsonl"
    seed = cfg.seed + 1000 * r
    env_py = {"PYTHONPATH": str(REPO)}

    def pgv(run_dir, gen_args, problems, ver_out, feedback=False):
        argv = [PY, str(SCRIPTS / "p2" / "pipeline_gen_verify.py"), "--run-dir", str(run_dir), "--gen-args", " ".join(map(str, gen_args)),
                "--problems", str(problems), "--verify-out", str(ver_out), "--workers", str(cfg.workers)]
        if feedback:
            argv.append("--feedback")
        return Cmd(argv, dict(env_py))

    st: list[Stage] = []

    # ---- select
    st.append(Stage("select", r, fn=lambda ctx: _do_select(ctx), params={
        "pool_sha": _sha_or_none(cfg.pool), "seed_ver": [_sha_or_none(p) for p in cfg.seed_ver] if r == cfg.start_round else "prev_round",
        "band": [cfg.band_lo, cfg.band_hi], "include_zero": cfg.include_zero, "max_zero": cfg.max_zero, "max_problems": cfg.max_problems,
        "seed": seed, "gen_cap": cfg.gen_cap},
        desc=f"band ({cfg.band_lo}, {cfg.band_hi}] on pass rates from "
             + (f"{[str(p) for p in cfg.seed_ver] or 'nothing (round 1 measures the whole pool)'}" if r == cfg.start_round else str(prev_table))
             + f" -> {sel_problems}",
        outputs=[sel_problems, L.sel / "selection.json"]))

    # ---- gen_verify (round-0 attempts)
    ga = ["--problems", sel_problems, "--out", gen_r0, "--model", gen_model, *mrev, "--n", cfg.n, "--cap-tokens", cfg.gen_cap,
          "--round", 0, "--tp", cfg.tp, "--seed", seed]
    if cfg.gen_temperature is not None:
        ga += ["--temperature", cfg.gen_temperature]  # part of the argv, hence of the stage identity and chain hash
    st.append(Stage("gen_verify", r, cmds=[pgv(L.gen / "run_r0", ga, sel_problems, ver_r0, feedback=cfg.corrections)], outputs=[gen_r0, ver_r0]))

    # ---- gen_verify_corr
    if cfg.corrections:
        gc = ["--problems", sel_problems, "--out", gen_r1, "--model", gen_model, *mrev, "--n", 1, "--cap-tokens", cfg.gen_cap,
              "--round", 1, "--tp", cfg.tp, "--seed", seed, "--corrections-from", ver_r0, gen_r0]
        if cfg.gen_temperature is not None:
            gc += ["--temperature", cfg.gen_temperature]
        st.append(Stage("gen_verify_corr", r, cmds=[pgv(L.gen / "run_r1", gc, sel_problems, ver_r1)], outputs=[gen_r1, ver_r1]))

    # ---- measure
    st.append(Stage("measure", r, fn=lambda ctx: _do_measure(ctx), params={"gen_cap": cfg.gen_cap},
                    desc=f"pass rates of this round's samples merged into the carried table -> {L.pass_rates}", outputs=[L.pass_rates]))

    # ---- sft_data
    gens = [gen_r0] + ([gen_r1] if cfg.corrections else [])
    vers = [ver_r0] + ([ver_r1] if cfg.corrections else [])
    sd = [PY, str(SCRIPTS / "p2" / "build_sft_data.py"), "--generations", *map(str, gens), "--verifications", *map(str, vers),
          "--problems", str(sel_problems), "--eval-problems", *map(str, list(cfg.eval_problems) + list(cfg.protected_problems)),
          "--cap-tokens", str(cfg.gen_cap), "--out-dir", str(L.data), *keep_rule_args(cfg), "--seed", str(cfg.seed)]
    st.append(Stage("sft_data", r, cmds=[Cmd(sd, dict(env_py))], outputs=[L.data / "manifest.json"], pre_clean=[L.data],
                    params={"min_rows": cfg.min_rows}))

    # ---- train
    tr = [PY, "-m", "torch.distributed.run", "--standalone", f"--nproc_per_node={cfg.nproc}", str(SCRIPTS / "sft.py"), str(cfg.train_config),
          f"--data.shard_dir={L.data / 'shards'}", f"--data.manifest_path={L.data / 'manifest.json'}", f"--paths.checkpoint_dir={L.sft}",
          f"--paths.model_snapshot={gen_model}"]
    if cfg.lr is not None:
        tr.append(f"--optimizer.lr={cfg.lr}")
    if cfg.epochs is not None:
        tr.append(f"--train.epochs={cfg.epochs}")
    st.append(Stage("train", r, cmds=[Cmd(tr, dict(env_py))], outputs=[L.sft]))

    # ---- export (merge LoRA into THIS round's input model; the script's base is overridden by env)
    ex = [PY, str(SCRIPTS / "qwen3_export_check.py"), "run", "--ckpt", LATEST + str(L.sft), "--out-dir", str(L.export), "--work", str(L.export_work),
          "--r", str(cfg.lora_r), "--alpha", str(cfg.lora_alpha), "--keep"]
    st.append(Stage("export", r, cmds=[Cmd(ex, {**env_py, "P2_EXPORT_BASE_SNAPSHOT": gen_model})], outputs=[L.export / "config.json"], pre_clean=[L.export]))

    # ---- average
    if cfg.avg_weight is not None:
        anchor = resolve_snapshot(cfg.base_model) if cfg.avg_anchor == "original" else gen_model
        av = [PY, "-m", "align.model.weight_average", "--base", anchor, "--ft", str(L.export), "--ft-weight", str(cfg.avg_weight), "--out", str(L.avg)]
        st.append(Stage("average", r, cmds=[Cmd(av, dict(env_py))], outputs=[L.avg / "weight_average.json"], pre_clean=[L.avg, Path(str(L.avg) + ".partial")]))

    # ---- eval
    ea = ["--problems", minif2f_file(cfg), "--out", ev_gen, "--model", L.model_out, *model_rev_args(cfg, str(L.model_out), f"r{r}"), "--n", en,
          "--cap-tokens", cfg.eval_cap, "--round", 0, "--tp", cfg.tp, "--seed", cfg.eval_seed]
    st.append(Stage("eval", r, cmds=[pgv(L.eval / "run", ea, minif2f_file(cfg), ev_ver)], outputs=[ev_gen, ev_ver]))

    # ---- report
    base_spec = cfg.baseline_run
    if base_spec is None and r > cfg.start_round:
        pg, pv = Layout(cfg, r - 1).eval_files()
        base_spec = f"prev:{pg}:{pv}"
    if base_spec is not None:
        rp = [PY, str(SCRIPTS / "p2" / "eval_report.py"), "--problems", *map(str, cfg.eval_problems), "--run", base_spec,
              "--run", f"r{r}:{ev_gen}:{ev_ver}", "--baseline", base_spec.split(":")[0], "--candidate", f"r{r}", "--out-dir", str(L.report)]
        st.append(Stage("report", r, cmds=[Cmd(rp, dict(env_py))], outputs=[L.report / "report.json"], pre_clean=[L.report]))

    # ---- manifest
    st.append(Stage("manifest", r, fn=lambda ctx: _do_manifest(ctx), desc=f"{L.manifest}: counts, eval, paired McNemar vs previous round, stop decision",
                    params={"stop": [cfg.min_gain, cfg.patience, cfg.max_regress, cfg.target_solved, cfg.rounds], "eval_n": en}, outputs=[L.manifest]))
    return st


def chain_stages(cfg: Cfg, upto_round: int) -> dict[int, list]:
    """Build every stage of rounds start..upto_round and assign chain hashes (round r+1's first stage chains
    off round r's last stage)."""
    out: dict[int, list] = {}
    prev = hashlib.sha256(json.dumps({"base": resolve_snapshot(cfg.base_model), "name": cfg.name}).encode()).hexdigest()
    for r in range(cfg.start_round, upto_round + 1):
        stages = build_stages(cfg, r)
        for s in stages:
            prev = hashlib.sha256((prev + s.identity()).encode()).hexdigest()
            s.chain = prev
        out[r] = stages
    return out


# =============================================================================== in-process stages


@dataclass
class Ctx:
    cfg: Cfg
    r: int
    L: Layout


def _write_json(path: Path, obj) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(obj, indent=1, sort_keys=True, default=str))
    os.replace(tmp, path)


def _do_select(ctx: Ctx) -> dict:
    cfg, L = ctx.cfg, ctx.L
    pool = C.load_problems(cfg.pool)
    if ctx.r == cfg.start_round:
        table = pass_rate_table(cfg.seed_ver, cfg.gen_cap) if cfg.seed_ver else {}
        src = [str(p) for p in cfg.seed_ver]
    else:
        table = json.loads(Layout(cfg, ctx.r - 1).pass_rates.read_text())["table"]
        src = [str(Layout(cfg, ctx.r - 1).pass_rates)]
    order, why, counts = select_by_band(table, list(pool), cfg.band_lo, cfg.band_hi, include_zero=cfg.include_zero, max_zero=cfg.max_zero,
                                        max_problems=cfg.max_problems, seed=cfg.seed + 1000 * ctx.r)
    if not order:
        raise DataStop(f"round {ctx.r}: no problems in band ({counts}); the pool is exhausted for this model")
    L.sel.mkdir(parents=True, exist_ok=True)
    C.write_jsonl_atomic(L.sel / "problems.jsonl", [pool[p] for p in order])
    _write_json(L.sel / "selection.json", {"round": ctx.r, "band": [cfg.band_lo, cfg.band_hi], "source": src, "counts": counts,
                                           "selected": len(order), "pool": len(pool), "reasons": why})
    return {"selected": len(order), "counts": counts}


def _do_measure(ctx: Ctx) -> dict:
    cfg, L = ctx.cfg, ctx.L
    prev = {} if ctx.r == cfg.start_round else json.loads(Layout(cfg, ctx.r - 1).pass_rates.read_text())["table"]
    if ctx.r == cfg.start_round and cfg.seed_ver:
        prev = pass_rate_table(cfg.seed_ver, cfg.gen_cap)
    new = pass_rate_table([L.gen / "ver_r0.jsonl"], cfg.gen_cap)
    table = dict(prev)
    for pid, m in new.items():
        table[pid] = {**m, "round": ctx.r}
    # a problem teaches something only if some sample was accepted and (unless --keep-easy) not all n>=8 were
    # (build_sft_data.EASY_MIN_SAMPLES): with none, sft_data would raise "nothing selected" -- say so as a data stop instead
    trainable = sum(1 for m in new.values() if m["c"] >= 1 and (cfg.keep_easy or m["n"] < 8 or m["c"] < m["n"]))
    _write_json(L.pass_rates, {"round": ctx.r, "n_measured_this_round": len(new), "n_trainable": trainable, "table": table})
    if trainable == 0:
        raise DataStop(f"round {ctx.r}: none of the {len(new)} sampled problems has both an accepted proof and a failure "
                       f"(all unsolved, or all solved {cfg.n}/{cfg.n}); there is nothing to distil")
    return {"measured": len(new), "trainable": trainable}


def _do_manifest(ctx: Ctx) -> dict:
    cfg, L, r = ctx.cfg, ctx.L, ctx.r
    from align.eval import proof_stats as S

    ids = list(C.load_problems(minif2f_file(cfg)))
    n_ev = cfg.eval_n_for(r)
    _, ev_ver = L.eval_files()
    ev = eval_summary(ev_ver, ids, n_ev, cfg.eval_cap)
    smap = solved_map(ev_ver, ids, n_ev, cfg.eval_cap)
    paired = {}
    if r > cfg.start_round:
        pm = json.loads(Layout(cfg, r - 1).manifest.read_text())
        if (pm["eval"]["n_samples_per_problem"], pm["eval"]["cap_tokens"]) == (n_ev, cfg.eval_cap):
            pg, pv = Layout(cfg, r - 1).eval_files()
            m = S.mcnemar_exact(solved_map(pv, ids, n_ev, cfg.eval_cap), smap)
            paired["vs_previous_round"] = {"only_prev": m.only_a, "only_this": m.only_b, "p_value_exact_two_sided": m.p_value}
        else:
            paired["vs_previous_round"] = "not comparable (eval n/cap differ)"
    sel = json.loads((L.sel / "selection.json").read_text())
    dm = json.loads((L.data / "manifest.json").read_text())
    rows = dm["rows"]["total"]
    ckpt = latest_checkpoint(L.sft)
    avg = json.loads((L.avg / "weight_average.json").read_text()) if cfg.avg_weight is not None else None
    hist = []
    for k in range(cfg.start_round, r):
        pm = json.loads(Layout(cfg, k).manifest.read_text())
        hist.append({"round": k, "n_eval": pm["eval"]["n_samples_per_problem"], "cap": pm["eval"]["cap_tokens"], "solved": pm["eval"]["solved_at_n"]})
    hist.append({"round": r, "n_eval": n_ev, "cap": cfg.eval_cap, "solved": ev["solved_at_n"]})
    dec = stop_decision(hist, rounds_done=r, max_rounds=cfg.rounds, min_gain=cfg.min_gain, patience=cfg.patience, max_regress=cfg.max_regress,
                        target_solved=cfg.target_solved)
    markers = {}
    for f in sorted(L.stages.glob("*.done.json")):
        mk = json.loads(f.read_text())
        markers[mk["stage"]] = {"seconds": mk.get("seconds"), "finished": mk.get("finished"), "chain": mk["chain"][:12]}
    man = {"round": r, "run_dir": cfg.run_dir, "model_in": model_in(cfg, r), "model_out": str(L.model_out),
           "selection": {"selected": sel["selected"], "pool": sel["pool"], "counts": sel["counts"], "band": sel["band"]},
           "sft_data": {"rows": rows, "problems_with_rows": dm["problems"]["with_rows"], "tokens": dm.get("tokens"), "params": dm["params"]},
           "train": {"checkpoint": str(ckpt) if ckpt else None, "lr": cfg.lr, "epochs": cfg.epochs},
           "average": avg and {k: avg[k] for k in ("ft_weight", "tensors", "shards", "peak_buffer_bytes")},
           "eval": ev, "paired": paired, "history": hist, "stop": dec, "stages": markers,
           "params": {k: v for k, v in cfg.__dict__.items() if k not in ("seed_ver", "redo", "dry")}, "written": time.strftime("%Y-%m-%d %H:%M:%S")}
    _write_json(L.manifest, man)
    return {"stop": dec}


# =============================================================================== running


class SubprocessRunner:
    """Runs one Cmd; its stdout/stderr go to `log`. SIGTERM/SIGINT are forwarded to the child's process group."""

    def __init__(self):
        self.sig = None

    def run(self, argv: list, env: dict, log: Path) -> int:
        log.parent.mkdir(parents=True, exist_ok=True)
        full = {**os.environ, **env}
        with open(log, "ab") as lf:
            lf.write(f"\n=== {time.strftime('%F %T')} {shlex.join(argv)}\n".encode())
            lf.flush()
            p = subprocess.Popen(argv, stdout=lf, stderr=subprocess.STDOUT, env=full, start_new_session=True, cwd=str(REPO))
            old = {s: signal.signal(s, lambda sg, _f: setattr(self, "sig", sg)) for s in (signal.SIGTERM, signal.SIGINT)}
            try:
                while p.poll() is None:
                    if self.sig is not None:
                        os.killpg(p.pid, signal.SIGTERM)
                        try:
                            p.wait(timeout=120)
                        except subprocess.TimeoutExpired:
                            os.killpg(p.pid, signal.SIGKILL)
                            p.wait()
                        return EXIT_SIGNAL
                    time.sleep(1.0)
            finally:
                for s, h in old.items():
                    signal.signal(s, h)
            return p.returncode


def expand_argv(argv: list) -> list:
    out = []
    for a in argv:
        if isinstance(a, str) and a.startswith(LATEST):
            ck = latest_checkpoint(Path(a[len(LATEST):]))
            if ck is None:
                raise DriverError(f"no checkpoint with optimizer_rank0.pt under {a[len(LATEST):]}: the train stage left nothing to export")
            out.append(str(ck))
        else:
            out.append(str(a))
    return out


def display_argv(argv: list) -> str:
    parts = []
    for a in argv:
        if isinstance(a, str) and a.startswith(LATEST):
            d = a[len(LATEST):]
            parts.append(f'"$(ls -d {d}/*/checkpoints/micro_???????? | sort -V | tail -1)"')
        else:
            parts.append(shlex.quote(str(a)))
    return " ".join(parts)


def marker_path(L: Layout, name: str) -> Path:
    return L.stages / f"{name}.done.json"


def stage_status(L: Layout, s: Stage) -> str:
    """done | pending | stale (marker exists for a different configuration)."""
    mp = marker_path(L, s.name)
    if not mp.exists():
        return "pending"
    mk = json.loads(mp.read_text())
    if mk.get("chain") != s.chain:
        return "stale"
    if not all(Path(o).exists() for o in s.outputs):
        return "pending"  # marker present, outputs gone: redo
    return "done"


def stage_artefacts(L: Layout, name: str) -> list:
    """Every path stage `name` leaves on disk, including the pipeline run dirs and the export scratch that the
    stage's declared outputs do not name. These are what a re-run would silently REUSE: sft.py resumes the newest
    checkpoint, lean_gen resumes any gen/ver key whose model_rev LABEL (constant for an exported model) matches."""
    g = L.gen
    return {
        "select": [L.sel],
        "gen_verify": [g / "gen_r0.jsonl", g / "ver_r0.jsonl", g / "run_r0"],
        "gen_verify_corr": [g / "gen_r1.jsonl", g / "ver_r1.jsonl", g / "run_r1"],
        "measure": [L.pass_rates],
        "sft_data": [L.data],
        "train": [L.sft],
        "export": [L.export, L.export_work],
        "average": [L.avg, Path(str(L.avg) + ".partial")],
        "eval": [L.eval],
        "report": [L.report],
        "manifest": [L.manifest],
    }[name]


def _remove(p: Path) -> None:
    if p.is_dir() and not p.is_symlink():
        shutil.rmtree(p)
    elif p.exists() or p.is_symlink():
        p.unlink()


def apply_redo(cfg: Cfg, redo: list) -> None:
    """Discard the stage, every later stage of the round, and (when the redone stage can change the round's output
    model or pass-rate table) every stage of every later round: their markers AND their artefacts. Markers alone are not
    enough: the children resume from whatever is on disk (newest checkpoint, keyed gen/ver rows), so a redo that left
    them would 'finish' instantly on the previous configuration's model and samples."""
    for spec in redo:
        rs, name = spec.split(":")
        r = int(rs)
        if name not in STAGE_ORDER:
            raise DriverError(f"--redo: unknown stage {name!r}; stages: {STAGE_ORDER}")
        L = Layout(cfg, r)
        for later in STAGE_ORDER[STAGE_ORDER.index(name):]:
            marker_path(L, later).unlink(missing_ok=True)
            for p in stage_artefacts(L, later):
                _remove(p)
        # round r+1 reads round r's model (export/average) and pass_rates (measure); eval/report/manifest feed nothing downstream
        feeds_next = STAGE_ORDER.index(name) <= STAGE_ORDER.index("average")
        for rr in range(r + 1, cfg.rounds + 1):  # later rounds chained off this one are stale too
            Lr = Layout(cfg, rr)
            for f in Lr.stages.glob("*.done.json") if Lr.stages.exists() else []:
                f.unlink()
            if feeds_next:
                for later in STAGE_ORDER:
                    for p in stage_artefacts(Lr, later):
                        _remove(p)


def generation_incomplete(gen_path, problems_path, n: int, cap_tokens: int) -> Optional[str]:
    """Why `gen_path` is not a complete round-0 sample set (None = complete).

    Two independent proofs, both required (audit U12-2 follow-up): every planned key (problem x sample_idx < n at the
    cap) is in the file, AND lean_gen's last manifest session says `complete` (an interrupted or crashed session
    can leave a file that looks finished to a reader that only checks that it exists)."""
    gen_path = Path(gen_path)
    want = set()
    with open(problems_path, encoding="utf-8") as f:
        for ln in f:
            if ln.strip():
                pid = json.loads(ln)["problem_id"]
                want.update((pid, i, 0, cap_tokens) for i in range(n))
    have = set()
    with open(gen_path, "rb") as f:
        for raw in f:
            if not raw.endswith(b"\n") or not raw.strip():
                continue  # torn tail of a killed writer
            r = json.loads(raw)
            have.add((r["problem_id"], int(r["sample_idx"]), int(r["round"]), int(r["cap_tokens"])))
    missing = sorted(want - have)
    if missing:
        return f"{gen_path} is missing {len(missing)} of {len(want)} planned samples (first: {missing[:3]})"
    man = Path(str(gen_path) + ".manifest.json")
    if not man.exists():
        return f"{man} does not exist: cannot confirm lean_gen finished"
    sessions = json.loads(man.read_text()).get("sessions") or []
    if not sessions or sessions[-1].get("status") != "complete":
        return f"lean_gen's last manifest session is {sessions[-1].get('status') if sessions else None!r}, not 'complete' ({man})"
    return None


def run_stage(cfg: Cfg, s: Stage, runner, say=print) -> str:
    L = Layout(cfg, s.r)
    status = stage_status(L, s)
    if status == "done":
        say(f"[round {s.r}] {s.name}: done (skipped)")
        return "skipped"
    if status == "stale":
        raise DriverError(f"[round {s.r}] {s.name}: a finished marker exists for a DIFFERENT configuration (command line or upstream "
                          f"stage changed). Refusing to mix them. Re-run with --redo {s.r}:{s.name} to discard it and every later stage "
                          f"of this round, or restore the old arguments.")
    t0 = time.time()
    say(f"[round {s.r}] {s.name}: running")
    L.stages.mkdir(parents=True, exist_ok=True)
    for p in s.pre_clean:
        if Path(p).exists():
            shutil.rmtree(p) if Path(p).is_dir() else Path(p).unlink()
    if s.fn is not None:
        s.fn(Ctx(cfg, s.r, L))
    for i, c in enumerate(s.cmds):
        rc = runner.run(expand_argv(c.argv), c.env, L.logs / f"{s.name}{'' if len(s.cmds) == 1 else f'_{i}'}.log")
        if rc == EXIT_SIGNAL:
            raise KeyboardInterrupt(f"{s.name} stopped by a signal")
        if rc != 0:
            raise DriverError(f"[round {s.r}] {s.name}: command {i} exited {rc} (log: {L.logs}); no marker written, re-run resumes it")
    missing = [str(o) for o in s.outputs if not Path(o).exists()]
    if missing:
        raise DriverError(f"[round {s.r}] {s.name}: finished but declared outputs are missing: {missing}")
    if s.name == "gen_verify":
        why = generation_incomplete(L.gen / "gen_r0.jsonl", L.sel / "problems.jsonl", cfg.n, cfg.gen_cap)
        if why:
            raise DriverError(f"[round {s.r}] gen_verify: generation is not complete, no marker written: {why}")
    if s.name == "sft_data":
        rows = json.loads((L.data / "manifest.json").read_text())["rows"]["total"]
        if rows < cfg.min_rows:
            raise DataStop(f"round {s.r}: only {rows} SFT rows (< --min-rows {cfg.min_rows}); the selected problems yielded too few accepted proofs")
    if s.name == "average" and cfg.prune_export and L.export.exists():
        shutil.rmtree(L.export)  # the averaged model is the round's output; the 16 GB export has served its purpose
    _write_json(marker_path(L, s.name), {"stage": s.name, "round": s.r, "chain": s.chain, "finished": time.strftime("%F %T"),
                                         "seconds": round(time.time() - t0, 1), "identity": json.loads(s.identity())})
    return "ran"


def run_all(cfg: Cfg, runner, say=print) -> int:
    chains = chain_stages(cfg, cfg.rounds)
    if cfg.redo:
        apply_redo(cfg, cfg.redo)
    try:
        for r in range(cfg.start_round, cfg.rounds + 1):
            err = split_problems(cfg)  # files can change while a round runs: check again before each round
            if err:
                say(f"[round {r}] REFUSED: {err}")
                return 3
            for s in chains[r]:
                run_stage(cfg, s, runner, say)
            man = json.loads(Layout(cfg, r).manifest.read_text())
            dec = man["stop"]
            say(f"[round {r}] solved@{man['eval']['n_samples_per_problem']} = {man['eval']['solved_at_n']}/{man['eval']['n_problems']} "
                f"(cap {man['eval']['cap_tokens']}); decision: {dec['reason']}")
            if dec["stop"]:
                say(f"STOP after round {r}: {dec['reason']} (best round: {dec['best_round']})")
                return EXIT_OK
    except DataStop as e:
        say(f"STOP (data): {e}")
        return EXIT_DATA_STOP
    except KeyboardInterrupt as e:
        say(f"interrupted: {e}")
        return EXIT_SIGNAL
    except (DriverError, C.ContractError) as e:
        say(f"FAILED: {e}")
        return EXIT_FAIL
    return EXIT_OK


def dry_run(cfg: Cfg, say=print) -> int:
    chains = chain_stages(cfg, cfg.rounds)
    say(f"# dry run: nothing is executed or written. run dir {cfg.run_dir}; rounds {cfg.start_round}..{cfg.rounds}")
    for r in range(cfg.start_round, cfg.rounds + 1):
        L = Layout(cfg, r)
        say(f"\n## round {r}: model in {model_in(cfg, r)} -> out {L.model_out}")
        for s in chains[r]:
            say(f"[{stage_status(L, s)}] {s.name}")
            if s.fn is not None:
                say(f"    # in-process: {s.desc}")
            for c in s.cmds:
                env = " ".join(f"{k}={shlex.quote(v)}" for k, v in sorted(c.env.items()))
                say(f"    {('env ' + env + ' ') if env else ''}{display_argv(c.argv)}")
    return EXIT_OK


# =============================================================================== CLI


def scored_splits(path: str) -> set:
    """The `split` values of the scored benchmark file. FAILS CLOSED: a missing, unreadable or
    malformed file, an empty file, or any record without a `split` raises ValueError -- an unknown
    split must never pass as "not test"."""
    try:
        with open(path, encoding="utf-8") as f:
            rows = [json.loads(line) for line in f if line.strip()]
    except (OSError, ValueError) as e:
        raise ValueError(f"cannot read the scored benchmark {path}: {e}") from e
    if not rows:
        raise ValueError(f"the scored benchmark {path} is empty")
    missing = sum(1 for r in rows if not isinstance(r, dict) or not r.get("split"))
    if missing:
        raise ValueError(f"{missing} record(s) in {path} have no `split`: cannot prove it is not a test split")
    return {r["split"] for r in rows}


def split_problems(cfg: Cfg) -> Optional[str]:
    """Split hygiene; returns an error message or None. The scored benchmark drives the stop rule, the
    best-round choice and the averaging check, so scoring a TEST split makes every later test number
    partly a selection artefact. Development runs score miniF2F-valid. --allow-test-selection is an
    escape hatch for one non-adaptive diagnostic round; the FINAL evaluation does not use this driver at
    all (docs/SOTA_PLAN.md: frozen checkpoint -> pipeline_gen_verify.py -> eval_report.py).
    Decontamination: every protected test benchmark must reach build_sft_data's exclusion set."""
    try:
        splits = scored_splits(str(cfg.eval_problems[0]))
    except ValueError as e:
        return str(e)
    if "test" in splits:
        if not cfg.allow_test_selection:
            return (f"--eval-problems[0] {cfg.eval_problems[0]} is a TEST split: rounds, the stop rule and averaging would "
                    "be selected on test. Score miniF2F-valid (scripts/p2/build_valid.py)")
        if cfg.rounds != 1 or cfg.target_solved is not None:
            return "--allow-test-selection is for ONE non-adaptive round: needs --rounds 1 and no --target-solved"
    guarded = set()
    for p in list(cfg.eval_problems) + list(cfg.protected_problems):
        try:
            guarded |= scored_splits(str(p))
        except ValueError as e:
            return str(e)
    if "test" not in guarded:
        return ("no TEST benchmark reaches the SFT decontamination gate: pass the protected test files "
                "(miniF2F-test, MOBench) via --protected-problems")
    return pool_informal_overlap(cfg)


def pool_informal_overlap(cfg: Cfg) -> Optional[str]:
    """build_sft_data excludes by exact statement_sha only, so 'the same informal problem, formalised differently' (the
    realistic leak, build_problems' informal 13-gram / 8-gram layers) must be absent from the pool already. It is applied
    to the pool against test+MOBench by build_problems and to the dev split by build_valid; this re-checks the files
    actually handed to this driver, so a stale pool or dev file cannot inflate the dev pass rate that drives the stop rule."""
    if not Path(cfg.pool).is_file():
        return None  # the missing pool is reported where it is read
    spec = importlib.util.spec_from_file_location("p2_build_problems_for_ei", Path(__file__).with_name("build_problems.py"))
    bp = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = bp
    try:
        spec.loader.exec_module(bp)
    finally:
        sys.modules.pop(spec.name, None)
    guarded = []
    for p in list(cfg.eval_problems) + list(cfg.protected_problems):
        guarded.extend(C.iter_jsonl(p))
    decon = bp.Decontaminator(guarded)
    hits = [(r["problem_id"], why) for r in C.iter_jsonl(cfg.pool) if (why := decon.check_informal(r.get("informal_prefix"))) is not None]
    if hits:
        return (f"{len(hits)} --pool problem(s) share informal statement text with --eval-problems/--protected-problems "
                f"(first: {hits[:3]}): training on them leaks the dev/test problems. Drop them from the pool, or rebuild the dev "
                "split with scripts/p2/build_valid.py (it drops valid problems that overlap the pool)")
    return None


def test_guard(ap, cfg: Cfg) -> None:
    err = split_problems(cfg)
    if err:
        ap.error(err)


def parse_args(argv=None) -> Cfg:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    g = ap.add_argument
    g("--run-dir", required=True)
    g("--pool", required=True, help="training pool Problem JSONL (split=train)")
    g("--eval-problems", nargs="+", required=True, help="eval Problem JSONL; the FIRST is the scored benchmark (miniF2F)")
    g("--base-model", required=True, help="HF snapshot dir (or its snapshots/ parent) of the starting model")
    g("--rounds", type=int, default=3)
    g("--start-round", type=int, default=1)
    g("--name", default="ei")
    g("--seed-ver", nargs="*", default=[], help="verification JSONL(s) measured at --gen-cap that seed the pass-rate table of round 1")
    g("--band-lo", type=float, default=0.0)
    g("--band-hi", type=float, default=0.75)
    g("--include-zero", action="store_true", help="also resample problems with measured pass rate 0")
    g("--max-zero", type=int, default=None)
    g("--max-problems", type=int, default=None)
    g("--seed", type=int, default=0)
    g("--n", type=int, default=8, help="samples per selected problem per round")
    g("--gen-cap", type=int, default=16384)
    g("--gen-temperature", type=float, default=None,
      help="lean_gen --temperature for the generation stages (omit: lean_gen's default 1.0, argv unchanged)")
    g("--tp", type=int, default=4)
    g("--workers", type=int, default=8)
    g("--corrections", action="store_true", help="also generate+train one Lean-feedback correction round")
    g("--max-per-problem", type=int, default=4)
    g("--max-correction-per-problem", type=int, default=1)
    g("--rank", choices=("order", "shortest"), default="shortest")
    g("--keep-easy", action="store_true")
    g("--fill-duplicates", action="store_true")
    g("--max-seq-len", type=int, default=18432)
    g("--min-rows", type=int, default=50)
    g("--train-config", default="configs/p2_sft.toml")
    g("--nproc", type=int, default=4)
    g("--lr", type=float, default=None)
    g("--epochs", type=int, default=None)
    g("--lora-r", type=int, default=64)
    g("--lora-alpha", type=float, default=128.0)
    g("--avg-weight", type=float, default=None, help="w in out=(1-w)*anchor+w*merged; omit for no averaging")
    g("--avg-anchor", choices=("original", "start"), default="original")
    g("--prune-export", action="store_true", help="delete the 16 GB merged export once the averaged model exists")
    g("--eval-n", type=int, default=8)
    g("--final-eval-n", type=int, default=None, help="n for the last round only (e.g. 32)")
    g("--eval-cap", type=int, default=30000)
    g("--eval-seed", type=int, default=0)
    g("--baseline-run", default=None, help="NAME:GEN:VER of the baseline eval (same n and cap)")
    g("--min-gain", type=int, default=2)
    g("--patience", type=int, default=2)
    g("--max-regress", type=int, default=3)
    g("--target-solved", type=int, default=None)
    g("--allow-test-selection", action="store_true",
      help="escape hatch: permit a TEST-split scored benchmark for ONE non-adaptive round (--rounds 1, no --target-solved); "
           "the final evaluation does not use this driver")
    g("--protected-problems", nargs="*", default=[],
      help="benchmark JSONL(s) never scored here but excluded from training data (miniF2F-test, MOBench, ...)")
    g("--redo", action="append", default=[], metavar="ROUND:STAGE")
    g("--dry-run", action="store_true")
    a = ap.parse_args(argv)
    d = vars(a)
    d["dry"] = d.pop("dry_run")
    cfg = Cfg(**d)
    if cfg.avg_weight is not None and not (0.0 < cfg.avg_weight <= 1.0):
        ap.error("--avg-weight must be in (0, 1]")
    test_guard(ap, cfg)
    if cfg.eval_cap < 30000:
        print(f"note: --eval-cap {cfg.eval_cap} < 30000: not comparable to the published pass@32 protocol", file=sys.stderr)
    return cfg


def main(argv=None, runner=None) -> int:
    cfg = parse_args(argv)
    if cfg.dry:
        return dry_run(cfg)
    return run_all(cfg, runner or SubprocessRunner())


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