align / code /scripts /p2 /backend_agreement.py
viai957's picture
Add files using upload-large-folder tool
63f41a6 verified
Raw History Blame Contribute Delete
19.1 kB
#!/usr/bin/env python
"""Agreement + throughput: our strict verifier vs the Kimina Lean server pre-filter.
hostguard.py run --cores 4 --mem-gb 24 --label agree --wait -- \\
.venv/bin/python scripts/p2/backend_agreement.py --out .tmp/kimina_agreement/run1
Item sets (all verified by BOTH backends on the same (formal statement, extracted Lean code)):
kg known-good proofs of the ceiling test (150 Lean-Workbook-proofs rows saved by the registered
ceiling run + 30 restated), statements derived exactly as `lean_ceiling_test.py` does
adv the 61 frozen adversarial cases of `lean_ceiling_cases.py` (3 bombs run on a 1 GiB / 30 s
pool for ours and on a second, small Kimina server for the candidate)
sample N random code-bearing records (finish_reason != length) of a P2 generation file
Phases (resumable, JSONL per phase): `ours` (StrictVerifier, the authority), `control` (OUR REPL pool,
compile-only: separates "Kimina server overhead" from "strict-tier overhead"), `kimina` (compile-only
through the server). They run strictly one after another: never two heavy jobs at once. The server(s)
must already be running (toolchains/work/kimina-run/start.sh); this script never starts or stops them.
Semantics of the comparison. Kimina is a PRE-FILTER: its verdict is `reject` / `maybe` / `error`, never
an accept. The "gated" outcome (accept iff strict accepts AND the pre-filter did not reject) is computed
from the two per-item records; the only accept that can exist is the strict verifier's.
"""
from __future__ import annotations
import argparse
import importlib.util
import json
import os
import random
import statistics
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO))
MAIN_REPO = Path("/apps/development/rlhf/align") # data lives in the main checkout's .tmp, not in worktrees
DEFAULT_GENS = MAIN_REPO / ".tmp/p2/r1/baseline/gen_minif2f_c40960.jsonl"
DEFAULT_PROBLEMS = MAIN_REPO / ".tmp/p2/problems/minif2f.jsonl"
DEFAULT_KG_RAW = MAIN_REPO / ".tmp/lean_ceiling/registered/known_good_raw.jsonl"
CORES = 4
def _load(name: str, path: Path):
spec = importlib.util.spec_from_file_location(name, path)
mod = importlib.util.module_from_spec(spec)
sys.modules[name] = mod
spec.loader.exec_module(mod)
return mod
LC = _load("lean_ceiling_test", REPO / "scripts" / "lean_ceiling_test.py")
CASES = LC.C
def pctl(xs, q):
s = sorted(xs)
return s[min(len(s) - 1, int(q * len(s)))] if s else float("nan")
# --------------------------------------------------------------------------- items
def build_items(args) -> list[dict]:
from align.verify.lean.reward import extract_lean_code
items: list[dict] = []
# known-good
cands = []
for line in open(args.kg_raw):
r = json.loads(line)
cands.append((r["problem_id"], "lean-workbook-proofs", r["full_proof"]))
for rs in CASES.RESTATED:
cands.append((rs["id"], "restated", CASES.HEADER + rs["code"]))
for pid, src, code in cands:
ds = LC.derive_statement(code)
if ds is None:
continue
items.append(dict(id=f"kg:{pid}", set="kg", sub=src, statement=ds[0], code=code, pool="main"))
# adversarial
for c in CASES.ADVERSARIAL:
stmt = CASES.HEADER + CASES.STATEMENTS[c["statement"]]
comp = c["completion"] if c["completion"] is not None else LC.as_completion(
(CASES.HEADER if c["raw_header"] is None else c["raw_header"]) + c["code"])
items.append(dict(id=f"adv:{c['id']}", set="adv", sub=c["category"], statement=stmt,
code=extract_lean_code(comp), pool=c["pool"]))
# sampled generations
probs = {}
for line in open(args.problems):
r = json.loads(line)
probs[r["problem_id"]] = r["formal_statement"]
pop = []
for line in open(args.generations):
r = json.loads(line)
if r["finish_reason"] == "length" or not (r["lean_code"] or "").strip() or r["problem_id"] not in probs:
continue
pop.append(r)
rng = random.Random(args.seed)
for r in rng.sample(pop, min(args.sample, len(pop))):
items.append(dict(id=f"gen:{r['problem_id']}#{r['sample_idx']}", set="sample", sub="minif2f",
statement=probs[r["problem_id"]], code=r["lean_code"], pool="main"))
items.append(dict(id="__meta__", set="meta", population=len(pop)))
return items
def done_ids(path: Path) -> dict:
out = {}
if path.exists():
for line in open(path):
try:
r = json.loads(line)
out[r["id"]] = r
except json.JSONDecodeError:
pass
return out
class Sink:
def __init__(self, path: Path):
self.f = open(path, "a")
self.lock = threading.Lock()
def write(self, rec: dict):
with self.lock:
self.f.write(json.dumps(rec, ensure_ascii=False, default=str) + "\n")
self.f.flush()
# --------------------------------------------------------------------------- phases
def phase_ours(items, out: Path, args) -> dict:
from align.verify.lean import ReplLimits, ReplPool, StrictConfig, StrictVerifier, find_toolchain
tc = find_toolchain()
sink = Sink(out / "ours.jsonl")
have = done_ids(out / "ours.jsonl")
walls = {}
for pool_name, lim, cfg in (("main", ReplLimits(), StrictConfig(proof_timeout_s=args.timeout)),
("bomb", ReplLimits(anon_cap_gb=1.0), StrictConfig(proof_timeout_s=30))):
todo = [it for it in items if it["set"] != "meta" and it["pool"] == pool_name and it["id"] not in have]
if not todo:
continue
n = args.workers if pool_name == "main" else 1
pool = ReplPool.from_toolchain(tc, n_workers=n, limits=lim, preflight=True,
scratch_dir=str(tc.root / "tmp" / f"agree_{pool_name}"))
try:
ver = StrictVerifier(pool, cfg)
t0 = time.monotonic()
def one(w, it):
r = ver._verify_on(w, (it["statement"], it["code"]))
sink.write(dict(id=it["id"], set=it["set"], accepted=bool(r.accepted), reasons=list(r.reasons),
stage=r.stage, elapsed_s=r.elapsed_s, undecided=bool(r.undecided)))
pool.map(one, todo)
walls[pool_name] = time.monotonic() - t0
finally:
pool.close()
return {"walls": walls, "workers": args.workers}
def phase_control(items, out: Path, args) -> dict:
"""OUR REPL pool, compile-only (no strictness): the engine-level control for the Kimina server."""
from align.verify.lean import ReplLimits, ReplPool, find_toolchain
tc = find_toolchain()
sink = Sink(out / "control.jsonl")
have = done_ids(out / "control.jsonl")
todo = [it for it in items if it["set"] in ("kg", "sample") and it["id"] not in have]
if not todo:
return {}
pool = ReplPool.from_toolchain(tc, n_workers=args.workers, limits=ReplLimits(), preflight=True,
scratch_dir=str(tc.root / "tmp" / "agree_control"))
try:
t0 = time.monotonic()
def one(w, it):
t = time.monotonic()
r = LC.compile_only(w, it["statement"], it["code"], args.timeout) if it["code"] else {"ok": False, "why": "no_code"}
sink.write(dict(id=it["id"], set=it["set"], ok=bool(r.get("ok")), clean=bool(r.get("clean")), why=r.get("why", ""),
elapsed_s=time.monotonic() - t))
pool.map(one, todo)
return {"wall": time.monotonic() - t0, "workers": args.workers}
finally:
pool.close()
def phase_kimina(items, out: Path, args) -> dict:
from align.verify_backends.kimina_server import KiminaClient, KiminaPrefilter
sink = Sink(out / "kimina.jsonl")
have = done_ids(out / "kimina.jsonl")
walls = {}
for pool_name, url, tmo, workers in (("main", args.kimina_url, args.timeout, args.workers),
("bomb", args.kimina_bomb_url, 30, 1)):
todo = [it for it in items if it["set"] != "meta" and it["pool"] == pool_name and it["id"] not in have]
if not todo:
continue
cl = KiminaClient(url)
if not cl.healthy():
raise SystemExit(f"Kimina server at {url} is not healthy; start it first")
pf = KiminaPrefilter(cl, tmo)
def one(it):
t = time.monotonic()
v = pf.check_code(it["code"])
sink.write(dict(id=it["id"], set=it["set"], verdict=v.verdict, reasons=list(v.reasons), detail=v.detail,
elapsed_s=time.monotonic() - t, server_s=v.elapsed_s))
t0 = time.monotonic()
with ThreadPoolExecutor(workers) as ex:
list(ex.map(one, todo))
walls[pool_name] = time.monotonic() - t0
return {"walls": walls, "workers": args.workers}
# --------------------------------------------------------------------------- report
def stats(xs):
return {"n": len(xs), "sum_s": round(sum(xs), 1), "mean": round(statistics.mean(xs), 3) if xs else None,
"p50": round(pctl(xs, .5), 3), "p90": round(pctl(xs, .9), 3), "p99": round(pctl(xs, .99), 3),
"max": round(max(xs), 3) if xs else None}
def report(items, out: Path, phase_meta: dict) -> dict:
ours, ctl, kim = done_ids(out / "ours.jsonl"), done_ids(out / "control.jsonl"), done_ids(out / "kimina.jsonl")
by_id = {it["id"]: it for it in items if it["set"] != "meta"}
rep: dict = {"phase_meta": phase_meta, "matrix": {}, "disagreements": [], "n_items": {}}
def okind(r): # ours: accept / reject / undecided
return "accept" if r["accepted"] else ("undecided" if r["undecided"] else "reject")
for s in ("kg", "adv", "sample", "ALL"):
ids = [i for i in by_id if (s == "ALL" or by_id[i]["set"] == s) and i in ours and i in kim]
m: dict = {}
for i in ids:
m.setdefault(okind(ours[i]), {}).setdefault(kim[i]["verdict"], 0)
m[okind(ours[i])][kim[i]["verdict"]] += 1
rep["matrix"][s] = {"n": len(ids), "rows(ours)xcols(kimina)": m}
# ceiling truth
kg_ids = [i for i in by_id if by_id[i]["set"] == "kg" and i in ours and i in kim]
adv_ids = [i for i in by_id if by_id[i]["set"] == "adv" and i in ours and i in kim]
rep["ceiling"] = {
"kg_ours_accept": sum(ours[i]["accepted"] for i in kg_ids), "kg_n": len(kg_ids),
"kg_kimina_maybe": sum(kim[i]["verdict"] == "maybe" for i in kg_ids),
"kg_kimina_reject": sum(kim[i]["verdict"] == "reject" for i in kg_ids),
"adv_ours_accept": sum(ours[i]["accepted"] for i in adv_ids), "adv_n": len(adv_ids),
"adv_kimina_maybe(compile-only would pass these on)": sum(kim[i]["verdict"] == "maybe" for i in adv_ids),
"adv_kimina_reject": sum(kim[i]["verdict"] == "reject" for i in adv_ids),
"gated_adv_accepts": sum(ours[i]["accepted"] and kim[i]["verdict"] != "reject" for i in adv_ids),
}
# disagreements: everything that is not (ours accept & kimina maybe) or (ours reject & kimina reject)
for i in by_id:
if i not in ours or i not in kim:
continue
o, k = ours[i], kim[i]
ok = okind(o)
kv = k["verdict"]
if (ok == "accept" and kv == "maybe") or (ok == "reject" and kv == "reject"):
continue
if ok == "accept" and kv == "reject":
klass = "PREFILTER_FALSE_REJECT (recall loss)"
elif ok == "reject" and kv == "maybe":
klass = "STRICT_ONLY_REJECT (expected: policy/audit layer)"
elif ok == "accept" and kv == "error":
klass = "PREFILTER_UNDECIDED_STRICT_ACCEPT"
elif ok == "undecided" or kv == "error":
klass = "UNDECIDED_ON_ONE_SIDE"
else:
klass = "OTHER"
rep["disagreements"].append(dict(id=i, set=by_id[i]["set"], klass=klass, ours=ok, ours_reasons=o["reasons"],
ours_s=round(o["elapsed_s"], 2), kimina=kv, kimina_reasons=k["reasons"],
kimina_detail=k.get("detail"), kimina_s=round(k["elapsed_s"], 2),
control=({"ok": ctl[i]["ok"], "clean": ctl[i]["clean"], "why": ctl[i]["why"]} if i in ctl else None)))
from collections import Counter
rep["disagreement_classes"] = dict(Counter(d["klass"] for d in rep["disagreements"]))
# throughput over the sample set (+ kg) : same items, three engines
for name, rows, key in (("ours_strict", ours, "elapsed_s"), ("ours_compile_only", ctl, "elapsed_s"), ("kimina_compile_only", kim, "elapsed_s")):
for s in ("sample", "kg"):
xs = [rows[i][key] for i in by_id if by_id[i]["set"] == s and i in rows]
rep.setdefault("latency_s", {}).setdefault(name, {})[s] = stats(xs)
# busy-time throughput on the SAME items (kg + sample): n / sum(per-item latency) = proofs per second per
# worker-core, independent of the heavy tail's wall-clock packing and of which extra items a phase ran
same = [i for i in by_id if by_id[i]["set"] in ("sample", "kg") and i in ours and i in kim and i in ctl]
rep["busy_throughput_proofs_per_s_per_core"] = {
"n": len(same),
**{nm: round(len(same) / sum(rows[i]["elapsed_s"] for i in same), 4)
for nm, rows in (("ours_strict", ours), ("ours_compile_only", ctl), ("kimina_compile_only", kim))}}
# gated pipeline, from the per-item records (never re-run)
samp = [i for i in by_id if by_id[i]["set"] in ("sample", "kg") and i in ours and i in kim]
if samp:
rep["gated"] = {
"n": len(samp),
"strict_accepts": sum(ours[i]["accepted"] for i in samp),
"gated_accepts": sum(ours[i]["accepted"] and kim[i]["verdict"] != "reject" for i in samp),
"prefilter_rejects": sum(kim[i]["verdict"] == "reject" for i in samp),
"strict_calls_saved": sum(kim[i]["verdict"] == "reject" for i in samp),
"sum_strict_s_saved": round(sum(ours[i]["elapsed_s"] for i in samp if kim[i]["verdict"] == "reject"), 1),
"sum_kimina_s_spent": round(sum(kim[i]["elapsed_s"] for i in samp), 1),
"sum_strict_s_total": round(sum(ours[i]["elapsed_s"] for i in samp), 1),
"sum_strict_s_on_survivors": round(sum(ours[i]["elapsed_s"] for i in samp if kim[i]["verdict"] != "reject"), 1),
}
# agreement on the sample set between the two COMPILE-ONLY engines (our REPL vs Kimina)
both = [i for i in by_id if i in ctl and i in kim and by_id[i]["set"] in ("sample", "kg")]
if both:
agree = sum((ctl[i]["clean"]) == (kim[i]["verdict"] == "maybe") for i in both if kim[i]["verdict"] != "error" and ctl[i]["why"] != "timeout")
n2 = sum(1 for i in both if kim[i]["verdict"] != "error" and ctl[i]["why"] != "timeout")
rep["engine_control"] = {"n_decided_by_both": n2, "clean_agree": agree,
"diffs": [dict(id=i, ours_compile_clean=ctl[i]["clean"], ours_why=ctl[i]["why"], kimina=kim[i]["verdict"], kimina_reasons=kim[i]["reasons"])
for i in both if kim[i]["verdict"] != "error" and ctl[i]["why"] != "timeout" and ctl[i]["clean"] != (kim[i]["verdict"] == "maybe")]}
return rep
def throughput(rep, items, out: Path, phase_meta: dict) -> None:
by_set = {}
for it in items:
if it["set"] != "meta":
by_set.setdefault(it["pool"], []).append(it["id"])
tp = {}
for name, f, wkey in (("ours_strict", "ours.jsonl", "ours"), ("ours_compile_only", "control.jsonl", "control"), ("kimina_compile_only", "kimina.jsonl", "kimina")):
m = phase_meta.get(wkey) or {}
wall = m.get("wall") if "wall" in m else (m.get("walls") or {}).get("main")
n_main = sum(1 for r in done_ids(out / f).values() if by_set.get("main") and r["id"] in set(by_set["main"]))
if wall:
tp[name] = {"n": n_main, "wall_s": round(wall, 1), "proofs_per_s": round(n_main / wall, 4),
"proofs_per_s_per_core": round(n_main / wall / CORES, 4), "cores": CORES}
rep["throughput"] = tp
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--out", required=True)
ap.add_argument("--phases", default="ours,control,kimina")
ap.add_argument("--workers", type=int, default=CORES)
ap.add_argument("--timeout", type=float, default=180.0)
ap.add_argument("--sample", type=int, default=500)
ap.add_argument("--seed", type=int, default=20261019)
ap.add_argument("--generations", default=str(DEFAULT_GENS))
ap.add_argument("--problems", default=str(DEFAULT_PROBLEMS))
ap.add_argument("--kg-raw", default=str(DEFAULT_KG_RAW))
ap.add_argument("--kimina-url", default="http://127.0.0.1:18000")
ap.add_argument("--kimina-bomb-url", default="http://127.0.0.1:18001")
ap.add_argument("--sets", default="kg,adv,sample", help="comma list; smoke runs use e.g. --sets sample --sample 6")
ap.add_argument("--report-only", action="store_true")
args = ap.parse_args()
if args.workers > CORES:
print(f"host rule: at most {CORES} cores/workers", file=sys.stderr)
return 2
out = Path(args.out)
out.mkdir(parents=True, exist_ok=True)
items = [i for i in build_items(args) if i["set"] == "meta" or i["set"] in args.sets.split(",")]
meta_it = [i for i in items if i["set"] == "meta"][0]
(out / "items_manifest.json").write_text(json.dumps(
{"n": len(items) - 1, "population_code_bearing": meta_it["population"], "seed": args.seed,
"sets": {s: sum(1 for i in items if i["set"] == s) for s in ("kg", "adv", "sample")},
"ids": [i["id"] for i in items if i["set"] != "meta"]}, indent=1))
meta_path = out / "phase_meta.json"
phase_meta = json.loads(meta_path.read_text()) if meta_path.exists() else {}
if not args.report_only:
for ph, fn in (("ours", phase_ours), ("control", phase_control), ("kimina", phase_kimina)):
if ph in args.phases.split(","):
print(f"== phase {ph} ==", flush=True)
t = time.monotonic()
m = fn(items, out, args)
if m:
phase_meta[ph] = m
meta_path.write_text(json.dumps(phase_meta, indent=1))
print(f" done in {time.monotonic() - t:.0f}s", flush=True)
rep = report(items, out, phase_meta)
throughput(rep, items, out, phase_meta)
(out / "report.json").write_text(json.dumps(rep, indent=1, default=str))
print(json.dumps({k: rep[k] for k in ("matrix", "ceiling", "disagreement_classes", "throughput", "gated") if k in rep}, indent=1))
return 0
if __name__ == "__main__":
sys.exit(main())