feat(token-id-capture): external-sink token capture (worker-owned staging, ledger, terminal attribution)

[pthombre]

Summary

Adds external-sink token capture: a mode where a framework inference worker (e.g. NeMo RL with TransferQueue) stages each model call's token delta durably in its own storage and acknowledges Gym with token-light commit coordinates. Gym's data plane never carries raw training tokens end-to-end — it keeps only a token-free custody ledger of hashes and metadata, and the framework reconstructs and cryptographically verifies the full training sequence at rollout end.

This builds directly on the existing token capture stack (capture middleware, CaptureContext, fingerprint-based lineage resolution, LineageStore, and witness-based terminal attribution) rather than replacing it: the legacy complete-record path (TokenSink/TokenSource, TokenCaptureStore) is unchanged and remains the default. External staging is a per-server opt-in (token_id_capture.external_staging: true).

Why

For RL training at scale, shipping full token arrays through the agent/env data plane is wasteful and couples Gym's storage to training infrastructure. With external staging, tokens live only where training consumes them (e.g. TransferQueue), while Gym retains just enough — chained digests and custody metadata — to prove, at rollout end, exactly which tokens in which order earned the reward.

How it works

1. Admission. Capture-prefixed requests are admitted by the existing capture middleware. resolve_parent performs strict tri-state admission against the rollout's ledger: a unique fingerprint match mints a token_in CaptureAdmission (parent id, prev_len, parent chain hash, staging chain); no model-authored context mints a text root; anything else writes a poison row and the call is unadmitted (fail closed).
2. Worker-side staging. The admission rides the engine-bound request as ng_capture. The worker runs RolloutTokenCapture: verifies its prompt begins with the authorized prefix (fetched from the external sink via the staging chain — large token arrays never cross the Gym wire), builds the mask/logprob-aligned delta, computes chain_hash_N = H(chain_hash_{N-1} || delta_N) plus a whole-sequence cumulative_hash, stages the StagedCallRecord to the StagingSink, and only then acknowledges with CommitCoords on the response.
3. Ledger commit. The vLLM server validates the coordinates against the active capture context, computes output/continuation fingerprints server-side, and appends a token-free custody row to the per-rollout ledger. All token/logprob transport fields are stripped before the response leaves the process. Any failure writes a poison row instead; an admitted call that returns without coordinates is poisoned by the middleware.
4. Rollout end. The framework reads the token-free manifest over a bearer-protected control route, attributes the terminal call (declared response id > served response id > content fingerprints, with corroboration; parent-link heuristic as the no-witness fallback), seals a RolloutReceipt, fetches the staged deltas from the StagingSource, and runs verify_and_linearize: every digest and chain hash is re-derived link by link, with a single terminal cumulative-hash anchor, producing the masked LinearizedRow for training. Any tampering, loss, or reordering of a delta breaks a hash and fails closed.

What's added

nemo_gym/token_id_capture/staging/ — portable wire/integrity contracts. Dependency-free by design (no fastapi/ray/torch/transfer_queue imports, test-enforced) so the framework side can consume it as a spec:

• digest.py — canonical encodings and the v2 staging digest: chain hashing, cumulative hashing, extras digests, delta construction with mask/logprob invariants.
• records.py — CaptureAdmission, StagedCallRecord/StagedCallSnapshot, CommitCoords, token-free CallRecord, RolloutManifest, RolloutReceipt; plus the centralized machine-readable reason-code vocabulary (poison and terminal-selection reasons) the framework switches on.
• capture.py / protocols.py — the worker-side lifecycle (RolloutTokenCapture: begin/complete/fail, single-shot claims, capture failures never fail the model completion) and the framework seams (StagingSink, StagingSource, CaptureAdapter, WeightVersionProvider, install_capture).
• rebuild.py — verify_and_linearize: receipt/version checks, manifest-vs-snapshot comparison, full digest recomputation, parent-graph validation, chain-hash re-derivation, routed-experts splicing.
• terminal.py — fail-closed heuristic terminal selection from parent links, for rollouts with no declared terminal and no witness.
• routes.py — compact routed-experts envelope (decode + token-count without materializing the tensor).
• conformance/ — installable golden vectors so an independent implementation can prove byte-identical digests at startup.

Capture ledger. The lineage store doubles as the rollout's capture ledger (CaptureLedger protocol): FileLineageStore gains a durable per-rollout *.lineage.jsonl of token-free custody rows (fsync'd, incrementally cached, sharing the token store's per-rollout lock and rollout-id validation), with ledger-row fallback resolution so chains continue without tokens ever crossing the wire.

Terminal attribution. The existing witness-corroboration engine (token_id_capture/terminal.py) is extended to serve both record families: token-free custody rows identify by (cumulative_hash, chain_hash, cum_len) and match content through recorded fingerprints, and a new authoritative declared_response_id witness carries the harness's declaration (a declared id the ledger cannot confirm attributes nothing and never falls back).

Model server integration (vllm_model). External-capture request shaping (admission attached, prefix arrays omitted in favor of staging-chain fetch), response finalization (coordinate validation, fingerprint stamping, ledger commit, transport-field stripping), and hard config guards (requires enabled, forbids rebuild_response, completions API, responses-native mode, n != 1, and streaming).

Agent plumbing (swe_agents). Rollout-scoped model base URLs threaded into the sandbox, an OpenHands sitecustomize overlay that routes the pinned in-container client through the capture-prefixed URL, the declared-terminal response id surfaced on verify, and routed-experts attachment from harness completions.

