| |
| """K5: inference-time levers at MATCHED total generated tokens per problem (SOTA plan v3.1 section 3; docs/K5_LEVERS.md). |
| |
| Three arms on the same base model (Goedel-Prover-V2-8B) and the same dev problems: |
| |
| restart independent whole-proof attempts (class A samples). |
| correct Goedel-format Lean-feedback correction (class B): attempt + up to 2 repairs per trajectory, cumulative |
| token ledger, repairs of EVERY failed root (so each trajectory has a complete cost record). |
| agentic (class C, exploratory) proof-bank-conditioned attempt, then R = 4 rounds in total of: APOLLO-style subgoal |
| repair (split the failing proof at its failing `have`s, prove each subgoal as its own problem, reassemble, |
| re-verify) alternating with compiler-feedback correction. Retrieval is from a decontaminated bank of |
| verified proofs of OTHER training problems. |
| |
| Matching. The unit of every arm is a TRAJECTORY (a restart sample; a root with its repair chain; an agentic chain with |
| its subgoal sub-generations) with a token cost: the completion tokens spent up to the end of the round in which it first |
| succeeded, or all of its tokens if it never did. Arms are generated generously (a pool of units per problem) and then |
| REPLAYED at a per-problem budget B_p: units are drawn in a random order (permutation-averaged, deterministic seed) and |
| run until the budget is spent or one succeeds; a unit that does not fit in the remaining budget ends the run (a real |
| run cannot know a unit's cost in advance, so it is clipped and fails). The estimate per problem is the fraction of |
| orderings that solve it, directly comparable across arms and testable with the registered paired sign-flip test. |
| `--budget restart_k:8` sets B_p to 8 mean restart attempts of that problem (the pass@8-equivalent spend); |
| `--budget fixed:N` is a scalar. A problem whose pool cannot fill the budget is reported `underfilled`, and |
| `status` says how many more units to generate: the budget equaliser never silently lets one arm spend less. |
| |
| Subcommands (`k5_levers.py <cmd> -h`): |
| plan print every stage (the exact command lines) for a configuration; starts nothing |
| run execute the stages with a resumable done-marker per stage (each inside `hostguard.py run`) |
| prep merge dev problems, build + decontaminate the proof bank, attach retrieved examples (CPU) |
| route agentic round k: choose subgoal vs feedback per failed parent, extract subgoals through Lean (CPU) |
| assemble agentic round k: splice proven subgoals back, write assembled candidates (CPU) |
| ledger per-arm unit tables from generation + verification files |
| report matched-budget comparison, paired tests, Holm over the K5 family |
| status under-filled problems per arm and how many units each needs |
| estimate box-hour estimate from measured rates |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import dataclasses |
| import hashlib |
| import importlib.util |
| import json |
| import math |
| import os |
| import re |
| import shlex |
| import subprocess |
| import sys |
| import time |
| from collections import Counter, defaultdict |
| from dataclasses import dataclass, field |
| from pathlib import Path |
| from typing import Iterable, Optional |
|
|
| import numpy as np |
|
|
| REPO = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(REPO)) |
|
|
|
|
| def _load(name: str, rel: str): |
| if name in sys.modules: |
| return sys.modules[name] |
| spec = importlib.util.spec_from_file_location(name, REPO / rel) |
| mod = importlib.util.module_from_spec(spec) |
| sys.modules[name] = mod |
| spec.loader.exec_module(mod) |
| return mod |
|
|
|
|
| C = _load("p2_common", "scripts/p2/p2_common.py") |
| LP = _load("k5_lean_prompts", "extras/vllm_rollout/lean_prompts.py") |
|
|
| |
| _VR = Path(os.environ.get("ALIGN_VENV_ROOT", REPO if (REPO / ".venv").exists() else "/apps/development/rlhf/align")) |
| PYR = str(_VR / ".venv-rollout" / "bin" / "python") |
| PY = str(_VR / ".venv" / "bin" / "python") |
| DEFAULT_MODEL = ( |
| "/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/" |
| "dfd02e6271a58375dfbf3ece0175277cf6b6a89a" |
| ) |
| ARMS = ("restart", "correct", "agentic") |
| ASM_BASE = 1_000_000 |
|
|
|
|
| |
|
|
|
|
| @dataclass(frozen=True) |
| class Unit: |
| """One trajectory of an arm. `tokens` is the cost the unit charges when it is run: the completion tokens of every |
| node up to the end of the round of its first success, or `tokens_total` when it never succeeded.""" |
|
|
| problem_id: str |
| arm: str |
| unit_id: str |
| tokens: int |
| solved: bool |
| tokens_total: int |
| by_kind: tuple = () |
|
|
|
|
| def node_tokens(g: dict) -> int: |
| t = g.get("trajectory") |
| if t and "completion_tokens" in t: |
| return int(t["completion_tokens"]) |
| ids = g.get("completion_ids") |
| if ids is None: |
| raise ValueError(f"{g.get('problem_id')}: generation record without completion_ids or a trajectory ledger") |
| return len(ids) |
|
|
|
|
| def node_kind(g: dict) -> str: |
| t = g.get("trajectory") or {} |
| return t.get("kind") or ("attempt" if int(g["round"]) == 0 else "repair") |
|
|
|
|
| def root_id(g: dict) -> str: |
| t = g.get("trajectory") |
| if t and "root_trajectory_id" in t: |
| return t["root_trajectory_id"] |
| if int(g["round"]) != 0: |
| raise ValueError(f"{g['problem_id']}: a round-{g['round']} record without a trajectory ledger cannot be assigned to a root") |
| return f"{g['problem_id']}|{int(g['sample_idx'])}" |
|
|
|
|
| def _join(gens: Iterable[dict], vers: Iterable[dict]) -> list[tuple[dict, dict]]: |
| vmap = {} |
| for v in vers: |
| vmap[C.gen_key(v)] = v |
| out, seen = [], set() |
| for g in gens: |
| k = C.gen_key(g) |
| if k in seen: |
| raise ValueError(f"duplicate generation key {k} across the input files") |
| seen.add(k) |
| v = vmap.get(k) |
| if v is None: |
| raise ValueError(f"generation {k} has no verification record; verify every node before building the ledger") |
| out.append((g, v)) |
| return out |
|
|
|
|
| def units_restart(gens: Iterable[dict], vers: Iterable[dict], arm: str = "restart") -> dict[str, list[Unit]]: |
| """Each round-0 sample is its own unit.""" |
| out: dict[str, list[Unit]] = defaultdict(list) |
| for g, v in _join(gens, vers): |
| if int(g["round"]) != 0: |
| continue |
| n = node_tokens(g) |
| out[g["problem_id"]].append(Unit(g["problem_id"], arm, root_id(g), n, bool(v["accepted"]), n, ((node_kind(g), n),))) |
| return dict(out) |
|
|
|
|
| def units_chains(gens: Iterable[dict], vers: Iterable[dict], arm: str) -> dict[str, list[Unit]]: |
| """One unit per root trajectory (class B and agentic). Nodes are grouped by `trajectory.root_trajectory_id`; the |
| cost through round r is the sum of ALL nodes of rounds <= r (branches included), so nothing generated is free.""" |
| groups: dict[str, list[tuple[dict, dict]]] = defaultdict(list) |
| for g, v in _join(gens, vers): |
| groups[root_id(g)].append((g, v)) |
| out: dict[str, list[Unit]] = defaultdict(list) |
| for rid, nodes in groups.items(): |
| if not any(int(g["round"]) == 0 for g, _ in nodes): |
| raise ValueError(f"{rid}: no round-0 node among the input files (the lineage is partial); pass every round's files") |
| pid = nodes[0][0]["problem_id"] |
| solved_round = min((int(g["round"]) for g, v in nodes if v["accepted"]), default=None) |
| charged = [(g, v) for g, v in nodes if solved_round is None or int(g["round"]) <= solved_round] |
| by_kind: Counter = Counter() |
| for g, _ in charged: |
| by_kind[node_kind(g)] += node_tokens(g) |
| total = sum(node_tokens(g) for g, _ in nodes) |
| out[pid].append(Unit(pid, arm, rid, sum(by_kind.values()), solved_round is not None, total, tuple(sorted(by_kind.items())))) |
| return dict(out) |
|
|
|
|
| def _seed_for(seed: int, pid: str) -> int: |
| return int.from_bytes(hashlib.sha256(f"{seed}|{pid}".encode()).digest()[:8], "big") % (2**32) |
|
|
|
|
| @dataclass |
| class Replay: |
| p_solve: float |
| spent: float |
| underfill: float |
| n_units: int |
|
|
|
|
| def replay_budget(units: list[Unit], budget: float, n_perm: int = 256, seed: int = 0, pid: str = "") -> Replay: |
| """Permutation-averaged solve probability of `units` at `budget` generated tokens (the equaliser's arithmetic). |
| |
| One ordering: units are run in order; a unit whose cost exceeds the remaining budget ends the ordering unsolved |
| (it would be clipped); a solved unit within budget ends it solved; an unsolved one that fits spends its tokens |
| and the next is tried. Spent tokens are what the ordering actually used.""" |
| if not units: |
| return Replay(0.0, 0.0, 1.0, 0) |
| rng = np.random.default_rng(_seed_for(seed, pid)) |
| cost = np.array([u.tokens for u in units], dtype=np.float64) |
| solved = np.array([u.solved for u in units], dtype=bool) |
| mean_cost = float(cost.mean()) |
| wins = under = 0 |
| spent_sum = 0.0 |
| n = len(units) |
| for _ in range(n_perm): |
| order = rng.permutation(n) |
| spent, won, exhausted = 0.0, False, True |
| for i in order: |
| if spent + cost[i] > budget: |
| exhausted = False |
| break |
| spent += cost[i] |
| if solved[i]: |
| won, exhausted = True, False |
| break |
| wins += won |
| spent_sum += spent |
| if exhausted and not won and budget - spent >= mean_cost: |
| under += 1 |
| return Replay(wins / n_perm, spent_sum / n_perm, under / n_perm, n) |
|
|
|
|
| def budget_table(spec: str, restart: dict[str, list[Unit]], problems: Iterable[str]) -> dict[str, float]: |
| """Per-problem budget B_p. `restart_k:K` = K x the mean tokens of that problem's restart attempts (the spend of a |
| pass@K run); `fixed:N` = N tokens everywhere.""" |
| kind, _, arg = spec.partition(":") |
| out = {} |
| if kind == "fixed": |
| return {p: float(arg) for p in problems} |
| if kind != "restart_k": |
| raise ValueError(f"unknown budget spec {spec!r} (restart_k:K or fixed:N)") |
| k = float(arg) |
| for p in problems: |
| us = restart.get(p) |
| if not us: |
| raise ValueError(f"problem {p}: no restart units to derive its budget from") |
| out[p] = k * sum(u.tokens_total for u in us) / len(us) |
| return out |
|
|
|
|
| def equalise(arms: dict[str, dict[str, list[Unit]]], budgets: dict[str, float], n_perm: int = 256, seed: int = 0, |
| fractions: tuple = (0.25, 0.5, 1.0)) -> dict: |
| """Replay every arm at every budget fraction. -> {arm: {fraction: {problem: Replay}}} plus summaries.""" |
| res: dict = {} |
| for arm, by_p in arms.items(): |
| res[arm] = {} |
| for f in fractions: |
| res[arm][f] = {p: replay_budget(by_p.get(p, []), b * f, n_perm, seed, p) for p, b in budgets.items()} |
| return res |
|
|
|
|
| def arm_summary(per_problem: dict[str, Replay], budgets: dict[str, float], fraction: float = 1.0) -> dict: |
| n = len(per_problem) |
| tot_b = sum(budgets[p] * fraction for p in per_problem) |
| return { |
| "n_problems": n, |
| "mean_p_solve": sum(r.p_solve for r in per_problem.values()) / max(n, 1), |
| "tokens_spent_per_problem": sum(r.spent for r in per_problem.values()) / max(n, 1), |
| "budget_per_problem": tot_b / max(n, 1), |
| "spend_ratio": sum(r.spent for r in per_problem.values()) / tot_b if tot_b else 0.0, |
| "underfilled_problems": sum(1 for r in per_problem.values() if r.underfill > 0.05), |
| "no_units": sum(1 for r in per_problem.values() if r.n_units == 0), |
| } |
|
|
|
|
| def units_needed(units: dict[str, list[Unit]], budgets: dict[str, float], per_problem: dict[str, Replay]) -> dict[str, int]: |
| """Additional units per under-filled problem: enough mean-cost units to cover the unspent budget, +1 slack.""" |
| need = {} |
| for p, r in per_problem.items(): |
| if r.underfill <= 0.05: |
| continue |
| us = units.get(p) or [] |
| mean = (sum(u.tokens for u in us) / len(us)) if us else max(budgets[p] / 4.0, 1.0) |
| need[p] = max(1, math.ceil(max(budgets[p] - r.spent, 0.0) / max(mean, 1.0))) + 1 |
| return need |
|
|
|
|
| |
|
|
|
|
| @dataclass |
| class BankEntry: |
| problem_id: str |
| statement: str |
| proof: str |
| informal: Optional[str] |
| tokens: int |
| code_sha: str |
| source_statement: str = "" |
|
|
|
|
| def split_example(lean_code: str) -> Optional[tuple[str, str]]: |
| """(`theorem ... := by`, body) of a verified `lean_code`, or None. Comments are masked for the search only.""" |
| from align.rl.envs.subgoal_repair import mask |
|
|
| mk = mask(lean_code) |
| m = re.search(r"\btheorem\b[\s\S]*?:=\s*by", mk) |
| if m is None: |
| return None |
| return lean_code[m.start():m.end()], lean_code[m.end():] |
|
|
|
|
| def eval_items_from_problems(rows: Iterable[dict]): |
| from align.data.decontam import EvalItem |
|
|
| return [EvalItem(f"{r.get('split', '')}:{r['problem_id']}", r.get("source", "eval"), r["formal_statement"], r.get("informal_prefix")) |
| for r in rows] |
|
|
|
|
| def decontaminate_bank(candidates: list[BankEntry], eval_rows: list[dict], ancestry: Optional[dict] = None) -> tuple[list[BankEntry], list[dict]]: |
| """Drop every candidate that hits any evaluation / query problem under the repo's layered detector |
| (`align.data.decontam`: normalised-text, signature and alpha hashes, MinHash-LSH, token n-gram containment, informal |
| text n-grams), or whose problem id IS an evaluation id, or whose ancestry (mutation parent) does. Returns |
| (kept, dropped-with-reason). Over-flagging is the policy: a false positive costs one bank entry.""" |
| from align.data.decontam import ContaminationLedger, Decontaminator |
|
|
| dec = Decontaminator(eval_items_from_problems(eval_rows)) |
| eval_ids = {r["problem_id"] for r in eval_rows} |
| led = ContaminationLedger() |
| for c in candidates: |
| led.add(c.problem_id, dec.check(c.source_statement or c.statement, c.informal), (ancestry or {}).get(c.problem_id)) |
| kept, dropped = [], [] |
| for c in candidates: |
| if c.problem_id in eval_ids: |
| dropped.append({"problem_id": c.problem_id, "why": "is_an_eval_problem"}) |
| continue |
| bad, why = led.resolve(c.problem_id) |
| if bad: |
| dropped.append({"problem_id": c.problem_id, "why": why, "hits": [dataclasses.asdict(h) for h in led.direct.get(c.problem_id, [])[:3]]}) |
| continue |
| kept.append(c) |
| return kept, dropped |
|
|
|
|
| def bank_candidates(problems: dict[str, dict], gens: Iterable[dict], vers: Iterable[dict], max_chars: int = 2500) -> list[BankEntry]: |
| """The shortest accepted proof per problem, as an example (statement + body).""" |
| best: dict[str, BankEntry] = {} |
| for g, v in _join(gens, vers): |
| if not v["accepted"] or not g.get("lean_code") or g["problem_id"] not in problems: |
| continue |
| ex = split_example(g["lean_code"]) |
| if ex is None or len(ex[0]) + len(ex[1]) > max_chars: |
| continue |
| pr = problems[g["problem_id"]] |
| e = BankEntry(g["problem_id"], ex[0], ex[1], pr.get("informal_prefix"), node_tokens(g), C.code_sha(g["lean_code"]), pr["formal_statement"]) |
| cur = best.get(g["problem_id"]) |
| if cur is None or len(e.proof) < len(cur.proof): |
| best[g["problem_id"]] = e |
| return [best[k] for k in sorted(best)] |
|
|
|
|
| class BankIndex: |
| """TF-IDF over statement tokens (identifiers kept: lemma names are the signal). Cosine retrieval, shorter proof breaks |
| ties. Pure numpy; the bank is thousands of entries.""" |
|
|
| def __init__(self, entries: list[BankEntry]): |
| from align.data.decontam import lean_tokens |
|
|
| self._tok = lambda s: lean_tokens(s, alpha=False) |
| self.entries = entries |
| docs = [Counter(self._tok(e.source_statement or e.statement)) for e in entries] |
| df = Counter(t for d in docs for t in d) |
| self.vocab = {t: i for i, t in enumerate(sorted(df))} |
| n = max(len(entries), 1) |
| self.idf = np.array([math.log((1 + n) / (1 + df[t])) + 1.0 for t in sorted(df)]) if df else np.zeros(0) |
| self.mat = np.zeros((len(entries), len(self.vocab)), dtype=np.float32) |
| for i, d in enumerate(docs): |
| self.mat[i] = self._vec(d) |
|
|
| def _vec(self, counts: Counter) -> np.ndarray: |
| v = np.zeros(len(self.vocab), dtype=np.float32) |
| for t, c in counts.items(): |
| j = self.vocab.get(t) |
| if j is not None: |
| v[j] = (1 + math.log(c)) * self.idf[j] |
| nrm = np.linalg.norm(v) |
| return v / nrm if nrm else v |
|
|
| def retrieve(self, statement: str, k: int = 3, exclude_ids: Iterable[str] = (), min_sim: float = 0.05) -> list[BankEntry]: |
| if not self.entries: |
| return [] |
| q = self._vec(Counter(self._tok(statement))) |
| sims = self.mat @ q |
| ex = set(exclude_ids) |
| order = sorted(range(len(self.entries)), key=lambda i: (-float(sims[i]), len(self.entries[i].proof))) |
| out = [] |
| for i in order: |
| if sims[i] < min_sim or self.entries[i].problem_id in ex: |
| continue |
| out.append(self.entries[i]) |
| if len(out) == k: |
| break |
| return out |
|
|
|
|
| def bank_examples(entries: list[BankEntry], max_chars: int = 6000) -> list[dict]: |
| out, used = [], 0 |
| for e in entries: |
| used += len(e.statement) + len(e.proof) |
| if used > max_chars: |
| break |
| out.append({"statement": e.statement, "proof": e.proof, "source": e.problem_id}) |
| return out |
|
|
|
|
| def attach_bank(problems: list[dict], index: BankIndex, k: int = 3) -> list[dict]: |
| """Problem rows with `bank_examples` (the `goedel_bank` prompt format reads them). A problem is never given its own |
| bank entry (`exclude_ids`).""" |
| out = [] |
| for p in problems: |
| ex = bank_examples(index.retrieve(p["formal_statement"], k, exclude_ids=[p["problem_id"]])) |
| out.append(dict(p, bank_examples=ex)) |
| return out |
|
|
|
|
| |
|
|
|
|
| def parent_info(g: dict, v: dict) -> Optional[dict]: |
| """Failed node with Lean feedback, or None (accepted, truncated, no code, no compile errors).""" |
| if v["accepted"] or not g.get("lean_code"): |
| return None |
| if not isinstance(v.get("repl_errors"), list) or not v["repl_errors"] or not isinstance(v.get("repl_code"), str): |
| return None |
| return {"gen": g, "ver": v} |
|
|
|
|
| def route_for_round(routes: str, round_k: int) -> str: |
| """`routes` = comma list for rounds 1..R-1, each `subgoal` (subgoal repair, falling back to feedback when nothing can |
| be extracted) or `feedback`. A design choice, not tuned: alternating heterogeneous moves.""" |
| rs = [x.strip() for x in routes.split(",")] |
| if not 1 <= round_k <= len(rs) or any(x not in ("subgoal", "feedback") for x in rs): |
| raise ValueError(f"routes {routes!r} has no valid entry for round {round_k}") |
| return rs[round_k - 1] |
|
|
|
|
| def sg_problem_id(parent_g: dict, j: int) -> str: |
| return f"{parent_g['problem_id']}::sg{int(parent_g['sample_idx'])}r{int(parent_g['round'])}g{j}" |
|
|
|
|
| def make_asm_nodes(plan, parent_g: dict, parent_v: dict, bodies: dict[int, str], sub_tokens: int, n_proven: int, round_k: int, cap: int): |
| """(gen record, ver record or None). The assembled candidate is a synthetic generation: its text is the code, its |
| completion_ids a single placeholder, and the real cost (the subgoal samples) sits in `trajectory.completion_tokens`. |
| The history `messages` are the parent's, so a following Goedel correction prompt reads the whole story. When the |
| candidate equals the parent's code (no subgoal proven) its verdict is copied, not re-verified.""" |
| from align.rl.envs.subgoal_repair import reassemble |
|
|
| code = reassemble(plan, bodies, restore_unproven=True) |
| pt = parent_g.get("trajectory") or {} |
| prior = int(pt.get("cum_completion_tokens", node_tokens(parent_g))) |
| sidx = ASM_BASE + int(parent_g["sample_idx"]) |
| rid = root_id(parent_g) |
| gen = { |
| "problem_id": parent_g["problem_id"], "sample_idx": sidx, "round": round_k, "cap_tokens": cap, |
| "prompt_ids": [0], "completion_ids": [0], "prompt_text": "", "completion_text": f"```lean4\n{code}\n```", |
| "finish_reason": "stop", "lean_code": code, "lean_code_error": None, "messages": parent_g["messages"], |
| "model_path": parent_g["model_path"], "model_rev": parent_g["model_rev"], |
| "sampling": dict(parent_g["sampling"], max_tokens=1), |
| "context_clipped": False, |
| "parent": {"problem_id": parent_g["problem_id"], "sample_idx": int(parent_g["sample_idx"]), "round": int(parent_g["round"])}, |
| "trajectory": { |
| "root_trajectory_id": rid, "problem_id": parent_g["problem_id"], |
| "root_sample_idx": int(pt.get("root_sample_idx", parent_g["sample_idx"])), "round": round_k, "kind": "assembled", |
| "parent_key": [parent_g["problem_id"], int(parent_g["sample_idx"]), int(parent_g["round"])], |
| "completion_tokens": int(sub_tokens), "prior_completion_tokens": prior, "cum_completion_tokens": prior + int(sub_tokens), |
| "truncated": False, "budget_clipped": False, "partial": False, |
| }, |
| "subgoal": {"gaps": len(plan.gaps), "proven": n_proven}, |
| } |
| ver = None |
| if code == parent_g["lean_code"]: |
| ver = dict(parent_v, sample_idx=sidx, round=round_k, cap_tokens=cap, code_sha256=C.code_sha(code)) |
| return gen, ver |
|
|
|
|
| |
|
|
|
|
| @dataclass |
| class Cfg: |
| run_dir: str |
| problems: list[str] |
| model: str = DEFAULT_MODEL |
| seed: int = 0 |
| cap: int = 40960 |
| traj_cap: int = 40000 |
| n_restart: int = 32 |
| n_agentic: int = 8 |
| rounds: int = 4 |
| routes: str = "subgoal,feedback,subgoal" |
| sub_n: int = 4 |
| sub_cap: int = 16384 |
| bank_k: int = 3 |
| budget: str = "restart_k:8" |
| workers: int = 8 |
| gpu_mem_util: float = 0.82 |
| bank_problems: Optional[str] = None |
| bank_gens: Optional[str] = None |
| bank_vers: Optional[str] = None |
| eval_problems: list = field(default_factory=list) |
| hostguard: bool = True |
| gpu_caps: str = "--cores 24 --mem-gb 100" |
| cpu_caps: str = "--cores 16 --mem-gb 64" |
|
|
| def p(self, *parts) -> str: |
| return str(Path(self.run_dir, *parts)) |
|
|
|
|
| @dataclass |
| class Stage: |
| name: str |
| argv: list |
| gpu: bool |
| outputs: list |
| skip_if_empty: Optional[str] = None |
|
|
| def marker(self, cfg: Cfg) -> Path: |
| return Path(cfg.p("done", self.name)) |
|
|
|
|
| def _hg(cfg: Cfg, st: Stage) -> list: |
| if not cfg.hostguard: |
| return st.argv |
| caps = (cfg.gpu_caps if st.gpu else cfg.cpu_caps).split() |
| return [PY, str(REPO / "scripts/hostguard.py"), "run", "--wait", *caps, "--label", f"k5-{st.name}", "--", *st.argv] |
|
|
|
|
| def _gen_args(cfg: Cfg, problems: str, out: str, n: int, rnd: int, cap: int, label: str, extra: list = ()) -> str: |
| a = ["--model", cfg.model, "--problems", problems, "--out", out, "--n", str(n), "--cap-tokens", str(cap), "--round", str(rnd), |
| "--tp", "4", "--seed", str(cfg.seed), "--max-num-seqs", "128", "--gpu-mem-util", str(cfg.gpu_mem_util), "--label", label, *extra] |
| return " ".join(shlex.quote(x) for x in a) |
|
|
|
|
| def _pgv(cfg: Cfg, name: str, problems: str, gen_out: str, ver_out: str, n: int, rnd: int, cap: int, extra: list = (), feedback: bool = True, |
| wait_marker: Optional[str] = None) -> Stage: |
| run = cfg.p("pipe", name) |
| argv = [PY, str(REPO / "scripts/p2/pipeline_gen_verify.py"), "--run-dir", run, "--problems", problems, "--verify-out", ver_out, |
| "--workers", str(cfg.workers), "--label", f"k5-{name}", "--gen-args", _gen_args(cfg, problems, gen_out, n, rnd, cap, f"k5-{name}", list(extra))] |
| if feedback: |
| argv.append("--feedback") |
| return Stage(name, argv, True, [gen_out, ver_out]) |
|
|
|
|
| def _self(cfg: Cfg, name: str, sub: str, *args: str, gpu: bool = False) -> Stage: |
| return Stage(name, [PY, str(REPO / "scripts/p2/k5_levers.py"), sub, *args], gpu, []) |
|
|
|
|
| def build_stages(cfg: Cfg) -> list[Stage]: |
| """Every stage of K5, in order. Pure: nothing is executed or created.""" |
| dev, dev_bank = cfg.p("dev.jsonl"), cfg.p("dev_bank.jsonl") |
| ss: list[Stage] = [] |
| prep = ["--run-dir", cfg.run_dir, "--problems", *cfg.problems, "--bank-k", str(cfg.bank_k)] |
| if cfg.bank_problems: |
| prep += ["--bank-problems", cfg.bank_problems, "--bank-gens", cfg.bank_gens, "--bank-vers", cfg.bank_vers] |
| if cfg.eval_problems: |
| prep += ["--eval-problems", *cfg.eval_problems] |
| ss.append(_self(cfg, "prep", "prep", *prep)) |
| |
| R0g, R0v = cfg.p("restart/gen_r0.jsonl"), cfg.p("restart/ver_r0.jsonl") |
| ss.append(_pgv(cfg, "restart_r0", dev, R0g, R0v, cfg.n_restart, 0, cfg.cap)) |
| |
| prev = [(R0v, R0g)] |
| for r in (1, 2): |
| g, v = cfg.p(f"correct/gen_r{r}.jsonl"), cfg.p(f"correct/ver_r{r}.jsonl") |
| extra = ["--trajectory-cap-tokens", str(cfg.traj_cap), "--repair-solved"] |
| for pv, pg in prev: |
| extra += ["--corrections-from", pv, pg] |
| ss.append(_pgv(cfg, f"correct_r{r}", dev, g, v, 1, r, cfg.cap, extra)) |
| prev = [(v, g)] + prev |
| |
| A0g, A0v = cfg.p("agentic/gen_r0.jsonl"), cfg.p("agentic/ver_r0.jsonl") |
| ss.append(_pgv(cfg, "agentic_r0", dev_bank, A0g, A0v, cfg.n_agentic, 0, cfg.cap, ["--prompt-format", "goedel_bank"])) |
| last = [(A0v, A0g)] |
| for r in range(1, cfg.rounds): |
| mode = route_for_round(cfg.routes, r) |
| rd = cfg.p(f"agentic/round{r}") |
| route_args = ["--run-dir", cfg.run_dir, "--round", str(r), "--mode", mode, "--workers", "2", "--bank-k", str(cfg.bank_k), "--sub-n", str(cfg.sub_n)] |
| for v, g in last: |
| route_args += ["--pairs", v, g] |
| ss.append(_self(cfg, f"agentic_route_r{r}", "route", *route_args)) |
| fg, fv = cfg.p(f"agentic/gen_r{r}.jsonl"), cfg.p(f"agentic/ver_r{r}.jsonl") |
| fbg, fbv = f"{rd}/fb_parents.gen.jsonl", f"{rd}/fb_parents.ver.jsonl" |
| nxt = [(fv, fg)] |
| if mode == "subgoal": |
| sgp, sgg, sgv = f"{rd}/sg_problems.jsonl", f"{rd}/sg_gen.jsonl", f"{rd}/sg_ver.jsonl" |
| st = _pgv(cfg, f"agentic_sg_r{r}", sgp, sgg, sgv, cfg.sub_n, 0, cfg.sub_cap, ["--prompt-format", "goedel_bank"], feedback=False) |
| st.skip_if_empty = sgp |
| ss.append(st) |
| extra = ["--trajectory-cap-tokens", str(cfg.traj_cap), "--repair-solved", "--corrections-from", fbv, fbg] |
| st = _pgv(cfg, f"agentic_fb_r{r}", dev_bank, fg, fv, 1, r, cfg.cap, extra) |
| st.skip_if_empty = fbg |
| ss.append(st) |
| if mode == "subgoal": |
| ag, av = f"{rd}/asm.gen.jsonl", f"{rd}/asm.ver.jsonl" |
| ss.append(_self(cfg, f"agentic_asm_r{r}", "assemble", "--run-dir", cfg.run_dir, "--round", str(r), "--cap", str(cfg.cap))) |
| st = Stage(f"agentic_asmverify_r{r}", [PY, str(REPO / "scripts/p2/verify_generations.py"), "--generations", ag, "--problems", dev, |
| "--out", av, "--workers", str(cfg.workers), "--budget-s", "300", "--feedback"], False, [av]) |
| st.skip_if_empty = ag |
| ss.append(st) |
| nxt.append((av, ag)) |
| last = nxt |
| ss.append(_self(cfg, "report", "report", "--run-dir", cfg.run_dir, "--budget", cfg.budget, "--rounds", str(cfg.rounds))) |
| return ss |
|
|
|
|
| class Runner: |
| def run(self, argv: list, log: Path) -> int: |
| log.parent.mkdir(parents=True, exist_ok=True) |
| with open(log, "ab") as lf: |
| return subprocess.call(argv, stdout=lf, stderr=subprocess.STDOUT, cwd=str(REPO)) |
|
|
|
|
| def run_stages(cfg: Cfg, stages: list[Stage], runner, say=print, only: Optional[set] = None) -> int: |
| for st in stages: |
| if only and st.name not in only: |
| continue |
| mk = st.marker(cfg) |
| if mk.exists(): |
| say(f"skip {st.name} (done)") |
| continue |
| if st.skip_if_empty and Path(st.skip_if_empty).exists() and not Path(st.skip_if_empty).read_text().strip(): |
| say(f"skip {st.name} (its input {st.skip_if_empty} is empty)") |
| mk.parent.mkdir(parents=True, exist_ok=True) |
| mk.write_text("empty") |
| continue |
| say(f"START {st.name}") |
| rc = runner.run(_hg(cfg, st), Path(cfg.p("logs", st.name + ".log"))) |
| say(f"END {st.name} rc={rc}") |
| if rc != 0: |
| say(f"STOP: {st.name} failed; re-running this command resumes (every stage is keyed and resumable)") |
| return rc |
| mk.parent.mkdir(parents=True, exist_ok=True) |
| mk.write_text(time.strftime("%F %T")) |
| return 0 |
|
|
|
|
| |
|
|
|
|
| MEASURED = { |
| "box_tok_per_h": 522.5 * 3600, |
| "mean_attempt_tokens": {"minif2f": 5685.0, "proofnet": 9000.0}, |
| "fail_rate": {"minif2f": 0.36, "proofnet": 0.75}, |
| } |
|
|
|
|
| def estimate(sets: dict[str, int], cfg: Cfg, m: dict = MEASURED, repair_cost: float = 1.15, repair_success: float = 0.12, |
| sub_share: float = 0.35) -> dict: |
| """Box-hour estimate per arm. `sets` = {"minif2f": N, "proofnet": N}. The repair multiplier covers the longer |
| (history-carrying) prompts; sub_share is the tokens a subgoal round spends relative to one attempt per failed parent.""" |
| out: dict = {"assumptions": {"repair_cost": repair_cost, "repair_success": repair_success, "sub_share": sub_share, **{"box_tok_per_h": m["box_tok_per_h"]}}} |
| tot: Counter = Counter() |
| for s, n in sets.items(): |
| a, f = m["mean_attempt_tokens"][s], m["fail_rate"][s] |
| restart = n * cfg.n_restart * a |
| f1 = f |
| f2 = f * (1 - repair_success) |
| correct = n * cfg.n_restart * a * repair_cost * (f1 + f2) |
| agent_root = n * cfg.n_agentic * a |
| agent = agent_root |
| fa = f * 0.85 |
| for r in range(1, cfg.rounds): |
| mode = route_for_round(cfg.routes, r) |
| k = n * cfg.n_agentic * fa |
| agent += k * a * (sub_share if mode == "subgoal" else repair_cost) |
| fa *= (1 - repair_success) |
| tot["restart"] += restart |
| tot["correct"] += correct |
| tot["agentic"] += agent |
| out["tokens"] = dict(tot) |
| out["box_hours"] = {k: v / m["box_tok_per_h"] for k, v in tot.items()} |
| out["box_hours"]["total"] = sum(out["box_hours"].values()) |
| return out |
|
|
|
|
| |
|
|
|
|
| def _read(path) -> list[dict]: |
| return list(C.iter_jsonl(path)) if path and Path(path).exists() else [] |
|
|
|
|
| def _problems(files: Iterable[str]) -> dict[str, dict]: |
| out: dict[str, dict] = {} |
| for f in files: |
| for pid, rec in C.load_problems(f).items(): |
| if pid in out and out[pid]["formal_statement"] != rec["formal_statement"]: |
| raise C.ContractError(f"problem {pid!r} appears in two files with different statements") |
| out[pid] = rec |
| return out |
|
|
|
|
| def _write(path, rows) -> int: |
| return C.write_jsonl_atomic(path, rows) |
|
|
|
|
| def cmd_prep(a) -> int: |
| rd = Path(a.run_dir) |
| probs = _problems(a.problems) |
| _write(rd / "dev.jsonl", probs.values()) |
| bank_entries: list[BankEntry] = [] |
| report: dict = {"dev_problems": len(probs)} |
| if a.bank_problems: |
| train = _problems([a.bank_problems]) |
| cands = bank_candidates(train, _read(a.bank_gens), _read(a.bank_vers)) |
| eval_rows = list(probs.values()) + [r for f in a.eval_problems for r in _read(f)] |
| bank_entries, dropped = decontaminate_bank(cands, eval_rows) |
| report.update(bank_candidates=len(cands), bank_kept=len(bank_entries), bank_dropped=len(dropped), |
| eval_items=len(eval_rows), dropped_reasons=dict(Counter(d["why"].split(":")[0] for d in dropped))) |
| _write(rd / "bank.jsonl", (dataclasses.asdict(e) for e in bank_entries)) |
| _write(rd / "bank_dropped.jsonl", dropped) |
| rows = attach_bank(list(probs.values()), BankIndex(bank_entries), a.bank_k) |
| _write(rd / "dev_bank.jsonl", rows) |
| report["with_examples"] = sum(1 for r in rows if r["bank_examples"]) |
| (rd / "prep.json").write_text(json.dumps(report, indent=1)) |
| print(json.dumps(report)) |
| return 0 |
|
|
|
|
| def _pairs(args) -> tuple[list[dict], list[dict]]: |
| gens, vers = [], [] |
| for v, g in args: |
| gens += _read(g) |
| vers += _read(v) |
| return gens, vers |
|
|
|
|
| def cmd_route(a) -> int: |
| """Agentic round k. Parents = failed, feedback-bearing nodes of round k-1 (every `--pairs` file). Subgoal rounds |
| plan each parent through Lean; parents with nothing extractable, and every parent of a feedback round, go to the |
| Goedel correction path via filtered copies of their (ver, gen) records.""" |
| from align.rl.envs import subgoal_repair as SR |
|
|
| rd = Path(a.run_dir, "agentic", f"round{a.round}") |
| rd.mkdir(parents=True, exist_ok=True) |
| probs = {r["problem_id"]: r for r in _read(Path(a.run_dir, "dev_bank.jsonl"))} |
| gens, vers = _pairs(a.pairs) |
| nodes = [(g, v) for g, v in _join(gens, vers) if int(g["round"]) == a.round - 1] |
| st = Counter(parents_total=len(nodes)) |
| fb, plans, sg_rows = [], [], [] |
| probe_pool = None |
| index = None |
| bank = _read(Path(a.run_dir, "bank.jsonl")) |
| if bank: |
| index = BankIndex([BankEntry(**b) for b in bank]) |
| for g, v in nodes: |
| info = parent_info(g, v) |
| if info is None: |
| st["no_feedback_or_solved"] += 1 |
| continue |
| route = "feedback" |
| if a.mode == "subgoal": |
| if probe_pool is None: |
| probe_pool = _open_pool(a.workers) |
| plan = _plan_with_pool(probe_pool, probs[g["problem_id"]]["formal_statement"], g, v, SR) |
| if plan is not None and plan.lemmas: |
| route = "subgoal" |
| rows = [] |
| for j, gap in enumerate(plan.gaps): |
| if gap.lemma is None: |
| continue |
| stmt = SR.lemma_problem_statement(plan, gap) |
| row = {"problem_id": sg_problem_id(g, j), "source": "k5_subgoal", "split": "dev", "formal_statement": stmt, |
| "informal_prefix": None, "statement_sha": C.statement_sha(stmt), "parent": [g["problem_id"], int(g["sample_idx"]), int(g["round"])], |
| "gap_index": j} |
| if index is not None: |
| row["bank_examples"] = bank_examples(index.retrieve(stmt, a.bank_k)) |
| rows.append(row) |
| sg_rows += rows |
| plans.append({"parent_key": [g["problem_id"], int(g["sample_idx"]), int(g["round"])], "plan": SR.plan_to_dict(plan), |
| "sg_ids": [r["problem_id"] for r in rows]}) |
| st[route] += 1 |
| if route == "feedback": |
| fb.append((g, v)) |
| _write(rd / "fb_parents.gen.jsonl", (g for g, _ in fb)) |
| _write(rd / "fb_parents.ver.jsonl", (v for _, v in fb)) |
| _write(rd / "sg_problems.jsonl", sg_rows) |
| _write(rd / "sg_plans.jsonl", plans) |
| (rd / "route.json").write_text(json.dumps(dict(st), indent=1)) |
| print(json.dumps(dict(st))) |
| if probe_pool is not None: |
| probe_pool.close() |
| return 0 |
|
|
|
|
| def _open_pool(workers: int): |
| from align.verify.lean import ReplPool, find_toolchain |
| from align.verify.lean.repl_pool import ReplLimits |
|
|
| tc = find_toolchain() |
| return ReplPool.from_toolchain(tc, n_workers=workers, limits=ReplLimits(), preflight=True, |
| scratch_dir=str(Path(os.environ.get("TMPDIR", str(REPO / ".tmp"))) / "k5_route_pool")) |
|
|
|
|
| def _plan_with_pool(pool, statement: str, g: dict, v: dict, SR): |
| """SubgoalPlan for a failed node, lemma statements checked to elaborate; None when Lean cannot run it.""" |
| from align.verify.lean.repl_pool import ReplError |
|
|
| def work(w, _): |
| try: |
| probe = SR.repl_probe(w, statement) |
| plan = SR.plan_subgoals(statement, g["lean_code"], v["repl_errors"], probe, parent_id=g["problem_id"], repl_code=v["repl_code"]) |
| for gap in plan.gaps: |
| if gap.lemma is not None and not SR.check_lemma_elaborates(w, plan, gap): |
| gap.lemma, gap.why_not = None, "lemma_does_not_elaborate" |
| return plan |
| except (ReplError, ValueError): |
| return None |
|
|
| r = pool.map(work, [None])[0] |
| return None if isinstance(r, BaseException) else r |
|
|
|
|
| def cmd_assemble(a) -> int: |
| from align.rl.envs import subgoal_repair as SR |
|
|
| rd = Path(a.run_dir, "agentic", f"round{a.round}") |
| plans = _read(rd / "sg_plans.jsonl") |
| sg_gens = _read(rd / "sg_gen.jsonl") |
| sg_vers = _read(rd / "sg_ver.jsonl") |
| ver_by = {C.gen_key(v): v for v in sg_vers} |
| by_pid: dict[str, list[tuple[dict, dict]]] = defaultdict(list) |
| for g in sg_gens: |
| v = ver_by.get(C.gen_key(g)) |
| if v is None: |
| raise ValueError(f"subgoal sample {C.gen_key(g)} is not verified: run the sg stage's verification to completion") |
| by_pid[g["problem_id"]].append((g, v)) |
| |
| parents = {} |
| for f in sorted(Path(a.run_dir, "agentic").glob("gen_r*.jsonl")) + sorted(Path(a.run_dir, "agentic").glob("round*/asm.gen.jsonl")): |
| for g in _read(f): |
| parents[(g["problem_id"], int(g["sample_idx"]), int(g["round"]))] = g |
| pvers = {} |
| for f in sorted(Path(a.run_dir, "agentic").glob("ver_r*.jsonl")) + sorted(Path(a.run_dir, "agentic").glob("round*/asm.ver.jsonl")): |
| for v in _read(f): |
| pvers[(v["problem_id"], int(v["sample_idx"]), int(v["round"]))] = v |
| gens, vers, st = [], [], Counter() |
| for pl in plans: |
| pk = tuple(pl["parent_key"]) |
| pg, pv = parents[pk], pvers[pk] |
| plan = SR.plan_from_dict(pl["plan"]) |
| bodies, sub_tokens = {}, 0 |
| for sid in pl["sg_ids"]: |
| sub_tokens += sum(node_tokens(g) for g, _ in by_pid.get(sid, [])) |
| for j, gap in enumerate(plan.gaps): |
| if gap.lemma is None: |
| continue |
| sid = sg_problem_id(pg, j) |
| ok = [g for g, v in by_pid.get(sid, []) if v["accepted"] and g.get("lean_code")] |
| ok.sort(key=lambda g: len(g["lean_code"])) |
| for g in ok: |
| body = SR.lemma_body(g["lean_code"], gap.lemma.name) |
| if body is not None: |
| bodies[j] = body |
| break |
| gen, ver = make_asm_nodes(plan, pg, pv, bodies, sub_tokens, len(bodies), a.round, a.cap) |
| gens.append(gen) |
| if ver is not None: |
| vers.append(ver) |
| st["assembled"] += 1 |
| st["gaps"] += len(plan.gaps) |
| st["proven"] += len(bodies) |
| st["unchanged_verdict_copied"] += ver is not None |
| _write(rd / "asm.gen.jsonl", gens) |
| if vers: |
| _write(rd / "asm.ver.jsonl", vers) |
| print(json.dumps(dict(st))) |
| return 0 |
|
|
|
|
| def _arm_files(run_dir: str, rounds: int) -> dict[str, tuple[list[str], list[str]]]: |
| rd = Path(run_dir) |
| ag = [rd / "agentic/gen_r0.jsonl"] |
| av = [rd / "agentic/ver_r0.jsonl"] |
| for r in range(1, rounds): |
| ag += [rd / f"agentic/gen_r{r}.jsonl", rd / f"agentic/round{r}/asm.gen.jsonl"] |
| av += [rd / f"agentic/ver_r{r}.jsonl", rd / f"agentic/round{r}/asm.ver.jsonl"] |
| files = { |
| "restart": ([rd / "restart/gen_r0.jsonl"], [rd / "restart/ver_r0.jsonl"]), |
| "correct": ([rd / "restart/gen_r0.jsonl", rd / "correct/gen_r1.jsonl", rd / "correct/gen_r2.jsonl"], |
| [rd / "restart/ver_r0.jsonl", rd / "correct/ver_r1.jsonl", rd / "correct/ver_r2.jsonl"]), |
| "agentic": (ag, av), |
| } |
| return {k: ([str(x) for x in g], [str(x) for x in v]) for k, (g, v) in files.items()} |
|
|
|
|
| def load_arms(run_dir: str, rounds: int = 4, arms: Iterable[str] = ARMS) -> dict[str, dict[str, list[Unit]]]: |
| out = {} |
| for arm, (gf, vf) in _arm_files(run_dir, rounds).items(): |
| if arm not in arms: |
| continue |
| if arm == "correct" and not Path(gf[1]).exists(): |
| continue |
| gens = [g for f in gf for g in _read(f)] |
| vers = [v for f in vf for v in _read(f)] |
| if not gens: |
| continue |
| out[arm] = units_restart(gens, vers) if arm == "restart" else units_chains(gens, vers, arm) |
| return out |
|
|
|
|
| def build_report(arms: dict[str, dict[str, list[Unit]]], budget_spec: str, n_perm: int = 256, seed: int = 0) -> dict: |
| from align.eval.paired import holm, paired_bootstrap_ci, sign_flip_test |
|
|
| if "restart" not in arms: |
| raise ValueError("the restart arm defines the budget and the pairing; run it first") |
| pids = sorted(arms["restart"]) |
| budgets = budget_table(budget_spec, arms["restart"], pids) |
| rep = equalise(arms, budgets, n_perm, seed) |
| out: dict = {"budget": budget_spec, "n_problems": len(pids), "n_perm": n_perm, "seed": seed, "arms": {}, "curves": {}} |
| for arm in arms: |
| out["arms"][arm] = arm_summary(rep[arm][1.0], budgets, 1.0) |
| out["curves"][arm] = {str(f): arm_summary(rep[arm][f], budgets, f)["mean_p_solve"] for f in rep[arm]} |
| fam = [("correct", "restart"), ("agentic", "restart"), ("agentic", "correct")] |
| tests, ps = [], [] |
| for x, y in fam: |
| if x in arms and y in arms: |
| d = [rep[x][1.0][p].p_solve - rep[y][1.0][p].p_solve for p in pids] |
| t = sign_flip_test(d, seed=seed) |
| lo, hi = paired_bootstrap_ci(d, seed=seed) |
| tests.append({"contrast": f"{x} - {y}", "mean_diff": t["observed"], "p": t["p"], "ci95": [lo, hi], "n": t["n"], "n_nonzero": t["n_nonzero"]}) |
| ps.append(t["p"]) |
| if ps: |
| h = holm(ps) |
| for t, adj, rej in zip(tests, h["adjusted"], h["reject"]): |
| t["p_holm"], t["reject_at_0.05"] = adj, rej |
| out["tests"] = tests |
| out["family"] = "K5 pairwise (Holm over the contrasts present); the training-seed null of K0 does not apply: no training here" |
| out["per_problem"] = {arm: {p: rep[arm][1.0][p].p_solve for p in pids} for arm in arms} |
| return out |
|
|
|
|
| def cmd_report(a) -> int: |
| arms = load_arms(a.run_dir, a.rounds) |
| rep = build_report(arms, a.budget, a.n_perm, a.seed) |
| Path(a.run_dir, "k5_report.json").write_text(json.dumps(rep, indent=1)) |
| for arm, s in rep["arms"].items(): |
| print(f"{arm:9s} p_solve@B={s['mean_p_solve']:.4f} spent/B={s['spend_ratio']:.3f} underfilled={s['underfilled_problems']}") |
| for t in rep["tests"]: |
| print(f"{t['contrast']:22s} diff={t['mean_diff']:+.4f} CI[{t['ci95'][0]:+.4f},{t['ci95'][1]:+.4f}] p_holm={t.get('p_holm', float('nan')):.4f}") |
| return 0 |
|
|
|
|
| def cmd_status(a) -> int: |
| arms = load_arms(a.run_dir, a.rounds) |
| pids = sorted(arms["restart"]) |
| budgets = budget_table(a.budget, arms["restart"], pids) |
| for arm, by_p in arms.items(): |
| rep = {p: replay_budget(by_p.get(p, []), budgets[p], a.n_perm, a.seed, p) for p in pids} |
| need = units_needed(by_p, budgets, rep) |
| print(json.dumps({"arm": arm, "underfilled": len(need), "extra_units_total": sum(need.values()), "worst": sorted(need.items(), key=lambda kv: -kv[1])[:5]})) |
| return 0 |
|
|
|
|
| def cmd_ledger(a) -> int: |
| arms = load_arms(a.run_dir, a.rounds) |
| for arm, by_p in arms.items(): |
| rows = [dataclasses.asdict(u) for us in by_p.values() for u in us] |
| _write(Path(a.run_dir, "ledger", f"units_{arm}.jsonl"), rows) |
| print(arm, len(rows), "units over", len(by_p), "problems") |
| return 0 |
|
|
|
|
| def cmd_estimate(a) -> int: |
| cfg = Cfg(run_dir="-", problems=[], n_restart=a.n_restart, n_agentic=a.n_agentic) |
| sets = {"minif2f": a.minif2f, "proofnet": a.proofnet} |
| print(json.dumps(estimate(sets, cfg), indent=1)) |
| return 0 |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| sub = ap.add_subparsers(dest="cmd", required=True) |
|
|
| def common(p, run=True): |
| if run: |
| p.add_argument("--run-dir", required=True) |
|
|
| def cfgargs(p): |
| d = Cfg(run_dir="", problems=[]) |
| common(p) |
| p.add_argument("--problems", nargs="+", required=True) |
| p.add_argument("--model", default=d.model) |
| p.add_argument("--seed", type=int, default=d.seed) |
| p.add_argument("--cap", type=int, default=d.cap) |
| p.add_argument("--traj-cap", type=int, default=d.traj_cap) |
| p.add_argument("--n-restart", type=int, default=d.n_restart) |
| p.add_argument("--n-agentic", type=int, default=d.n_agentic) |
| p.add_argument("--rounds", type=int, default=d.rounds) |
| p.add_argument("--routes", default=d.routes) |
| p.add_argument("--sub-n", type=int, default=d.sub_n) |
| p.add_argument("--sub-cap", type=int, default=d.sub_cap) |
| p.add_argument("--bank-k", type=int, default=d.bank_k) |
| p.add_argument("--budget", default=d.budget) |
| p.add_argument("--workers", type=int, default=d.workers) |
| p.add_argument("--bank-problems") |
| p.add_argument("--bank-gens") |
| p.add_argument("--bank-vers") |
| p.add_argument("--eval-problems", nargs="*", default=[]) |
| p.add_argument("--no-hostguard", dest="hostguard", action="store_false") |
|
|
| p = sub.add_parser("plan", help="print the stages and exact command lines") |
| cfgargs(p) |
| p = sub.add_parser("run", help="execute the stages (resumable)") |
| cfgargs(p) |
| p.add_argument("--only", nargs="*", default=None) |
|
|
| p = sub.add_parser("prep") |
| common(p) |
| p.add_argument("--problems", nargs="+", required=True) |
| p.add_argument("--bank-problems") |
| p.add_argument("--bank-gens") |
| p.add_argument("--bank-vers") |
| p.add_argument("--eval-problems", nargs="*", default=[]) |
| p.add_argument("--bank-k", type=int, default=3) |
|
|
| p = sub.add_parser("route") |
| common(p) |
| p.add_argument("--round", type=int, required=True) |
| p.add_argument("--mode", choices=("subgoal", "feedback"), required=True) |
| p.add_argument("--pairs", nargs=2, action="append", metavar=("VER", "GEN"), required=True) |
| p.add_argument("--workers", type=int, default=2) |
| p.add_argument("--bank-k", type=int, default=3) |
| p.add_argument("--sub-n", type=int, default=4) |
|
|
| p = sub.add_parser("assemble") |
| common(p) |
| p.add_argument("--round", type=int, required=True) |
| p.add_argument("--cap", type=int, default=40960) |
|
|
| for name, fn in (("report", cmd_report), ("status", cmd_status), ("ledger", cmd_ledger)): |
| p = sub.add_parser(name) |
| common(p) |
| p.add_argument("--budget", default="restart_k:8") |
| p.add_argument("--rounds", type=int, default=4) |
| p.add_argument("--n-perm", type=int, default=256) |
| p.add_argument("--seed", type=int, default=0) |
|
|
| p = sub.add_parser("estimate") |
| p.add_argument("--minif2f", type=int, default=226) |
| p.add_argument("--proofnet", type=int, default=185) |
| p.add_argument("--n-restart", type=int, default=32) |
| p.add_argument("--n-agentic", type=int, default=8) |
|
|
| a = ap.parse_args(argv) |
| if a.cmd in ("plan", "run"): |
| cfg = Cfg(run_dir=a.run_dir, problems=a.problems, model=a.model, seed=a.seed, cap=a.cap, traj_cap=a.traj_cap, n_restart=a.n_restart, |
| n_agentic=a.n_agentic, rounds=a.rounds, routes=a.routes, sub_n=a.sub_n, sub_cap=a.sub_cap, bank_k=a.bank_k, budget=a.budget, |
| workers=a.workers, bank_problems=a.bank_problems, bank_gens=a.bank_gens, bank_vers=a.bank_vers, |
| eval_problems=a.eval_problems, hostguard=a.hostguard) |
| stages = build_stages(cfg) |
| if a.cmd == "plan": |
| for st in stages: |
| print(f"# {st.name} [{'GPU' if st.gpu else 'CPU'}]\n{' '.join(shlex.quote(x) for x in _hg(cfg, st))}\n") |
| return 0 |
| return run_stages(cfg, stages, Runner(), only=set(a.only) if a.only else None) |
| return {"prep": cmd_prep, "route": cmd_route, "assemble": cmd_assemble, "report": cmd_report, "status": cmd_status, |
| "ledger": cmd_ledger, "estimate": cmd_estimate}[a.cmd](a) |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
