Blame for code/extras/vllm_rollout/lean_gen.py

"""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  # noqa: E402  (sibling module, stdlib only)

DEFAULT_MODEL = (
    "/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/"
    "dfd02e6271a58375dfbf3ece0175277cf6b6a89a"
)
TEMPERATURE = 1.0  # Goedel inference.py --temp default and the dpskcot setting
TOP_P = 0.95  # hard-coded in Goedel inference.py
Key = tuple  # (problem_id, sample_idx, round, cap_tokens)


# --------------------------------------------------------------------------
# pure logic: seeds, keys, jsonl, planning
# --------------------------------------------------------------------------


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
        # walk back to the previous newline
        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  # 0 if there was no newline at all
        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  # drop anything after the last newline
    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:  # streamed: the pool file is >1 GB
        lines = enumerate(f, 1)
        for ln, raw in lines:
            if not raw.endswith(b"\n"):
                break  # torn tail of a killed writer (the repair truncates it); not a record
            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


# lean_prompts.py module bytes before the prompt-format / temperature work (commits d956862, 901a82c). Goedel
# prompts rendered by the current module are pinned byte-identical to that version by a test
# (tests/test_gen_temperature_kimina.py), so goedel sessions recorded under it stay resumable.
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":  # probe the example block too
            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 []  # nothing on disk was sampled under the old module
    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  # {"problem_id","sample_idx","round"}
    parent_ctx: Optional[dict] = None  # messages / completion_text / compile_code / repl_errors
    resample: bool = False  # fresh initial-style prompt at round >= 1 (no-feedback parents)
    trajectory: Optional[dict] = None  # ledger inputs; None = a root attempt (round 0)


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  # a record without its token ids cannot be counted: lower bound, flagged
    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  # the trajectory token cap, not the context window or --cap-tokens, set max_tokens


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:
        # Goedel inference.py drops prompts longer than 3/4 of max_model_len
        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  # absent for goedel: those records stay byte-identical to before the flag
    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:
            # one write call per group: a kill leaves either a whole group or a torn tail, which
            # repair_torn_tail removes; never an interleaving
            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)  # vLLM 0.26 MaxLoRARanks


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)


# --------------------------------------------------------------------------
# engine side (lazy third-party imports)
# --------------------------------------------------------------------------


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")  # flashinfer's JIT ninja build fails on this box
    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()
    # lock FIRST: with --wait-lock we may sleep for hours behind another job, and a preflight taken
    # before that sleep describes a host that no longer exists
    if inherit_pid is not None:  # a supervisor holds the lock for the whole job (scripts/p2/launch_baseline.py)
        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:
        # explicit outcome, not a silent drop: every failed parent with nothing to show is listed
        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

    # materialise lazily: tokenising 10^5 long correction prompts up front would stall the start
    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)

    # --- bring the engine up -------------------------------------------------
    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:  # informational only; vLLM also logs it
        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()
    # tokenisation self-check: vLLM's own text path must yield OUR prompt ids
    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"  # overwritten on a clean exit; a raise leaves it
    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())