Control plane. A bearer-protected manifest route (/training-token-capture/control/rollouts/{id}/manifest) plus RolloutControlClient; the token is read from $NEMO_GYM_TOKEN_CAPTURE_CONTROL_TOKEN (name configurable) and never serialized into run config.

Configuration

token_id_capture:
  enabled: true
  external_staging: true          # opt-in; forbids rebuild_response
  control_auth_token_env: NEMO_GYM_TOKEN_CAPTURE_CONTROL_TOKEN

No new runtime dependencies; TransferQueue is integrated purely by protocol (no transfer_queue import anywhere in Gym).

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[pthombre]

/claude review

[pthombre]

/ok to test c5dccdb

[pthombre]

/ok to test 7dc7ed4

[claude]

SHIP WITH CARE

Reviewed the full 40-file diff. This adds an opt-in "framework-owned token staging" path for training-time token capture (token_id_capture.external_staging), the worker/gate wire contracts, a capture-ledger surface on FileLineageStore/InMemoryLineageStore, a bearer-protected read-only manifest control route, and the SWE-agent plumbing to route sandboxed agent traffic through the capture-prefixed URL. The design is fail-closed throughout (unresolved parent → poison, not silent new-root; missing commit coords → poison; sink/adapter exceptions poison capture without failing the model completion), the digests are domain-separated and versioned, and test coverage is genuinely thorough (ledger, admission, attribution, rebuild, worker, staging core). Async hygiene is clean — no httpx.AsyncClient, no ray.get(), blocking ledger I/O is offloaded via asyncio.to_thread under flock, and the control client goes through Gym's global aiohttp request().

The subsystem itself is gated behind config that defaults off and validates its prerequisites eagerly (external_staging requires enabled, non-InMemoryLineageStore, CaptureLedger, rebuild_response=false, a control bearer, and non-streaming), so the risk to existing runs is low.

One finding worth confirming before merge (inline on responses_converter.py:790):

• RISK — chat_completion_to_response now reuses chat_completion.id as the Responses id on the default path (not just capture). This flips the public id from resp_* to the backend's chatcmpl-* for all chat-backed Responses traffic and makes uniqueness backend-dependent. If the terminal-attribution join only needs this under external_staging, gate it; otherwise call it out as a deliberate API contract change. Everything else in the capture path is correctly scoped behind the opt-in flag.

Non-blocking notes:

• pyproject.toml package-data / exclude additions and the .secrets.baseline golden-vector entries look consistent with the new files.
• The OpenHands sitecustomize.py overlay monkeypatches ServerClient.request in the pinned-Gym venv; it's idempotent (guard attr) and scoped to the capture URL, which is a reasonable approach given OpenHands pins its own nemo-gym.

No async-hang, verifier-scoring, or broken-public-API blockers found. The public base-class surface is extended (new optional BaseRunRequest.capture_rollout_id alias, new CaptureLedger protocol) without removing or renaming existing members.

[claude]

SHIP WITH CARE — no blockers found.

Reviewed the full diff (42 files) and the key correctness paths in the head tree. This is a large but carefully-engineered, fail-closed subsystem for framework-owned token staging.

What makes this low-risk to merge:

• Fully gated behind the new token_id_capture.external_staging flag (default false). Config validators require enabled=true + rebuild_response=false, and VLLMModel.model_post_init rejects incompatible modes (use_completions_api, is_responses_native, return_token_id_information). Existing runs are untouched.
• Async hygiene holds: no httpx/ray.get() in async paths; FileLineageStore file I/O is dispatched via asyncio.to_thread; RolloutControlClient routes through nemo_gym.server_utils.request(). New deps are stdlib only (hashlib, struct, base64).
• Correctness is fail-closed end to end: unresolved/ambiguous parents, missing commit coords, admission/identity divergence, and custody rows without a served response_id all write a poison row (or abstain) rather than emitting a silently-wrong training row. Digest chaining (compute_chain_hash/cumulative_hash), verify_and_linearize, and multi-witness terminal attribution are invariant-checked with strong dedicated tests.
• _strip_capture_transport_fields runs in a finally, so token IDs / logprobs / routed_experts never leak onto the agent hop.

RISK — test coverage of the model-server finalization seam. VLLMModel._finalize_external_capture / _apply_external_capture / _strip_capture_transport_fields (responses_api_models/vllm_model/app.py:1031-1150) are the glue that turns a worker CommitCoords ack into the custody ledger row that becomes a training row — and that strips custody fields before the response reaches the agent. This path has no direct unit test in the diff (the ledger/sink/attribution primitives it calls are well covered, but not the assembly: coords→record(...), the missing-response_id fail-closed branch, and the double-poison-vs-mark_external_staging_committed ordering). Given the >=96% coverage bar and that this seam manufactures training-row data, a focused test (happy commit, capture_failed, missing coords, missing envelope id, coords/admission divergence) would close the gap. Author's call whether to defer to integration coverage.

Per repo policy: this needs a real rollout with a model inspected end-to-end before relying on it — green unit tests alone aren't sufficient for a capture/training-signal path.

[pthombre]

/ok to test 7f2f5ac

[ananthsub]

this lock covers the external sink.stage() call, not only the single-completion state transition.

this means concurrent requests serialize on the capture host. for an out-of-tree sink implementation, can the protocol require thread-safe sinks and move stage() outside of this global lock? or could we use a per-rollout/keyed synchronization instead?

[pthombre]

I think for the TQ write this lock should be safe to remove

[pthombre]

/ok to test b781987

[pthombre]

/ok to test b09d85f