| |
| """Verify P2 Generation JSONL with the strict Lean verifier -> resumable Verification JSONL. |
| |
| .venv/bin/python scripts/p2/verify_generations.py \\ |
| --generations gens_r0.jsonl --problems minif2f.jsonl --out ver_r0.jsonl --workers 8 |
| |
| Output: one line per generation, keyed by (problem_id, sample_idx, round, cap_tokens): |
| |
| {problem_id, sample_idx, round, cap_tokens, accepted, reasons, verify_s, budget_s, |
| verifier: {pins, config, workers}, code_sha256, stage, diagnostics[, repl_errors, repl_code]} |
| |
| Properties (each one is a defect the first version of such a script has): |
| |
| * RESUMABLE. The file is append-only and fsynced per record; a rerun skips every key that is |
| already there. A line cut in half by SIGKILL is dropped on restart (`repair_torn_tail`); a |
| malformed line elsewhere raises (appends only ever damage the tail). |
| * NO STALE VERDICTS. Every record carries `code_sha256` of the code that was verified. A |
| resumed run whose generation file was regenerated (same keys, different samples) finds a |
| different hash and RAISES instead of attaching the old verdict to new text. |
| * NO LEAN FOR NOTHING. `finish_reason == "length"` (a truncated proof is not a proof) and |
| code-less samples are recorded as rejected with reasons ["no_code"] without calling Lean. |
| * A DEAD VERIFIER IS NOT A ZERO. Results whose reason is an infrastructure death (crash, |
| protocol, memory, output_too_large) are held back until their chunk is complete; if more |
| than half of a chunk of >= 8 died that way the run aborts WITHOUT writing them, so a |
| resume retries them. Otherwise they are written (a proof that kills its REPL scores 0). |
| * HOST RULES. >2 workers => exclusive heavy-job flock (waited for, not skipped) and the |
| host preflight; workers <= 8; scratch under /apps. |
| * `--feedback` additionally stores, for samples rejected with `compile_error`, `repl_errors` |
| (every Lean error: pos / endPos / data) and `repl_code` (the exact text the positions index: |
| Goedel's `compile_view(lean_code)`), which `extras/vllm_rollout/lean_gen.py` needs to build |
| correction prompts and which `StrictResult.diagnostics` (3 errors, 240 chars, no positions) |
| cannot supply. Rejections that are not compile errors (banned construct, statement |
| mismatch, timeout...) carry no Lean error list; the generator skips and counts them. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib.util |
| import json |
| import logging |
| import os |
| import sys |
| import threading |
| import time |
| from dataclasses import asdict |
| 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() |
| log = logging.getLogger("p2.verify") |
|
|
| INFRA_REASONS = ("crash", "protocol", "memory", "output_too_large") |
| MAX_WORKERS = 8 |
|
|
|
|
| class StaleVerification(C.ContractError): |
| """An existing verification record belongs to different code than the generation now holds.""" |
|
|
|
|
| class VerifierDied(RuntimeError): |
| """More than half of a chunk died for infrastructure reasons; nothing from it was written.""" |
|
|
|
|
| class _NullWriter: |
| def append(self, rec) -> None: |
| pass |
|
|
| def close(self) -> None: |
| pass |
|
|
|
|
| class Result: |
| """What the verify callback returns: the minimum the record needs.""" |
|
|
| __slots__ = ("accepted", "reasons", "elapsed_s", "stage", "diagnostics", "feedback") |
|
|
| def __init__(self, accepted, reasons, elapsed_s, stage="", diagnostics=None, feedback=None): |
| self.accepted = bool(accepted) |
| self.reasons = list(reasons) |
| self.elapsed_s = float(elapsed_s) |
| self.stage = stage |
| self.diagnostics = diagnostics or {} |
| self.feedback = feedback |
|
|
| @property |
| def infra(self) -> bool: |
| return any(r in INFRA_REASONS for r in self.reasons) |
|
|
|
|
| |
|
|
|
|
| def load_existing(path: Path) -> dict[tuple, dict]: |
| out: dict[tuple, dict] = {} |
| if not path.exists(): |
| return out |
| for rec in C.iter_jsonl(path): |
| for k in (*C.KEY_FIELDS, "accepted", "reasons", "verify_s", "budget_s", "verifier", "code_sha256"): |
| if k not in rec: |
| raise C.ContractError(f"{path}: verification record lacks {k!r}: {str(rec)[:200]}") |
| key = C.gen_key(rec) |
| if key in out: |
| raise C.ContractError(f"{path}: duplicate verification key {key} (the file is append-only)") |
| out[key] = rec |
| return out |
|
|
|
|
| def iter_generations(paths): |
| seen: set[tuple] = set() |
| for p in paths: |
| for rec in C.iter_jsonl(p): |
| C.validate_generation(rec) |
| key = C.gen_key(rec) |
| if key in seen: |
| raise C.ContractError(f"{p}: duplicate generation key {key} across the input files") |
| seen.add(key) |
| yield rec |
|
|
|
|
| def make_record(gen: dict, res: Result, budget_s: float, verifier_meta: dict) -> dict: |
| rec = { |
| "problem_id": gen["problem_id"], |
| "sample_idx": gen["sample_idx"], |
| "round": gen["round"], |
| "cap_tokens": gen["cap_tokens"], |
| "accepted": res.accepted, |
| "reasons": res.reasons, |
| "verify_s": round(res.elapsed_s, 4), |
| "budget_s": budget_s, |
| "verifier": verifier_meta, |
| "code_sha256": C.code_sha(gen["lean_code"]), |
| "stage": res.stage, |
| "diagnostics": res.diagnostics, |
| } |
| if res.feedback is not None: |
| rec.update(res.feedback) |
| return rec |
|
|
|
|
| def no_code_cause(gen: dict) -> Optional[str]: |
| if gen["finish_reason"] == "length": |
| return "finish_reason=length" |
| if gen["lean_code"] is None or not gen["lean_code"].strip(): |
| return "no_extracted_code" |
| return None |
|
|
|
|
| def verify_all( |
| gen_paths, |
| problems: dict[str, dict], |
| out_path: Path, |
| batch_fn: Callable[[list[tuple[str, str]], Callable[[int, Result], None]], None], |
| budget_s: float, |
| verifier_meta: dict, |
| chunk: int = 128, |
| limit: Optional[int] = None, |
| dry_run: bool = False, |
| ) -> dict: |
| """Verify every generation not yet in `out_path`. `batch_fn(items, on_result)` must call |
| `on_result(index, Result)` once per item (from any thread). Returns run statistics.""" |
| out_path = Path(out_path) |
| torn = C.repair_torn_tail(out_path) |
| existing = load_existing(out_path) |
| stats = {"generations": 0, "already_done": 0, "no_code": 0, "verified": 0, "accepted": 0, "torn_bytes_removed": torn} |
| pending: list[dict] = [] |
| writer = _NullWriter() if dry_run else C.AppendJsonl(out_path) |
| try: |
| for gen in iter_generations(gen_paths): |
| stats["generations"] += 1 |
| if gen["problem_id"] not in problems: |
| raise C.ContractError(f"generation {C.gen_key(gen)}: problem_id not in the problems file(s)") |
| key = C.gen_key(gen) |
| if key in existing: |
| if existing[key]["code_sha256"] != C.code_sha(gen["lean_code"]): |
| raise StaleVerification( |
| f"{key}: the stored verdict was computed for different code (sha " |
| f"{existing[key]['code_sha256']} vs {C.code_sha(gen['lean_code'])}); the generation file was " |
| "regenerated under the same keys. Use a fresh --out, never reuse verdicts." |
| ) |
| stats["already_done"] += 1 |
| continue |
| cause = no_code_cause(gen) |
| if cause is not None: |
| rec = make_record(gen, Result(False, ["no_code"], 0.0, "skipped", {"cause": cause}), budget_s, verifier_meta) |
| writer.append(rec) |
| existing[key] = rec |
| stats["no_code"] += 1 |
| continue |
| pending.append(gen) |
| if limit is not None and len(pending) >= limit: |
| break |
| log.info("%d generations: %d already verified, %d no_code written, %d to verify", stats["generations"], |
| stats["already_done"], stats["no_code"], len(pending)) |
| stats["pending"] = len(pending) |
| t0 = time.monotonic() |
| for lo in range(0, 0 if dry_run else len(pending), chunk): |
| part = pending[lo : lo + chunk] |
| items = [(problems[g["problem_id"]]["formal_statement"], g["lean_code"]) for g in part] |
| held: list[tuple[int, Result]] = [] |
| lock = threading.Lock() |
| results: dict[int, Result] = {} |
|
|
| def on_result(i: int, res: Result, part=part) -> None: |
| with lock: |
| results[i] = res |
| if res.infra: |
| held.append((i, res)) |
| return |
| writer.append(make_record(part[i], res, budget_s, verifier_meta)) |
|
|
| batch_fn(items, on_result) |
| missing = [i for i in range(len(part)) if i not in results] |
| if missing: |
| raise RuntimeError(f"batch_fn returned without a result for items {missing[:5]}") |
| if len(part) >= 8 and len(held) / len(part) > 0.5: |
| raise VerifierDied( |
| f"{len(held)}/{len(part)} verifications of chunk {lo // chunk} died for infrastructure reasons " |
| f"({sorted({r for _, x in held for r in x.reasons})}); not writing them, so a resume retries. " |
| "Either the pool is dead or the policy is emitting REPL-killing proofs: look before rerunning." |
| ) |
| for i, res in held: |
| writer.append(make_record(part[i], res, budget_s, verifier_meta)) |
| for res in results.values(): |
| stats["verified"] += 1 |
| stats["accepted"] += int(res.accepted) |
| log.info("verified %d/%d (%.0fs elapsed, %d accepted so far)", min(lo + chunk, len(pending)), len(pending), |
| time.monotonic() - t0, stats["accepted"]) |
| finally: |
| writer.close() |
| return stats |
|
|
|
|
| |
|
|
|
|
| def collect_feedback(worker, statement: str, code: str, timeout_s: float = 120.0) -> Optional[dict]: |
| """Full Lean error list for `code`, in the form the correction-round prompt needs. |
| |
| Reproduces Goedel's compile step: `compile_view(code)` (import / set_option / open lines |
| dropped) is run in an environment that already holds the statement's header, and the |
| REPL's error dicts are kept as `repl_errors` with the exact text they index as `repl_code` |
| (`extras/vllm_rollout/lean_gen.py` joins on these two keys). `StrictResult.diagnostics` |
| cannot serve: it keeps 3 errors, 240 characters each, without positions. |
| |
| Returns None when the statement/code cannot be run at all (nothing to show a model).""" |
| 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 or lex(code).broken: |
| return None |
| view = C.compile_view(code) |
| try: |
| env = worker.header_env(ref.header_text, "") |
| resp = worker.command(view, env, timeout_s) |
| except ReplError as e: |
| return {"repl_infra": e.code} |
|
|
| def pos(p): |
| return None if not p else {"line": int(p["line"]), "column": int(p["column"])} |
|
|
| errors = [ |
| {"pos": pos(m.get("pos")), "endPos": pos(m.get("endPos")), "data": m.get("data", "")} |
| for m in resp.get("messages", []) |
| if isinstance(m, dict) and m.get("severity") == "error" |
| ] |
| return {"repl_errors": errors, "repl_code": view} |
|
|
|
|
| def build_real(workers: int, budget_s: float, feedback: bool): |
| from align.verify.lean import ReplPool, StrictConfig, StrictVerifier, find_toolchain |
| from align.verify.lean.repl_pool import ReplLimits |
|
|
| tc = find_toolchain() |
| cfg = StrictConfig(proof_timeout_s=budget_s) |
| pool = ReplPool.from_toolchain( |
| tc, n_workers=workers, limits=ReplLimits(), preflight=True, |
| scratch_dir=str(Path(os.environ.get("TMPDIR", str(REPO / ".tmp"))) / "p2_verify_pool"), |
| ) |
| verifier = StrictVerifier(pool, cfg) |
| cfgd = asdict(cfg) |
| cfgd["axiom_whitelist"] = sorted(cfgd["axiom_whitelist"]) |
| meta = {"pins": tc.pins, "config": cfgd, "workers": pool.n_workers} |
|
|
| def batch_fn(items, on_result): |
| def one(w, ix_item): |
| ix, (stmt, code) = ix_item |
| sr = verifier._verify_on(w, (stmt, code)) |
| fb = None |
| if feedback and not sr.accepted and sr.reasons and sr.reasons[0] == "compile_error": |
| fb = collect_feedback(w, stmt, code) |
| on_result(ix, Result(sr.accepted, sr.reasons, sr.elapsed_s, sr.stage, _jsonable(sr.diagnostics), fb)) |
| return None |
|
|
| out = pool.map(one, list(enumerate(items))) |
| for o in out: |
| if isinstance(o, BaseException): |
| raise o |
|
|
| return batch_fn, meta, pool |
|
|
|
|
| def _jsonable(d: dict) -> dict: |
| return json.loads(json.dumps(d, default=str)) |
|
|
|
|
| |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--generations", nargs="+", required=True) |
| ap.add_argument("--problems", nargs="+", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--workers", type=int, default=8) |
| ap.add_argument("--inherit-lock", type=int, default=None, metavar="PARENT_PID", |
| help="the heavy-job lock is already held by this (parent) pid; verify that and do not take it again") |
| ap.add_argument("--budget-s", type=float, default=180.0, help="per-proof wall budget (compile + audit); recorded in every row") |
| ap.add_argument("--feedback", action="store_true", help="store full Lean errors for compile_error rejections") |
| ap.add_argument("--chunk", type=int, default=128) |
| ap.add_argument("--limit", type=int, default=None, help="verify at most this many pending samples (smoke tests)") |
| ap.add_argument("--dry-run", action="store_true", help="count what would be verified; start no Lean") |
| a = ap.parse_args(argv) |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") |
| if not 1 <= a.workers <= MAX_WORKERS: |
| raise SystemExit(f"--workers must be in [1, {MAX_WORKERS}]") |
| out = Path(a.out).resolve() |
| if not str(out).startswith("/apps/"): |
| raise SystemExit(f"--out {out} must be on /apps (never / or /tmp)") |
|
|
| problems: dict[str, dict] = {} |
| for p in a.problems: |
| for pid, rec in C.load_problems(p).items(): |
| if pid in problems and problems[pid]["formal_statement"] != rec["formal_statement"]: |
| raise C.ContractError(f"problem {pid!r} appears in two problem files with different statements") |
| problems[pid] = rec |
|
|
| if a.dry_run: |
| print(json.dumps(verify_all(a.generations, problems, out, None, a.budget_s, {}, a.chunk, a.limit, dry_run=True))) |
| return 0 |
|
|
| from align.utils import guard |
|
|
| if a.inherit_lock is not None: |
| guard.inherit_heavy_job_lock(a.inherit_lock) |
| elif a.workers > 2: |
| log.info("waiting for the heavy-job lock (Lean pool with %d workers)", a.workers) |
| guard.acquire_heavy_job_lock("p2-verify-generations", blocking=True) |
| guard.preflight_host_check(strict=True) |
| batch_fn, meta, pool = build_real(a.workers, a.budget_s, a.feedback) |
| stats = verify_all(a.generations, problems, out, batch_fn, a.budget_s, meta, a.chunk, a.limit) |
| stats["pool"] = asdict(pool.stats) |
| print(json.dumps(stats)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
