# vLLM rollout (optional extra)

This directory is **not** part of the `align` package and nothing under `align/`
imports it (`tests/test_dependency_firewall.py::test_extras_are_not_reachable_from_core`).
The default GRPO rollout is the dependency-free native sampler,
`align/rollout/native.py`. Use this only when generation throughput is the
bottleneck and a dedicated engine GPU is available.

## Runs in its own venv

`rollout_worker.py` and `rollout_pilot.py` run under `.venv-rollout`
(vllm 0.26.0, ray, pyarrow, torch cu130), **never** the `align` venv, whose
only dependencies are torch, triton, numpy and safetensors. Build it with
`uv venv .venv-rollout` plus `torch==2.11.0 triton==3.6.0 vllm==0.26.0 ray pyarrow`.

## How it talks to the trainer

Only through `align/rollout/protocol.py` (stdlib-only by contract): one JSON
file per message in a spool directory, atomically written, with a heartbeat
file for liveness. The worker puts `rollout_worker.py`'s repo root on
`sys.path` and imports that single module; `import align` and
`import align.rollout` are lazy / stdlib-only so no `align` dependency reaches
this venv.

Trainer side:

```python
from align.rollout.client import make_rollout_client, spawn_worker

proc = spawn_worker(spool, gpu=3, log_path=".../rollout_worker.log")
client = make_rollout_client("spool", spool=spool, timeout_s=3600)
```

`spawn_worker` launches `extras/vllm_rollout/rollout_worker.py` with the
required environment (flashinfer sampler off, caches under `/apps`, torchrun
variables scrubbed, own process group so teardown reaps the engine core).

## Logprob basis caveat

The worker requests `SamplingParams(logprobs=0)`, which in vLLM are the
engine's **raw** logprobs (before temperature). The native backend returns
`log_softmax(logits / T)` untruncated, which is what the GRPO module recomputes
teacher-forced. The two coincide at `temperature == 1.0` (the stage-b setting).
For `T != 1` the spool backend's logprobs differ from the training-side basis
by the temperature rescaling, so rho will not be ~1 on-policy there; use
`T = 1.0` with this worker, or set vLLM's `logprobs_mode` to a processed mode
and accept that top-p/top-k then enter the logprob too.

Weight sync is filesystem-based: a changed `weights_dir` tears the engine down
and rebuilds it (~60 s), so sync every `sync_every` optimizer steps and let
`use_rho_correction=True` absorb the off-policyness.

## Lean proving generation driver (experiment P2): `lean_gen.py` + `lean_prompts.py`

Generates Goedel-Prover-V2 proof attempts with vLLM and writes the shared Generation
record JSONL (one line per sample). Runs in `.venv-rollout`; `lean_prompts.py` and the
pure half of `lean_gen.py` are stdlib-only at import and are unit-tested from the align
venv (`tests/test_lean_gen_contract.py`, 36 tests, 15/15 hand mutants killed).

```bash
cd /apps/development/rlhf/align && export TMPDIR=$PWD/.tmp/p2/<you>
# round 0
.venv-rollout/bin/python extras/vllm_rollout/lean_gen.py --problems problems.jsonl \
    --out gen_r0.jsonl --n 4 --cap-tokens 16384 --tp 4 --seed 0 --wait-lock
# correction round 1 (parents = failed round-0 samples; needs a verification JSONL, see below)
.venv-rollout/bin/python extras/vllm_rollout/lean_gen.py --problems problems.jsonl \
    --out gen_r1.jsonl --n 1 --cap-tokens 16384 --round 1 \
    --corrections-from verif_r0.jsonl gen_r0.jsonl --wait-lock
# plan + prompts only, no GPU / lock
... --dry-run
```

Rerunning the same command resumes. `<out>.manifest.json` has one `sessions[]` entry per
launch (model rev, vllm/torch versions, args, load/generate wall, tokens, finish counts, KV
blocks, peak GPU memory); `<out>.skipped.jsonl` lists correction prompts that did not fit.

**Fidelity.** `python extras/vllm_rollout/lean_prompts.py fidelity --model <snapshot> --out r.json`
(rollout venv) imports the reference `src/utils.py` and asserts identical prompt strings and
token ids on all 244 minif2f problems, identical `extrac_code` + `problem_check` on 63 synthetic
outputs, `get_error_str` / `generate_correction_prompt` on 60 random error sets, and
`compile.py:handle` vs `compile_view`.

**What a record means.**
* `cap_tokens` is the label every reported number carries. `sampling.max_tokens` is what vLLM
  got: `min(cap, max_model_len - len(prompt))`. They differ only when the context window ends the
  sample (`context_clipped: true`); at `--cap-tokens 40960 --max-model-len 40960` every sample is
  clipped by its prompt length (~250-600 tokens). `--max-model-len` above the model's 40960
  position limit raises.
* `completion_ids` are vLLM's `token_ids`: they end with the stop token (`<|im_end|>`) when
  `finish_reason == "stop"`. Train on them as they are; do not re-tokenize `completion_text`.
