feat(token-id-capture): worker-locus staging, gate custody, and control plane - #2278
Draft
pthombre wants to merge 60 commits into
Draft
feat(token-id-capture): worker-locus staging, gate custody, and control plane#2278pthombre wants to merge 60 commits into
pthombre wants to merge 60 commits into
Conversation
Contributor
Author
|
Companion NeMo-RL PR (submodule pin + TransferQueue StagingSink/StagingSource + finalizer): NVIDIA-NeMo/RL#3455 /ok to test b605153 |
pthombre
added a commit
to NVIDIA-NeMo/RL
that referenced
this pull request
Jul 31, 2026
Token-in/token-out capture for NeMo-Gym rollouts in the async SingleController GRPO path, rebuilt on the nano SWE recipe branch (3fcc696, main + the TQ-honoured SWE recipe) per docs/design-docs/tq-gym-tokidcap-migration.md §9b.1a. The Gym submodule is pinned to pthombre/tq-tokidcap-capture (b6051536 = upstream tokidcap stack top 81ac2736/#2182 + the 7-commit gate/staging series; companion Gym PR NVIDIA-NeMo/Gym#2278). This squashes the capture series (S1-S5 primitives/worker/gate/receipts/ verification + the stack migration) into one commit ported onto the main-lineage SC: the granular history lives on the prior branch head 624bb27. Highlights: - TQTokenSink/TQTokenSource: per-call token deltas staged from the vLLM worker straight into TransferQueue, durable before the response releases; BlackboxFinalizer rebuilds digest-verified always-N training rows via Gym's terminal-aware linearize over run_builder. - Identity: rollout ids ride the run body as the opaque _ng_rollout_id key (agents stamp /ng-rollout/<id>; all agent impls unmodified); the middleware-minted model_call_id is the call id; TQ sample ids ARE the capture keys. - Gate hosting config: LineageIndex capacity derived from the training config; per-run control-plane bearer token; hard per-call control deadlines (gate death surfaces as failed dispatches + placeholders, not a silent retry stall); base capture dir under the run's log dir. - Receipt-mode dispatch in RolloutManager (_generate_and_finalize: reserve-with-ids, seal -> receipt, finalize_group, commit_finalized, always-N with placeholders), weight-version fan-out to workers, gate metrics logging (token_in_rate, fallback-by-cause). - Port adaptations onto main's structures: streaming run_rollouts (receipt postprocess branch + picklable aiohttp error re-raise), commit()'s pre-write and evicted-during-write guards keep main's semantics, legacy failure path keeps main's remove_group. Dormant by default: everything is behind token_capture.enabled=false. Working log: docs/design-docs/tq-gym-tokidcap-migration-log.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
ananthsub
force-pushed
the
ananthsub/tokidcap/read-route
branch
3 times, most recently
from
August 3, 2026 15:49
30e059b to
26f902e
Compare
pthombre
marked this pull request as draft
August 3, 2026 21:44
ananthsub
force-pushed
the
ananthsub/tokidcap/read-route
branch
6 times, most recently
from
August 5, 2026 16:00
6bad3e2 to
79a1883
Compare
pthombre
pushed a commit
that referenced
this pull request
Aug 6, 2026
…ng harnesses OpenHands (and opencode) compose fresh model calls through NemoGymClient, which posts to a fixed url_path via ServerClient — the /ng-rollout/<id> correlation prefix cannot ride that hop, so every call arrived anonymous, nothing staged, and capture rollouts sealed with empty manifests (job 5764598; same finding as RL docs/design-docs/ tq-gym-gate-openhands-attribution.md on the prior stack). Two-sided fix along the doc's option-2 'path -> metadata -> marker' order: - _CaptureMiddleware peeks unprefixed observed-path requests for metadata.ng_rollout_id (body replayed verbatim downstream); the path prefix stays authoritative when present. - The SWE harness snapshots the /run rollout_context into SWEBenchWrapperInstanceConfig.ng_rollout_id (contextvars do not survive the runner_ray_remote hop) and the OpenHands processor forwards it via llm.model.completion_kwargs.metadata, which nv-OpenHands merges into every assembled request (llm.py:210-226 @ 5f01800). Belongs in PR #2278 (the run-wide opt-in alone cannot correlate harnesses whose clients rebuild requests). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@oci-hsg-cs-001-vscode-02.cm.cluster>
pthombre
added a commit
to NVIDIA-NeMo/RL
that referenced
this pull request
Aug 12, 2026
Token-in/token-out capture for NeMo-Gym rollouts in the async SingleController GRPO path, rebuilt on the nano SWE recipe branch (3fcc696, main + the TQ-honoured SWE recipe) per docs/design-docs/tq-gym-tokidcap-migration.md §9b.1a. The Gym submodule is pinned to pthombre/tq-tokidcap-capture (b6051536 = upstream tokidcap stack top 81ac2736/#2182 + the 7-commit gate/staging series; companion Gym PR NVIDIA-NeMo/Gym#2278). This squashes the capture series (S1-S5 primitives/worker/gate/receipts/ verification + the stack migration) into one commit ported onto the main-lineage SC: the granular history lives on the prior branch head 624bb27. Highlights: - TQTokenSink/TQTokenSource: per-call token deltas staged from the vLLM worker straight into TransferQueue, durable before the response releases; BlackboxFinalizer rebuilds digest-verified always-N training rows via Gym's terminal-aware linearize over run_builder. - Identity: rollout ids ride the run body as the opaque _ng_rollout_id key (agents stamp /ng-rollout/<id>; all agent impls unmodified); the middleware-minted model_call_id is the call id; TQ sample ids ARE the capture keys. - Gate hosting config: LineageIndex capacity derived from the training config; per-run control-plane bearer token; hard per-call control deadlines (gate death surfaces as failed dispatches + placeholders, not a silent retry stall); base capture dir under the run's log dir. - Receipt-mode dispatch in RolloutManager (_generate_and_finalize: reserve-with-ids, seal -> receipt, finalize_group, commit_finalized, always-N with placeholders), weight-version fan-out to workers, gate metrics logging (token_in_rate, fallback-by-cause). - Port adaptations onto main's structures: streaming run_rollouts (receipt postprocess branch + picklable aiohttp error re-raise), commit()'s pre-write and evicted-during-write guards keep main's semantics, legacy failure path keeps main's remove_group. Dormant by default: everything is behind token_capture.enabled=false. Working log: docs/design-docs/tq-gym-tokidcap-migration-log.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
pthombre
force-pushed
the
pthombre/tq-tokidcap-capture
branch
from
August 12, 2026 22:27
9181f5d to
cbe9399
Compare
pthombre
pushed a commit
that referenced
this pull request
Aug 12, 2026
…ng harnesses OpenHands (and opencode) compose fresh model calls through NemoGymClient, which posts to a fixed url_path via ServerClient — the /ng-rollout/<id> correlation prefix cannot ride that hop, so every call arrived anonymous, nothing staged, and capture rollouts sealed with empty manifests (job 5764598; same finding as RL docs/design-docs/ tq-gym-gate-openhands-attribution.md on the prior stack). Two-sided fix along the doc's option-2 'path -> metadata -> marker' order: - _CaptureMiddleware peeks unprefixed observed-path requests for metadata.ng_rollout_id (body replayed verbatim downstream); the path prefix stays authoritative when present. - The SWE harness snapshots the /run rollout_context into SWEBenchWrapperInstanceConfig.ng_rollout_id (contextvars do not survive the runner_ray_remote hop) and the OpenHands processor forwards it via llm.model.completion_kwargs.metadata, which nv-OpenHands merges into every assembled request (llm.py:210-226 @ 5f01800). Belongs in PR #2278 (the run-wide opt-in alone cannot correlate harnesses whose clients rebuild requests). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@oci-hsg-cs-001-vscode-02.cm.cluster>
pthombre
added a commit
to NVIDIA-NeMo/RL
that referenced
this pull request
Aug 12, 2026
Token-in/token-out capture for NeMo-Gym rollouts in the async SingleController GRPO path, rebuilt on the nano SWE recipe branch (3fcc696, main + the TQ-honoured SWE recipe) per docs/design-docs/tq-gym-tokidcap-migration.md §9b.1a. The Gym submodule is pinned to pthombre/tq-tokidcap-capture (b6051536 = upstream tokidcap stack top 81ac2736/#2182 + the 7-commit gate/staging series; companion Gym PR NVIDIA-NeMo/Gym#2278). This squashes the capture series (S1-S5 primitives/worker/gate/receipts/ verification + the stack migration) into one commit ported onto the main-lineage SC: the granular history lives on the prior branch head 624bb27. Highlights: - TQTokenSink/TQTokenSource: per-call token deltas staged from the vLLM worker straight into TransferQueue, durable before the response releases; BlackboxFinalizer rebuilds digest-verified always-N training rows via Gym's terminal-aware linearize over run_builder. - Identity: rollout ids ride the run body as the opaque _ng_rollout_id key (agents stamp /ng-rollout/<id>; all agent impls unmodified); the middleware-minted model_call_id is the call id; TQ sample ids ARE the capture keys. - Gate hosting config: LineageIndex capacity derived from the training config; per-run control-plane bearer token; hard per-call control deadlines (gate death surfaces as failed dispatches + placeholders, not a silent retry stall); base capture dir under the run's log dir. - Receipt-mode dispatch in RolloutManager (_generate_and_finalize: reserve-with-ids, seal -> receipt, finalize_group, commit_finalized, always-N with placeholders), weight-version fan-out to workers, gate metrics logging (token_in_rate, fallback-by-cause). - Port adaptations onto main's structures: streaming run_rollouts (receipt postprocess branch + picklable aiohttp error re-raise), commit()'s pre-write and evicted-during-write guards keep main's semantics, legacy failure path keeps main's remove_group. Dormant by default: everything is behind token_capture.enabled=false. Working log: docs/design-docs/tq-gym-tokidcap-migration-log.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
An agent harness that drives its own model calls hands back a transcript with no
token ids, because the wire formats it speaks have no field for them. RL trains on
token ids, and re-tokenizing the returned text gives a sequence that differs from
what the policy sampled by an unmeasured amount.
The ids still exist inside the model server, for the moment before it converts the
response to the harness's dialect and synthesizes a stream. Capture takes them
there, keyed to the rollout that produced them, and writes one TokenEntry per
model call.
Calls are correlated to a rollout by the /ng-rollout/<id> path prefix already on
main for evaluation capture, so this adds no second correlation scheme. The
agent-side gate now serves both consumers, and a per-agent token_id_capture flag
scopes which agents participate; native agents leave it off because they carry
token ids on their own response items.
The capture key is derived from a run request's task and rollout indices, which
assumes each dispatch gets a distinct pair. A caller that restarts numbering per
dispatch produces a repeated id, so two dispatches share one key and their calls
stitch into one trajectory. An explicit _ng_rollout_id on the run body replaces
the derivation, with the attempt suffix still applied on top. An id that would not
survive the path segment is refused rather than rewritten, and the id pattern is
defined once so the body check and the middleware cannot disagree.
Settings live in one `token_id_capture` block rather than as flat keys, and it
names where records go:
token_id_capture:
enabled: true
dir: /tmp/ng_tokcap
sink: my_pkg.sinks:MyDataPlaneSink
`sink` is constructed once per server process at app startup. That matters at
num_workers > 1: uvicorn is handed an app string and workers=N and spawns those
workers, re-importing the app module rather than inheriting the launcher's memory,
so a sink installed programmatically by a launcher does not exist in any worker.
Measured, capture then falls back to the file store, or writes nothing at all when
no directory is set, and logs no error either way. install_token_sink remains for
programmatic use under the same constraint. The validator refuses combinations
that would silently capture nothing: settings with `enabled: false`, a sink beside
a directory, an unknown key, and a sink that cannot report a lost call.
TokenSink and TokenSource are protocols in a module that imports no web framework,
cluster runtime or tensor library, so an inference worker can write into its own
data plane without pulling in the server stack.
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
…no token ids capture_tokens returned quietly when a response carried no token ids, so a rollout that lost a call looked identical to a complete one. The builder reads the gap between one call's tokens and the next call's prompt as tool output, which closes the chain over the hole: the missing call's generated tokens are delivered inside the next prompt at mask 0, and tokens the policy sampled train as if the environment had written them. Mark the rollout instead, on the same path an exception already takes. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Separate rollout correlation from capture intent and expose durable paired transport contracts so external frameworks can integrate without silent partial training data. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Use the current FastAPI lifespan API so configured capture transports are closed without breaking server startup. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Make agent selection explicit without changing evaluation defaults, and avoid rescanning growing token payloads on every durable write. Rename snapshot and URL contracts so their lifecycle and training purpose are clear. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Keep comments and docstrings aligned with static agent selection, the training-specific route marker, and frozen source snapshots. Use short standalone sentences throughout the capture path. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Keep remaining protocol and storage comments to one complete thought per line. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Keep a frozen state tombstone after conditional drop so a late writer from the retired attempt cannot recreate its records. Explicit pre-dispatch cleanup starts the next attempt. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Remove framework source construction from Gym configuration so consumers create and inject sources in their own process. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Keep token_id_capture out of server discovery so env prefetch does not treat run-wide capture settings as a server. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Pranav Thombre <pthombre@oci-hsg-cs-001-vscode-02.cm.cluster>
pthombre
force-pushed
the
pthombre/tq-tokidcap-capture
branch
from
August 19, 2026 17:00
96a3cb0 to
1b34208
Compare
pthombre
changed the base branch from
ananthsub/tokidcap/read-route
to
ananthsub/tokidcap/docs-extras
August 19, 2026 17:01
Promote LineageStore into the per-rollout append-only capture ledger and
delete the RolloutCaptureGate/GateStateStore pair entirely:
- New CaptureLedger protocol: record() custody columns (CallRecord fields +
logical_request_id + staging_digest), record_failure() poison rows,
token-free manifest(), has_rows(). Implemented in FileLineageStore and
InMemoryLineageStore (the latter rejected in external mode: eviction
breaks completeness).
- Gate-free tri-state admission in resolve_parent(): unique verified match
-> token_in; empty fingerprint or seeded history on an empty ledger ->
text root; anything else -> poison row, never a silent root fallback.
- Gate-free commit in vllm_model: reconstruct cumulative tokens from
CommitCoords and publish one ledger row; failures poison.
- Control plane shrinks to one read-only route:
GET /training-token-capture/rollouts/{id}/manifest.
- token_id_capture.external_staging replaces the gate config block; leftover
gate keys fail validation loudly (extra=forbid).
- Data-capability authorization dropped with the gate; the SWE sandbox
capture URL prefix now keys off agent-level token_id_capture enablement.
- gate.py, gate_store.py, gate control routes, and gate tests deleted; the
surviving invariants are re-expressed in test_token_capture_ledger.py.
Signed-off-by: Prashant Thombre <pthombre@nvidia.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Prashant Thombre <pthombre@nvidia.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Prashant Thombre <pthombre@nvidia.com>
ananthsub
force-pushed
the
ananthsub/tokidcap/docs-extras
branch
2 times, most recently
from
August 20, 2026 21:50
f9222b4 to
2432ffb
Compare
The middleware stamps admitted_at once on the CaptureContext when it admits the call; the commit hook threads it through the lineage stores' custody columns into the token-free manifest as CallRecord.admitted_at. Stamping at admission (not at record time) keeps retried commits byte-identical, which the stores' idempotency-by-equality requires. Rows written before the column carry None and still validate. The column exists so heuristic terminal selection can order candidate roots for harnesses that do not declare a terminal response id. Signed-off-by: Prashant Thombre <pthombre@nvidia.com>
…llouts A harness that reports the response id it kept gives receipt assembly an authoritative terminal row; one that reports nothing previously masked every rollout. select_terminal_call infers the terminal from the manifest's explicit parent links alone: the earliest-admitted root (an extended root beats an abandoned sibling), then a walk that eliminates abandoned retries (childless child loses to an extended sibling). Selection is token-free and fail-closed: a retry of the final call, divergent extended branches, orphaned or cyclic rows all select nothing and surface a reason the caller records as the failure reason. The verifier still checks every digest on the chosen chain downstream. RolloutReceipt.terminal_selection records which path chose the terminal (declared vs heuristic) so consumers can meter heuristic reliance. The README now documents the ledger (the gate it described was removed) and the declared > heuristic > mask precedence. Signed-off-by: Prashant Thombre <pthombre@nvidia.com>
ananthsub
force-pushed
the
ananthsub/tokidcap/docs-extras
branch
10 times, most recently
from
August 21, 2026 17:49
77bfef1 to
6e92c53
Compare
Signed-off-by: Prashant Thombre <pthombre@nvidia.com>
Replace cumulative_token_ids in external (custody) lineage rows with a pair of worker-computed hashes: chain_hash = H(domain ‖ parent_chain_hash ‖ token_ids_delta) cumulative_hash = hash_token_ids(prompt + generated) CommitCoords.token_ids_delta is removed (hard drop; gate and worker must deploy together). The gate records token-free external rows and verify_and_linearize checks both hashes during the terminal-chain walk. Old rows without hashes skip verification; legacy external rows without chain_hash are rejected at admission (fail closed). Signed-off-by: Prashant Thombre <pthombre@nvidia.com>
Staged, fail-closed terminal selection: declared response id, scored response envelope id, and content fingerprints each independently name a manifest row; agreeing witnesses attribute, disagreement attributes nothing, heuristic parent-link selection remains the no-witness fallback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Prashant Thombre <pthombre@nvidia.com>
ananthsub
force-pushed
the
ananthsub/tokidcap/docs-extras
branch
2 times, most recently
from
August 25, 2026 18:23
c162b8a to
67c1693
Compare
ananthsub
force-pushed
the
ananthsub/tokidcap/docs-extras
branch
3 times, most recently
from
September 1, 2026 18:05
0159069 to
654b421
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds the worker-owned token-capture custody layer on top of the rebased token-capture stack.
This PR is based on #2349 (
ananthsub/tokidcap/docs-extras,d2123272), the current top of the stack that begins at #2124. The PR contains only the five-commit delta from that stack head.Changes
Commit series
feat(token-id-capture): define staging integrity contractfeat(token-id-capture): expose worker-owned staging hooksfeat(token-id-capture): verify receipts and rebuild terminal ancestryfeat(token-id-capture): add authenticated rollout gate lifecyclefeat(token-capture): complete multi-worker custody rebaseStack
Testing
Not rerun as part of this PR history update.