feat(sc): gate-authoritative token capture via an external TransferQueue sink

[pthombre]

What this PR adds

Gate-authoritative token capture for NeMo-Gym rollouts on the SingleController path: instead of round-tripping token ids through the environment (token echo) or re-tokenizing text, the vLLM worker stages each model call's exact token delta into TransferQueue (the external sink), NeMo-Gym's policy model server custodies the lineage in a capture ledger, and a finalizer actor pool assembles verified receipts into training rows.

Enabled by token_capture.enabled=true (NeMo-Gym + vLLM async engine only, validated loudly at setup). The legacy token-echo path is unchanged and remains the default.

How it works (rollout → training row)

1. Staging (external sink). The vLLM async worker hosts Gym's capture core in-process. For every model call it stages the call's token delta (prompt/generation coordinates, logprobs, and optionally routed experts) into a TQ staging partition, keyed by a per-rollout id (/ng-rollout/<id>) that rides the Gym run body — the TQ sample id is the capture key end to end (nemo_rl/data_plane/tq_token_sink.py, vllm_worker_async.py).
2. Custody (capture ledger). Gym's policy model server records admission → staged-coordinates → terminal attribution as a token-free custody ledger (chained digest). Uncommitted off-chain calls poison only their own rollout, and undeclared rollouts get witness-based / heuristic terminal selection rather than silent acceptance.
3. Receipt assembly. After a rollout finishes, NeMo-RL fetches the manifest over the ledger control plane (bearer-token authenticated) and assembles a receipt: the verified prefix chain naming exactly which staged rows constitute the trajectory (nemo_rl/environments/nemo_gym.py, nemo_rl/experience/blackbox_finalizer.py).
4. Finalization (actor pool). The rollout pump hands a metadata-only FinalizationRequest (group id, rollout ids, receipts, rewards — never token tensors) to a finalizer actor pool, which publishes canonical training rows into the TQ training partition and commits the replay-buffer slot (nemo_rl/experience/finalizer_actor.py, route_plan.py, single_controller.py actor dispatch path).
5. Training. The train pump consumes committed groups as usual. Capture placeholders carry sample_mask=0, and the advantage estimator now takes a real valid_mask instead of hardwired all-ones (advantage_estimator.py).

Supporting changes

