feat(token-id-capture): capture training tokens from external harnesses

[ananthsub]

Records the exact prompt token IDs, generated token IDs, and generation log probabilities for model calls made by an external agent harness.

Control and data flow

sequenceDiagram
    participant RC as Rollout collection
    participant AS as Agent server
    participant H as External harness
    participant MW as Model middleware
    participant MS as Model route
    participant TS as TokenSink

    RC->>AS: /run with rollout identity
    AS->>AS: select static agent capture or all_agents
    AS->>H: model URL with /ng-rollout/id/training-token-capture
    H->>MW: model request
    MW->>MW: mint model_call_id and CaptureContext
    MW->>MS: normalized model request
    MS->>MS: assemble exact token arrays
    MS->>TS: await put(TokenEntry)
    alt capture fails or arrays are incomplete
        MS->>TS: await mark_incomplete(rollout_id, model_call_id)
    end
    MS-->>H: dialect response or stream

Loading

Capture happens before dialect conversion or stream synthesis can discard token fields. The write is awaited before the model response returns.

Summary

• Separates run-level infrastructure enablement from static per-agent capture selection; all_agents selects every configured agent for a training run.
• Encodes selected training intent in the explicit /training-token-capture path, independently of rollout correlation and evaluation observability.
• Defines framework-neutral async TokenSink and TokenSource contracts with durable incomplete state, atomic frozen snapshots, lifecycle hooks, and versioned conditional retirement.
• Keeps a frozen tombstone after retirement so a late writer cannot resurrect a consumed attempt; explicit pre-dispatch cleanup starts the next attempt.
• Uses durable call digests and tail reconciliation so the common append path does not reread or reserialize prior token arrays.
• Rejects malformed or conflicting token records and closes owned endpoints through the app lifespan.

Gym imports no training-framework data plane. A framework configures sink and lineage proxy factories in Gym model-server workers and constructs its source in the rollout-consumer process.

Stack base. Followed by #2125.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[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.

[cmunley1]

/claude review

[claude]

SHIP WITH CARE

Training-token capture: a new leaf package (nemo_gym.token_id_capture) that records prompt/generation token ids + logprobs per model call into a rollout-keyed store, gated independently from eval capture. Reviewed the correctness-critical paths (verifier/scorer analog here is the capture→readback→train chain, since a hole in it silently corrupts training data).

What I checked and where it holds up:

• Durability before return. capture_tokens awaits sink.put on the model server response path, so the entry is fsynced+flock'd before the model call returns — a post-rollout reader in another process never races a partial file. Tested (test_a_record_is_readable_as_soon_as_put_returns, test_concurrent_appends_to_one_rollout_stay_intact).
• Incomplete = loud, not silent. The worst outcome — a lost/token-less call that leaves the rollout looking complete, so its generated tokens leak into the next prompt at mask 0 — is guarded on every path: write failure, build failure, and a response with no token ids all mark_incomplete. A sink lacking mark_incomplete is logged at error, not swallowed. Directly tested across store, installed-sink, and configured-sink routes.
• Schema versioning refuses newer records rather than decoding them clean under extra="allow". Tested end-to-end through the store read.
• Async hygiene. Blocking file IO offloaded via asyncio.to_thread; no httpx, no ray.get in async paths; sink ContextVar reset in a finally. The leaf-package import boundary (no fastapi/ray/torch) is enforced by a subprocess test.
• Explicit rollout id round-trips through the path-prefix charset (writer/reader agree), with malformed ids refused rather than sanitized. Tested on both sides.
• num_workers > 1 correctly handled: sink built per-worker at app startup, with a spawn-context test pinning why install_token_sink alone doesn't survive.

Watch items (not blockers):

• Multi-carrier assumption. extract_token_fields/strip_token_fields assume exactly one output item carries token arrays and take the last one. True for today's vLLM single-choice path, but a future backend emitting token ids on multiple output items (multi-message / interleaved tool calls) would silently drop all but the last item's tokens into the record — a correctness hole with no error. Fine now; worth an assert or a follow-up if output shapes broaden.
• One dead-code note inline (_request_messages).

Public API is additive only (new optional kwargs with defaults, new per-agent token_id_capture: bool = False); no signature breaks. Test coverage is genuinely thorough and asserts real behavior. Merge is a deliberate go — the async-durability and incomplete-marking invariants are the load-bearing ones and they hold.

[ananthsub]

/claude review

[claude]

SHIP WITH CARE — training-token capture infrastructure. Correct-by-construction on the failure paths that matter for RLHF, but one write-path efficiency concern and a couple of NOTEs.

The design is sound and the review-critical property — an incomplete rollout must never look complete — is defended in depth: capture_tokens/commit_entry never raise into the model call, a missing token payload marks the rollout incomplete (_capture_missing), a sink lacking mark_incomplete is logged at ERROR rather than silently swallowed, and the schema version refuses newer records instead of decoding them under extra="allow". The store fsyncs + flocks per file and put() is awaited (durable before the call returns), so a post-rollout reader never races a partial file. Async hygiene checks out: blocking store ops go through asyncio.to_thread, no ray.get() in async, no httpx. The leaf-package import test enforces the fastapi/ray/torch-free boundary the framework-worker path depends on. Test coverage is genuinely thorough — streaming Anthropic end-to-end, external-staging mode, spawned-worker sink isolation, schema round-trips, concurrent appends.

RISK (inline, store.py:151): append() re-reads and re-serializes every prior entry of the rollout on every call → O(N²·T) over a long trajectory, which is exactly the multi-turn external-harness case the two new configs enable. Doesn't corrupt data; degrades write throughput superlinearly. Fix: a seen model_call_id set in state instead of rescanning the file.

NOTE — request-body buffering under token-capture-only mode. When self._store is None but capture_wanted (training capture on, eval capture off), the middleware forwards without buffering the request body (base_responses_api_model.py:1131). That's correct for capture — the sink reads the response — but confirm no eval-only consumer relies on the request bytes in this mode. The tests cover the token side; they don't assert the eval record is intentionally absent here beyond test_tokens_captured_even_when_eval_capture_disabled, which looks right.

NOTE — TOKEN_ID_CAPTURE_KEY_NAME (global_config.py:101) is unused. The per-agent flag is read via config.token_id_capture and the block via TOKEN_ID_CAPTURE_BLOCK; this second constant with the same string value has no referent. Dead constant, linter's call, but it invites future confusion between the block key and the per-agent key since they collide on the literal "token_id_capture".

No blockers. The correctness-critical paths for training data integrity are well-handled.

[ananthsub]

/claude review

[claude]

Review incomplete: the PR head advanced during review (analysis ran against f480877826ff9dbda89a44411ca608a5443ab9b3, but the pre-publish revision check found the head is now 6aca3bf814baf9fc839609dd6577d886450dd4c7). Inline findings anchored to the old commit could land on the wrong lines, so none were posted. Please re-run the review against the current head.

[ananthsub]

/claude review