code/scripts/lean_ceiling_test.py · viai957/align at main

align / code /scripts /lean_ceiling_test.py
Add files using upload-large-folder tool
bcc8d35 verified 15 days ago
34.4 kB
"""Ceiling test + measurements for the strict Lean verifier (Track L / P1).
House rule (paid for twice: a perfect policy scored 0/100 through a broken extractor, and
mprm's PPO runs measured the grader): ceiling-test every instrument before trusting any
number it produces. For a VERIFIER the ceiling has two sides, and both are checked:
  * known-good  : proofs that independently compile under the pinned toolchain must be
                  ACCEPTED (>= 98%): a strict gate that rejects honest proofs starves RL.
  * adversarial : hand-built exploits must ALL be REJECTED (0% accepted): a lax gate is
                  exactly what a policy under RL finds.
The criteria, the case tables (frozen, SHA-256 pinned) and the pass thresholds are
pre-registered in docs/LEAN_VERIFIER.md BEFORE the registered run. `--verify-prereg`
checks that the tables in this checkout hash to the registered value; the run refuses to
start otherwise (override: --allow-unregistered, recorded in the output).
Known-good sources
  1. Goedel-LM/Lean-workbook-proofs (public, proofs verified by Goedel-Prover-V1 under
     the same Lean 4.9 / Mathlib environment as DeepSeek-Prover-V1.5 and
     Goedel-Prover-V2): 6 contiguous slices of 25 rows at fixed offsets, fetched over
     HTTP from the HF datasets-server (never the full dataset).
  2. 30 hand-restated Mathlib-style theorems with proofs (scripts/lean_ceiling_cases.py).
  Each is first run through an INDEPENDENT compile-only pass (same REPL, no strictness at
  all: no error message, no `sorry` warning). Items that fail it are version-incompatible
  with the pin and are excluded and listed; the strict verdict is scored on the rest.
Host rules (/apps/CLAUDE.md): `apply_resource_guards(heavy=True)` runs the preflight and
takes the exclusive heavy-job flock; scratch lives under align/.tmp/; nothing touches / or
/tmp; every blocking call has a timeout; the process has a global wall-clock watchdog.
    python scripts/lean_ceiling_test.py --workers 8
    python scripts/lean_ceiling_test.py --verify-prereg
    python scripts/lean_ceiling_test.py --measure-only
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import re
import statistics
import subprocess
import sys
import threading
import time
import urllib.parse
import urllib.request
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO))

import importlib.util  # noqa: E402

# Loaded by path (not `import lean_ceiling_cases`): scripts/ is not a package, and the
# dependency firewall (rightly) rejects an import it cannot place in stdlib/allowed/align.
_spec = importlib.util.spec_from_file_location("lean_ceiling_cases", Path(__file__).with_name("lean_ceiling_cases.py"))
C = importlib.util.module_from_spec(_spec)
sys.modules["lean_ceiling_cases"] = C
_spec.loader.exec_module(C)

DOC = REPO / "docs" / "LEAN_VERIFIER.md"
DATASET = "Goedel-LM/Lean-workbook-proofs"
SLICE_OFFSETS = (0, 5000, 10000, 15000, 20000, 25000)
SLICE_LEN = 25
KNOWN_GOOD_FLOOR = 0.98
MIN_KNOWN_GOOD = 50
MIN_ADVERSARIAL = 30


# ---------------------------------------------------------------------------- prereg


def registered_hash() -> str | None:
    if not DOC.exists():
        return None
    m = re.search(r"PREREG_TABLE_SHA256:\s*`?([0-9a-f]{64})`?", DOC.read_text())
    return m.group(1) if m else None


def verify_prereg() -> bool:
    reg, cur = registered_hash(), C.table_hash()
    print(f"registered : {reg}\ncurrent    : {cur}\nmatch      : {reg == cur}")
    return reg == cur


# ---------------------------------------------------------------------------- verdict


def empty_known_good_entry(candidates: int, excluded: int) -> dict:
    """Report entry for a run in which NO known-good proof reached the strict stage (toolchain broken,
    REPL crashing, header mismatch: every candidate failed the independent compile-only filter).
    That is the absence of the ceiling's first side, i.e. a FAILURE: a verifier that rejects
    everything would otherwise sail through the adversarial side (accepted == 0 is trivially true)."""
    return {"n": 0, "accepted": 0, "rate": float("nan"), "floor": KNOWN_GOOD_FLOOR, "pass": False,
            "why": f"no known-good proof survived the compile-only filter ({excluded}/{candidates} excluded)",
            "rejected": []}


def ceiling_ok(report: dict, only: str | None) -> bool:
    """Exit-status verdict. Both sides must be PRESENT and passing in a full run; `--only` runs need
    the side they ran. A missing side is never a pass (the old `all(v.get('pass', True) ...)` read
    absence as success)."""
    need = {None: ("known_good", "adversarial"), "good": ("known_good",), "adversarial": ("adversarial",)}[only]
    if any(not isinstance(report.get(k), dict) or not report[k].get("pass") for k in need):
        return False
    return all(v.get("pass", True) for v in report.values() if isinstance(v, dict))


# ---------------------------------------------------------------------------- statistics


def wilson(k: int, n: int, z: float = 1.96) -> tuple[float, float]:
    """Wilson score interval for a binomial proportion (exact-ish at k=n or k=0)."""
    if n == 0:
        return (float("nan"), float("nan"))
    p = k / n
    d = 1 + z * z / n
    c = (p + z * z / (2 * n)) / d
    h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
    return (max(0.0, c - h), min(1.0, c + h))


def pctl(xs: list[float], q: float) -> float:
    if not xs:
        return float("nan")
    s = sorted(xs)
    return s[min(len(s) - 1, int(q * len(s)))]


# ---------------------------------------------------------------------------- data


def http_json(url: str, timeout: float = 60.0, retries: int = 4) -> dict:
    last = None
    for i in range(retries):
        try:
            req = urllib.request.Request(url, headers={"User-Agent": "align-lean-ceiling/1"})
            with urllib.request.urlopen(req, timeout=timeout) as r:
                return json.loads(r.read().decode("utf-8"))
        except Exception as e:  # noqa: BLE001
            last = e
            time.sleep(2 * (i + 1))
    raise RuntimeError(f"GET {url} failed: {last!r}")


def fetch_known_good(out_dir: Path) -> tuple[list[dict], dict]:
    rows: list[dict] = []
    for off in SLICE_OFFSETS:
        q = urllib.parse.urlencode(
            {"dataset": DATASET, "config": "default", "split": "train", "offset": off, "length": SLICE_LEN}
        )
        d = http_json(f"https://datasets-server.huggingface.co/rows?{q}")
        for r in d["rows"]:
            rows.append({"problem_id": r["row"]["problem_id"], "full_proof": r["row"]["full_proof"], "row_idx": r["row_idx"]})
    raw = out_dir / "known_good_raw.jsonl"
    with open(raw, "w") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    meta = {"dataset": DATASET, "offsets": list(SLICE_OFFSETS), "slice_len": SLICE_LEN, "n_rows": len(rows),
            "raw_sha256": hashlib.sha256(raw.read_bytes()).hexdigest()}
    try:
        info = http_json(f"https://huggingface.co/api/datasets/{DATASET}")
        meta["dataset_revision"] = info.get("sha")
    except Exception as e:  # noqa: BLE001
        meta["dataset_revision"] = f"unavailable: {e!r}"
    return rows, meta


def derive_statement(code: str):
    """(statement, target_name) from a full proof file: header + the LAST theorem/lemma's
    signature (cut at the signature-ending `:=`, via the verifier's own scanner) + `sorry`."""
    from align.verify.lean.strict import find_sig_end, lex

    lx = lex(code)
    toks = lx.toks
    if lx.broken:
        return None
    first = next((i for i, t in enumerate(toks) if t.kind == "id" and t.text in
                  ("theorem", "lemma", "def", "abbrev", "instance", "example", "structure", "inductive")), None)
    thm = [i for i, t in enumerate(toks) if t.kind == "id" and t.text in ("theorem", "lemma")]
    if first is None or not thm:
        return None
    k = thm[-1]
    e = find_sig_end(toks, k + 2)
    if e is None or k + 1 >= len(toks):
        return None
    name = toks[k + 1].text
    header = code[: toks[first - 1].end] if first > 0 else ""
    sig = code[toks[k + 2].start : toks[e].start]
    return header + "\n\n" + f"theorem {name} {sig} := by\n  sorry\n", name


def as_completion(code: str) -> str:
    return "I will state a plan and then give the proof.\n\n```lean4\n" + code.rstrip("\n") + "\n```\n"


# ---------------------------------------------------------------------------- compile-only


def compile_only(w, statement: str, code: str, timeout: float = 120.0) -> dict:
    """The NAIVE grader: elaborate the candidate and look at nothing but error messages.
    Used (a) as the independent reference compile that defines "known-good", and (b) to show
    how many adversarial cases a compile-only reward would pay."""
    from align.verify.lean.repl_pool import ReplError
    from align.verify.lean.strict import lex, parse_statement

    ref = parse_statement(statement)
    if not ref.ok:
        return {"ok": False, "why": "unsupported_statement"}
    lx = lex(code)
    nh = len(ref.header_toks)
    if lx.broken or tuple(t.text for t in lx.toks[:nh]) != ref.header_toks:
        return {"ok": False, "why": "header"}
    body = code[lx.toks[nh - 1].end :] if nh else code
    try:
        env = w.header_env(ref.header_text)
        resp = w.command(body, env, timeout)
    except ReplError as e:
        return {"ok": False, "why": e.code}
    msgs = [m for m in resp.get("messages", []) if isinstance(m, dict)]
    err = [m for m in msgs if m.get("severity") == "error"]
    sorry = bool(resp.get("sorries")) or any("sorry" in str(m.get("data", "")) for m in msgs)
    return {"ok": not err, "clean": (not err) and (not sorry), "why": "error" if err else ("sorry" if sorry else "")}


def naive_fresh(w, code: str, timeout: float = 120.0) -> dict:
    """Compile-only on the candidate's OWN header (fresh import, no cache) -- for the
    adversarial cases whose header is the attack. Each case gets a FRESH process: a second
    `import Mathlib` in a live REPL costs +3.9GiB anonymous (see `split_imports`)."""
    from align.verify.lean.repl_pool import ReplError

    w.discard()
    try:
        proc = w.ensure()
        resp = proc.request({"cmd": code}, timeout)
    except ReplError as e:
        w.discard()
        return {"ok": False, "why": e.code}
    w.discard()
    err = [m for m in resp.get("messages", []) if isinstance(m, dict) and m.get("severity") == "error"]
    return {"ok": not err, "why": "error" if err else ""}


# ---------------------------------------------------------------------------- measurement


def dir_bytes(p: Path) -> int:
    out = subprocess.run(["du", "-sb", str(p)], capture_output=True, text=True, timeout=300)
    return int(out.stdout.split()[0]) if out.returncode == 0 and out.stdout else -1


def evict_page_cache(tc) -> dict:
    """Drop the page cache for every .olean/.ilean under the install (POSIX_FADV_DONTNEED
    works unprivileged for clean, unmapped pages) so the next import is a COLD read."""
    n = tot = 0
    roots = [tc.workspace / ".lake" / "packages", tc.root / "elan" / "toolchains"]
    for r in roots:
        for dp, _dn, fns in os.walk(r):
            for fn in fns:
                if fn.endswith((".olean", ".ilean", ".so")):
                    fp = os.path.join(dp, fn)
                    try:
                        fd = os.open(fp, os.O_RDONLY)
                        try:
                            tot += os.fstat(fd).st_size
                            os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
                            n += 1
                        finally:
                            os.close(fd)
                    except OSError:
                        pass
    return {"files": n, "bytes": tot}


def cached_fraction(tc) -> float | None:
    """Fraction of the Mathlib .olean bytes resident in page cache, via `fincore` if present."""
    import shutil

    f = shutil.which("fincore")
    if not f:
        return None
    lib = tc.workspace / ".lake" / "packages" / "mathlib" / ".lake" / "build" / "lib" / "Mathlib"
    files = [str(p) for p in lib.rglob("*.olean")][:4000]
    try:
        out = subprocess.run([f, "--bytes", "--raw", "--noheadings", *files], capture_output=True, text=True, timeout=120)
        res = tot = 0
        for line in out.stdout.splitlines():
            parts = line.split()
            if len(parts) >= 3 and parts[0].isdigit():
                res += int(parts[0])
                tot += int(parts[1])
        return res / tot if tot else None
    except Exception:  # noqa: BLE001
        return None


def measure_install(tc, workers: int, out: dict) -> None:
    from align.verify.lean.repl_pool import ReplLimits, ReplPool, ReplProcess

    m: dict = {}
    m["disk_bytes"] = {
        "total": dir_bytes(tc.root),
        "elan_toolchain": dir_bytes(tc.root / "elan"),
        "workspace": dir_bytes(tc.workspace),
        "mathlib_package": dir_bytes(tc.workspace / ".lake" / "packages" / "mathlib"),
        "repl": dir_bytes(tc.root / "repl"),
    }
    hdr = C.HEADER.rstrip("\n")

    def one_import() -> tuple[float, dict | None]:
        p = ReplProcess(tc.argv(), tc.child_env, cwd=tc.cwd, limits=ReplLimits())
        p.start()
        t = time.monotonic()
        r = p.request({"cmd": hdr}, 900)
        dt = time.monotonic() - t
        assert "env" in r and not r.get("messages"), r
        mem = p.memory()
        p.kill()
        return dt, mem

    ev = evict_page_cache(tc)
    m["evicted"] = ev
    m["cached_fraction_after_evict"] = cached_fraction(tc)
    cold, cold_mem = one_import()
    m["cached_fraction_after_cold_import"] = cached_fraction(tc)
    warm = [one_import()[0] for _ in range(3)]
    m["import_s"] = {"cold": cold, "warm": warm, "warm_median": statistics.median(warm)}
    m["single_repl_after_import_mib"] = cold_mem

    # N workers concurrently warm: per-process PSS (RSS double-counts the shared oleans)
    pool = ReplPool.from_toolchain(tc, n_workers=workers, scratch_dir=str(tc.root / "tmp" / "measure"))
    try:
        t = time.monotonic()
        errs = pool.warm(hdr)
        m["pool_warm_s"] = time.monotonic() - t
        assert not any(errs), errs
        rep = pool.memory_report()
        m["n_workers"] = workers
        m["per_worker_mib"] = rep
        pss = [r["pss"] for r in rep if r]
        m["pss_mib"] = {"min": min(pss), "mean": statistics.mean(pss), "max": max(pss), "sum": sum(pss)}
        m["pss_anon_mib"] = {"mean": statistics.mean(r["pss_anon"] for r in rep if r)}
        m["rss_sum_mib"] = sum(r["rss"] for r in rep if r)
    finally:
        pool.close()
    out["install"] = m


# ---------------------------------------------------------------------------- confusion


def confusion(good_acc: list[bool], bad_acc: list[bool]) -> dict:
    tp = sum(good_acc)
    fn = len(good_acc) - tp
    fp = sum(bad_acc)
    tn = len(bad_acc) - fp
    return {"TP": tp, "FN": fn, "FP": fp, "TN": tn, "n_good": len(good_acc), "n_bad": len(bad_acc)}


def render_confusion(c: dict) -> str:
    return (
        "                      verdict ACCEPT   verdict REJECT   total\n"
        f"truth GOOD proofs     {c['TP']:>14}   {c['FN']:>14}   {c['n_good']:>5}\n"
        f"truth BAD  proofs     {c['FP']:>14}   {c['TN']:>14}   {c['n_bad']:>5}\n"
    )


# ---------------------------------------------------------------------------- main


def sample_memory(pool, stop: threading.Event, rows: list) -> None:
    while not stop.is_set():
        try:
            rep = pool.memory_report()
            with open("/proc/meminfo") as f:
                avail = next(int(l.split()[1]) for l in f if l.startswith("MemAvailable:")) / 1024**2
            rows.append({"t": time.time(), "workers": rep, "mem_available_gb": avail})
        except Exception:  # noqa: BLE001
            pass
        stop.wait(0.5)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--workers", type=int, default=None)
    ap.add_argument("--out", default=None)
    ap.add_argument("--verify-prereg", action="store_true")
    ap.add_argument("--allow-unregistered", action="store_true")
    ap.add_argument("--measure-only", action="store_true")
    ap.add_argument("--skip-measure", action="store_true")
    ap.add_argument("--max-wall-s", type=float, default=3600.0)
    ap.add_argument("--no-semantic-pass", action="store_true")
    ap.add_argument("--no-naive-pass", action="store_true")
    ap.add_argument("--only", choices=("good", "adversarial"), default=None)
    ap.add_argument("--exploratory", action="store_true",
                    help="debug run: <=2 workers, no heavy-job lock (a registered run always takes it)")
    ap.add_argument("--lock-wait-s", type=float, default=7200.0)
    ap.add_argument("--no-kernel-recheck", action="store_true",
                    help="known-good pass without the kernel recheck (default: on, as in the shipped config)")
    ap.add_argument("--kg-limit", type=int, default=0,
                    help="exploratory only: keep the first N rows of each fetched slice (0 = all)")
    args = ap.parse_args()

    if args.verify_prereg:
        return 0 if verify_prereg() else 2
    registered = registered_hash() == C.table_hash()
    if not registered and not args.allow_unregistered and not args.measure_only:
        print("REFUSING: case tables do not match the registered hash (see --verify-prereg); "
              "pass --allow-unregistered for an exploratory run", file=sys.stderr)
        return 2

    os.environ.setdefault("HF_HOME", "/apps/development/rlhf/hf_cache")
    os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
    from align.utils.guard import apply_resource_guards

    apply_resource_guards(cpu_threads=8, heavy=False, label="lean_ceiling_test")
    if not args.exploratory:
        # A Lean pool is a heavy job (research report, section 7A): same preflight + the same
        # exclusive flock as the trainers. If a GPU job holds it, WAIT -- never overlap.
        from align.exit_codes import PreflightRefused
        from align.utils.guard import acquire_heavy_job_lock, preflight_host_check

        t_wait = time.monotonic()
        while True:
            try:
                preflight_host_check(strict=True)
                acquire_heavy_job_lock("lean_ceiling_test")
                break
            except PreflightRefused as e:
                if time.monotonic() - t_wait > args.lock_wait_s:
                    print(f"gave up waiting for the host: {e}", file=sys.stderr)
                    return 83
                print(f"host busy ({str(e)[:120]}...); retrying in 30s", flush=True)
                time.sleep(30)
    elif (args.workers or 2) > 2:
        print("--exploratory caps workers at 2", file=sys.stderr)
        args.workers = 2

    from align.verify.lean import (
        LeanAccuracy, ReplLimits, ReplPool, StrictConfig, StrictVerifier, default_workers, find_toolchain,
    )
    from align.verify.lean.repl_pool import meminfo_available_gb, plan_workers

    tc = find_toolchain()
    limits = ReplLimits()
    n = args.workers or default_workers()
    n = plan_workers(n, limits.anon_cap_gb, 4.0, meminfo_available_gb())
    stamp = time.strftime("%Y%m%d_%H%M%S")
    out_dir = Path(args.out or REPO / ".tmp" / "lean_ceiling" / stamp)
    out_dir.mkdir(parents=True, exist_ok=True)
    pools: list = []

    def watchdog() -> None:
        time.sleep(args.max_wall_s)
        print(f"WATCHDOG: {args.max_wall_s}s exceeded; killing pools", file=sys.stderr, flush=True)
        for p in pools:
            try:
                p.close()
            except Exception:  # noqa: BLE001
                pass
        os._exit(124)

    threading.Thread(target=watchdog, daemon=True).start()

    def on_signal(signum, _frame):
        print(f"signal {signum}: closing pools", file=sys.stderr, flush=True)
        for p in pools:
            try:
                p.close()
            except Exception:  # noqa: BLE001
                pass
        os._exit(128 + signum)

    import signal

    signal.signal(signal.SIGTERM, on_signal)
    signal.signal(signal.SIGINT, on_signal)

    result: dict = {
        "started": stamp, "toolchain": tc.describe(), "pins": tc.pins, "workers": n,
        "limits": limits.__dict__, "prereg_registered": registered, "table_sha256": C.table_hash(),
        "cores": os.cpu_count(),
    }

    def new_pool(workers: int, lim: ReplLimits = limits, label: str = "pool") -> ReplPool:
        p = ReplPool.from_toolchain(tc, n_workers=workers, limits=lim, scratch_dir=str(tc.root / "tmp" / label))
        pools.append(p)
        return p

    # -------------------------------------------------- install measurements
    if not args.skip_measure and args.only is None:
        print("== measuring install ==", flush=True)
        measure_install(tc, n, result)
        print(json.dumps(result["install"], indent=1, default=str)[:3000], flush=True)
        if args.measure_only:
            (out_dir / "result.json").write_text(json.dumps(result, indent=1, default=str))
            return 0

    # -------------------------------------------------- known-good
    good_items: list[dict] = []
    if args.only in (None, "good"):
        print("== known-good: fetching ==", flush=True)
        rows, meta = fetch_known_good(out_dir)
        if args.kg_limit:
            if not args.exploratory:
                print("--kg-limit is for exploratory runs only", file=sys.stderr)
                return 2
            rows = [r for i, r in enumerate(rows) if i % SLICE_LEN < args.kg_limit]
            meta["kg_limit"] = args.kg_limit
        result["known_good_source"] = meta
        cands = []
        for r in rows:
            code = r["full_proof"]
            ds = derive_statement(code)
            if ds is None:
                cands.append({"id": r["problem_id"], "source": "lean-workbook-proofs", "code": code, "statement": None})
            else:
                cands.append({"id": r["problem_id"], "source": "lean-workbook-proofs", "code": code, "statement": ds[0]})
        for rs in C.RESTATED:
            code = C.HEADER + rs["code"]
            ds = derive_statement(code)
            cands.append({"id": rs["id"], "source": "restated", "code": code, "statement": ds[0] if ds else None})
        result["known_good_candidates"] = len(cands)

        print("== known-good: independent compile-only filter ==", flush=True)
        pool = new_pool(n, label="kg_compile")
        t0 = time.monotonic()
        refs = pool.map(
            lambda w, it: compile_only(w, it["statement"], it["code"]) if it["statement"] else {"ok": False, "clean": False, "why": "no_statement"},
            cands,
        )
        pool.close()
        excluded = []
        for it, r in zip(cands, refs):
            if isinstance(r, Exception):
                r = {"ok": False, "clean": False, "why": repr(r)}
            it["ref"] = r
            if r.get("clean"):
                good_items.append(it)
            else:
                excluded.append({"id": it["id"], "source": it["source"], "why": r.get("why")})
        result["known_good_excluded"] = excluded
        result["known_good_filter_s"] = time.monotonic() - t0
        print(f"compile-clean: {len(good_items)}/{len(cands)}  excluded: {len(excluded)}", flush=True)

    # -------------------------------------------------- known-good: strict
    good_results: list[dict] = []
    if good_items:
        print("== known-good: strict verification ==", flush=True)
        pool = new_pool(n, label="kg_strict")
        acc = LeanAccuracy(StrictVerifier(pool, StrictConfig(kernel_recheck=not args.no_kernel_recheck)))
        from align.verify.lean.strict import AUDIT_PRELUDE, static_check

        plan, _, _ = static_check(good_items[0]["statement"], good_items[0]["code"])
        t = time.monotonic()
        pool.warm(plan.header if plan else C.HEADER, AUDIT_PRELUDE)
        result["strict_warm_s"] = time.monotonic() - t
        result["warm_memory_mib"] = pool.memory_report()
        stop, mem_rows = threading.Event(), []
        th = threading.Thread(target=sample_memory, args=(pool, stop, mem_rows), daemon=True)
        th.start()
        t = time.monotonic()
        def prog(ix, r):
            print(f"  [{ix:3d}] {good_items[ix]['id']:34s} {'ACCEPT' if r.accepted else 'REJECT'} {r.elapsed_s:7.2f}s {list(r.reasons)}", flush=True)

        accs, res = acc.accuracy_batch([g["statement"] for g in good_items], [as_completion(g["code"]) for g in good_items], progress=prog)
        wall = time.monotonic() - t
        stop.set()
        th.join(5)
        for g, a, r in zip(good_items, accs, res):
            good_results.append({"id": g["id"], "source": g["source"], "accepted": bool(a), "reasons": list(r.reasons),
                                 "diag": r.diagnostics, "elapsed_s": r.elapsed_s})
        lat = [r.elapsed_s for r in res]
        result["throughput"] = {
            "n_proofs": len(res), "wall_s": wall, "proofs_per_s": len(res) / wall, "workers": n,
            "latency_s": {"p50": pctl(lat, 0.5), "p90": pctl(lat, 0.9), "p99": pctl(lat, 0.99), "max": max(lat), "mean": statistics.mean(lat)},
            "note": "unique statement per proof (no reference-env reuse across proofs); warm workers",
        }
        # per-worker PSS during the run
        per_w: dict[int, list] = {}
        for row in mem_rows:
            for i, m in enumerate(row["workers"]):
                if m:
                    per_w.setdefault(i, []).append(m)
        result["pss_during_run_mib"] = {
            str(i): {"pss_mean": statistics.mean(m["pss"] for m in ms), "pss_max": max(m["pss"] for m in ms),
                     "pss_anon_max": max(m["pss_anon"] for m in ms), "rss_max": max(m["rss"] for m in ms),
                     "samples": len(ms)} for i, ms in sorted(per_w.items())
        }
        result["host_during_run"] = {"min_mem_available_gb": min((r["mem_available_gb"] for r in mem_rows), default=None)}
        result["pool_stats"] = pool.stats.__dict__
        # repeat-G run: GRPO shape (same statement x8, reference env reused) on the 12
        # MEDIAN-latency statements of the pass above (the heavy tail would swamp a 96-proof run)
        order = sorted(range(len(good_items)), key=lambda i: res[i].elapsed_s)
        mid = max(0, len(order) // 2 - 6)
        rep = [good_items[i] for i in order[mid : mid + 12]] * 8
        t = time.monotonic()
        accs2, _ = acc.accuracy_batch([g["statement"] for g in rep], [as_completion(g["code"]) for g in rep])
        w2 = time.monotonic() - t
        result["throughput_grpo_shape"] = {"n_proofs": len(rep), "wall_s": w2, "proofs_per_s": len(rep) / w2, "accepted": int(sum(accs2)),
                                           "shape": "12 median-latency statements x 8 identical samples"}
        pool.close()

    # -------------------------------------------------- adversarial
    adv_results: list[dict] = []
    if args.only in (None, "adversarial"):
        print("== adversarial: strict ==", flush=True)

        def completion_for(case) -> tuple[str, str | None]:
            stmt = C.HEADER + C.STATEMENTS[case["statement"]]
            if case["completion"] is not None:
                return stmt, case["completion"]
            hdr = C.HEADER if case["raw_header"] is None else case["raw_header"]
            return stmt, as_completion(hdr + case["code"])

        main_cases = [c for c in C.ADVERSARIAL if c["pool"] == "main"]
        bomb_cases = [c for c in C.ADVERSARIAL if c["pool"] == "bomb"]
        pool = new_pool(n, label="adv_strict")
        acc = LeanAccuracy(StrictVerifier(pool, StrictConfig(proof_timeout_s=120)), min_batch=10**9)
        pairs = [completion_for(c) for c in main_cases]
        _, rs = acc.accuracy_batch([p[0] for p in pairs], [p[1] for p in pairs])
        pool.close()
        by_id = {}
        for c, r in zip(main_cases, rs):
            by_id[c["id"]] = {"id": c["id"], "category": c["category"], "accepted": r.accepted, "reasons": list(r.reasons),
                              "stage": r.stage, "diag": r.diagnostics, "elapsed_s": r.elapsed_s, "pool": "main"}
        # bombs: 1 worker, 1 GiB anon cap, 30 s budget so they terminate promptly (same code paths)
        blim = ReplLimits(anon_cap_gb=1.0)
        bpool = new_pool(1, blim, label="adv_bomb")
        bacc = LeanAccuracy(StrictVerifier(bpool, StrictConfig(proof_timeout_s=30)), min_batch=10**9)
        for c in bomb_cases:
            stmt, comp = completion_for(c)
            sc = bacc.score(stmt, comp)
            by_id[c["id"]] = {"id": c["id"], "category": c["category"], "accepted": sc.result.accepted,
                              "reasons": list(sc.result.reasons), "stage": sc.result.stage, "diag": sc.result.diagnostics,
                              "elapsed_s": sc.result.elapsed_s, "pool": "bomb(1w,1GiB,30s)"}
        bpool.close()
        adv_results = [by_id[c["id"]] for c in C.ADVERSARIAL]

        # naive compile-only on each case's own text (what a compile-only reward would pay)
        if not args.no_naive_pass:
            print("== adversarial: naive compile-only (liveness) ==", flush=True)
            npool = new_pool(n, label="adv_naive")
            live_cases = [c for c in main_cases if c["completion"] is None and c["code"]]
            texts = [(c["raw_header"] if c["raw_header"] is not None else C.HEADER) + c["code"] for c in live_cases]
            nv = npool.map(lambda w, t: naive_fresh(w, t, 120.0), texts)
            npool.close()
            naive = {c["id"]: (v if not isinstance(v, Exception) else {"ok": False, "why": repr(v)}) for c, v in zip(live_cases, nv)}
            for r in adv_results:
                r["naive_compile_ok"] = naive.get(r["id"], {}).get("ok")

        # semantic layer alone (lexical gate OFF): how much does the Lean-side audit catch by itself?
        if not args.no_semantic_pass:
            print("== adversarial: semantic layer only (lexical gate off) ==", flush=True)
            spool = new_pool(n, label="adv_sem")
            sv = StrictVerifier(spool, StrictConfig(lexical_gate=False, kernel_recheck=True, proof_timeout_s=120))
            sem_cases = [c for c in main_cases if c["completion"] is None and c["code"]]
            from align.verify.lean.reward import extract_lean_code

            sem_items = []
            for c in sem_cases:
                stmt, comp = completion_for(c)
                sem_items.append((stmt, extract_lean_code(comp)))
            sres = sv.verify_batch(sem_items)
            spool.close()
            sem = {c["id"]: (r.accepted, list(r.reasons)) for c, r in zip(sem_cases, sres)}
            for r in adv_results:
                if r["id"] in sem:
                    r["semantic_only_accepted"], r["semantic_only_reasons"] = sem[r["id"]]

    # -------------------------------------------------- report
    report = {}
    if not good_results and args.only in (None, "good"):
        report["known_good"] = empty_known_good_entry(
            result.get("known_good_candidates", 0), len(result.get("known_good_excluded", []))
        )
    if good_results:
        acc_good = [g["accepted"] for g in good_results]
        lo, hi = wilson(sum(acc_good), len(acc_good))
        report["known_good"] = {"n": len(acc_good), "accepted": sum(acc_good), "rate": sum(acc_good) / len(acc_good),
                                "wilson95": [lo, hi], "floor": KNOWN_GOOD_FLOOR,
                                "pass": (sum(acc_good) / len(acc_good) >= KNOWN_GOOD_FLOOR and len(acc_good) >= MIN_KNOWN_GOOD),
                                "rejected": [g for g in good_results if not g["accepted"]]}
    if adv_results:
        acc_bad = [a["accepted"] for a in adv_results]
        by_cat: dict[str, list] = {}
        for a in adv_results:
            by_cat.setdefault(a["category"], []).append(a["accepted"])
        report["adversarial"] = {"n": len(acc_bad), "accepted": sum(acc_bad), "pass": (sum(acc_bad) == 0 and len(acc_bad) >= MIN_ADVERSARIAL),
                                 "by_category": {k: {"n": len(v), "accepted": sum(v)} for k, v in sorted(by_cat.items())},
                                 "accepted_cases": [a for a in adv_results if a["accepted"]]}
        live = [a for a in adv_results if a.get("naive_compile_ok")]
        report["adversarial"]["naive_compile_only_would_accept"] = len(live)
        if any("semantic_only_accepted" in a for a in adv_results):
            sem_acc = [a for a in adv_results if a.get("semantic_only_accepted")]
            report["adversarial"]["semantic_only_would_accept"] = [a["id"] for a in sem_acc]
    if good_results and adv_results:
        report["confusion"] = confusion([g["accepted"] for g in good_results], [a["accepted"] for a in adv_results])
    result["report"] = report
    result["good_results"] = good_results
    result["adversarial_results"] = adv_results
    (out_dir / "result.json").write_text(json.dumps(result, indent=1, default=str, ensure_ascii=False))
    print("\n================ RESULT ================")
    if "confusion" in report:
        print(render_confusion(report["confusion"]))
    if report.get("known_good", {}).get("n"):
        k = report["known_good"]
        print(f"known-good accepted {k['accepted']}/{k['n']} = {k['rate']:.4f}  Wilson95 [{k['wilson95'][0]:.4f}, {k['wilson95'][1]:.4f}]  floor {k['floor']}  PASS={k['pass']}")
        for g in k["rejected"]:
            print("   REJECTED known-good:", g["id"], g["reasons"], g["diag"])
    if "adversarial" in report:
        a = report["adversarial"]
        print(f"adversarial accepted {a['accepted']}/{a['n']}  PASS={a['pass']}  (naive compile-only would accept {a['naive_compile_only_would_accept']})")
        for c in a["accepted_cases"]:
            print("   ACCEPTED adversarial:", c["id"], c["reasons"])
    if "throughput" in result:
        print("throughput:", json.dumps(result["throughput"]))
    print(f"results: {out_dir/'result.json'}")
    if "known_good" in report and not report["known_good"]["n"]:
        print(f"known-good: NO EVIDENCE -- {report['known_good']['why']}  PASS=False")
    return 0 if ceiling_ok(report, args.only) else 1


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