feat(sc): gate-authoritative token capture on the Gym tokidcap stack

docs/guides/nano-swe-token-capture.md

@@ -0,0 +1,227 @@
+# Nano SWE RL with Gate-Authoritative Token Capture
+
+A reproducible 6-node recipe that runs agentic SWE RL on
+Nemotron-3-Nano-30B-A3B with **exact-token capture** enabled: the vLLM worker
+stages each model call's token delta durably into the TransferQueue data
+plane, the Gym gate serves verified prefix token ids back on every follow-up
+call (token-in), and the trainer consumes rows rebuilt from the staged deltas.
+No token echo over HTTP, no re-tokenization of agent history — the tokens the
+engine sampled are byte-for-byte the tokens the trainer sees.
+
+It builds directly on the [Nano SWE TransferQueue
+recipe](nano-swe-transferqueue.md); read that first for the cluster shape,
+`swe_nano.env` setup, and the SingleController constraints. This guide covers
+only what token capture adds.
+
+## Verified result
+
+Run on 6 GB200 NVL72 nodes:
+
+| | |
+|---|---|
+| Entrypoint | `examples/run_grpo_single_controller.py` |
+| Config | `examples/configs/ultra/nano_swe_teacher_sc.yaml` |
+| Model | `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16` |
+| Shape | train 4 nodes (TP2·PP2·CP4, EP1) / gen 2 nodes (vLLM TP2 → 4 engines) |
+| Capture flags | `token_capture.enabled=true`, `NG_TIC_FP_CANONICAL=1` |
+| Outcome | 15/15 steps, `token_in_rate ≈ 0.9999`, zero `empty_manifest`, zero `capture_failed`, `valid_seqs = GBS` |
+| Launched via | `swe_nano_sc_capture.sh` (batch) / `swe_nano_sc_capture_interactive.sh` (attach + iterate) |
+
+## How it works
+
+Legacy SWE runs recover training tokens by echoing token ids in every model
+response and re-tokenizing the agent's rendered history each turn. Both are
+lossy for a reasoning model — the chat template strips `<think>` content from
+history, and re-tokenizing an assistant turn can split differently than the
+tokens the model actually sampled. Token capture replaces that path.
+
+### Anatomy of one rollout
+
+The call flow for a single SWE rollout, end to end:
+
+1. **Dispatch.** The `SingleControllerActor`'s rollout pump pulls a prompt
+   group from the dataset and hands it to the NemoGym environment actor,
+   which POSTs `/run` to the `swe_agents_train` Gym server. The run body
+   carries a fresh `ng_rollout_id`; the gate **registers** the rollout
+   (control plane) before any model call happens.
+2. **Sandbox + agent.** The SWE harness materializes the SWE-bench instance
+   in a sandbox and starts the OpenHands agent. The rollout id is snapshotted
+   into the agent's config so every LLM request it makes can be attributed
+   back to this rollout (`metadata.ng_rollout_id` on the request body).
+3. **Agent turn loop.** Each turn, OpenHands POSTs `/v1/chat/completions`
+   with the *full rendered history* to the policy model server, which runs
+   in **gate mode**:
+   - the gate admits the call (rejects unregistered rollouts), fingerprints
+     the incoming history, and resolves which previously served call this
+     one continues (its *parent*) in the lineage index;
+   - on a match it attaches the parent's **exact cumulative token ids**
+     (`required_prefix_token_ids`) plus a capture context
+     (`rollout_id`, `call_id`, `parent_call_id`, `prev_len`), and forwards
+     the request to a vLLM engine (rollout→engine affinity keeps a
+     rollout's calls on the engine holding its KV cache);
+   - no match (a true root, or a rewritten/condensed history) starts a new
+     chain — a fallback, never an error.
+4. **Generate + stage.** The vLLM worker splices the supplied prefix
+   verbatim, renders only the new tail through the chat template, and
+   generates. Before acknowledging the call it **stages the call's token
+   delta** — `rendered_prompt[prev_len:] + generated` ids, a loss mask
+   (0.0 on carried prompt, 1.0 on generated), and per-token logprobs — to
+   the TransferQueue staging partition (synchronous `tq_put`: bytes are
+   durable before the response leaves the worker). The response back to the
+   gate carries only text plus token-light `CommitCoords` (~4 B/token).
+5. **Commit.** The gate ingests the coords into its lineage index (this
+   *is* the authoritative commit — the next turn's parent resolution
+   depends on it), strips them, and returns a plain OpenAI-shaped
+   completion to the agent. Steps 3-5 repeat for every tool call the agent
+   makes (tool execution happens agent-side between turns).
+6. **Verify + seal.** When the agent finishes, the harness runs the
+   SWE-bench verifier to score the patch (reward 0/1). The `/run` response
+   returns the reward plus a **token-free `RolloutReceipt`**: the ordered
+   manifest of committed calls (call ids, staging keys, digests, weight
+   versions) and a `terminal_call_id` naming the chain that is the
+   training row.
+7. **Finalize.** Trainer-side, the `BlackboxFinalizer` fetches the staged
+   deltas the manifest names from TQ, re-verifies each (digest, lengths,
+   mask shape, weight-version tags), linearizes the terminal chain into one
+   exact token row, and publishes it to the training partition — a rejected
+   rollout becomes a masked placeholder so the GRPO group keeps its shape.
+   Staged rows are cleared after publish.
+8. **Train.** Once a global batch of rows is buffered, the SC takes an
+   optimizer step and syncs weights to the engines; the bumped
+   `weight_version` is stamped on subsequent calls so refit boundaries are
+   visible in the data.
+
+The token path in that flow, compressed:
+
+```
+agent (nv-OpenHands)                     gate (vllm_model, gate mode)
+  /run body carries ng_rollout_id  ───►  registers the rollout, admits calls,
+                                         fingerprints incoming history →
+                                         resolves the parent call → sends the
+                                         parent's exact prefix token ids
+                                              │  required_prefix_token_ids
+                                              ▼
+vLLM worker: splices the prefix verbatim, renders only the new tail,
+  generates, then STAGES the call's token delta (ids + mask + logprobs)
+  to TransferQueue — a synchronous tq_put, durable before the call is
+  acked — and returns token-light CommitCoords on the response
+                                              │  coords (≈4 B/token)
+                                              ▼
+gate ingests coords into its lineage index; when the rollout ends it
+returns a token-free RolloutReceipt (a manifest of call_ids and staging
+keys) on the /run response
+                                              ▼
+finalizer (trainer side): fetches the staged deltas by key, verifies
+digests, rebuilds the exact training row, publishes it, clears staging
+```
+
+The heavy bytes (token arrays, logprobs) move exactly once, worker→TQ,
+node-locally. The gate hop and the `/run` response stay token-light.
+
+The pieces, by repo:
+
+| Component | Where |
+|---|---|
+| Gate mode, prefix serving, lineage | Gym `responses_api_models/vllm_model/app.py` + `nemo_gym/token_id_capture/gate.py` |
+| Rollout attribution (`ng_rollout_id`) | Gym capture middleware + `swe_agents` harness forwarding |
+| Wire schema (deltas, coords, receipts) | Gym `nemo_gym/token_id_capture/staging/records.py` |
+| Worker-side capture + prefix splice | `nemo_rl/models/generation/vllm/vllm_worker_async.py` |
+| Staging sink/source over TransferQueue | `nemo_rl/data_plane/tq_token_sink.py` |
+| Receipt → training row | `nemo_rl/experience/blackbox_finalizer.py` |
+
+## Quick start
+
+Batch, one command from a networked shell at the repo root:
+
+```bash
+DRY_RUN=0 SC_EXP_NAME=my-capture-run NG_TIC_FP_CANONICAL=1 \
+  WALLTIME=3:59:00 bash swe_nano_sc_capture.sh
+```
+
+Without `DRY_RUN=0` it prints the resolved driver command and exits, which is
+worth doing once. Interactive mode — allocate once, keep Ray up, run the
+driver by hand, edit, re-run — follows the same attach/run-cmd workflow as
+the TransferQueue recipe:
+
+```bash
+NG_TIC_FP_CANONICAL=1 bash swe_nano_sc_capture_interactive.sh
+```
+
+The non-capture baseline remains `swe_nano_sc.sh` /
+`swe_nano_sc_interactive.sh`. Batch and interactive build byte-for-byte the
+same driver command.
+
+`WALLTIME` must be a Slurm time string (`3:59:00`). A bare `3h` fails with
+`sbatch: error: Invalid --time specification` printed at the very *end* of
+the launch output — easy to miss. Budget ~1 h 40 of setup (venv rebuild +
+checkpoint load) before the first step.
+
+## What the capture launchers add, and why
+
+Every line of the capture posture in `swe_nano_sc_capture*.sh` exists because
+its absence broke a run:
+
+| Setting | Without it |
+|---|---|
+| `token_capture.enabled=true` | Capture never engages. The `token_capture` block (with `enabled: false` and all defaults documented) lives in `nano_swe_teacher_sc.yaml`; the launcher flips the one flag. |
+| `NG_TIC_FP_CANONICAL=1` | Token-in silently degrades to ~0: reasoning models echo history with `<think>` blocks stripped, the gate's fingerprint of the served turn never matches, and every call falls back to text mode. With it, `token_in_rate ≈ 0.9999`. |
+| `NRL_DRIVER_PYTHONPATH=/opt/nemo-rl/3rdparty/Gym-workspace/Gym` | Driver `ModuleNotFoundError: nemo_gym` — the driver imports the staging record schema, and the baked driver venv has no nemo_gym. |
+| `NRL_DRIVER_PIP_INSTALL=orjson` | Driver `ModuleNotFoundError: orjson` — Gym's `token_id_capture/__init__` eagerly imports the store. |
+| `VllmAsyncGenerationWorker` in `NRL_FORCE_REBUILD_VENVS_LIST` | Worker `ModuleNotFoundError: orjson` — venv caching is spec-unaware and silently reuses a non-capture worker venv built by an earlier job. |
+| capture env set *after* sourcing `swe_nano.env` | `swe_nano.env` exports `NRL_FORCE_REBUILD_VENVS_LIST` unconditionally and clobbers an env-prefix value — which is why these are dedicated launchers rather than an env prefix on `swe_nano_sc.sh`. |
+| `CALL_TIMING=0` (optional) | Per-call latency JSONL is on by default in these launchers (`NRL_CALL_TIMING_DIR`/`NG_CALL_TIMING_DIR`); set 0 to disable. All probes are env-gated and dormant without the dir. |
+
+## Verifying capture is really engaged
+
+Config echo is not evidence. Check, in order:
+
+1. **Gate metrics in the SC step output.** Grade from the SC worker `.out`
+   under `<job>-logs/ray/session_*/logs/worker-*-<pid>.out` — the driver
+   log's actor-stdout forwarding is not reliable and can make a healthy run
+   look stalled:
+
+   ```
+   gate_metrics={'registered': 128.0, 'token_in': 7042.0,
+                 'fallback_no_match': 0.0, 'capture_failed': 0.0, ...}
+   ```
+
+   `token_in / (token_in + fallback_*)` should be ≥ 0.99. A rate near 0 with
+   large `fallback_no_match` means canonical fingerprints are off (see
+   above). `unattributed_calls` > 0 means the harness is not forwarding
+   `ng_rollout_id` (attribution fix missing from the Gym pin).
+
+2. **No finalize rejections.** `finalize: ... rejected (empty_manifest)` on
+   every rollout means calls reach the worker without gate context — the run
+   generates but nothing is trainable.
+
+3. **TQ staging traffic**: `PUT_DATA` on the staging partition fires per
+   model call (tens of thousands per run), not just per training batch.
+
+4. **Training equivalence**: `token_mult_prob_error` should sit near 1.0
+   (max ≲ 3), `gen_kl_error` in the same band as a non-capture run (~0.004).
+
+## Known limits
+
+- **Grade runs from the SC worker `.out` or W&B, never the driver log** —
+  Ray's driver-log stdout forwarding dropped entire actors in testing (runs
+  looked stalled while training normally).
+- **Receipt-mode W&B rollout metrics are not yet comparable to legacy**:
+  `gen_tokens_per_sample` counts the carried prompt tail and
+  `truncation_rate` is constant on the capture arm.
+- **Weight-version mixing** across spliced chains at a refit boundary is
+  tag-checked per call; a per-row guard in the finalizer is a pending
+  hardening item.
+- **Router replay (R3)**: the capture path stages routed experts beside the
+  token delta (`routed_experts_delta` extras) so the per-token arrays never
+  transit HTTP — but vLLM 0.20.0's engine currently crashes with
+  `enable_return_routed_experts=true` on this model; blocked on an engine
+  fix.
+
+## Related
+
+- [Nano SWE with TransferQueue](nano-swe-transferqueue.md) — base recipe,
+  cluster shape, SingleController constraints
+- [Router Replay](router-replay.md) — R3 background and trainer-side replay
+- `nemo_rl/data_plane/tq_token_sink.py` — the staging sink/source over TQ
+- `nemo_rl/experience/blackbox_finalizer.py` — receipt → training row
+- Gym `nemo_gym/token_id_capture/staging/records.py` — the wire schema