* `messages` are the PROMPT messages (what the chat template rendered; no completion), so a
  correction round's history is `parent.messages + [assistant: parent.completion_text, user: errors]`
  exactly as Goedel builds it. `lean_code` is Goedel's `full_code` or null, with the reason in
  `lean_code_error` (extra field).
* Sampling: temperature 1.0, top_p 0.95, no top_k (Goedel's `SamplingParams`; the model's
  `generation_config.json` says top_k=20 / T=0.6 but explicit params override it). One seed per
  sample, derived from (seed, problem, sample_idx, round); not the cap.

**Correction rounds need two extra fields in the VERIFICATION record** (the shared contract does
not carry them): `repl_errors` (list of Lean REPL messages with `severity == "error"`, each
`{pos:{line,column}, endPos:{line,column}|null, data}`) and `repl_code` (the exact text those
positions index into; for the Goedel flow that is `lean_prompts.compile_view(lean_code)`, i.e. the
code with `import` / `set_option` / `open` lines removed). Without them `<error>` markers land on
the wrong lines, so a parent lacking them is skipped and counted (`no_repl_errors`), never guessed.

**Throughput/memory choices.** `gpu_memory_utilization=0.90` (20.7 GB of the A10G's 22.5 GiB;
measured peak `nvidia-smi` use 21.7-22.5 GB, i.e. 0.5-1.3 GB of slack for NCCL/allocator; 0.92 was
not tried), `max_num_seqs=128`, prefix caching on, CUDA graphs on, FLASH_ATTN backend (works on
sm_86), `VLLM_USE_FLASHINFER_SAMPLER=0`. KV pool: 28,533 blocks x 16 = 456,528 tokens (TP=4: each
GPU holds 2 of 8 KV heads, 36 KB/token). The driver steps the engine itself (`add_request`/`step`)
with a refill window instead of `LLM.generate` chunks, so a 16k-token straggler does not idle the
GPUs. Measured on Goedel-Prover-V2-8B, TP=4, seeded sampling (minif2f in file order, T=1.0, top_p=0.95):

| run | cap | samples | mean / median / p90 completion tokens | finish=length | end-to-end tok/s | tok/s while >=64 running |
|---|---|---|---|---|---|---|
| 32 problems x n=4, max_num_seqs 128 | 16384 | 128 | 5066 / 3242 / 13993 | 10/128 (7.8%) | 918 | 1602 |
| 64 problems x n=4, max_num_seqs 256 | 16384 | 256 | 4996 / 3033 / 13641 | 21/256 (8.2%) | 1043 | 1363 (KV pool 100% full at peak) |
| 8 problems x n=2, max_num_seqs 128 | 40960 (context-clipped, max_tokens ~40.6-40.8k) | 16 | 4450 / 2662 / 7339 | 0/16 | 256 (tail of 1 sequence) | n/a (max 16 running) |

Per-sequence decode speed is ~25 tok/s at 85-128 concurrent sequences and ~74 tok/s alone: TP=4
over PCIe (no NVLink) is allreduce-bound, so aggregate throughput saturates near 1.4-2.2k tok/s
beyond ~100 concurrent sequences, and 256 sequences only fill the KV pool. Keep `--max-num-seqs 128`.

### Class-B trajectories, exit status, un-merged adapter (lane V2)

* `--trajectory-cap-tokens 40000` enforces the class-B cumulative cap; every record carries a
  `trajectory` ledger (root id, round, cumulative completion tokens, truncation, budget clipping);
  failed parents with no compiler error are `skipped_no_feedback` in `<out>.skipped.jsonl`
  (or re-sampled with `--no-feedback-policy resample`). Goedel's rule (repair only problems
  without an accepted sample) is the default; `--repair-solved` lifts it. Pass every earlier
  round's `VERIF GEN` pair in `--corrections-from` so round-2+ lineage is complete (otherwise the
  entry is `partial` and a cap refuses). `build_trajectory_ledger()` aggregates a final state per root.
* Exit status: 0 only for a complete sample set, 143 on SIGTERM/SIGINT. `pipeline_gen_verify.py`
  also checks the manifest (exit 4 if not `complete`) and reaps the generator's process group,
  sweeps `/dev/shm` and waits for the GPUs to drain after every generator exit.
* Un-merged adapter: `--lora-adapter <peft dir> [--lora-name adapter]` serves it via vLLM
  `enable_lora`/`LoRARequest` (no bf16 merge). Adapter sha256 goes in the manifest, each record and
  `model_rev` (`<base>+lora-<sha12>`). GPU smoke (run after the chains exit, takes the heavy lock):

```bash
.venv-rollout/bin/python extras/vllm_rollout/lean_gen.py --problems <5 problems>.jsonl \
    --out .tmp/v2_smoke/gen_lora.jsonl --n 2 --cap-tokens 2048 --tp 1 --max-model-len 8192 \
    --lora-adapter <adapter_dir> --wait-lock
# expect: manifest sessions[-1].adapter.sha256 set, vllm engine log shows LoRA enabled, every
# record has "adapter", exit 0; then rerun the same command (must resume, 0 new samples).
```