feat(telemetry): add nemo-lens instrumentation and efficiency buckets

[rrs45]

What does this PR do ?

Add Nemo Lens instrumentation covering basic components along with efficiency categorization

Usage

Emitting OpenTelemetry data

Telemetry is opt-in and a zero-cost no-op unless (a) nemo-lens is installed and (b) it's enabled. No code changes are needed to get spans/metrics out of the built-in GRPO/algorithm loops.

1. Install the extra

uv sync --extra telemetry        # or: uv pip install 'nemo-lens[sdk]'

2a. Enable via env vars (quickest)

export NEMO_RL_OTEL_ENABLED=1
export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317   # your OTLP collector/backend
export NEMO_RL_OTEL_SPAN_GROUPS=default                    # default | per_step | all
uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml

2b. Or via the telemetry: config block

# examples/configs/grpo_math_1B.yaml
telemetry:
  enabled: true
  span_groups: default          # default | per_step | all
  export_strategy: single_rank  # single_rank | all_ranks | first_rank_per_node
  exporter: otlp                # otlp | console  (console = print spans to stdout, no backend)

Raw NEMO_RL_OTEL_* / OTEL_EXPORTER_OTLP_* env vars always override the YAML.

3. No backend handy? Print spans to the console

NEMO_RL_OTEL_ENABLED=1 NEMO_RL_OTEL_EXPORTER=console \
  uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml

4. Instrumenting your own code

Import the wrappers from nemo_rl.telemetry.instrumentation (not raw nemo-lens) so leaf spans automatically get the rl.bucket efficiency tag:

from nemo_rl.telemetry.instrumentation import managed_span, trace_fn
from nemo_rl.telemetry.span_groups import RLSpanGroup
# Context manager
with managed_span(RLSpanGroup.GENERATION, "rl.vllm.generate", **{"rl.backend": "vllm"}):
    outputs = model.generate(batch)
# Or as a decorator
@trace_fn(RLSpanGroup.POLICY_UPDATE, "rl.grpo.policy_update")
def policy_update(...):
    ...

Driver/worker lifecycle is already wired in the examples/run_<algo>.py entrypoints: init_telemetry_driver(config, algorithm="grpo") runs before init_ray() (so NEMO_RL_OTEL_* propagates to every Ray worker via the runtime env), and shutdown_telemetry() runs at the end.

What you get

• Spans per phase (rl.grpo.step, rl.vllm.generate, rl.grpo.policy_update, …), with leaf spans tagged rl.bucket ∈ {productive, overhead, idle, wasted} for offline goodput rollups.
• Metrics (rl.*) mirrored from the training Logger.
  See docs/observability/ for the full config reference, span-group catalog, and OTLP/collector setup.

Before your PR is "Ready for review"

Pre checks:

• [✔ ] Make sure you read and followed Contributor guidelines
• [✔] Did you write any new necessary tests?
• [✔] Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• [✔] Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

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

[saumishr]

Code review via review-pr skill. Solid opt-in telemetry design (fallbacks, env-before-Ray, span-group gating, good docs/tests). Three follow-ups below for worker init, shutdown flushing, and idle-bucket wiring.

[saumishr]

@rrs45

Telemetry coverage on v2 pathways

Follow-up from review: the instrumentation in this PR covers the legacy algorithm loops and standard examples/run_*.py entrypoints well, but the newer v2 training paths are largely uncovered.

Covered

• Legacy sync GRPO (grpo_train, data_plane.enabled=false) — step / rollout / reward / logprob / policy_update / checkpoint / evaluate spans
• Legacy async GRPO (async_grpo_train) — training-step phases (idle/generation gaps remain; see inline comment on instrumentation.py)
• PPO, DPO, SFT, RM, distillation — same driver-side span pattern
• examples/run_grpo.py, run_ppo.py, etc. — init_telemetry_driver() before init_ray(), shutdown_telemetry() at end
• vLLM generation (VllmGeneration.generate) — rl.vllm.generate spans + token metrics
• Logger — rl.* metric tee on any path using Logger.log_metrics

Running with policy.dtensor_cfg._v2: true through the legacy trainers still gets driver-side phase spans; worker-internal timing (e.g. DTensorPolicyWorkerV2) is not instrumented.

Not covered

Path	Gap
TQ / data-plane sync GRPO (data_plane.enabled=true → grpo_train_sync in grpo_sync.py)	No phase spans — PR does not touch grpo_sync.py. run_grpo.py still calls init_telemetry_driver, so metrics tee works, but no step/rollout/policy_update spans.
Single-controller async GRPO (run_grpo_single_controller.py → SingleControllerActor)	Not wired — no driver init/shutdown, no spans in single_controller.py.
Worker-side export (any path)	init_telemetry_worker() is never called from Ray actors — export_strategy has no effect on workers yet.
Idle/wasted buckets (async)	EFFICIENCY_CATEGORY_BUCKET is defined but not hooked to timer.time("idle/buffer_starvation") etc.

