Skip to content

feat(token-id-capture): worker-locus staging, gate custody, and control plane - #2278

Draft
pthombre wants to merge 60 commits into
mainfrom
pthombre/tq-tokidcap-capture
Draft

feat(token-id-capture): worker-locus staging, gate custody, and control plane#2278
pthombre wants to merge 60 commits into
mainfrom
pthombre/tq-tokidcap-capture

Conversation

@pthombre

@pthombre pthombre commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

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

  • Defines versioned staging records, strict integrity digests, extras validation, and routed-expert envelopes.
  • Exposes engine-neutral worker-owned staging hooks and a dependency-light vLLM adapter.
  • Verifies sealed receipts and rebuilds only the declared terminal ancestry before publication.
  • Adds an authenticated, retry-safe rollout gate with process-shared state, data capabilities, tombstones, TTL cleanup, and metrics.
  • Integrates the gate with Gym's model middleware, the vLLM proxy path, rollout correlation, and SWE harness credential routing.
  • Keeps token IDs, log probabilities, masks, routed-expert data, and commit coordinates off the agent-facing response.

Commit series

  1. feat(token-id-capture): define staging integrity contract
  2. feat(token-id-capture): expose worker-owned staging hooks
  3. feat(token-id-capture): verify receipts and rebuild terminal ancestry
  4. feat(token-id-capture): add authenticated rollout gate lifecycle
  5. feat(token-capture): complete multi-worker custody rebase

Stack

Testing

Not rerun as part of this PR history update.

@copy-pr-bot

copy-pr-bot Bot commented Jul 31, 2026

Copy link
Copy Markdown

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

Copy link
Copy Markdown
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
ananthsub force-pushed the ananthsub/tokidcap/read-route branch 3 times, most recently from 30e059b to 26f902e Compare August 3, 2026 15:49
@pthombre
pthombre marked this pull request as draft August 3, 2026 21:44
@ananthsub
ananthsub force-pushed the ananthsub/tokidcap/read-route branch 6 times, most recently from 6bad3e2 to 79a1883 Compare August 5, 2026 16:00
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
pthombre force-pushed the pthombre/tq-tokidcap-capture branch from 9181f5d to cbe9399 Compare August 12, 2026 22:27
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
pthombre force-pushed the pthombre/tq-tokidcap-capture branch from 96a3cb0 to 1b34208 Compare August 19, 2026 17:00
@pthombre
pthombre changed the base branch from ananthsub/tokidcap/read-route to ananthsub/tokidcap/docs-extras August 19, 2026 17:01
Pranav Thombre and others added 3 commits August 19, 2026 14:38
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
ananthsub force-pushed the ananthsub/tokidcap/docs-extras branch 2 times, most recently from f9222b4 to 2432ffb Compare August 20, 2026 21:50
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
ananthsub force-pushed the ananthsub/tokidcap/docs-extras branch 10 times, most recently from 77bfef1 to 6e92c53 Compare August 21, 2026 17:49
pthombre and others added 3 commits August 23, 2026 02:21
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
ananthsub force-pushed the ananthsub/tokidcap/docs-extras branch 2 times, most recently from c162b8a to 67c1693 Compare August 25, 2026 18:23
@ananthsub
ananthsub force-pushed the ananthsub/tokidcap/docs-extras branch 3 times, most recently from 0159069 to 654b421 Compare September 1, 2026 18:05
Base automatically changed from ananthsub/tokidcap/docs-extras to main September 1, 2026 19:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants