feat(telemetry)!: instrument the single-controller path and propagate traces - #4052
Open
rrs45 wants to merge 21 commits into
Open
feat(telemetry)!: instrument the single-controller path and propagate traces#4052rrs45 wants to merge 21 commits into
rrs45 wants to merge 21 commits into
Conversation
…vLLM engine metrics Three blind spots made a NeMo-RL trace hard to read. The single-controller path emitted no spans at all, so its steps were one opaque block; everything before the first step (Ray init, worker construction, the initial buffer fill) was invisible, though it is often where a run's first minutes go; and the vLLM engine's own counters never left the worker. Adds the rl.sc.* phase spans plus data-plane spans for transfer-queue traffic, an rl.startup umbrella with per-phase rl.setup.* spans and a matching rl.setup.duration metric, and tees the engine's token, sequence length and request-outcome deltas alongside the training scalars. Worker model loads get spans too. The driver can only see setup() as one block because the worker builds run concurrently in threads OTel context does not reach, so rl.policy.load_model / rl.value.load_model are opened in the workers themselves, carrying rl.backend to compare megatron against dtensor. Loading and sharding a large checkpoint is routinely the longest phase of a startup and was previously unattributed time. Two of the new spans are emitted once per prompt rather than once per step, so a 10k-prompt rollout gets ~20k of them where every other group emits a fixed handful. They are gated behind their own per_prompt group, absent from per_step so that preset's cost keeps scaling with steps rather than with dataset size. The rollout's data-plane put is one of the two, and the client that emits it is shared with the batch stages, so the group comes from a per_prompt_scope contextvar the rollout path enters rather than from the op name -- the same shape bucket_scope already uses for reclassification. Umbrella groups now say so at the call site: RLSpanGroup.U_* aliases with umbrella_span/umbrella_trace_fn, so a reader can tell without a lookup whether a span feeds the goodput rollup. Concurrent work stays unbucketed -- its durations sum past wall time and cannot be summed into a bucket. Signed-off-by: Raj Singh <rajsin@nvidia.com>
The `rl` metric group is declared through lens's consumer-driven registry (`MetricSpec` / `register_metric_group` / `record_metrics`), which landed after v0.2.0. The pin was still on v0.2.0, so `ensure_metric_group_registered` took its fail-soft path and no RL metric was ever exported. megatron-core pins v0.2.0 in its own [tool.uv.sources] and the workspace resolves a single nemo-lens, so the two git URLs are a hard resolution conflict; an override-dependencies entry reconciles them. That is safe because the only removal in the refactor was the shipped `instruments.rl` module, which megatron-core never imported. Signed-off-by: Raj Singh <rajsin@nvidia.com>
…st the span tables `distillation.validate` drives generation through the same `run_multi_turn_rollout` path as training, so without a `bucket_scope(Bucket.OVERHEAD)` its generate spans stay `productive` and a validation pass reads as goodput. grpo and ppo already scope it; distillation was the odd one out. The span-group table also named six spans no site emits — `save_checkpoint`, `collect_rollouts`, `compute_logprobs`, `compute_rewards`, `compute_advantages`, `policy_update` — all renamed when the span tails were aligned to their timer keys. The drift guard only enforces emitted <= documented, and these are all `rl.<algo>.*` forms its regex skips, so nothing caught them. vllm-tracing.md named the same stale rollout parent. Signed-off-by: Raj Singh <rajsin@nvidia.com>
…PO rollouts as ppo `record_inference_metrics` records `gen_ai.server.request.duration` only when `request_duration_s` is passed, and the vLLM call omitted it — so the metric that metrics.md documents was never emitted. Time the driver-side `generate` / `generate_text` call and pass it. The async trajectory collector hardcoded `rl.grpo.generation`, but `async_ppo_train` builds the same collector, so async PPO rollouts were labelled grpo. Pick the name in the branch that already discriminates the two master configs. Both names stay spelled out so they remain greppable from the emit site. Signed-off-by: Raj Singh <rajsin@nvidia.com>
…gating Bumps nemo-lens to b71263c (11 commits on from the previous pin), which carries three breaking changes to the API NeMo-RL consumes. `SpanRegistry` replaces the shipped `SpanGroup` class. Lens now names no span groups of its own -- a consuming library registers what it emits under its own namespace -- so `RLSpanGroup` becomes a bag of `str` constants that `register_span_groups()` declares at import, and it declares the eight groups lens used to ship (`job`, `step`, ...) alongside the RL-specific ones. The names are unchanged, so every existing `span_groups` value and all ~220 call sites keep working; `managed_span` and `trace_fn` take the group as a plain string, which is what makes that possible. `all` is no longer a preset NeMo-RL owns: lens reserves it and resolves it as a wildcard over the live registry, which is what keeps it correct when Megatron registers alongside. An unrecognised spec entry is no longer fatal. Lens returns it as `pending`, because a registry is per-process while a spec is job-wide, so the driver reports a typo as a warning naming the registered groups rather than raising. `rank` / `world_size` leave `setup_telemetry` and become the `nv.dl.rank` / `nv.dl.world_size` resource attributes. Rank-based export gating goes away with the `strategies.py` and `sampling.py` modules lens deleted: `export_strategy`, `export_rank`, `export_sample_rate` and `sampler_enabled` are removed from `TelemetryConfig` and from the worker env projection, along with `init_telemetry_worker`'s `always_export` and the `_unrank` / `_always_export` helpers that existed to neutralise those filters for singleton processes. Every process that enables telemetry now exports and labels itself by rank; narrowing a fleet down is a collector-side filter on `nv.dl.rank`, or `telemetry.enabled: false` on the ranks that should stay quiet. A stale `export_strategy` in an existing YAML still parses and is not projected anywhere, so it reaches nothing rather than erroring. Docs and tests follow, including guards that the removed fields cannot quietly return and that registration really is an import side effect. Signed-off-by: Raj Singh <rajsin@nvidia.com>
…workers Ray does not carry OTel context across an actor boundary, so every span a worker opened started its own trace -- a recent run's 1,200 data-plane spans and the entire Gym leg were unreachable from the job trace. Adds a dispatch/receive pair: dispatch_with_trace_context and trace_context_kwargs at the call site, @accepts_trace_context on the worker method. Wired through the TQPolicy, TQValue and teacher presharded methods and NemoGym.run_rollouts, plus aiohttp client instrumentation so the HTTP hop into the Gym service keeps the trace. Ray validates remote signatures after inspect.unwrap, which hides a wrapper's **kwargs, so the decorator also advertises the carrier on __signature__; without that every Gym rollout raised TypeError once telemetry was enabled. A dispatch to an undecorated method retries without the carrier rather than failing the run. Every span emitted per step now shares one trace. The six traced_worker_init constructor spans still root their own, since they open during actor construction before any call exists to carry a carrier. Also: - Bound shutdown_telemetry to its own timeout_ms. Nothing below honours it: the SDK's force_flush accepts timeout_millis and drains the queue anyway, and the provider.shutdown() calls after it take no timeout at all. Against an unreachable collector, 3k buffered spans took 36s under a 5s budget -- enough to blow the 15s ray.get async GRPO uses to save the collector's last rollout spans, killing the collector and losing them instead. - Drop four span groups that were registered and bucketed but never emitted. reference_policy was the misleading one: it sat in per_step, so the preset advertised spans it could not deliver, and the work it names is already one policy_and_reference_logprobs span under logprob. A drift guard now asserts registered <= emitted so this cannot return. - Add job to per_step, giving the preset a run-scoped span for current_trace_carrier to hand the trajectory collector. - Gate Gym's aiohttp instrumentation on per_prompt rather than rollout. Its volume is prompts x turns x tool calls, which is exactly what per_step's cost-scales-with-steps contract excludes. - Stop doing telemetry setup ahead of the span-group gate in the data-plane client: 1.82us -> 0.045us per op with telemetry off. - Shield the two async shutdown awaits so a cancellation during cleanup cannot replace the exception that ended the run. Signed-off-by: Raj Singh <rajsin@nvidia.com>
A backend collecting several stages of one model's lifecycle (pretrain -> SFT -> RL) needs to select the RL stage without matching on service names, which differ per launcher. Constant rather than configurable: a process running this package is the RL stage. Seeded into both the driver and the worker attribute builders. A resource is per process and the two paths never see each other, so tagging only the driver would leave every policy, value, vLLM and NeMo-Gym worker out of that view. Seeded first in the worker dict so an explicit resource_attributes argument still wins, which the nemo_gym actor relies on. Lens owns no semconv constant for this name, so it stays an RL-side literal alongside rl.algorithm and nemo.precision. Signed-off-by: Raj Singh <rajsin@nvidia.com>
… snapshot
Three defects found in self-review, all in code this branch added.
efficiency_span("idle/buffer_starvation") wrapped the whole batch-selection
block rather than the wait inside it. The block evicts stale groups, which
clears their rows through the data plane and so opens an overhead-bucketed
rl.data_plane.clear inside an idle-bucketed parent -- and a rollup that sums
durations by bucket has no notion of nesting, so it counts that interval twice.
The span's own docstring forbids exactly this. Two more consequences of the same
scope: the `continue` sits inside the `with`, so during starvation the span
reopened every 5ms (~200/s, and efficiency ships in the per_step preset), and
the healthy path -- buffer full, batch returned at once -- reported as
starvation. Now wraps only the sleep, as the sync analog in grpo.py does. The
enclosing exposed_generation timer keeps its scope: it predates this branch and
feeds a different breakdown, and the single controller logs no efficiency/*
scalars for the span to disagree with.
snapshot_step_metrics called _get_raw_spec_counters unguarded. This branch
broadened that reader from grepping for spec_decode to walking every series the
engine exposes, and its sibling get_step_metrics was given a try/except whose
comment says an exception there "would end the run" -- but the snapshot half
reads the same version-dependent surface, and vLLM's get_metrics_snapshot()
raises AssertionError on a metric type it does not recognise. The callers do not
defend it either: grpo.py is bare and the single controller catches only
RayActorError. Left as None on failure so the paired read returns {} rather than
delta-ing against a stale baseline.
SingleControllerActor initialised telemetry with no resource_attributes. It is
built directly rather than by RayWorkerGroup, so nothing sets NRL_WORKER_GROUP
and its resource carried no rl.worker_group -- reporting rank 0 of 1 with no
group name, which is exactly what the launcher driver reports, leaving every
rl.sc.* span indistinguishable from the driver's. Named explicitly, as NemoGym
already does for the same reason.
Signed-off-by: Raj Singh <rajsin@nvidia.com>
The driver warned about an unresolvable telemetry.span_groups entry itself, and
got the scope wrong in both directions: lens resolves the spec against every
namespace registered in the process, so an entry is pending only when nothing at
all owns it -- yet the message said the names "match nothing NeMo-RL registers"
and listed only RLSpanGroup.ALL_GROUPS, omitting whatever Megatron registered
and calling a real group unregistered. Lens already emits an accurate version
from set_span_group_spec inside the setup_telemetry call just below, naming the
registered groups, presets and namespaces, so one typo produced two warnings
with the wrong one first. Dropped ours. span_groups is still imported here for
the registration side effect, which is what setup_telemetry resolves against.
ensure_metric_group_registered caught every exception and reported one cause, "a
lens build with the metric registry is required". Lens raises ValueError from
register_metric_group for a duplicate key, an empty spec list, or a second
registration of the group -- all our own declaration bugs, none of them a
dependency problem, and all of them sent the reader to the wrong file. Split so
ImportError/AttributeError keeps the version message and ValueError names the
group and quotes the cause. Both are permanent for the process, so a new flag
stops the attempt repeating: the tee runs once per log_metrics, and the old code
rebuilt all 17 MetricSpec objects every step for a group that would never
register, with warn_once demoting the repeats to debug so nothing showed it. The
telemetry conftest resets the flag too, or one test forcing a failure would
disable the tee for every test after it.
_get_raw_spec_counters imported four names from vllm.utils inside the function
body while the same module is imported at module top, with no circular-import or
optional-dependency reason. Its docstring also still promised
AssertionError "if called before vLLM engine is initialized"; there is no
assert, it returns {}.
Signed-off-by: Raj Singh <rajsin@nvidia.com>
Three places -- metrics.md, extending.md and the _TeedMetric docstring -- said a drift test failed the build when a declared logger key stopped being emitted. No such test existed: test_source_drift guarded span names, span groups, efficiency categories and carrier dispatch, and never read logger_key. That claim is the whole justification for keeping the logger key and the OTel name in one row, and it is the one failure mode nothing at runtime can catch -- the tee reads the key with .get, so a miss is indistinguishable from a step that did not report the value, and the result is a gauge reporting a flat line rather than an error. An earlier design had exactly this bug in three entries. Direction is declared <= emitted. Declarations are parsed out of metrics.py for the reason the module header gives, and emit sites are matched in the two shapes these scalars are actually produced in: a literal key in a dict display, and a subscript assignment onto an accumulating dict. Deliberately not "every string constant in the file", which would also match a key named only in a docstring and let a rename that updated the prose but not the code pass. Both halves carry a non-empty assertion so neither a moved declaration nor a stale matcher can turn this into a tautology. Passes on all 14 declared keys, and mutation-checked to fail on a renamed emit site. Signed-off-by: Raj Singh <rajsin@nvidia.com>
…laims Filtering instructions that return nothing are worse than no instructions. configuration.md said to "filter by run_id in your backend" and metrics.md said every rl.* point carries the run_id resource attribute. Lens emits it as nemo.run.id -- run_id is the config field you set, not the key your backend indexes -- so a user pasting run_id="abc" into Grafana gets zero rows with nothing to explain them. Corrected in five files, the contrast between the two spellings called out where the field is documented, and nemo.run.id added to the resource-attribute table, which omitted it entirely. The rest: - extending.md's worked example opened rl.sc.generate_and_push with U_ROLLOUT. The span is emitted under U_PER_PROMPT, and span-groups.md documents it that way in two places, so copying the example would move it into per_step -- the dataset-scaling regression per_prompt exists to prevent. The same section counted six U_ aliases; there are seven, and the missing one is U_PER_PROMPT. - span-groups.md still gated emission on "the rank is exporting" in two places. Rank-based export gating was removed earlier on this branch. - metrics.md said rl.efficiency.pct is tagged window="step" unconditionally. The tee derives the window from the denominator and defaults to run; async PPO passes no per-step wall time, so its points really are window="run", and by the doc's own argument that ratio climbs toward 100% as the run lengthens whatever the idle time does. Now says to read the attribute, and the run row lists pct. - Eight span attributes were documented nowhere: the five rl.data_plane.* tags set on every transfer-queue op, plus rl.target_step, rl.critic_epochs and rl.ppo_epoch from the single-controller instrumentation. The op vocabulary listed four of eleven. - The telemetry README advertised a throughput metric no MetricSpec declares, and claimed telemetry no-ops when nemo-lens is absent. Lens is a base dependency and init_telemetry_driver imports it above the enabled check, so a missing lens raises ImportError even with telemetry off; six launchers repeated the claim in a comment. index.md listed two of five metric families. Also reworded the exposed_generation sentence, which described the starvation span's old scope. Signed-off-by: Raj Singh <rajsin@nvidia.com>
Both would send a reader looking for something that is not there. span-groups.md offered "enable per_prompt and drop data_plane" as a way to keep the rollout view while shedding the ~20k per-rollout spans. It does nothing: inside a rollout MetricsDataPlaneClient gates the put on U_PER_PROMPT, not DATA_PLANE, so per_prompt alone still emits both spans and dropping data_plane only removes the batch-shaped ones. The same page already said as much twelve lines earlier and again in the next subsection, so the page contradicted itself twice. vllm-tracing.md steered users to the vllm/* metrics for queue time and preemptions, in the caveat and again in the layer-picking summary. Neither is teed: _KEPT_COUNTER_NAMES keeps the token counters, the length histograms and request_success, and nothing else -- num_preemptions_total is literally the "not kept" fixture in test_vllm_utils. Says so now, and points at Layer 2 as the only way to see either today. Signed-off-by: Raj Singh <rajsin@nvidia.com>
github-actions
Bot
added
the
Documentation
Collaborator
Author
|
/ok to test 493d853 |
rrs45
changed the title
Catches the telemetry branch up to main, which had moved 22 commits and tripped pr-branch-up-to-date-check (max 10 behind). Eleven files conflicted. Resolutions, in the order a reviewer would want them: - single_controller.py: main split _dispatch_one_prompt into a token-capture branch (generate_for_finalization + the finalizer actor pool, #3837) and moved the legacy generate_and_push loop into an else. Git interleaved our per-prompt span into the new branch, so the function was restored to main's exactly and the per_prompt_scope() + rl.sc.generate_and_push umbrella re-applied to the legacy loop only. The token-capture branch is left uninstrumented: it commits through the finalizer pool rather than here, so its attempt boundary is a different shape than this span describes. Recorded in the coverage-gap table in docs/observability/span-groups.md. - run_grpo_single_controller.py: keeps our startup_span / setup_span phases and adopts main's VLM processor (#4009), with the processor-aware tokenizer construction moved inside setup_span ("tokenizer") and processor= threaded to setup_single_controller. - run_grpo.py: adopts main's make_policy_factory() helper in place of our inline factory selection, keeping setup_span("workers"). - nemo_gym.py: keeps @accepts_trace_context and the run_rollouts / _stream_rollouts span split, alongside main's ledger control plane; both signatures widen to main's 4-tuple yield. - rollout_manager.py: keeps dispatch_with_trace_context and unpacks main's added resolved_agent_ref. - factory.py: adopts main's LocalDataPlaneConfig-aware observability lookup, keeping the telemetry_enabled_in_env() arm that installs the wrapper for its spans. - vllm_worker.py: keeps umbrella_trace_fn(U_MODEL_INIT) on _load_model beside main's _refit_with_reload_api_enabled. - config.py, worker_mixin.py, virtual_cluster.py, vllm_generation.py: both sides added adjacent imports or fields; all kept. The textually merged uv.lock was corrupt (a missing source field on opentelemetry-instrumentation-aiohttp-client, which matched more than one package), so it was regenerated from main's with uv 0.11.28 to match CI's revision 3. The result is 52 insertions: the nemo-lens rev bump and the three otel aiohttp packages, nothing else. Both uv lock --check runs pass. Verified: all 54 changed Python files compile, ruff check and format are clean, and the seven test_source_drift.py guards pass -- including the carrier guard, which confirms run_rollouts is still dispatched with trace context after the signature change. Signed-off-by: Raj Singh <rajsin@nvidia.com>
github-actions
Bot
removed
the
CI:L1
Picks up periodic rollout checkpointing (#3924) and the partial-rerun routing fix (#4048), which landed while the previous merge was being verified. One conflict, in the critic training block: main now raises _optimizer_commit_in_progress before the value_training timer so a periodic snapshot cannot land mid-update. Kept, with our rl.sc.value_training span re-applied around the timer and the flag left where main set it -- outside the span, since it guards the whole irreversible update rather than the measured region. Signed-off-by: Raj Singh <rajsin@nvidia.com>
Collaborator
Author
|
/ok to test de7ab13 |
Collaborator
Author
|
/ok to test de7ab13 |
…gates to run_rollouts became a thin umbrella_span wrapper that delegates the streaming to a new _stream_rollouts, so the span covers the whole generator without indenting the body. The mock self this test drives run_rollouts with predates that split and defines only the attributes the old body touched, so the delegation raised AttributeError and the test streamed nothing. Point the mock at the real helper rather than reverting the split: the test already stands in for self attribute by attribute, so supplying one more is consistent with how it is written. Signed-off-by: Raj Singh <rajsin@nvidia.com>
The docs claimed vLLM's OTLP span exporter is gRPC-only and so cannot reach an http/protobuf endpoint. It is not: get_span_exporter defaults to grpc but supports http/protobuf, selected by OTEL_EXPORTER_OTLP_TRACES_PROTOCOL. Readers were being sent to stand up a gRPC collector they may not need. The real trap is narrower, so document that instead: vLLM reads only the traces-specific protocol var, never the generic OTEL_EXPORTER_OTLP_PROTOCOL lens honours, so an http/protobuf run gets vLLM speaking gRPC at an HTTP port. Also document that the vllm_native_tracing switch governs engine spans alone. vLLM's worker processes call maybe_init_worker_tracer unconditionally and gate purely on OTEL_EXPORTER_OTLP_TRACES_ENDPOINT being present in the environment, which init_ray spreads cluster-wide -- so exporting it in a job script turns on per-request worker spans with the opt-in still false. Signed-off-by: Raj Singh <rajsin@nvidia.com>
Collaborator
Author
|
/ok to test bec17bd |
Catches up with 5 commits; only #3730 (colocated MInf) conflicted, in single_controller.py (3 hunks) and test_single_controller_actor.py (1). Resolutions: - __init__: main moved validate_single_controller_config out of the actor into single_controller_utils/setup.py, so its call is dropped here. Kept our init_telemetry_worker + self._tracer setup, which main does not touch. - Refit: main relocated the whole block below the step bookkeeping and split it in two -- a colocated engine bound for a checkpoint now skips _sync_weights, offloads, and is woken after the save instead. Took main's structure and put efficiency_span("idle/refit_bubble") on the real weight sync only. The deferred branch is left unbucketed because it syncs nothing, and the wake that stands in for it runs after the save; bucketing the no-op would report a bubble whose duration says nothing about how long generation served stale weights. Non-colocated runs never reach that branch, so their refit attribution is unchanged. - Checkpointing: re-applied the rl.sc.checkpointing span onto main's new will_save_checkpoint condition, which replaced the inline enabled/period test. - Tests: kept our test_tracer_is_declared_on_the_class; dropped test_rejects_multiple_optimizer_steps_per_rl_step, which main deleted here because the validation it covers moved to test_setup.py. No uv.lock, pyproject or submodule changes. ruff and the 7 telemetry drift guards pass. Signed-off-by: Raj Singh <rajsin@nvidia.com>
Collaborator
Author
|
/ok to test 130756c |
The previous commit fixed one of three tests that call run_rollouts unbound with a hand-rolled self; CI surfaced them one at a time because the job stops at the first failure. Bind the helper on the other two rather than wait for another round trip, and the file already uses this idiom for _postprocess_nemo_gym_to_nemo_rl_result. Checked exhaustively: no test calling modified_class.run_rollouts is left without _stream_rollouts. Signed-off-by: Raj Singh <rajsin@nvidia.com>
Collaborator
Author
|
/ok to test 45c2002 |
On the single-controller path a logical train step is three separate driver->worker dispatches, and only the middle one carried the caller's trace context. begin/finish/abort_train_step_presharded neither accepted a carrier nor were dispatched with one, so mcore's megatron.grad_sync.* spans -- which finish_train_step reaches through finalize_model_grads -- started a new trace per rank per step instead of nesting under rl.sc.policy_optimizer_step. Wire all three the way train_microbatch_presharded already is: @accepts_trace_context on the worker entrypoint, **trace_context_kwargs() on the run_all_workers_single_data fanout. test_every_context_accepting_method_is_dispatched_with_a_carrier could not see this: it starts from the decorated set, so an entrypoint wired on neither side is invisible to it. Add a guard over the presharded family that requires both halves, and correct the trace_context_kwargs docstring example, which passed its data positionally to a method that asserts it was given no positional arguments. Signed-off-by: Raj Singh <rajsin@nvidia.com>
Collaborator
Author
|
/ok to test 7bf922c |
Conflict in teacher_worker_group.py: main added the opd_full_payload kwargs to the same common_kwargs dict this branch added the trace carrier to. Kept both, with the carrier spread last as at every other call site. Signed-off-by: Raj Singh <rajsin@nvidia.com>
Collaborator
Author
|
/ok to test 88cb61d |
rrs45
added
CI:L2
Collaborator
Author
|
/ok to test 246b047 |
rrs45
added
CI:Lfast
Collaborator
Author
|
/ok to test 246b047 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 revb71263cforSpanRegistryand 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.