feat(telemetry)!: instrument the single-controller path and propagate traces - #4052
Conversation
|
/ok to test 493d853 |
|
/ok to test de7ab13 |
|
/ok to test de7ab13 |
|
/ok to test bec17bd |
|
/ok to test 130756c |
|
/ok to test 45c2002 |
|
/ok to test 7bf922c |
|
/ok to test 88cb61d |
|
/ok to test 246b047 |
|
/ok to test 246b047 |
|
/ok to test e2235e6 CI:Lfast |
1 similar comment
|
/ok to test e2235e6 CI:Lfast |
|
/ok to test e2235e6 |
terrykong
left a comment
There was a problem hiding this comment.
Thanks for turning this around so quickly. I checked each reply against the code at e2235e6, and ran the fixes where I could (the Gym cancel path, the telemetry-only wrapper, the stage per algorithm, the counter filter, the real-Ray tests).
Of the 29 threads, 21 are fixed; I replied with the commit and resolved them. The other 8 have a follow-up reply in this review: 7 are partly fixed, and 1 is fixed with a small docs leftover.
Before this can merge:
- Four unit tests fail at this head (three in
test_metrics.py, one intest_vllm_utils.py). CI has not shown them because the unit lanes have not run yet. - CI Lint fails on the return type of
streaming_umbrella_span. - The new Ray test's fixture shuts down the shared session cluster.
metric_utils.py's shared key constants are not used by the producers yet (reply on the_TRAIN_SCALARSthread).
In the PR description: the nemo-lens rev is still given as b71263c (the pin is b0f977d).
Generated by Claude Code
Signed-off-by: Raj Singh <rajsin@nvidia.com>
… traces Open job, step, and rollout spans on the single-controller path, add startup and vLLM engine metrics, and carry W3C trace context across the Ray hops into the policy, value, and NeMo-Gym workers so one training step reads as a single trace. Signed-off-by: Raj Singh <rajsin@nvidia.com>
Signed-off-by: Raj Singh <rajsin@nvidia.com>
…lity stack Signed-off-by: Raj Singh <rajsin@nvidia.com>
startup_span opened a span and yielded nothing, so `with startup_span() as span:` bound None even with the setup group enabled. It now yields the span, matching setup_span. Also from review: - shutdown_telemetry narrows its except to (ImportError, AttributeError) and warns instead of logging at debug, matching ensure_metric_group_registered. - Annotate tracer as Optional[Tracer] on the span helpers, type the trace_fn / umbrella_trace_fn decorators via the existing _F, widen safe_set_span_attributes to dict[str, Any], and spell _metric_specs' real return type. - SingleControllerActor uses the handle init_telemetry_worker returns rather than discarding it and re-reading the global, matching trajectory_collector. Signed-off-by: Raj Singh <rajsin@nvidia.com>
Moves both the `[tool.uv.sources]` pin and the matching `override-dependencies` entry from b71263cc to b0f977d4, picking up two upstream commits: - `feat: Add single trace id for the entire job run based on config` (#61) adds `NemoLensConfig.single_trace_id` (default False, opt in via `ENABLE_SINGLE_TRACEID`) and a `DeterministicTraceIdGenerator`. The provider only selects it when the flag is set and a job seed exists, so our runs keep the existing `SeedIndependentIdGenerator`. - `refactor(semconv): organize attribute definitions and resource mappings` (#66) splits `semconv.py` into a package. `semconv/__init__.py` star-imports `attributes.py`, so `NV_DL_RANK` / `NV_DL_WORLD_SIZE` stay importable from `nemo.lens.semconv` as before. No module we import was touched other than those, so every lens call site keeps its signature. Relocked with the uv version pinned in docker/Dockerfile (0.11.28) to keep the lockfile at `revision = 3`; the re-resolution also consolidates gitpython onto 3.1.62 for the sglang extra combination, which previously resolved to 3.1.59. Signed-off-by: Raj Singh <rajsin@nvidia.com>
The single-controller train pump retries every 5ms while the replay buffer is empty, and wrapped each sleep in its own `rl.idle.buffer_starvation` span. A startup stall therefore arrived as thousands of near-identical spans -- ~10.5k for a minute of waiting -- which buries the trace it was meant to explain. Async GRPO shares the category name but polls at 0.5s, so it needs no coalescing and keeps wrapping its bare sleep. `start_efficiency_span` returns a started span the caller ends by hand, since a wait that spans loop iterations cannot be a `with`. It pointedly does not attach the span to the context: held open across `await` points, a current span is copied into every task created during the wait, which would reparent unrelated rollout work under an idle span. It returns None when the group is off, so the `is not None` guards a hand-managed span needs are also the telemetry-off no-op. The pump opens the span on the first starved poll, closes it when a batch becomes selectable, and closes it again after the loop for the two exits that continue the run (the `break` on a target that fell to what is already dispatched, and the condition going false mid-wait). The exits that end the run are deliberately uncovered, since dropping the span costs nothing there; a try/finally would have re-indented 383 lines of loop body to say the same thing. Widening the span past the sleeps is only sound because a *starved* poll never reaches the data plane: evict() returns 0 before calling remove() when nothing is stale, and select() gives up in _finalize_selection before claiming. So the span still has no bucketed children for a rollup summing by rl.bucket to count twice. The new rl.idle.polls attribute is what keeps the wider duration readable -- ten seconds is two thousand clean polls or two hundred whose selection ran long, opposite diagnoses. The bucket and category rules move into a helper shared with efficiency_span so the unbucketed-category list cannot drift between the two ways of opening one of these spans. Signed-off-by: Raj Singh <rajsin@nvidia.com>
The stub restated the adapter's constructor keywords, so the `checkpointing` argument the factory now passes made every test using it fail on an unexpected keyword. Which keywords reach the adapter is already checked against the adapter itself, so the stub takes whatever it is given instead of keeping a copy that goes stale. Signed-off-by: Raj Singh <rajsin@nvidia.com>
Move the teed rows out of nemo_rl/telemetry/metrics.py and next to the code that produces their keys, so a rename touches one place: the training scalars to algorithms/metric_utils.py, kl_penalty/approx_entropy to the loss that computes them, critic/loss to ppo.py, and the vllm/* family to a torch-free module beside compute_engine_step_metrics. Telemetry keeps the tee and the registration. nemo_rl/telemetry/vocabulary.py holds what both sides need and imports nothing, which is what lets a torch-heavy owner declare a row and lets telemetry stay importable without the training stack. The registry key is derived from the series name rather than typed a second time, and the init/total category, the timing/setup prefix and the setup field-naming rule each get a single home. ALL_GROUPS and UMBRELLA_GROUPS are read off RLSpanGroup instead of restated, so a group added to the class cannot go unregistered. Signed-off-by: Raj Singh <rajsin@nvidia.com>
…ollout path - accepts_trace_context attached the caller's context across the generator's yield. Ray abandons a cancelled streaming call without closing the generator, so the token was detached later by the garbage collector in a different context and OTel logged 'Failed to detach context' twice per cancelled shard. Attach around each __anext__ instead, and give NemoGym the same shape through streaming_umbrella_span. - rl.gym.run_rollouts is dispatched once per prompt by the single controller, so it now takes U_PER_PROMPT there and keeps U_ROLLOUT on the sync batch path, restoring per_step's promise to scale with steps. - telemetry alone installed the data-plane wrapper, and is_metrics_client answered on type, so observability.enabled: false still paid for per-worker stat polling and a payload walk per put. - nv.dl.campaign.stage was RL for every process, SFT runs included. - the vLLM step-metric reads swallowed RayActorError, so the single controller's handlers never fired; the two methods also shared one warn key. - the carrier-refusal warning keyed on a repr for .options() handles, so it repeated per dispatch and never named the method. Signed-off-by: Raj Singh <rajsin@nvidia.com>
…histogram - rl.sc.generate_and_push built its attribute dict before checking whether PER_PROMPT was on, once per prompt. Gate first and fall back to a shared NO_SPAN, now exported from instrumentation so the data plane and the single controller hold one instance. per_prompt_scope() is still entered either way, since the data-plane put inside reads it. - the batch generate() wall clock went to gen_ai.server.request.duration, which semconv defines per inference request; one call here covers a batch across every data-parallel shard. It gets rl.vllm.batch.duration instead, declared next to the code that records it. - startup in run_grpo_single_controller ran outside the try, so a setup failure skipped shutdown_telemetry() and dropped the buffered spans that would have explained it. - run_grpo's startup timer labels did not match their span names. - docs: adding a metric now means declaring it next to its producer; drop the claims that telemetry.enabled is per-rank, that lens names the rl.* attributes, and that the engine series need _total aliases. Signed-off-by: Raj Singh <rajsin@nvidia.com>
…passes Each algorithm's validate() repeated the same four things: resolve a tracer, open an U_EVALUATE umbrella, build the span name, and scope the enclosed generation to OVERHEAD because validation tokens are scored and discarded. Three of the six had already dropped the bucket scope, so their validation passes read as goodput. Folding the pairing into the helper makes that property part of what an eval span *is*. Signed-off-by: Raj Singh <rajsin@nvidia.com>
…M gaps - the carrier tests all ran against fakes, which cannot catch the thing that actually broke: Ray validates .remote() arguments on the caller, against a signature it gets by unwrapping. Put the methods on a real actor, across all three shapes plus an .options() handle, and assert each runs under the caller's trace. - a counter that went backwards, which is what a mid-step engine restart looks like. - the engine series names are copies of vLLM's; check them against vLLM's own declarations so a rename fails a test rather than dropping a dashboard line. Signed-off-by: Raj Singh <rajsin@nvidia.com>
Nineteen comment blocks ran five lines or longer inside a function body, several of them twelve to nineteen. Keep the constraint each one records and drop the argument around it; the durable rationale for the span groups already lives in the RLSpanGroup docstrings, which is where a reader looking for it will go. Signed-off-by: Raj Singh <rajsin@nvidia.com>
Use the metric key constants at the sites that build the logged dicts instead of retyping the strings, across grpo, grpo_sync, the single controller utils and the two policy workers. Enter per_prompt_scope() around the token-capture finalizer dispatch, so the gym and data-plane spans under that branch are gated on PER_PROMPT like the sibling branch already is. Widen the streaming umbrella helper's yield to ContextManager[Any]: use_span yields the Span, which failed pyrefly's invalid-yield check. Add vocabulary.py and metric_names.py to pyrefly.toml. Both type-check clean, so the CI whitelist step fails while they are missing from it. Guard the four tests that build a lens handle with requires_lens, and drop test_run_window_categories_match_the_efficiency_summary: utils.py now imports the category set rather than restating it, so there is no second copy left to drift. Reword the docs and the drift-guard docstring that called rl.vllm.generate worker-side; it is driver-side on the sync path, and the async path has no generation span yet. Signed-off-by: Raj Singh <rajsin@nvidia.com>
|
/ok to test 5b4217c |
|
Thanks for the detailed pass — the reproductions and the suggested one-liners made this quick to act on. All 13 open threads are fixed in 5b4217c (rebased onto The five merge blockers1. Four failing unit tests. All fixed. In
In 2. CI Lint. Fixed, and green on this head: run 35940079731. The yield is Worth flagging: fixing that exposed the step immediately after it. 3. The Ray fixture. The module-scoped 4. 5. PR description rev. Corrected to The other eight threadsThe docs and comment leftovers are all done: The token-capture branch enters
Registration now closes: On the comment-length thread I took all nine suggested one-liners as written, then re-ran an AST scan restricted to PR-added lines (comment runs over 4 lines whose first line is inside a function body, intersected with the diff against the merge-base). It found four more the table had not listed — One more fix, not from your reviewFour tests this PR added in VerificationLocally: The macOS limitation from the description still applies — the lockfile is Linux-only, so anything needing a real lens handle, |
`_stream_rows` now passes `per_prompt=in_per_prompt_scope()` when it dispatches `run_rollouts`, so the fakes standing in for that actor method reject the call with a TypeError. Seven doubles across four files take the kwarg with the same `False` default the real method has. The eighth, in test_run_async_nemo_gym_rollout_streams_complete_prompt_groups, is left alone: it stands in for the sharded dispatch in `rollouts.py`, which passes its arguments positionally and never sends `per_prompt`. Only one of these failed CI, because `addopts` carries `-x` and the run stopped at the first failure. The other three files had not been reached. Signed-off-by: Raj Singh <rajsin@nvidia.com>
|
/ok to test b3552d8 |
Brings #4052 (single-controller telemetry; moves nemo-lens to b0f977d and adds the aiohttp extra), #4252 (reward-model environment expected score, which fixes `L0_Unit_Tests_Environments` on the previous merge), #4248, #4234, #4254 and #4171. Only uv.lock conflicted: regenerated from this branch's lock; nemo-lens moves to main's rev and the three opentelemetry-instrumentation packages its aiohttp extra needs are added. Nothing else moved, and nothing resolves below main except the deliberate numpy / llguidance / flash-attn pins. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
What does this PR do ?
Extends NeMo-RL's OpenTelemetry/nemo-lens layer to cover the async single-controller path, propagates trace context across Ray actor boundaries so a run reads as one trace, and adopts nemo-lens's
SpanRegistryand consumer-driven metric registry.Before this, the sync GRPO path was instrumented and everything else was not. A recent run's 1,200 data-plane spans and the entire NeMo-Gym leg were unreachable from the job trace, because Ray does not carry OTel context across an actor boundary and every worker span started its own trace.
What lands here:
rl.sc.*spans covering the actor's phases (step, logprobs, value inference, advantage, training, optimizer step, checkpointing) plus the per-prompt dispatch span. The run lives insideSingleControllerActor, so that is the process that opens the job span.dispatch_with_trace_context/@accepts_trace_contextpair wired through theTQPolicy,TQValueand teacher presharded entrypoints andNemoGym.run_rollouts, plus aiohttp client instrumentation so the HTTP hop into the Gym service keeps the trace. Ray validates remote signatures afterinspect.unwrap, which hides a wrapper's**kwargs, so the decorator also advertises the carrier on__signature__— Ray hits the same problem with its own_ray_trace_ctxand fixes it the same way.rl.startupoverinit_ray()andsetup(), withrl.setup.<phase>beneath, so the startup phases arrive as one waterfall instead of unrelated root traces.0.0, so a vLLM rename leaves a gap in the dashboard instead of a plausible-looking zero.telemetry.vllm_native_tracing, defaultfalse) — one span per request, so it is a debugging tool you switch on for a few steps, not something to leave on.nv.dl.campaign.stage="RL"on every process, so a backend collecting several stages of a model's lifecycle (pretrain → SFT → RL) can select this stage without matching on service names.Breaking change
TelemetryConfiglosesexport_strategy,export_rank,export_sample_rateandsampler_enabled, following nemo-lens deleting the rank-gating machinery they drove. Every process that enables telemetry now exports and labels itself by rank. Narrowing a fleet down is now a collector-side filter onnv.dl.rank, ortelemetry.enabled: falseon the ranks that should stay quiet.A stale key in an existing YAML still parses (
TelemetryConfigallows extras) and is not projected into the worker environment, so it reaches nothing rather than erroring. Migration note indocs/observability/configuration.md.Issues
None.
Usage
Off by default. Enable per run:
export OTEL_EXPORTER_OTLP_ENDPOINT=http://collector:4318 uv run examples/run_grpo_single_controller.py --config examples/configs/grpo_math_1B.yamlThen filter by the
nemo.run.idresource attribute in your backend to isolate the run. Adding a span to a new phase:See docs/observability/ — configuration, span groups, metrics, extending, and vLLM tracing.
Before your PR is "Ready for review"
Pre checks:
On the unit-test box: the suite could not run locally —
uv run --group test pytestrefuses on macOS arm64 because the lockfile's supported environments are Linux-only. Locally verified instead:ruff checkandruff formatclean on all changed files,py_compileon all 54 changed Python files, and the seventest_source_drift.pyguards run green (by putting the pinned nemo-lenssrctree onsys.path). The rest of the telemetry suite is on CI.Additional Information
Test coverage. ~2,600 lines of new unit tests across
tests/unit/telemetry/,tests/unit/data_plane/,tests/unit/models/generation/,tests/unit/distributed/andtests/unit/environments/.test_source_drift.pyis worth a look during review: it parses the sources and fails the build when a declaration drifts from the call sites that use it — every emitted span name must be documented, every registered span group must have an emitter, every efficiency timer must be declared, every@accepts_trace_contextmethod must be dispatched with a carrier, and every teed logger key must be emitted somewhere. Those are the failure modes nothing at runtime can catch, since a missing metric is indistinguishable from a step that did not report one.Dependency.
nemo-lensmoves to revb0f977dforSpanRegistryand the metric registry, via[tool.uv.sources]plus a[tool.uv] override-dependenciesentry — megatron-core pins v0.2.0 in its own sources and the workspace resolves a singlenemo-lens, so without the override the two git URLs are a hard conflict. Drop the override once Megatron-LM bumps to the same or newer rev.uv.lockchanges are 52 insertions: thenemo-lensrev bump plus the three OpenTelemetry aiohttp packages the extra pulls in, and nothing else.Known gaps, documented rather than hidden:
async_ppo_traininppo.pyis still uninstrumented (timer-only). Pre-existing, untouched here.docs/observability/span-groups.md.token_capturefinalizers configured, the single controller's per-prompt dispatch commits through the finalizer pool rather thangenerate_and_push, and that branch opens norl.sc.generate_and_pushspan. It landed on main after this branch was written; recorded in the coverage-gap table and left for a follow-up.idle/validationisidleas a metric andoverheadon the spans covering the same seconds. Both are true of different fleets; attributing it properly needs per-fleet accounting.Performance. Every instrumentation site gates on its span group before allocating. The data-plane wrapper — the most frequent call site in the repo, once per prompt on the rollout path — checks
is_span_group_enabledbefore building the span name or the attribute dict, so a run with telemetry off pays nothing.