feat(telemetry): instrument NeMo-RL with OpenTelemetry via nemo-lens

[ahmadki]

What does this PR do ?

Adds optional OpenTelemetry instrumentation (distributed traces + metrics) across the NeMo-RL training loops, vLLM generation, and metrics logging, built on the shared nemo-lens library. Telemetry is fully optional — when nemo-lens is not installed (or telemetry.enabled is false) every instrumentation site degrades to a ~0-cost no-op via nemo_rl/telemetry/_fallbacks.py.

Highlights:

• Telemetry module (nemo_rl/telemetry/): TelemetryConfig (the telemetry: block of MasterConfig), the driver/worker lifecycle (init_telemetry_driver runs before init_ray() and exports NEMO_RL_OTEL_* env so Ray workers inherit it; raw env overrides YAML), RLSpanGroup presets, and the metrics tee.
• Traces: group-gated spans on GRPO (sync + async), PPO, SFT, DPO, RM and distillation — step / data_processing / rollout / generation / logprob / reward / advantage / policy(+value)_update / checkpoint / evaluate, plus a job span per trainer. Spans wrap the existing timer.time() blocks with no body re-indentation.
• vLLM: driver-side rl.vllm.generate / generate_text spans (nested under rollouts) + gen_ai.* token/latency metrics, plus opt-in vLLM-native OTLP engine tracing (NEMO_RL_OTEL_VLLM_NATIVE_TRACING).
• Metrics: Logger.log_metrics tees the standard train-prefix scalars (reward, loss, grad_norm, learning_rate, throughput, …) into rl.* metrics.
• Deps / image: new optional telemetry extra (nemo-lens[sdk]); the release Dockerfiles sync it into the driver venv so images ship telemetry-ready.
• Tests + docs: tests/unit/telemetry/ and a 7-page docs/observability/ guide (+ nemo_rl/telemetry/README.md).

Backend-agnostic: exports OTLP to any OTLP-compatible backend or an OpenTelemetry Collector via the standard OTEL_EXPORTER_OTLP_* env vars; console/JSON output via NEMO_RL_OTEL_EXPORTER=console.

Issues

N/A

Usage

Enable telemetry and point it at a backend (or the console). No code changes to your run — just env or the telemetry: config block:

# Console (offline; prints spans + metrics as JSON to stdout)
NEMO_RL_OTEL_ENABLED=1 NEMO_RL_OTEL_EXPORTER=console NEMO_RL_OTEL_SPAN_GROUPS=per_step \
  uv run examples/run_grpo.py --config examples/configs/grpo_smoke.yaml

# Any OTLP backend / collector
export NEMO_RL_OTEL_ENABLED=1
export NEMO_RL_OTEL_SPAN_GROUPS=per_step
export OTEL_EXPORTER_OTLP_ENDPOINT=<your-otlp-endpoint>
uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml

Or in the run config:

telemetry:
  enabled: true
  span_groups: per_step   # default | per_step | all
  exporter: otlp          # otlp | console

You then get rl.grpo.* spans (step, collect_rollouts, compute_logprobs, …), rl.vllm.generate spans, and rl.* metrics (rl.reward.mean, rl.policy_loss, rl.grad_norm, rl.tokens_per_sec, …).

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? (tests/unit/telemetry/ — span-group presets/resolution, lens-absent fallbacks, config→env translation, driver setup + resource attrs, metrics mapping, and in-memory span/metric assertions)
• Did you run the unit tests and functional tests locally? (telemetry unit tests pass locally; functional/e2e still to run in CI)
• Did you add or update any necessary documentation? (docs/observability/ 7-page guide + docs/index.md toctree + nemo_rl/telemetry/README.md)

Additional Information

• Optional by design: only opentelemetry-api is imported at module load; the SDK + exporters install via the telemetry extra / the release image. No hard runtime dependency, no import-time cost when disabled (gate check only).
• Config: NEMO_RL_OTEL_* env vars (with NEMO_LENS_* fallback) plus the telemetry: YAML block; raw env always overrides YAML. Endpoint / protocol / headers are the standard OTEL_EXPORTER_OTLP_* vars, so any backend works.
• Companion change: the grad_norm / learning_rate / tokens_per_sec metrics rely on additive gauges added to nemo-lens record_rl_metrics (separate nemo-lens change); without them those specific metrics are skipped and everything else is unaffected.
• Scope notes: DPO and RM only have the shared loop phases (no rollout/generation spans, since those loops have no such phase); dataset download/setup is not instrumented. vLLM native tracing is gRPC-only (needs a gRPC OTLP endpoint / collector), so it's opt-in and off by default.

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

[rrs45]

Is this missed?

[rrs45]

This may not be the right place for this metric since for async, the rollout is occurring earlier in trajectory_collector, here it is mainly collecting the buffer and not generation

[rrs45]

Is this missed?