[macandro96]

What does this PR do?

Adds batch age to the rollout trace: the token-weighted number of policy updates between generating a token and training on it, per training step.

Stacked on #5048 (the rollout event log).

How it works

1. Per group: the reassembler returns each trained model call's ID, the weight version it began on and the tokens it generated. At publish, the controller writes a group_published event that sums tokens by version.
2. Per step: step_trained also records the policy version the step updated.
3. Exact at refits: vLLM keeps running calls across a refit, so a call's later tokens come from the new weights. While generation is paused for a refit, the vLLM worker reports how many tokens each running captured call has generated so far, written as a refit_cut event. group_published lists those calls individually so the analysis can split them; all other calls arrive summed by version.
4. Offline: batch_age_by_step in rollout_checkpointing/trace.py computes each step's batch age; python -m nemo_rl.algorithms.rollout_checkpointing.trace <log_dir>/rollout_trace.jsonl prints it.

The two commits separate the per-version count (exact except for calls that span a refit) from the exact refit split.

Notes for review

• vLLM internals: the running token count comes from vLLM's output-processor request state, which has no public API (external_req_id, detokenizer.num_output_tokens(), as in vLLM 0.29). An engine without that state reports nothing, and those calls count under the version they began on. Other backends inherit a default that reports nothing.
• Request-to-call link: a context variable set when capture begins a call, read where the same request task submits it to the engine.
• Overlap with feat(checkpoint): recover active vLLM generation prefixes #4506: the second commit touches vllm_worker_async.py, the generation interface and vllm_generation.py, which feat(checkpoint): recover active vLLM generation prefixes #4506 also changes; whichever lands second needs a small rebase there.
• Nothing Gym stores or a checkpoint saves changes.

Testing

• Unit: controller and analysis tests (split across two refits, an interrupted refit, token weighting, the refit and publish events), two vLLM wrapper tests (run in vLLM 0.29's venv), and the 20-turn capture pipeline test checks the token counts. 2116 passed (default lane), 121 passed (Gym lane). ruff and pyrefly clean.
• GPU: a Gym run with continuous dispatch traced 6 steps (batch age 0.0, 1.0, 1.09, 1.27, 1.28, 1.15). One refit paused a call before its first token; the trace counted its 60 tokens under the new weights ({v1: 60, v0: 30}), where counting by start version gives {v0: 90}.

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