| """vLLM generation driver for Lean proving (experiment P2): Goedel prompts in, Generation records out. |
| |
| OPTIONAL EXTRA -- NOT part of the `align` package. Runs under `.venv-rollout` |
| (vllm 0.26.0, python3.12); never the `align` venv. All third-party imports (vllm, |
| transformers) are lazy inside functions, so the pure planning / resume / record logic |
| below is importable -- and tested -- without them (tests/test_lean_gen_contract.py). |
| |
| .venv-rollout/bin/python extras/vllm_rollout/lean_gen.py \\ |
| --problems P.jsonl --out gen_r0.jsonl --n 4 --cap-tokens 16384 --round 0 --tp 4 |
| |
| # correction round k: parents are round k-1 samples + their verification records |
| ... --round 1 --n 1 --corrections-from verif_r0.jsonl gen_r0.jsonl --out gen_r1.jsonl |
| |
| Design decisions, each tied to a defect it prevents: |
| |
| * Honest caps. `cap_tokens` is the nominal cap every reported number is labelled with; |
| `sampling.max_tokens` is what vLLM was actually given, `min(cap, max_model_len - |
| len(prompt))`. They differ only when the context window, not the cap, ends a sample, and |
| then the record says `context_clipped: true`. Without this a "40960 cap" run would |
| silently be a "40610 cap" run (prompt tokens come out of the same 40960 window). |
| * Resumable and kill-safe. Output is append-only JSONL. A sample is keyed |
| (problem_id, sample_idx, round, cap_tokens); keys already on disk are skipped. Records |
| are fsync'd in groups (`--flush-every` / `--flush-secs`); SIGKILL loses the unflushed |
| group plus whatever is in flight (<= `--inflight` samples, never the finished ones). A |
| torn last line (kill mid-write) is truncated back to the last newline on startup rather |
| than parsed or appended after -- appending after a torn line would glue two records. |
| A resume is only into a file this invocation could have written: records already on disk must |
| carry this run's `model_rev`, the seed `derive_seed` gives for their key, and the registered |
| temperature / top_p, and the manifest's earlier sessions must have used this `lean_prompts.py`; |
| otherwise it RAISES (`resume_conflicts`, `prompt_template_conflicts`) instead of skipping the |
| keys and appending a second policy's samples to the first one's file. |
| * No chunk barrier. vLLM `LLM.generate` returns only when its slowest sample finishes, and |
| sample lengths here span ~1k..16k+ tokens, so a batch barrier idles the GPUs on the |
| longest tail. This driver instead drives `llm.llm_engine.add_request/step` with a refill |
| window, so finished samples are written immediately and the batch stays full. |
| * Per-sample seeds. n is expanded client-side (n=1 per request) so every sample has its |
| own seed in its own record. Seeds are derived from (base seed, problem, sample, round) |
| and NOT the cap, so the same sample at cap 16384 and cap 40960 shares its seed. |
| Bitwise replay is not promised (batch composition changes float reduction order); the |
| seed makes the sampling randomness a recorded input, not a reproducibility claim. |
| * No re-tokenisation downstream. `prompt_ids` come from the model's tokenizer on the chat |
| template text; `completion_ids` are vLLM's `token_ids`, which INCLUDE the terminating |
| stop token (<|im_end|>) when `finish_reason == "stop"`; the SFT shard builder trains on |
| exactly these ids. |
| * Class-B trajectory ledger (docs/SOTA_PLAN.md v3 section 1). Every record carries |
| `trajectory`: {root_trajectory_id "<problem>|<root sample_idx>", root_sample_idx, round, |
| kind attempt|repair|resampled_no_feedback, completion_tokens, prior_completion_tokens, |
| cum_completion_tokens (the whole lineage so far), truncated, budget_clipped, parent_key, |
| partial}. `--trajectory-cap-tokens N` (class B: 40000) enforces a cumulative completion-token |
| cap per trajectory: a repair's `max_tokens` is clipped to the remaining budget |
| (`budget_clipped: true`) and one with none left is NOT generated but listed in |
| `<out>.skipped.jsonl` as `skipped_trajectory_cap`. A failed parent with no compiler error to |
| show (truncated at the cap, no `repl_errors`, or an empty list) is never silently dropped: |
| `--no-feedback-policy skip` (default) writes it to `<out>.skipped.jsonl` as |
| `skipped_no_feedback` with its reason; `resample` draws fresh round-0-style attempts for it |
| (`kind: resampled_no_feedback`, same root, tokens still accumulate). |
| Selection rule: Goedel's, the DEFAULT -- only failed samples of problems with NO accepted |
| sample in round k-1 are repaired (`load_parents`); `--repair-solved` also repairs failed samples |
| of already-solved problems (needed to account every trajectory, not to match Goedel). |
| Backward compatible: pre-ledger files have no `trajectory`; lineage is reconstructed from the |
| `parent` chain when the earlier rounds' files are passed in `--corrections-from`, otherwise the |
| entry is `partial` (a lower bound) and enforcing a cap on it refuses. Resume only keys on |
| (problem, sample, round, cap), so old files resume unchanged. |
| * Un-merged LoRA serving. `--lora-adapter DIR` (a PEFT dir: adapter_config.json + |
| adapter_model.safetensors) serves the adapter through vLLM `enable_lora` / `LoRARequest` |
| on the base `--model`, instead of the bf16 merge that drops small delta elements. The adapter's |
| sha256 is in the manifest session, in every record (`adapter`), and in `model_rev` |
| (`<base>+lora-<sha12>`), so a resume cannot mix it with the base file. GPU smoke command: |
| see extras/vllm_rollout/README.md. |
| * Exit status. 0 only when every planned sample was generated; SIGTERM/SIGINT exit 143 so a |
| supervisor can never mistake a partial sample set for a finished one (U12-2). |
| * Unsupported configuration raises: max_model_len above the model's position limit, |
| tp above the visible GPUs, a round-0 prompt that cannot fit, missing verification |
| records for a parent, and so on. The only skips are correction prompts that cannot fit |
| the context and parents with nothing to correct; both are counted in the manifest and |
| listed in `<out>.skipped.jsonl`. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import dataclasses |
| import hashlib |
| import json |
| import os |
| import signal |
| import subprocess |
| import sys |
| import threading |
| import time |
| from typing import Iterator, Optional |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| REPO_ROOT = os.path.dirname(os.path.dirname(HERE)) |
| sys.path.insert(0, HERE) |
|
|
| import lean_prompts as LP |
|
|
| DEFAULT_MODEL = ( |
| "/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/" |
| "dfd02e6271a58375dfbf3ece0175277cf6b6a89a" |
| ) |
| TEMPERATURE = 1.0 |
| TOP_P = 0.95 |
| Key = tuple |
|
|
|
|
| |
| |
| |
|
|
|
|
| def derive_seed(base_seed: int, problem_id: str, sample_idx: int, round_: int) -> int: |
| """Per-sample seed in [0, 2**31). Independent of the cap on purpose (see module doc).""" |
| h = hashlib.sha256(f"{base_seed}|{problem_id}|{sample_idx}|{round_}".encode()).digest() |
| return int.from_bytes(h[:8], "big") % (2**31) |
|
|
|
|
| def rec_key(rec: dict) -> Key: |
| return (rec["problem_id"], int(rec["sample_idx"]), int(rec["round"]), int(rec["cap_tokens"])) |
|
|
|
|
| def repair_torn_tail(path: str) -> int: |
| """Truncate a JSONL file back to its last complete line. Returns bytes removed. |
| |
| A kill between `write` and `fsync` can leave a partial last line. Reading it would |
| raise (or worse, parse as a different record); appending after it would fuse two |
| records into one unparseable line and lose the NEW record too. |
| """ |
| if not os.path.exists(path): |
| return 0 |
| size = os.path.getsize(path) |
| if size == 0: |
| return 0 |
| with open(path, "rb+") as f: |
| f.seek(-1, os.SEEK_END) |
| if f.read(1) == b"\n": |
| return 0 |
| |
| pos = size |
| chunk = 1 << 16 |
| last_nl = -1 |
| while pos > 0 and last_nl < 0: |
| start = max(0, pos - chunk) |
| f.seek(start) |
| buf = f.read(pos - start) |
| i = buf.rfind(b"\n") |
| if i >= 0: |
| last_nl = start + i |
| pos = start |
| keep = last_nl + 1 |
| f.truncate(keep) |
| f.flush() |
| os.fsync(f.fileno()) |
| return size - keep |
|
|
|
|
| def read_jsonl(path: str) -> Iterator[dict]: |
| with open(path, encoding="utf-8") as f: |
| for ln, line in enumerate(f, 1): |
| if not line.strip(): |
| continue |
| try: |
| yield json.loads(line) |
| except json.JSONDecodeError as e: |
| raise ValueError(f"{path}:{ln}: not valid JSON ({e})") from e |
|
|
|
|
| def load_done_keys(out_path: str) -> set: |
| if not os.path.exists(out_path): |
| return set() |
| return {rec_key(r) for r in read_jsonl(out_path)} |
|
|
|
|
| def resume_state(out_path: str, repair: bool = True) -> tuple[set, int]: |
| """(done_keys, torn_bytes_removed) for an output file that a killed run may have left torn. |
| |
| The repair MUST precede the read: found by the real SIGKILL test, where the first |
| version read the file first and died on its own torn last line, i.e. the run could not |
| be resumed from exactly the crash the resume logic exists for. `repair=False` (dry run) |
| must not touch the file, so it ignores an unterminated last line instead. |
| """ |
| if repair: |
| torn = repair_torn_tail(out_path) |
| return load_done_keys(out_path), torn |
| if not os.path.exists(out_path): |
| return set(), 0 |
| with open(out_path, "rb") as f: |
| data = f.read() |
| cut = data.rfind(b"\n") + 1 |
| keys = {rec_key(json.loads(line)) for line in data[:cut].decode("utf-8").splitlines() if line.strip()} |
| return keys, len(data) - cut |
|
|
|
|
| def resume_conflicts(out_path: str, model_rev: str, base_seed: int, limit: int = 5, |
| adapter_sha: Optional[str] = None, temperature: float = TEMPERATURE, top_p: float = TOP_P, |
| prompt_format: str = "goedel") -> list[str]: |
| """Records already in `out_path` that THIS invocation could not have produced (empty list = resumable). |
| |
| THE MIXED-OUTPUT DEFECT (R2 review). A sample is keyed (problem, sample, round, cap) and nothing |
| else, so resuming into a file written by a DIFFERENT model (the post-SFT run pointed at the |
| baseline's --out), a different `--model-rev`, or a different `--seed` silently skips every key that |
| is already there and appends the rest: one output file, two policies, each row honestly labelled |
| with its own `model_rev` and nothing complaining. A baseline/candidate comparison built on it |
| compares the model with a blend of itself. Same call, same answer on a clean resume: the check is |
| over the records' own provenance (`model_rev`, `sampling.seed`), which a faithful resume reproduces |
| exactly (`derive_seed` is a pure function of (seed, problem, sample, round)). |
| """ |
| bad: list[str] = [] |
| n_bad = 0 |
| if not os.path.exists(out_path): |
| return bad |
| with open(out_path, "rb") as f: |
| lines = enumerate(f, 1) |
| for ln, raw in lines: |
| if not raw.endswith(b"\n"): |
| break |
| if not raw.strip(): |
| continue |
| r = json.loads(raw) |
| want = derive_seed(base_seed, r["problem_id"], int(r["sample_idx"]), int(r["round"])) |
| why = None |
| if r.get("model_rev") != model_rev: |
| why = f"model_rev {r.get('model_rev')!r} != {model_rev!r}" |
| elif r.get("sampling", {}).get("seed") != want: |
| why = f"sampling.seed {r.get('sampling', {}).get('seed')!r} != {want} (--seed {base_seed})" |
| elif (r.get("adapter") or {}).get("sha256") != adapter_sha: |
| why = f"adapter {(r.get('adapter') or {}).get('sha256')!r} != {adapter_sha!r}" |
| elif r.get("sampling", {}).get("temperature") != temperature or r.get("sampling", {}).get("top_p") != top_p: |
| why = f"sampling {r.get('sampling')!r} differs from temperature={temperature}, top_p={top_p}" |
| elif r.get("prompt_format", "goedel") != prompt_format: |
| why = f"prompt_format {r.get('prompt_format', 'goedel')!r} != {prompt_format!r}" |
| if why: |
| n_bad += 1 |
| if len(bad) < limit: |
| bad.append(f"line {ln} {rec_key(r)}: {why}") |
| if n_bad > len(bad): |
| bad.append(f"... and {n_bad - len(bad)} more") |
| return bad |
|
|
|
|
| |
| |
| |
| LEGACY_GOEDEL_PROMPTS_SHA256 = frozenset({"700ab6c7fc1f1cb342ed2969bd5ddbcf1ffe0f8e070b9332fb1247fad57677af"}) |
|
|
| _PROBE_PROBLEMS = ( |
| {"problem_id": "probe_a", "formal_statement": "theorem probe_a (x : ℕ) (h : x = 2) : x + x = 4 := by sorry", |
| "informal_prefix": "/-- Show that x + x = 4 when x = 2. -/\n"}, |
| {"problem_id": "probe_b", "formal_statement": "theorem probe_b (a b : ℝ) (h₀ : 0 < a) (h₁ : a < b) : a ^ 2 < b ^ 2 := by", |
| "informal_prefix": None}, |
| ) |
|
|
|
|
| def prompt_render_sha(prompt_format: str, tokenizer) -> str: |
| """sha256 of the prompts this invocation would actually SEND (rendered through the tokenizer's chat |
| template) for fixed probe problems, plus one goedel correction prompt. Unlike the module-bytes hash, |
| an edit to lean_prompts.py that leaves the rendered bytes unchanged does not block a resume, and an |
| edit that changes them does.""" |
| h = hashlib.sha256(prompt_format.encode()) |
| for prob in _PROBE_PROBLEMS: |
| if prompt_format == "goedel_bank": |
| prob = dict(prob, bank_examples=[{"statement": "theorem probe_ex : 1 = 1 := by", "proof": " rfl"}]) |
| text, _ = LP.initial_prompt_for(prompt_format, prob, tokenizer) |
| h.update(text.encode()) |
| if prompt_format == "goedel": |
| p0 = _PROBE_PROBLEMS[0] |
| _, msgs = LP.initial_prompt_for("goedel", p0, tokenizer) |
| text, _ = LP.correction_prompt(p0["formal_statement"], msgs, "```lean4\nexact rfl\n```", "exact rfl", |
| [{"severity": "error", "pos": {"line": 1, "column": 0}, |
| "endPos": {"line": 1, "column": 9}, "data": "probe error"}], |
| tokenizer, 1) |
| h.update(text.encode()) |
| return h.hexdigest() |
|
|
|
|
| def prompt_template_conflicts(manifest_path: str, render_sha256: str, out_path: str, prompt_format: str = "goedel") -> list[str]: |
| """Earlier sessions of this output file that sampled under DIFFERENT prompt bytes. |
| |
| The policy is conditioned on the rendered prompt bytes; an edit between two sessions that changes them |
| leaves a file whose halves answer different prompts while every key looks consistent. Sessions record |
| `prompt_render_sha256`. Older sessions recorded only the module bytes (`lean_prompts_sha256`); they are |
| compatible iff that sha is the pinned legacy goedel module and this run is goedel.""" |
| if not os.path.exists(manifest_path) or not os.path.exists(out_path) or os.path.getsize(out_path) == 0: |
| return [] |
| with open(manifest_path, encoding="utf-8") as f: |
| doc = json.load(f) |
| bad = set() |
| for s in doc.get("sessions", []): |
| rs = s.get("prompt_render_sha256") |
| if rs is not None: |
| if rs != render_sha256: |
| bad.add(f"render:{rs}") |
| elif not (prompt_format == "goedel" and s.get("lean_prompts_sha256") in LEGACY_GOEDEL_PROMPTS_SHA256): |
| bad.add(f"module:{s.get('lean_prompts_sha256')}") |
| return sorted(bad) |
|
|
|
|
| PROBLEM_FIELDS = ("problem_id", "source", "split", "formal_statement", "informal_prefix", "statement_sha") |
|
|
|
|
| def load_problems(path: str, limit: Optional[int] = None) -> list[dict]: |
| rows = [] |
| seen = set() |
| for r in read_jsonl(path): |
| missing = [k for k in PROBLEM_FIELDS if k not in r] |
| if missing: |
| raise ValueError(f"{path}: problem {r.get('problem_id')!r} lacks fields {missing}") |
| if r["problem_id"] in seen: |
| raise ValueError(f"{path}: duplicate problem_id {r['problem_id']!r}") |
| seen.add(r["problem_id"]) |
| rows.append(r) |
| if limit is not None and len(rows) >= limit: |
| break |
| if not rows: |
| raise ValueError(f"{path}: no problems") |
| return rows |
|
|
|
|
| @dataclasses.dataclass |
| class Spec: |
| """One sample to generate. `parent_ctx` is set only for correction rounds.""" |
|
|
| problem: dict |
| sample_idx: int |
| round: int |
| parent: Optional[dict] = None |
| parent_ctx: Optional[dict] = None |
| resample: bool = False |
| trajectory: Optional[dict] = None |
|
|
|
|
| def plan_initial(problems: list[dict], n: int, round_: int = 0) -> Iterator[Spec]: |
| for p in problems: |
| for i in range(n): |
| yield Spec(p, i, round_) |
|
|
|
|
| def record_trajectory(g: dict, index: dict) -> tuple[int, int, bool]: |
| """(root_sample_idx, cumulative completion tokens through `g`, partial) for a generation record. |
| |
| Trusts the record's own `trajectory` ledger when present; otherwise rebuilds it from the |
| `parent` chain through `index` ((problem, sample_idx, round) -> record, every round the caller |
| passed). A chain that leaves the provided files is `partial`: the root is the furthest ancestor |
| named and the token count a lower bound -- reported, never guessed.""" |
| t = g.get("trajectory") |
| if t and "cum_completion_tokens" in t: |
| return int(t["root_sample_idx"]), int(t["cum_completion_tokens"]), bool(t.get("partial", False)) |
| ids = g.get("completion_ids") |
| n = len(ids) if ids is not None else 0 |
| miss = ids is None |
| if int(g["round"]) == 0: |
| return int(g["sample_idx"]), n, miss |
| par = g.get("parent") |
| if not par: |
| return int(g["sample_idx"]), n, True |
| nxt = index.get((par["problem_id"], int(par["sample_idx"]), int(par["round"]))) |
| if nxt is None: |
| return int(par["sample_idx"]), n, True |
| root, cum, partial = record_trajectory(nxt, index) |
| return root, cum + n, partial or miss |
|
|
|
|
| def load_parents_ex(pairs: list[tuple[str, str]], round_k: int, repair_solved: bool = False) -> tuple[list[dict], dict, list[dict]]: |
| """Join verification and generation records of round k-1 -> parent contexts. |
| |
| Returns (parents, stats, no_feedback). Mirrors Goedel's `load_data_for_correction` (the DEFAULT): |
| * only samples that FAILED, and only for problems with NO accepted sample in the |
| previous round, are corrected (`repair_solved=True` lifts the second condition); |
| * every failed sample is a parent (not just one per problem). |
| Failed parents with nothing to show (no `repl_errors`/`repl_code` keys, an empty list, or a |
| sample truncated at its cap) are returned in `no_feedback` with a `reason` instead of being |
| dropped: Goedel would send an empty "error message:" block; here the caller decides |
| (`--no-feedback-policy`). Each parent / no-feedback entry carries `traj` = (root_sample_idx, |
| cumulative completion tokens, partial) for the class-B ledger. |
| """ |
| gens: dict = {} |
| vers: dict = {} |
| index: dict = {} |
| for vpath, gpath in pairs: |
| for g in read_jsonl(gpath): |
| index[(g["problem_id"], int(g["sample_idx"]), int(g["round"]))] = g |
| if int(g["round"]) == round_k - 1: |
| gens[rec_key(g)] = g |
| for v in read_jsonl(vpath): |
| if int(v["round"]) == round_k - 1: |
| vers[rec_key(v)] = v |
| if not gens: |
| raise ValueError(f"no generation records with round == {round_k - 1} in {pairs}") |
| unverified = sorted(set(gens) - set(vers)) |
| if unverified: |
| raise ValueError( |
| f"{len(unverified)} round-{round_k - 1} generation records have no verification record " |
| f"(first: {unverified[0]}); verify every parent before correcting" |
| ) |
| solved = {k[0] for k, v in vers.items() if v["accepted"] and k in gens} |
| stats = {"parents_total": len(gens), "problems_solved_prev_round": len(solved), |
| "skipped_solved_problem": 0, "skipped_accepted": 0, "no_repl_errors": 0, "no_error_messages": 0, |
| "no_feedback_truncated": 0} |
| parents, no_feedback = [], [] |
| for k in sorted(gens): |
| g, v = gens[k], vers[k] |
| if v["accepted"]: |
| stats["skipped_accepted"] += 1 |
| continue |
| if k[0] in solved and not repair_solved: |
| stats["skipped_solved_problem"] += 1 |
| continue |
| traj = record_trajectory(g, index) |
| has_keys = isinstance(v.get("repl_errors"), list) and isinstance(v.get("repl_code"), str) |
| if not has_keys or not v["repl_errors"]: |
| stats["no_repl_errors" if not has_keys else "no_error_messages"] += 1 |
| if g.get("finish_reason") == "length": |
| stats["no_feedback_truncated"] += 1 |
| reason = "truncated" |
| else: |
| reason = "no_repl_keys" if not has_keys else "no_error_messages" |
| no_feedback.append({"gen": g, "ver": v, "reason": reason, "traj": traj}) |
| continue |
| parents.append({"gen": g, "ver": v, "traj": traj}) |
| return parents, stats, no_feedback |
|
|
|
|
| def load_parents(pairs: list[tuple[str, str]], round_k: int) -> tuple[list[dict], dict]: |
| """Back-compatible (parents, stats) view of `load_parents_ex` with Goedel's selection rule.""" |
| parents, stats, _ = load_parents_ex(pairs, round_k) |
| return parents, stats |
|
|
|
|
| def _parent_ref(g: dict) -> dict: |
| return {"problem_id": g["problem_id"], "sample_idx": int(g["sample_idx"]), "round": int(g["round"])} |
|
|
|
|
| def _traj_inputs(traj: tuple, g: dict, kind: str) -> dict: |
| root, cum, partial = traj |
| return {"root_sample_idx": root, "prior_completion_tokens": cum, "partial": partial, "kind": kind, |
| "parent_key": [g["problem_id"], int(g["sample_idx"]), int(g["round"])]} |
|
|
|
|
| def plan_corrections(problems: list[dict], parents: list[dict], n: int, round_k: int) -> Iterator[Spec]: |
| """n correction samples per failed parent; sample_idx = parent.sample_idx * n + i is unique |
| within (problem, round) and stable under resume.""" |
| by_id = {p["problem_id"]: p for p in problems} |
| for par in parents: |
| g, v = par["gen"], par["ver"] |
| prob = by_id.get(g["problem_id"]) |
| if prob is None: |
| raise ValueError(f"parent problem {g['problem_id']!r} is not in --problems") |
| ctx = { |
| "messages": g["messages"], |
| "completion_text": g["completion_text"], |
| "compile_code": v["repl_code"], |
| "repl_errors": v["repl_errors"], |
| } |
| for i in range(n): |
| yield Spec(prob, int(g["sample_idx"]) * n + i, round_k, _parent_ref(g), ctx, |
| trajectory=_traj_inputs(par["traj"], g, "repair")) |
|
|
|
|
| def plan_resamples(problems: list[dict], no_feedback: list[dict], n: int, round_k: int) -> Iterator[Spec]: |
| """`--no-feedback-policy resample`: n fresh round-0-style attempts for each failed parent that has |
| no compiler error to show. Same root trajectory; the failed parent's tokens stay on the ledger.""" |
| by_id = {p["problem_id"]: p for p in problems} |
| for nf in no_feedback: |
| g = nf["gen"] |
| prob = by_id.get(g["problem_id"]) |
| if prob is None: |
| raise ValueError(f"parent problem {g['problem_id']!r} is not in --problems") |
| for i in range(n): |
| yield Spec(prob, int(g["sample_idx"]) * n + i, round_k, _parent_ref(g), None, resample=True, |
| trajectory=_traj_inputs(nf["traj"], g, "resampled_no_feedback")) |
|
|
|
|
| def no_feedback_entry(nf: dict, round_k: int, policy: str, cap: Optional[int] = None) -> dict: |
| """The `<out>.skipped.jsonl` line for a failed parent with no compiler error (never a silent drop).""" |
| g = nf["gen"] |
| root, cum, partial = nf["traj"] |
| return {"key": [g["problem_id"], int(g["sample_idx"]), round_k, cap], |
| "outcome": "skipped_no_feedback" if policy == "skip" else "resampled_no_feedback", |
| "reason": nf["reason"], "root_trajectory_id": f"{g['problem_id']}|{root}", |
| "cum_completion_tokens": cum, "partial": partial, "round": round_k} |
|
|
|
|
| def skip_outcome(why: str) -> str: |
| if why.startswith("skipped_trajectory_cap"): |
| return "skipped_trajectory_cap" |
| if why.startswith("prompt_"): |
| return "skipped_prompt_too_long" |
| return why |
|
|
|
|
| def append_skipped(path: str, entry: dict) -> bool: |
| """Append `entry` to the skipped ledger unless an identical line exists (a resume replans the same skips).""" |
| line = json.dumps(entry, sort_keys=True) |
| if os.path.exists(path): |
| with open(path, encoding="utf-8") as f: |
| for old in f: |
| try: |
| if json.dumps(json.loads(old), sort_keys=True) == line: |
| return False |
| except json.JSONDecodeError: |
| continue |
| with open(path, "a", encoding="utf-8") as f: |
| f.write(line + "\n") |
| return True |
|
|
|
|
| @dataclasses.dataclass |
| class Prepared: |
| spec: Spec |
| key: Key |
| prompt_text: str |
| prompt_ids: list |
| messages: list |
| max_tokens: int |
| seed: int |
| context_clipped: bool |
| budget_clipped: bool = False |
|
|
|
|
| def prepare(spec: Spec, tokenizer, cap: int, max_model_len: int, base_seed: int, error_thres: bool = True, |
| traj_cap: Optional[int] = None, prompt_format: str = "goedel"): |
| """Build the prompt and decide the per-request budget. |
| |
| Returns (Prepared, None) or (None, skip_reason). A round-0 prompt that does not fit is a |
| configuration error and raises; a correction prompt that does not fit (parent output + |
| errors can exceed the window) is skipped, with the reason recorded by the caller. |
| """ |
| fs = spec.problem["formal_statement"] |
| remaining = None |
| if traj_cap is not None and spec.trajectory is not None: |
| if spec.trajectory.get("partial"): |
| raise ValueError(f"{spec.problem['problem_id']}: the trajectory lineage is incomplete (pre-ledger files); " |
| "pass every earlier round's generation file in --corrections-from to enforce " |
| "--trajectory-cap-tokens") |
| remaining = traj_cap - int(spec.trajectory["prior_completion_tokens"]) |
| if remaining < 1: |
| return None, (f"skipped_trajectory_cap: {spec.trajectory['prior_completion_tokens']} completion tokens " |
| f"already spent of {traj_cap}") |
| if spec.round == 0 or spec.resample: |
| text, messages = (LP.initial_prompt(fs, tokenizer) if prompt_format == "goedel" |
| else LP.initial_prompt_for(prompt_format, spec.problem, tokenizer)) |
| else: |
| if prompt_format != "goedel": |
| raise ValueError(f"--prompt-format {prompt_format} has no published correction prompt; round >= 1 needs goedel") |
| c = spec.parent_ctx |
| text, messages = LP.correction_prompt( |
| fs, c["messages"], c["completion_text"], c["compile_code"], c["repl_errors"], |
| tokenizer, spec.round, error_thres=error_thres) |
| ids = list(tokenizer.encode(text)) |
| room = max_model_len - len(ids) |
| if len(ids) > max_model_len * 3 // 4 or room < 1: |
| |
| if spec.round == 0: |
| raise ValueError(f"round-0 prompt of {spec.problem['problem_id']} is {len(ids)} tokens; " |
| f"max_model_len={max_model_len} cannot hold it") |
| return None, f"prompt_{len(ids)}_tokens_over_3/4_of_max_model_len_{max_model_len}" |
| limit = cap if remaining is None else min(cap, remaining) |
| mt = min(limit, room) |
| return Prepared(spec, (spec.problem["problem_id"], spec.sample_idx, spec.round, cap), text, ids, messages, |
| mt, derive_seed(base_seed, spec.problem["problem_id"], spec.sample_idx, spec.round), |
| room < limit, remaining is not None and remaining < min(cap, room)), None |
|
|
|
|
| def build_record(pr: Prepared, completion_ids: list, completion_text: str, finish_reason: str, cap: int, |
| model_path: str, model_rev: str, extra: Optional[dict] = None, temperature: float = TEMPERATURE, |
| top_p: float = TOP_P, prompt_format: str = "goedel") -> dict: |
| if finish_reason not in ("stop", "length"): |
| raise ValueError(f"unsupported finish_reason {finish_reason!r} (contract: 'stop'|'length')") |
| if finish_reason == "length" and len(completion_ids) != pr.max_tokens: |
| raise ValueError(f"{pr.key}: finish_reason=length with {len(completion_ids)} tokens != max_tokens {pr.max_tokens}") |
| ti = pr.spec.trajectory or {"root_sample_idx": pr.spec.sample_idx, "prior_completion_tokens": 0, |
| "kind": "attempt", "partial": False, "parent_key": None} |
| n_comp = len(completion_ids) |
| traj = { |
| "root_trajectory_id": f"{pr.spec.problem['problem_id']}|{ti['root_sample_idx']}", |
| "problem_id": pr.spec.problem["problem_id"], "root_sample_idx": int(ti["root_sample_idx"]), |
| "round": pr.spec.round, "kind": ti["kind"], "parent_key": ti["parent_key"], |
| "completion_tokens": n_comp, "prior_completion_tokens": int(ti["prior_completion_tokens"]), |
| "cum_completion_tokens": int(ti["prior_completion_tokens"]) + n_comp, |
| "truncated": finish_reason == "length", "budget_clipped": pr.budget_clipped, "partial": bool(ti["partial"]), |
| } |
| code, err = LP.lean_code_for_record(pr.spec.problem["formal_statement"], completion_text, prompt_format) |
| rec = { |
| "problem_id": pr.spec.problem["problem_id"], |
| "sample_idx": pr.spec.sample_idx, |
| "round": pr.spec.round, |
| "cap_tokens": cap, |
| "prompt_ids": pr.prompt_ids, |
| "completion_ids": completion_ids, |
| "prompt_text": pr.prompt_text, |
| "completion_text": completion_text, |
| "finish_reason": finish_reason, |
| "lean_code": code, |
| "lean_code_error": err, |
| "messages": pr.messages, |
| "model_path": model_path, |
| "model_rev": model_rev, |
| "sampling": {"temperature": temperature, "top_p": top_p, "top_k": 0, "n": 1, |
| "max_tokens": pr.max_tokens, "seed": pr.seed}, |
| "context_clipped": pr.context_clipped, |
| "parent": pr.spec.parent, |
| "trajectory": traj, |
| } |
| if prompt_format != "goedel": |
| rec["prompt_format"] = prompt_format |
| if extra: |
| rec.update(extra) |
| return rec |
|
|
|
|
| def sampled_logprobs_extra(c) -> dict: |
| """`{"sampling_logprobs": [...]}`: the logprob of each SAMPLED token, from vLLM's `logprobs=0` output |
| (a list of {token_id: Logprob} per position). Raises if any position lacks the sampled token, so a record |
| never carries a misaligned behaviour distribution.""" |
| toks = list(c.token_ids) |
| lps = c.logprobs |
| if lps is None or len(lps) != len(toks): |
| raise RuntimeError(f"engine returned {None if lps is None else len(lps)} logprob positions for {len(toks)} tokens") |
| out = [] |
| for t, d in zip(toks, lps): |
| if t not in d: |
| raise RuntimeError(f"sampled token {t} missing from its logprob entry") |
| out.append(float(d[t].logprob)) |
| return {"sampling_logprobs": out} |
|
|
|
|
| def model_rev_of(model_path: str, override: Optional[str] = None) -> str: |
| p = os.path.realpath(model_path) |
| if override: |
| return override |
| if os.path.basename(os.path.dirname(p)) == "snapshots": |
| return os.path.basename(p) |
| raise ValueError(f"cannot infer a revision from {model_path!r} (not an HF snapshot dir); pass --model-rev") |
|
|
|
|
| class AppendWriter: |
| """Append-only JSONL with grouped fsync. `add` buffers; `flush` writes + fsyncs the group.""" |
|
|
| def __init__(self, path: str, flush_every: int = 8, flush_secs: float = 5.0): |
| self.torn_bytes_removed = repair_torn_tail(path) |
| self.f = open(path, "ab") |
| self.flush_every, self.flush_secs = flush_every, flush_secs |
| self.buf: list[bytes] = [] |
| self.last = time.monotonic() |
| self.n_written = 0 |
|
|
| def add(self, rec: dict) -> None: |
| self.buf.append((json.dumps(rec, ensure_ascii=False, separators=(",", ":")) + "\n").encode("utf-8")) |
| if len(self.buf) >= self.flush_every or time.monotonic() - self.last >= self.flush_secs: |
| self.flush() |
|
|
| def flush(self) -> None: |
| if self.buf: |
| |
| |
| self.f.write(b"".join(self.buf)) |
| self.f.flush() |
| os.fsync(self.f.fileno()) |
| self.n_written += len(self.buf) |
| self.buf = [] |
| self.last = time.monotonic() |
|
|
| def close(self) -> None: |
| self.flush() |
| self.f.close() |
|
|
|
|
| def build_trajectory_ledger(gens: list[dict], vers: list[dict], skipped: list[dict]) -> dict: |
| """Per root trajectory: rounds, cumulative completion tokens, and a final state. Pure. |
| |
| final = accepted_round<r> (earliest accepted attempt on the trajectory) | the skip outcome |
| (`skipped_no_feedback`, `skipped_trajectory_cap`, ...) | `unrepaired` (failed, no later record).""" |
| index = {(g["problem_id"], int(g["sample_idx"]), int(g["round"])): g for g in gens} |
| accepted = {rec_key(v): bool(v["accepted"]) for v in vers} |
| led: dict = {} |
| for g in gens: |
| root, cum, partial = record_trajectory(g, index) |
| tid = f"{g['problem_id']}|{root}" |
| e = led.setdefault(tid, {"rounds_seen": set(), "cum_completion_tokens": 0, "accepted_round": None, |
| "skips": [], "partial": False}) |
| e["rounds_seen"].add(int(g["round"])) |
| e["cum_completion_tokens"] = max(e["cum_completion_tokens"], cum) |
| e["partial"] |= partial |
| if accepted.get(rec_key(g)): |
| r = int(g["round"]) |
| e["accepted_round"] = r if e["accepted_round"] is None else min(e["accepted_round"], r) |
| for sk in skipped: |
| tid = sk.get("root_trajectory_id") |
| if tid in led: |
| led[tid]["skips"].append(sk.get("outcome")) |
| led[tid]["cum_completion_tokens"] = max(led[tid]["cum_completion_tokens"], int(sk.get("cum_completion_tokens", 0))) |
| out = {} |
| for tid, e in led.items(): |
| if e["accepted_round"] is not None: |
| final = f"accepted_round{e['accepted_round']}" |
| elif e["skips"]: |
| final = e["skips"][-1] |
| else: |
| final = "unrepaired" |
| out[tid] = {"final": final, "rounds": len(e["rounds_seen"]), "cum_completion_tokens": e["cum_completion_tokens"], |
| "partial": e["partial"]} |
| return out |
|
|
|
|
| LORA_RANKS = (1, 8, 16, 32, 64, 128, 256, 320, 512) |
|
|
|
|
| def _sha256_file(path: str) -> str: |
| h = hashlib.sha256() |
| with open(path, "rb") as f: |
| for chunk in iter(lambda: f.read(1 << 20), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def adapter_provenance(path: str) -> dict: |
| """Validate a PEFT adapter dir and hash it (config + weights): the identity recorded for an un-merged run.""" |
| cfg_p = os.path.join(path, "adapter_config.json") |
| if not os.path.isfile(cfg_p): |
| raise ValueError(f"{path}: no adapter_config.json (not a PEFT adapter directory)") |
| wfile = next((f for f in ("adapter_model.safetensors", "adapter_model.bin") |
| if os.path.isfile(os.path.join(path, f))), None) |
| if wfile is None: |
| raise ValueError(f"{path}: no adapter_model.safetensors / adapter_model.bin") |
| with open(cfg_p, encoding="utf-8") as f: |
| cfg = json.load(f) |
| if "r" not in cfg: |
| raise ValueError(f"{cfg_p}: no LoRA rank `r`") |
| wsha, csha = _sha256_file(os.path.join(path, wfile)), _sha256_file(cfg_p) |
| return {"path": os.path.abspath(path), "weights_file": wfile, "weights_sha256": wsha, "config_sha256": csha, |
| "sha256": hashlib.sha256((csha + wsha).encode()).hexdigest(), "rank": int(cfg["r"]), |
| "lora_alpha": cfg.get("lora_alpha"), "target_modules": cfg.get("target_modules"), |
| "base_model_name_or_path": cfg.get("base_model_name_or_path")} |
|
|
|
|
| def lora_engine_kwargs(pv: dict) -> dict: |
| """`LLM(...)` kwargs that enable one adapter slot of at least the adapter's rank.""" |
| ranks = [r for r in LORA_RANKS if r >= pv["rank"]] |
| if not ranks: |
| raise ValueError(f"adapter rank {pv['rank']} exceeds vLLM's largest supported max_lora_rank {LORA_RANKS[-1]}") |
| return {"enable_lora": True, "max_loras": 1, "max_lora_rank": ranks[0]} |
|
|
|
|
| def rev_with_adapter(base_rev: str, pv: Optional[dict], explicit: bool) -> str: |
| """model_rev that cannot collide with the base model's: `<base>+lora-<sha12>` (unless the user fixed one).""" |
| return base_rev if (pv is None or explicit) else f"{base_rev}+lora-{pv['sha256'][:12]}" |
|
|
|
|
| def make_lora_request(pv: Optional[dict], name: str, request_cls, int_id: int = 1): |
| """`vllm.lora.request.LoRARequest(name, id, path)`; the class is injected so this stays importable without vllm.""" |
| if pv is None: |
| return None |
| return request_cls(lora_name=name, lora_int_id=int_id, lora_path=pv["path"]) |
|
|
|
|
| def adapter_record_extra(pv: Optional[dict], name: str) -> Optional[dict]: |
| if pv is None: |
| return None |
| return {"adapter": {"name": name, "sha256": pv["sha256"], "weights_sha256": pv["weights_sha256"], |
| "config_sha256": pv["config_sha256"], "rank": pv["rank"], "path": pv["path"], |
| "serving": "unmerged_vllm_lora"}} |
|
|
|
|
| def final_status(stop_flag: bool, inflight_left: int, exhausted: bool) -> str: |
| """`interrupted` only if a stop request actually left planned work undone.""" |
| return "interrupted" if stop_flag and (inflight_left > 0 or not exhausted) else "complete" |
|
|
|
|
| def exit_code_for(status: str) -> int: |
| """0 only for a complete sample set. 143 = 128 + SIGTERM for an interruption (U12-2).""" |
| return {"complete": 0, "interrupted": 143}.get(status, 1) |
|
|
|
|
| def write_manifest(path: str, session: dict) -> None: |
| """sessions[] list: a resumed run appends a session instead of overwriting history.""" |
| doc = {"sessions": []} |
| if os.path.exists(path): |
| with open(path, encoding="utf-8") as f: |
| doc = json.load(f) |
| sessions = [s for s in doc["sessions"] if s.get("session_id") != session["session_id"]] |
| sessions.append(session) |
| doc["sessions"] = sessions |
| tmp = path + ".tmp" |
| with open(tmp, "w", encoding="utf-8") as f: |
| json.dump(doc, f, indent=1, ensure_ascii=False) |
| f.flush() |
| os.fsync(f.fileno()) |
| os.replace(tmp, path) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def prepare_env() -> None: |
| """Environment that must be right BEFORE `import vllm` (see rollout_worker.py).""" |
| d = os.environ.setdefault |
| d("HF_HOME", "/apps/development/rlhf/hf_cache") |
| d("HF_HUB_CACHE", "/apps/development/rlhf/hf_cache/hub") |
| d("HF_HUB_DISABLE_XET", "1") |
| d("TMPDIR", os.path.join(REPO_ROOT, ".tmp")) |
| d("XDG_CACHE_HOME", os.path.join(REPO_ROOT, ".cache")) |
| d("TRITON_CACHE_DIR", os.path.join(REPO_ROOT, ".cache", "triton")) |
| d("TORCHINDUCTOR_CACHE_DIR", os.path.join(REPO_ROOT, ".cache", "torchinductor")) |
| d("VLLM_USE_FLASHINFER_SAMPLER", "0") |
| d("TOKENIZERS_PARALLELISM", "false") |
| d("PYTHONUNBUFFERED", "1") |
| for var in ("TMPDIR", "XDG_CACHE_HOME", "TRITON_CACHE_DIR", "TORCHINDUCTOR_CACHE_DIR", "HF_HOME", "HF_HUB_CACHE"): |
| if not os.path.realpath(os.environ[var]).startswith("/apps/"): |
| raise RuntimeError(f"{var}={os.environ[var]!r} is not under /apps; / has 13 GB and /tmp is noexec tmpfs") |
| os.makedirs(os.environ[var], exist_ok=True) |
|
|
|
|
| def preflight_and_lock(label: str, wait: bool, inherit_pid: int | None = None) -> dict: |
| """Host preflight (>=45 GB MemAvailable, load < 2x cores, disk queue < 64) + the repo |
| heavy-job flock. `align.utils.guard` is stdlib-only at import (align/__init__ is lazy).""" |
| sys.path.insert(0, REPO_ROOT) |
| from align.utils import guard |
|
|
| t0 = time.monotonic() |
| |
| |
| if inherit_pid is not None: |
| guard.inherit_heavy_job_lock(inherit_pid) |
| else: |
| guard.acquire_heavy_job_lock(label, blocking=wait) |
| st = guard.preflight_host_check() |
| return {"host_state": st, "lock_wait_s": round(time.monotonic() - t0, 1)} |
|
|
|
|
| class GpuMonitor(threading.Thread): |
| """Peak per-GPU memory.used from nvidia-smi (what the driver sees, not torch's view).""" |
|
|
| def __init__(self, period: float = 10.0): |
| super().__init__(daemon=True) |
| self.period, self.peak, self.stop_ev = period, {}, threading.Event() |
|
|
| def run(self) -> None: |
| while not self.stop_ev.is_set(): |
| try: |
| out = subprocess.run( |
| ["nvidia-smi", "--query-gpu=index,memory.used", "--format=csv,noheader,nounits"], |
| capture_output=True, text=True, timeout=20).stdout |
| for line in out.strip().splitlines(): |
| i, m = (x.strip() for x in line.split(",")) |
| self.peak[i] = max(self.peak.get(i, 0), int(m)) |
| except Exception: |
| pass |
| self.stop_ev.wait(self.period) |
|
|
|
|
| def main(argv: Optional[list] = None) -> int: |
| ap = argparse.ArgumentParser(description="vLLM generation driver for Lean proving (P2).") |
| ap.add_argument("--problems", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--model", default=DEFAULT_MODEL) |
| ap.add_argument("--model-rev", default=None) |
| ap.add_argument("--n", type=int, required=True, help="samples per problem (round 0) or per failed parent (round>=1)") |
| ap.add_argument("--cap-tokens", type=int, required=True) |
| ap.add_argument("--round", type=int, default=0) |
| ap.add_argument("--corrections-from", nargs=2, action="append", metavar=("VERIF_JSONL", "GEN_JSONL"), default=[]) |
| ap.add_argument("--tp", type=int, default=4) |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--max-model-len", type=int, default=40960) |
| ap.add_argument("--gpu-mem-util", type=float, default=0.90) |
| ap.add_argument("--max-num-seqs", type=int, default=128) |
| ap.add_argument("--inflight", type=int, default=None, help="refill window (default max_num_seqs + 32)") |
| ap.add_argument("--flush-every", type=int, default=8) |
| ap.add_argument("--flush-secs", type=float, default=5.0) |
| ap.add_argument("--limit-problems", type=int, default=None) |
| ap.add_argument("--verify-tokenization", type=int, default=4, |
| help="send N prompts to vLLM as TEXT and assert its token ids equal ours") |
| ap.add_argument("--trajectory-cap-tokens", type=int, default=None, metavar="N", |
| help="cumulative completion-token cap per trajectory (class B: 40000); a repair is clipped to the " |
| "remaining budget, one with none left is skipped as skipped_trajectory_cap") |
| ap.add_argument("--no-feedback-policy", choices=("skip", "resample"), default="skip", |
| help="failed parents with no compiler error: record as skipped_no_feedback (default) or draw " |
| "fresh round-0-style attempts on the same trajectory") |
| ap.add_argument("--repair-solved", action="store_true", |
| help="also repair failed samples of problems that already have an accepted sample " |
| "(default = Goedel's rule: repair only unsolved problems)") |
| ap.add_argument("--lora-adapter", default=None, metavar="DIR", |
| help="serve this PEFT adapter UN-MERGED through vLLM enable_lora on top of --model") |
| ap.add_argument("--lora-name", default="adapter") |
| ap.add_argument("--no-error-thres", action="store_true", help="Goedel --error_thres False (all errors, full spans)") |
| ap.add_argument("--wait-lock", action="store_true", help="block on the heavy-job lock instead of refusing") |
| 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("--dry-run", action="store_true", help="plan + build prompts only; no GPU, no lock") |
| ap.add_argument("--label", default="p2-lean-gen") |
| ap.add_argument("--temperature", type=float, default=TEMPERATURE, |
| help=f"sampling temperature (default {TEMPERATURE}: Goedel protocol; Kimina card recommends 0.6)") |
| ap.add_argument("--top-p", type=float, default=TOP_P, help=f"nucleus top_p (default {TOP_P})") |
| ap.add_argument("--prompt-format", choices=LP.PROMPT_FORMATS, default="goedel", |
| help="goedel (default, byte-identical prompts) or kimina (docs/KIMINA_BASE.md; round 0 / resample only)") |
| ap.add_argument("--record-logprobs", action="store_true", |
| help="RL rollouts (scripts/p2/rl_lean.py): store the sampler's logprob of every sampled token as " |
| "`sampling_logprobs` (the BEHAVIOUR policy). Meaningful only with top_p == 1.0 and top_k " |
| "disabled; the RL loop refuses records that say otherwise.") |
| a = ap.parse_args(argv) |
| if not a.temperature > 0: |
| ap.error("--temperature must be > 0") |
| if not 0 < a.top_p <= 1: |
| ap.error("--top-p must be in (0, 1]") |
| if a.prompt_format != "goedel" and a.round >= 1: |
| ap.error(f"--prompt-format {a.prompt_format} has no correction prompt; use --round 0") |
|
|
| if a.round == 0 and a.corrections_from: |
| ap.error("--corrections-from needs --round >= 1") |
| if a.round >= 1 and not a.corrections_from: |
| ap.error("--round >= 1 needs --corrections-from VERIF GEN") |
| if a.n < 1 or a.cap_tokens < 1 or a.tp < 1: |
| ap.error("--n, --cap-tokens, --tp must be positive") |
| if a.trajectory_cap_tokens is not None and a.trajectory_cap_tokens < 1: |
| ap.error("--trajectory-cap-tokens must be positive") |
| if a.round == 0 and (a.repair_solved or a.no_feedback_policy != "skip" or a.trajectory_cap_tokens is not None): |
| ap.error("--repair-solved / --no-feedback-policy / --trajectory-cap-tokens only apply to --round >= 1") |
| adapter = adapter_provenance(a.lora_adapter) if a.lora_adapter else None |
| lora_kw = lora_engine_kwargs(adapter) if adapter else {} |
| rev = rev_with_adapter(model_rev_of(a.model, a.model_rev), adapter, explicit=bool(a.model_rev)) |
| cfg = json.load(open(os.path.join(a.model, "config.json"))) |
| if a.max_model_len > cfg["max_position_embeddings"]: |
| raise ValueError(f"--max-model-len {a.max_model_len} exceeds the model's max_position_embeddings " |
| f"{cfg['max_position_embeddings']} (no rope scaling configured); refusing") |
| inflight = a.inflight or a.max_num_seqs + 32 |
|
|
| prepare_env() |
| problems = load_problems(a.problems, a.limit_problems) |
| parent_stats = None |
| no_feedback: list = [] |
| skipped_path = a.out + ".skipped.jsonl" |
| if a.round == 0: |
| specs: Iterator[Spec] = plan_initial(problems, a.n, 0) |
| else: |
| parents, parent_stats, no_feedback = load_parents_ex( |
| [tuple(p) for p in a.corrections_from], a.round, repair_solved=a.repair_solved) |
| specs = plan_corrections(problems, parents, a.n, a.round) |
| if a.no_feedback_policy == "resample": |
| specs = _chain(list(specs), plan_resamples(problems, no_feedback, a.n, a.round)) |
|
|
| os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True) |
| done, torn_at_start = resume_state(a.out, repair=not a.dry_run) |
| conflicts = resume_conflicts(a.out, rev, a.seed, adapter_sha=adapter["sha256"] if adapter else None, |
| temperature=a.temperature, top_p=a.top_p, prompt_format=a.prompt_format) |
| if conflicts: |
| raise ValueError( |
| f"--out {a.out} already holds records this invocation (--model-rev {rev}, --seed {a.seed}) did not " |
| "produce; resuming would mix two policies / seed streams in one file. Use a fresh --out. First: " |
| + "; ".join(conflicts)) |
| from transformers import AutoTokenizer |
|
|
| tok = AutoTokenizer.from_pretrained(a.model) |
|
|
| prompts_sha = hashlib.sha256(open(os.path.join(HERE, "lean_prompts.py"), "rb").read()).hexdigest() |
| render_sha = prompt_render_sha(a.prompt_format, tok) |
| stale_prompts = prompt_template_conflicts(a.out + ".manifest.json", render_sha, a.out, a.prompt_format) |
| if stale_prompts: |
| raise ValueError(f"--out {a.out} holds samples written under different prompt bytes " |
| f"({stale_prompts}); the current rendering is {render_sha}. Use a fresh --out.") |
|
|
| session = { |
| "session_id": f"{int(time.time())}-{os.getpid()}", |
| "started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), |
| "ended_at": None, "status": "running", |
| "argv": sys.argv, "args": vars(a), |
| "model_path": a.model, "model_rev": rev, |
| "lean_prompts_sha256": prompts_sha, |
| "prompt_render_sha256": render_sha, |
| "temperature": a.temperature, "top_p": a.top_p, "prompt_format": a.prompt_format, |
| "prompt_format_sha256": hashlib.sha256(f"{prompts_sha}|{a.prompt_format}".encode()).hexdigest(), |
| "parent_stats": parent_stats, "already_done_at_start": len(done), |
| "adapter": adapter, "no_feedback_parents": len(no_feedback), |
| "no_feedback_policy": a.no_feedback_policy if a.round else None, |
| "trajectory_cap_tokens": a.trajectory_cap_tokens, |
| } |
| manifest_path = a.out + ".manifest.json" |
| if a.round and a.no_feedback_policy == "skip" and not a.dry_run: |
| |
| for nf in no_feedback: |
| if append_skipped(skipped_path, no_feedback_entry(nf, a.round, "skip", a.cap_tokens)): |
| session["skipped_no_feedback"] = session.get("skipped_no_feedback", 0) + 1 |
|
|
| |
| def todo() -> Iterator[Prepared]: |
| for s in specs: |
| key = (s.problem["problem_id"], s.sample_idx, s.round, a.cap_tokens) |
| if key in done: |
| session["skipped_already_done"] = session.get("skipped_already_done", 0) + 1 |
| continue |
| pr, why = prepare(s, tok, a.cap_tokens, a.max_model_len, a.seed, not a.no_error_thres, |
| traj_cap=a.trajectory_cap_tokens, prompt_format=a.prompt_format) |
| if pr is None: |
| outcome = skip_outcome(why) |
| session[outcome] = session.get(outcome, 0) + 1 |
| if not a.dry_run: |
| ent = {"key": list(key), "reason": why, "outcome": outcome} |
| if s.trajectory: |
| ent["root_trajectory_id"] = f"{s.problem['problem_id']}|{s.trajectory['root_sample_idx']}" |
| ent["cum_completion_tokens"] = s.trajectory["prior_completion_tokens"] |
| append_skipped(skipped_path, ent) |
| continue |
| session["budget_clipped"] = session.get("budget_clipped", 0) + int(pr.budget_clipped) |
| yield pr |
|
|
| if a.dry_run: |
| n = clipped = 0 |
| lens = [] |
| for pr in todo(): |
| n += 1 |
| clipped += pr.context_clipped |
| lens.append(len(pr.prompt_ids)) |
| session.update(status="dry-run", planned=n, context_clipped=clipped, |
| prompt_tokens_max=max(lens, default=0), prompt_tokens_mean=sum(lens) / max(1, len(lens))) |
| print(json.dumps(session, ensure_ascii=False, indent=1, default=str)) |
| return 0 |
|
|
| session.update(preflight_and_lock(a.label, a.wait_lock, a.inherit_lock)) |
| write_manifest(manifest_path, session) |
|
|
| |
| import vllm |
| from vllm import LLM, SamplingParams |
| from vllm.sampling_params import RequestOutputKind |
| import torch |
|
|
| if a.tp > torch.cuda.device_count(): |
| raise RuntimeError(f"--tp {a.tp} but only {torch.cuda.device_count()} CUDA devices visible") |
| session.update(vllm_version=vllm.__version__, torch_version=torch.__version__) |
| mon = GpuMonitor() |
| mon.start() |
| t_load = time.perf_counter() |
| llm = LLM(model=a.model, tokenizer=a.model, dtype="bfloat16", tensor_parallel_size=a.tp, |
| enable_prefix_caching=True, gpu_memory_utilization=a.gpu_mem_util, max_model_len=a.max_model_len, |
| max_num_seqs=a.max_num_seqs, seed=a.seed, disable_log_stats=False, **lora_kw) |
| lora_req = None |
| if adapter: |
| from vllm.lora.request import LoRARequest |
|
|
| lora_req = make_lora_request(adapter, a.lora_name, LoRARequest) |
| rec_extra = adapter_record_extra(adapter, a.lora_name) |
| if a.record_logprobs and a.top_p != 1.0: |
| raise ValueError("--record-logprobs needs --top-p 1.0: vLLM returns raw logprobs, which equal the sampling " |
| "(behaviour) distribution only without nucleus truncation") |
| session["engine_load_s"] = round(time.perf_counter() - t_load, 1) |
| eng = llm.llm_engine |
| try: |
| cc = eng.vllm_config.cache_config |
| session["kv_cache"] = {"num_gpu_blocks": cc.num_gpu_blocks, "block_size": cc.block_size, |
| "tokens": (cc.num_gpu_blocks or 0) * cc.block_size, |
| "gpu_memory_utilization": cc.gpu_memory_utilization} |
| except Exception as e: |
| session["kv_cache"] = {"error": repr(e)} |
|
|
| stop = {"flag": False} |
|
|
| def _sig(signum, _frm): |
| stop["flag"] = True |
|
|
| signal.signal(signal.SIGTERM, _sig) |
| signal.signal(signal.SIGINT, _sig) |
|
|
| writer = AppendWriter(a.out, a.flush_every, a.flush_secs) |
| session["torn_bytes_removed"] = torn_at_start + writer.torn_bytes_removed |
| gen_iter = todo() |
| |
| if a.verify_tokenization: |
| probe = [] |
| for pr in gen_iter: |
| probe.append(pr) |
| if len(probe) >= a.verify_tokenization: |
| break |
| gen_iter = _chain(probe, gen_iter) |
| ok = _verify_tokenization(eng, probe, SamplingParams) |
| session["tokenization_check"] = {"n": len(probe), "ok": ok} |
| if not ok: |
| raise RuntimeError("vLLM's tokenisation of the prompt TEXT differs from tokenizer.encode(); " |
| "prompt_ids in records would not be what the model saw") |
|
|
| inflight_map: dict[str, Prepared] = {} |
| n_done = n_tokens = n_len = 0 |
| t0 = time.perf_counter() |
| last_log = t0 |
| exhausted = False |
| status = "crashed" |
| try: |
| while True: |
| while not stop["flag"] and not exhausted and len(inflight_map) < inflight: |
| pr = next(gen_iter, None) |
| if pr is None: |
| exhausted = True |
| break |
| sp = SamplingParams(n=1, temperature=a.temperature, top_p=a.top_p, max_tokens=pr.max_tokens, |
| seed=pr.seed, output_kind=RequestOutputKind.FINAL_ONLY, |
| **({"logprobs": 0} if a.record_logprobs else {})) |
| rid = f"{pr.key[0]}|{pr.key[1]}|{pr.key[2]}|{pr.key[3]}" |
| eng.add_request(rid, {"prompt_token_ids": pr.prompt_ids}, sp, lora_request=lora_req) |
| inflight_map[rid] = pr |
| if not inflight_map: |
| break |
| if stop["flag"]: |
| eng.abort_request(list(inflight_map)) |
| break |
| for o in eng.step(): |
| if not o.finished: |
| continue |
| pr = inflight_map.pop(_match_rid(o.request_id, inflight_map)) |
| c = o.outputs[0] |
| if list(o.prompt_token_ids) != pr.prompt_ids: |
| raise RuntimeError(f"{pr.key}: engine prompt ids differ from the recorded prompt ids") |
| writer.add(build_record(pr, list(c.token_ids), c.text, c.finish_reason, a.cap_tokens, a.model, rev, |
| extra=({**(rec_extra or {}), **sampled_logprobs_extra(c)} if a.record_logprobs |
| else rec_extra), |
| temperature=a.temperature, top_p=a.top_p, prompt_format=a.prompt_format)) |
| n_done += 1 |
| n_tokens += len(c.token_ids) |
| n_len += c.finish_reason == "length" |
| now = time.perf_counter() |
| if now - last_log > 30: |
| print(f"[lean_gen] done={n_done} inflight={len(inflight_map)} tokens={n_tokens} " |
| f"avg_tok/s={n_tokens / (now - t0):.0f}", flush=True) |
| last_log = now |
| status = final_status(stop["flag"], len(inflight_map), exhausted) |
| finally: |
| writer.close() |
| mon.stop_ev.set() |
| wall = time.perf_counter() - t0 |
| session.update( |
| ended_at=time.strftime("%Y-%m-%dT%H:%M:%S%z"), status=status, |
| generate_wall_s=round(wall, 1), samples_written=n_done, tokens_generated=n_tokens, |
| finish_length=n_len, tokens_per_s_end_to_end=round(n_tokens / wall, 1) if wall else None, |
| gpu_mem_used_peak_mib=mon.peak, aborted_inflight=len(inflight_map), |
| ) |
| write_manifest(manifest_path, session) |
| print(json.dumps({k: session[k] for k in ("status", "samples_written", "tokens_generated", "generate_wall_s", |
| "tokens_per_s_end_to_end")})) |
| return exit_code_for(status) |
|
|
|
|
| def _chain(head: list, tail: Iterator): |
| yield from head |
| yield from tail |
|
|
|
|
| def _match_rid(rid: str, inflight: dict) -> str: |
| """vLLM may decorate request ids; accept an exact match or a unique prefix match.""" |
| if rid in inflight: |
| return rid |
| cand = [k for k in inflight if rid.startswith(k) or k.startswith(rid)] |
| if len(cand) != 1: |
| raise RuntimeError(f"cannot match engine request id {rid!r} to an in-flight request ({len(cand)} candidates)") |
| return cand[0] |
|
|
|
|
| def _verify_tokenization(eng, probe: list, SamplingParams) -> bool: |
| from vllm.sampling_params import RequestOutputKind |
|
|
| want = {} |
| for i, pr in enumerate(probe): |
| rid = f"tokcheck-{i}" |
| eng.add_request(rid, pr.prompt_text, SamplingParams(max_tokens=1, temperature=0.0, |
| output_kind=RequestOutputKind.FINAL_ONLY)) |
| want[rid] = pr.prompt_ids |
| ok = True |
| left = set(want) |
| while left: |
| for o in eng.step(): |
| if o.finished: |
| rid = _match_rid(o.request_id, {k: 0 for k in left}) |
| left.discard(rid) |
| ok &= list(o.prompt_token_ids) == want[rid] |
| return ok |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
