How the step number now reaches every span in a training step

NVIDIA-NeMo/RL #4607 (nemo-lens bump, rl.iteration on every span, shared nv.dl.* attributes). All permalinks pin the PR head 2f6b955.

Before
Only the step span and the training span carried the step number. Generation, reward, checkpointing and the other spans inside a step carried none, so "what else happened in step 412" could not be asked.
This PR
Sets the step number once per step, and every span the driver opens inside that step gets it.
Status
New feature, under review. The design works. The review asks for doc fixes and one bug fix (see the second page).

Two pages on this PR, one question each. This one: what the step number change does.

Background. NeMo-RL sends spans through the nemo-lens library. To find all the work in one step you filter on a tag, e.g. rl.iteration = 412. That filter only finds spans that carry the tag. Dotted words have a short explanation: hover or tap them.

Each training loop now opens its step inside iteration_scope(step) (for GRPO: grpo.py:3231). It puts {rl.iteration: 412} in a ContextVar (lens helpers.py:37), and a lens span processor copies it onto every span as the span starts (providers.py:42-48). No call site has to pass the number. A new span added inside the step later is covered too.

Which spans in GRPO step 412 carry rl.iteration

Filter: rl.iteration = 412 — which spans match? before after (this PR) driver process, inside iteration_scope(412) rl.grpo.step 412 (passed by hand)412 rl.grpo.generation —412 rl.grpo.reward_calculation —412 rl.grpo.policy_training 412 (passed by hand)412 rl.grpo.checkpointing —412 rl.grpo.evaluate rl.step=412 (other key)412 …and every other driver span opened inside the step (logprobs, advantages, data processing) —412 those spans outside the scope rl.checkpoint.finalize (new) did not exist412 (set explicitly) runs on a background thread that outlives the step, so the step is read from tmp_step_412 worker / Gym / collector spans —— (other process) matches the filter does not match no such span

A span that sets rl.iteration itself keeps its own value. Validation uses this: it also runs once before training and reports 0 there. The finalize span uses it too, because it runs on its own thread, and a raw thread starts with an empty ContextVar.

What else the PR adds

ChangeBeforeAfter
Shared nv.dl.* resource attributes2 (rank, world size)17: topology, batch sizes, optimizer, step target, library versions, run uuid. The same keys a Megatron run reports.
rl.checkpoint.finalize spanCheckpoint span ended when the loop was released, while async writers were still goingA second span that ends when tmp_step_N is renamed, i.e. when the checkpoint can be resumed from
rl.host.memory.* gaugesHost memory reached W&B onlyThe driver reads each node's Ray metrics page every 10 s and records used/total bytes per host
Renamed keys (breaking)rl.step, dl.tensor_parallel.size, dl.pipeline_parallel.sizerl.iteration, nv.dl.topology.size.tp, nv.dl.topology.size.pp
What to watch. Dashboards that filter on rl.step or dl.*_parallel.size must be updated. Four loops never open a step scope (grpo_sync, async PPO, xtoken distillation, MPO), so their spans still carry no rl.iteration, even though the docs say every loop does. The new rl.checkpoint.async_save tag on the finalize span is always true, which is the subject of the second page.

🟠 review comment — span-groups.md:149 (loops without the scope)