Bottom line: legacy in-memory GRPO gets the full span catalog; TQ sync gets metrics-only; SingleController gets essentially nothing today. Worth tracking as follow-up if v2 paths are the primary deployment target.

[rrs45]

/ok to test 784dce1

[rrs45]

/ok to test f8cafca

[copy-pr-bot]

/ok to test f8cafca

@rrs45, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[rrs45]

/ok to test f8cafca

[rrs45]

/ok to test 74b5dad

[rrs45]

/ok to test 74b5dad

[rrs45]

/ok to test aad00c1

[saumishr]

LGTM! Discussed offline that the SingleController path will be instrumented as a follow up.

[terrykong]

Review: telemetry instrumentation

Thanks for this — it's careful work, and a few things made reviewing it much cheaper. Every nemo-lens API call checks out against the pinned 0.1.0 wheel (I unpacked it and diffed the signatures). The "Not yet wired" table in nemo_rl/telemetry/README.md is an honest, specific gap list of the kind most PRs don't ship. And the AST drift-guard in tests/unit/telemetry/conftest.py, which keeps the restated category lists in lockstep with algorithms/utils.py, is the right answer to a duplication you can't remove — two of the comments below just ask you to point that same tool at two more places.

There are a lot of comments, so here is the shape of them. Most cluster into four themes and several are one-liners that fall out of a single decision:

1. Make nemo-lens a base dependency. If it's always installed, the nine try/except ImportError guards around it have no job, and ~163 lines of parallel no-op implementation go away — including a 46-line stub SpanGroup that reimplements lens's own presets. telemetry.enabled becomes the single switch. The comments on pyproject.toml, _fallbacks.py, the six guard sites, and docker/Dockerfile are all one decision.

2. Make the config the interface. All 12 telemetry: fields are projected 1:1 into NEMO_RL_OTEL_*, so the env layer carries nothing the config doesn't. Documenting both costs provenance — a run's settings should be recoverable from its config, not from someone's shell. Keep the standard OTEL_* vars; they're the OTel interop surface. Related: three constrained fields are typed as bare str, so nothing rejects a bad value at load time.

3. Don't map names centrally — standardize them at the source. _RL_OTEL_METRIC_MAP declares sets of metric names to be equivalent, and it has already drifted inside this PR: 3 of its 8 gauges can never receive a value, and 10 of its 16 candidate names are emitted nowhere. Same theme, different place: 39 of 45 timer.time() / span-name pairs disagree, which means correlating a slow span with its timing metric needs a mapping that exists only in the reader's head.

4. Close two drift gaps with the tool you already have. 47 span-name literals are restated by hand in a 332-line doc table with nothing checking they agree; and an idle/* timer added at a call site but not declared in WALL_CLOCK_EFFICIENCY_CATEGORIES is silently dropped — its time then counts as productive, so the efficiency number reads higher than reality.

Two are straight correctness bugs rather than shape: rl.efficiency.pct is tagged window="step" unconditionally even when the value is run-cumulative, and vllm_native_tracing escapes the enabled master switch because it's exported to env before the switch is checked.

On performance evidence — I don't think an overhead A/B is warranted, and I'd rather say so than leave it hanging. All 84 instrumentation sites are per-step or coarser; there's no per-token, per-microbatch or per-request span anywhere, and the comment at trajectory_collector.py:605-612 shows per-sample spans were explicitly considered and rejected. With telemetry off the residual cost is a dict copy and a ContextVar.get() against a step measured in seconds. The one thing worth adding to the PR description is the new teardown cost, which is bounded but real: up to 15s of ray.get plus a 5s force-flush.

Three things I checked and found clean, so they don't become churn later: the _fallbacks.py no-ops faithfully match the real lens contract (upstream managed_span also yields None for a disabled group, and no call site binds the yielded span); the Ray env propagation genuinely works (virtual_cluster.py:306 snapshots all of os.environ); and tests/unit/telemetry/ is collected by the L0_Unit_Tests_Other shard with nemo-lens actually present in the image, so the lens-installed branches do run in CI.

Caveat on this review's evidence: it was produced on macOS with no GPU against a Linux-only uv.lock, so no tests and no linter were run. Everything is from reading source, plus the unpacked nemo-lens 0.1.0 wheel and the pinned OpenTelemetry SDK. Anything depending on runtime behaviour should be treated as unverified.

[rrs45]

/ok to test 3e9783c

[rrs45]

/ok to test 53282e5

[rrs45]

/ok to test 0fc4d28