• TQReplayBuffer token-capture mode: reserve(rollout_ids=…), commit_finalized, abort, staging-aware remove (nemo_rl/algorithms/async_utils/replay_buffer.py).
• Routed-experts transport: nrlre1 envelope decode in the sink, int16 fields broadcast as int32 (NCCL has no Short), all-placeholder groups dropped when routed dims are unknown (data_plane/schema.py, preshard.py, worker_mixin.py).
• token_capture config block with setup-time validation of the supported matrix; derived fields (control-plane auth token, capture dir) filled at setup (single_controller_utils/config.py, setup.py).
• Build: Gym becomes an editable path dependency instead of a uv workspace member — both projects define a vllm extra with intentionally different server stacks, and uv applies a requested extra to every workspace member, making the RL-vLLM + base-Gym worker environment unselectable otherwise (pyproject.toml, regenerated uv.lock).
• Gym submodule pinned to fd5e84d6 (feat(token-id-capture): external-sink token capture (worker-owned staging, ledger, terminal attribution) Gym#2872, the tokidcap capture stack, merged into Gym main 2026-09-04).

Tests

• New unit suites: token sink, blackbox finalizer, finalizer actor + lifecycle, route plan, worker route assembly, TQ policy routes, nemo-gym token capture, vLLM capture hosting, advantage validity, plus TQ replay-buffer capture-mode coverage.
• L1 functional gains a token-capture arm: the SC+Gym smoke rerun with ++token_capture.enabled=true.
• Legacy-path invariance: all existing SC/experience/data-plane suites pass unchanged (RAY_ADDRESS=local run on a CPU node; final revalidation in flight).

Not in this PR

The nano SWE launch recipe (configs, launchers) and the recipe guides are kept out to scope this PR to core logic. History note: this supersedes #3456 (same feature, rebased onto current main with a clean history; the TQ actor runtime-env pin and the mtp refit fix from that branch were dropped as both are superseded upstream).

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

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: b3e70a0 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: f130e85 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[pthombre]

/ok to test 14dde62

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 14dde62 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 65fa9b0 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[pthombre]

/ok to test 82689d1

[ZhiyuLi-Nvidia]

Review finding: advantage-estimator valid_mask is only honored by GRPO.

[ZhiyuLi-Nvidia]

_cleanup_consumed_metas blocks the controller loop — rely on TQ's native async clear API

The end-of-step batching is a good change and should stay; it is fewer round trips than the per-chunk clears it replaced. The issue is only which TQ entry point it goes through.

_cleanup_consumed_metas                 single_controller.py:1314  (train pump, :2299)
└ self._call_dp("clear_samples")        :1057-1065 → local client → INLINE, no yield
  └ TQAdapter.clear_samples             adapters/transfer_queue.py:990-1015  (sync def)
    └ tq.kv_clear                       interface.py:558-579     ← sync veneer
      └ _run_coroutine(async_kv_clear(keys, partition_id))
        │                               client.py:1238-1241 → run_coroutine_threadsafe
        │                               future.result()   ← no timeout, not cancellable
        └ async_kv_clear                interface.py:933-959  ← the coroutine, already
          ├ async_kv_retrieve_meta      client.py:898
          ├ async_clear_samples         client.py:542
          └ storage_manager.clear_data
              simple → gather over SUs  simple_storage_manager.py:486   truly async
              KV     → NO await at all  managers/base.py:803-813        blocks the loop

_call_dp resolves to a local DataPlaneClient, so method(**kwargs) runs inline (:1063) with no suspension point, and the thread then blocks on future.result(). Called once per step from the train pump, that parks the rollout pump, the fleet-health probe and the finalizer pool — including _stall_watchdog_pump, which is a task on the same loop (:498) and so cannot report the stall it exists to catch.

Recommendation: call tq.async_kv_clear (interface.py:933) instead of tq.kv_clear. As the stack shows, kv_clear is _run_coroutine(async_kv_clear(...)) — the native API is the same coroutine without the thread hop and the un-cancellable wait. Same round trips; the loop keeps turning.

# nemo_rl/data_plane/adapters/transfer_queue.py
    async def async_clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None:
        """Async twin of clear_samples: awaits TQ's native coroutine directly."""
        self._mark_data_operation_started()
        await tq.async_kv_clear(keys=list(sample_ids), partition_id=partition_id)

# single_controller.py, _cleanup_consumed_metas
        await self._call_dp("async_clear_samples", sample_ids=unique_ids, partition_id=partition_id)

No plumbing needed — _call_dp (:1064) and call_data_plane (data_plane/async_utils.py:54) already await a returned coroutine.

Two things to settle first:

• Safe on simple only today. KV-backend clear_data is an async def with no await in it (managers/base.py:803-813), so on mooncake_cpu the blocking batch_remove would land on SingleController's loop instead of TQ's. Needs the upstream await asyncio.to_thread(...) fix before that backend can use this path.
• Cancellation becomes real. async_clear_samples marks-clearing before it clears (client.py:566-572), so a cancelled clear leaves rows unreadable but unfreed. Shield it, or accept that deliberately.

Interim one-liner if you'd rather not add the adapter's first async method in this PR: call_data_plane(self._dp_client, "clear_samples", offload_sync=True, ...) at both call sites — keeps the block off the loop on every backend, no upstream change.

Separately: _clear_data_plane_samples (:1068) now has no production callers — only tests/unit/single_controller/test_checkpointing.py:1192,1214 exercise it. Delete and repoint those tests, or reuse it here with a partition_id parameter.

[pthombre]

/ok to test 4e50901

[pthombre]

Thanks for the thorough review — went through all 21 threads. Summary below, grouped by outcome.

Applied

• _FORBIDDEN_RPC_KEYS gap — added logprobs to the denylist and a comment explaining the list's purpose/maintenance rule. (4222c0c46)
• RPC dataclass field tripwire — added test_rpc_dataclass_fields_are_classified and test_every_forbidden_key_is_rejected. (ce8604092)
• token_capture missing from exemplars — added the block to both grpo_math_1B_megatron_single_controller.yaml and ppo_math_1B_megatron_single_controller.yaml. (f98536a28)
• Gym pin is an open PR head — Gym fix(megatron): honor policy.logprob_chunk_size in the training loss path #2872 has since merged; re-pinned to the resulting fd5e84d6 commit on Gym main, relocked uv.lock, and corrected the PR description's stale d9ad9d2e/perf: enable MoE GroupedGEMM for MoE models #2278 reference. (6e5bddf41)
• Capture suite silently skips — added test_gym_staging_package_is_importable_in_the_nemo_gym_lane as a tripwire, and wired up a real test_tq_sink_source_passes_gym_golden_vectors calling Gym's assert_golden_vectors() (chose to wire up real conformance rather than just fix the docstring). (a361ca83a)
• Wrong valid_mask passed to the baseline — single_controller.py now passes final_sample_mask instead of sample_mask; docstring updated to match. This fixes the leak in both capture and non-capture modes (traced through — mask_sample/overlong_filtering/seq-logprob-error-masked rows were voting in siblings' baselines regardless of capture). (ad3a4722d)
• Functional test doesn't verify capture ran — added the finalize/total_ms > 0 guard, gated on token_capture.enabled=true being passed. (22df1ff4e)
• Deferred-mode route fallback invisible — fixed the missing per-step reset on _route_fallback_counts (was lifetime-cumulative), plumbed it through the existing worker→driver metrics pipeline, and surfaced it as routed_experts_deferred_fallback/<reason>. (68d84f409)
• Dead on_capture_failure knob — dropped the field rather than wiring it up (traced the alternative — reusing the existing group-drop machinery — and decided the smaller, contained fix was the right call here); fixed the tq_token_sink.py docstring that falsely claimed it was live. (f7e8a4114)
• train/reward skewed by placeholders — now a sample_mask-weighted mean instead of a bare mean, threaded through single_controller.py → reduce_advantage_pump_metrics. (49c366f85)
• Routed-experts backfill invisible — backfill_missing_routed_experts now returns per-message-log backfill counts, merged into the existing VIOLATION_TAG_KEYS/per-row-tag metrics pipeline (num_routed_experts_backfilled, routed_experts_backfilled_rate) rather than inventing new plumbing. (47b0ab3b2)
• mixed_weight_version_policy fires on the normal case, AI-1 — dropped the field entirely; group_min_wv/max_staleness_versions is now the sole staleness policy. (8a8f52d87)
• Same thread, AI-2 — moved the min_valid_fraction_per_group decision out of the finalizer and into the controller's existing replacement-seeking path (_take_replacement), so a low-valid-fraction group now gets a chance at a spare prompt the same way generation-side drops do, and the drop now counts against max_skipped_prompts/max_consecutive_dropped_prompts via a new record_finalizer_dropped_prompt(). FinalizedGroup gained valid_row_count/total_row_count; the finalizer no longer makes this policy call itself. (bc399288a)
• Stale § N design-doc references — dropped the 4 decorative ones, replaced the 2 substantive ones (setup.py's MVP-matrix comment, the sink test's conformance-split comment) with the actual content. (3086d388e)
• STAGING_FIELDS provenance undocumented — added a comment naming the source (StagedCallRecord), the maintenance rule, and the conformance-check blind spot for optional fields. (bundled in f7e8a4114)
• No per-reason capture-failure metric — added finalize/capture_failure_reason_<bucket>_count, bucketed by the static prefix before : to avoid a metric-name explosion from the templated reason strings. (933972108)
• max_skipped_prompts no-op under capture, AI-1 — added a warnings.warn at setup, same shape as the three existing token_capture cross-checks. AI-2 (dedup the two retry-loop scaffolds) is a follow-up per your own note, not touched here. (46afc518e)
• Suggestion: drop ++ override — applied as suggested, now that token_capture is in the exemplar. (c5654441a)
• overlong_filtering silently no-ops under capture — went with Option A: truncated is no longer carried from the dispatcher (it never could compute it correctly there); the finalizer now computes it from each row's rebuilt length against max_seq_len, threaded through FinalizerActorConfig. FinalizationRequest.truncated dropped from the RPC contract accordingly. (127164647)
• Hand-copied UNCOMMITTED_CALL_REASON / hand-copied terminal_selection methods — AI-1: now imports the constant from nemo_gym.token_id_capture instead of retyping it. AI-2: the terminal-selection metric breakdown now derives its method list from RolloutReceipt.model_fields["terminal_selection"].annotation via get_args, instead of a hand-copied tuple. (adaefd03b, 127164647)
• BlackboxFinalizer/FinalizerActor naming — went with the whole-family rename rather than just the one class: BlackboxFinalizer→RolloutReassembler, FinalizerActor→RolloutReassemblerActor, FinalizationRequest→ReassemblyRequest, FinalizerActorConfig→RolloutReassemblerActorConfig, create_finalizer_actors→create_rollout_reassembler_actors, num_finalizer_workers→num_reassembler_workers, plus the two file renames. (2b0957573)

Acknowledged, not changed

• FinalizerPool/TokenCaptureHost extraction — agreed this is the right shape eventually, but flagged as "a follow-up rather than something to hold this PR on," so leaving it for a subsequent PR.
• Barrier/event-loop landed already and replay_buffer thread — no action needed; the latter is @macandro96's to confirm.

All commits signed off, pushed to this branch. Running the full unit suite (default + --nemo-gym-only lanes) now to validate the full pass before this is ready for another look.

[pthombre]

/ok to test f6305dc

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: f6305dc (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 37cbd31 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[pthombre]

/ok to test 37cbd31

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 0e28303 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 30d300c (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[pthombre]

/ok to test 30d300c

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 9286c3e (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: f54bbb1 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[pthombre]

/ok to test f54bbb1