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

docker/Dockerfile

@@ -190,7 +190,7 @@ fi
 uv sync --link-mode symlink --locked --extra mcore --no-install-project
 uv sync --link-mode symlink --locked --extra automodel --no-install-project
 uv sync --link-mode symlink --locked --extra modelopt --no-install-project
-uv sync --link-mode symlink --locked --all-groups --no-install-project
+uv sync --link-mode symlink --locked --all-groups --extra telemetry --no-install-project
 
 # Remove the aiohttp in this uv cache dir to fully address CVE GHSA-mqqc-3gqh-h2x8
 # The ray install will include the older aiohttp version in its cache

docker/Dockerfile.ngc_pytorch

@@ -120,7 +120,7 @@ uv sync --link-mode symlink --locked --inexact --extra vllm --no-install-project
 uv sync --link-mode symlink --locked --inexact --extra mcore --no-install-project $UV_NO_INSTALL_PACKAGES
 uv sync --link-mode symlink --locked --inexact --extra automodel --no-install-project $UV_NO_INSTALL_PACKAGES
 uv sync --link-mode symlink --locked --inexact --extra modelopt --no-install-project $UV_NO_INSTALL_PACKAGES
-uv sync --link-mode symlink --locked --inexact --all-groups --no-install-project $UV_NO_INSTALL_PACKAGES
+uv sync --link-mode symlink --locked --inexact --all-groups --extra telemetry --no-install-project $UV_NO_INSTALL_PACKAGES
 EOF
 
 ENV NEMO_RL_VENV_DIR=/opt/ray_venvs

docs/index.md

@@ -184,6 +184,13 @@ Deep dive into NeMo RL's architecture, APIs, and design decisions for scalable R
 Tools and techniques for debugging distributed Ray applications and RL training runs.
 :::
 
+:::{grid-item-card} {octicon}`graph` Observability
+:link: observability/index
+:link-type: doc
+
+OpenTelemetry traces and `rl.*` metrics via nemo-lens: span groups, configuration, vLLM tracing, and an OTLP export stack.
+:::
+
 :::{grid-item-card} {octicon}`zap` FP8 Quantization
 :link: fp8
 :link-type: doc
@@ -312,6 +319,12 @@ fp8.md
 guides/use-custom-vllm.md
 ```
 
+```{toctree}
+:caption: Observability
+
+observability/index.md
+```
+
 ```{toctree}
 :caption: Design Docs
 

docs/observability/configuration.md

@@ -0,0 +1,138 @@
+# Configuration
+
+Telemetry can be configured two ways, which compose:
+
+1. A `telemetry:` block in your run config (YAML).
+2. `NEMO_RL_OTEL_*` and standard `OTEL_*` environment variables.
+
+**Raw environment variables always win over the YAML block.** On the driver, the `telemetry:` block is translated into `NEMO_RL_OTEL_*` env vars with `os.environ.setdefault` *before* `init_ray()` — so anything already present in the environment is left untouched, and the resulting environment is snapshotted into the Ray `runtime_env` and inherited by every worker.
+
+## The `telemetry:` config block
+
+`telemetry:` is an optional top-level field of every algorithm's `MasterConfig`. It is **documented here, not baked into the exemplar configs** — add it to your own run config, or configure purely via env vars.
+
+```yaml
+telemetry:
+  enabled: false              # master switch; when false, every site is a ~0-cost no-op
+  service_name: nemo-rl       # service.name reported to the backend
+  span_groups: default        # preset (default | per_step | all) or a comma-separated group list
+  export_strategy: single_rank # single_rank | all_ranks | sampled | first_rank_per_node
+  export_rank: -1             # for single_rank: which rank exports (-1 = last rank)
+  traces_enabled: true        # emit trace spans
+  metrics_enabled: true       # emit the rl.* metric instruments
+  logs_enabled: false         # bridge Python logging to OTel logs (trace-correlated)
+  exporter: otlp              # otlp | console
+  vllm_native_tracing: false  # opt in to vLLM's own OTLP tracing (gRPC-only — see vllm-tracing.md)
+```
+
+The defaults above are the field defaults of `TelemetryConfig` (`nemo_rl/telemetry/config.py`). The endpoint, headers, and protocol are **not** in this block — they come from the standard `OTEL_EXPORTER_OTLP_*` env vars (see below).
+
+The driver always exports (it hosts the training loop and the metrics logger); `export_strategy` / `export_rank` govern the Ray **worker** ranks.
+
+## NeMo-RL environment variables
+
+Each `NEMO_RL_OTEL_*` variable maps onto a `NemoLensConfig` field. Lens reads `NEMO_RL_OTEL_<KEY>` first and falls back to `NEMO_LENS_<KEY>`, so you can set a shared `NEMO_LENS_*` default and override it per-run with the RL-scoped prefix.
+
+| Variable | Maps to | Default |
+|---|---|---|
+| `NEMO_RL_OTEL_ENABLED` | `enabled` | `0` |
+| `NEMO_RL_OTEL_SPAN_GROUPS` | `span_groups` | `default` |
+| `NEMO_RL_OTEL_EXPORT_STRATEGY` | `export_strategy` | `single_rank` |
+| `NEMO_RL_OTEL_EXPORT_RANK` | `export_rank` | `-1` |
+| `NEMO_RL_OTEL_TRACES_ENABLED` | `traces_enabled` | `1` |
+| `NEMO_RL_OTEL_METRICS_ENABLED` | `metrics_enabled` | `1` |
+| `NEMO_RL_OTEL_LOGS_ENABLED` | `logs_enabled` | `0` |
+| `NEMO_RL_OTEL_EXPORTER` | `exporter` | `otlp` |
+| `NEMO_RL_OTEL_VLLM_NATIVE_TRACING` | `vllm_native_tracing` | `0` |
+| `NEMO_RL_OTEL_RUN_ID` | run identifier | (auto) |
+| `NEMO_RL_OTEL_USER_ID` | optional user/team label | (empty) |
+
+`service_name` maps onto the standard `OTEL_SERVICE_NAME` (lens reads it unprefixed).
+
+For the full config model, field semantics, and validation rules, see [lens: configuration](https://github.com/NVIDIA-NeMo/Lens).
+
+## Standard OTel SDK variables
+
+Endpoint, protocol, and headers are honoured by the OTel SDK directly:
+
+| Variable | Example |
+|---|---|
+| `OTEL_SERVICE_NAME` | `nemo-rl` |
+| `OTEL_EXPORTER_OTLP_ENDPOINT` | `http://localhost:4317` |
+| `OTEL_EXPORTER_OTLP_PROTOCOL` | `grpc` or `http/protobuf` |
+| `OTEL_EXPORTER_OTLP_HEADERS` | `<header>=<value>,<header>=<value>` (e.g. auth headers your backend requires) |
+
+Pick the protocol to match your backend: a local collector or Jaeger typically speaks gRPC on `:4317`; a direct-to-SaaS OTLP endpoint typically speaks `http/protobuf` on `:443`. See [Observability Stack](observability-stack.md).
+
+## Export strategy
+
+`export_strategy` controls which **worker** ranks actually send telemetry:
+
+- `single_rank` (default) — only the rank named by `export_rank` (`-1` = last rank).
+- `all_ranks` — every worker exports.
+- `sampled` / `first_rank_per_node` — sample a subset.
+
+The driver is independent of this — it always exports. Non-exporting ranks get an empty (`frozenset()`) span-group set, so `is_span_group_enabled()` is `False` everywhere and no span objects are created at all. See [lens: sampling](https://github.com/NVIDIA-NeMo/Lens) for the detailed semantics.
+
+## Run identification
+
+Every run gets a `run_id` that flows to all backends as a resource attribute and is shared by the driver and every worker.
+
+**Priority order:**
+
+1. `NEMO_RL_OTEL_RUN_ID` (explicit, highest priority).
+2. `SLURM_JOB_ID` (auto-detected on SLURM clusters).
+3. Auto-generated 12-character hex id (fallback).
+
+The `run_id` is written to the environment on the driver **before** `init_ray()`, so every worker inherits the same value and correlates to the same run. This is also how vLLM's native spans are correlated back to the RL run — see [vLLM Tracing](vllm-tracing.md).
+
+Filter by `run_id` in your backend to isolate a specific run.
+
+## Resource attributes
+
+`init_telemetry_driver` sets stable-for-the-run values on the OTel `Resource`, so they appear on every span/metric as backend "Process" tags:
+
+| Attribute | Source |
+|---|---|
+| `rl.algorithm` | the `algorithm="<algo>"` passed to `init_telemetry_driver` |
+| `rl.model` | `policy.model_name` |
+| `nemo.precision` | `policy.precision` |
+| `dl.tensor_parallel.size` | `policy.megatron_cfg` / `dtensor_cfg` TP size |
+| `dl.pipeline_parallel.size` | `policy.megatron_cfg` PP size |
+| `dl.rank`, `dl.world_size` | set automatically by lens |
+
+Attribute construction is best-effort: a missing config key simply omits that attribute; it never raises. Plus auto-detected host / GPU / SLURM / Kubernetes attributes from lens's resource detection.
+
+## Typical configurations
+
+### Console exporter (no backend)
+
+```bash
+export NEMO_RL_OTEL_ENABLED=1
+export NEMO_RL_OTEL_EXPORTER=console
+uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml
+```
+
+Spans and metrics print to stdout — a quick dry run with no backend to stand up.
+
+### Direct to an OTLP backend (http/protobuf)
+
+```bash
+export NEMO_RL_OTEL_ENABLED=1
+export OTEL_EXPORTER_OTLP_ENDPOINT=https://<your-otlp-endpoint>:443
+export OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf
+export OTEL_EXPORTER_OTLP_HEADERS="<header>=<value>"   # any auth headers your backend requires
+uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml
+```
+
+See [Observability Stack](observability-stack.md) for the full backend-export setup.
+
+### Per-step granularity
+
+```bash
+export NEMO_RL_OTEL_ENABLED=1
+export NEMO_RL_OTEL_SPAN_GROUPS=per_step
+export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317
+```
+
+`per_step` makes each training step its own root trace (rollout, generation, reward, advantage, policy update). See [Span Groups](span-groups.md).

docs/observability/extending.md

@@ -0,0 +1,114 @@
+# Extending Instrumentation
+
+To add new spans or metrics to NeMo-RL code, use the instrumentation primitives from nemo-lens (`managed_span`, `trace_fn`, `span_cm`). The primitives themselves are documented in [lens: instrumentation](https://github.com/NVIDIA-NeMo/Lens); this page covers NeMo-RL conventions.
+
+```{tip}
+If you work in this repo with Claude Code, the `add-span-group` skill (new span group), the `new-instrument` lens skill (new `rl.*` metric), and the `instrumentation-site-helper` agent (new span/metric site) automate the steps below and keep the cross-repo fallback contract in sync. They are optional — everything here can be done by hand.
+```
+
+## The import / fallback pattern
+
+Every lens import in NeMo-RL code must go through `nemo_rl.telemetry._fallbacks`, so the code runs unchanged when nemo-lens is not installed:
+
+```python
+from nemo_rl.telemetry._fallbacks import managed_span, trace_fn
+from nemo_rl.telemetry.setup import get_telemetry
+from nemo_rl.telemetry.span_groups import RLSpanGroup
+```
+
+`_fallbacks.py` re-exports the real nemo-lens implementations when it is installed, and provides identical no-op stubs when it is not. Never import from `nemo.lens.*` directly in algorithm code. See [lens: optional dependency](https://github.com/NVIDIA-NeMo/Lens).
+
+## Adding a span
+
+### Decorator — `trace_fn`
+
+For a whole function (this is how `rl.vllm.generate` and the `rl.<algo>.job` spans are done):
+
+```python
+@trace_fn(RLSpanGroup.GENERATION, "rl.vllm.generate")
+def generate(self, ...):
+    ...
+```
+
+### Group-gated block — `managed_span`
+
+For a hot path where you want minimal cost when the group is disabled:
+
+```python
+with managed_span(RLSpanGroup.ROLLOUT, "rl.grpo.collect_rollouts",
+                  **{"rl.iteration": iteration}) as span:
+    result = collect()
+    if span is not None:
+        span.set_attribute("rl.num_generations_per_prompt", n)
+```
+
+`managed_span` yields `None` when the group is disabled; the body still runs, so guard attribute-setting with `if span is not None`.
+
+### Always-on block — `span_cm`
+
+`span_cm` always creates a span when telemetry is active (no group gate) — for cold, top-level paths only:
+
+```python
+telemetry = get_telemetry()
+if telemetry is not None:
+    with span_cm("rl.grpo.job", tracer=telemetry.tracer):
+        ...
+```
+
+## Naming conventions
+
+| Kind | Convention | Example |
+|---|---|---|
+| Span name | `rl.<algorithm>.<operation>` | `rl.grpo.collect_rollouts` |
+| Span tag | `rl.<attr>` categorical | `rl.iteration`, `rl.backend` |
+| Resource attribute | `rl.<attr>` / shared `dl.<attr>` | `rl.model`, `dl.tensor_parallel.size` |
+| Metric name | `rl.<subsystem>.<metric>` (application scope) | `rl.reward.mean` |
+
+Metric names use the **application scope** (`rl.*`) — never `dl.*`. Attribute names shared across consumers use the constants in `nemo.lens.semconv`; RL-specific short strings are fine hard-coded.
+
+## Choosing a span group
+
+Pick from `RLSpanGroup` before inventing a new one:
+
+- Once per run (setup/whole-job)? → `job`
+- Once per training step? → `step`
+- Rollout collection? → `rollout`; generation? → `generation`
+- Log-probs? → `logprob` (or `reference_policy` for the reference model)
+- Reward / advantage / policy update? → `reward` / `advantage` / `policy_update`
+- Checkpoint / eval? → `checkpoint` / `evaluate`
+
+## Adding a new span group
+
+If nothing fits, add a group to `RLSpanGroup` in `nemo_rl/telemetry/span_groups.py`:
+
+1. Add the constant, add it to `ALL_GROUPS`, and slot it into the right preset(s) in `_PRESETS`. Decide per preset: `default` is coarse (rarely add here); `per_step` for per-step spans; `all` always includes it.
+2. **Update the fallback stub** in the same file — the stub `SpanGroup` used when nemo-lens is absent must keep the same constants and presets in lockstep.
+3. Document the new group in [Span Groups](span-groups.md).
+
+The `add-span-group` skill walks these steps and keeps the base-class contract (shared with lens and the other consumers) consistent.
+
+## Adding a metric
+
+The `rl.*` gauges are populated by teeing `Logger.log_metrics` (see [Metrics](metrics.md)) — not by scattering `record_rl_metrics()` calls. So there are two cases:
+
+- **The scalar already flows through `Logger.log_metrics`** under a `train` prefix but isn't teed. Add a candidate key (or a new field) to `_RL_OTEL_METRIC_MAP` in `nemo_rl/telemetry/metrics.py`, and add the matching field to `record_rl_metrics` in lens's `nemo.lens.instruments.rl`.
+- **You need a brand-new instrument** (a new counter/gauge/histogram, or a value that doesn't go through the Logger). Add it to `nemo.lens.instruments.rl` following the per-Meter `WeakKeyDictionary` caching pattern, then record it from the driver via `telemetry.meter`. The `new-instrument` lens skill covers this.
+
+Keep `rl.<subsystem>.<metric>` naming and record only non-`None` values. See [lens: metrics](https://github.com/NVIDIA-NeMo/Lens).
+
+## Testing new instrumentation
+
+NeMo-RL telemetry tests live under `tests/` and use lens's in-memory exporter fixtures (global OTel state reset per test). When adding a span:
+
+1. Assert the span is emitted when its group is enabled and absent when disabled.
+2. Assert on span name, tags, and parent relationships.
+
+For a pure metrics-tee change, `map_rl_metrics` in `nemo_rl/telemetry/metrics.py` is a pure function — unit-test the key mapping directly with no OTel setup.
+
+## When not to add instrumentation
+
+- Inside a tight inner loop (per-token) — even a gated `managed_span`'s frozenset lookup adds up.
+- On high-cardinality attributes (raw prompts, tensor shapes) — cardinality explosion at the backend.
+- As a replacement for logging — structured logs belong in logs (correlate via the log bridge, `NEMO_RL_OTEL_LOGS_ENABLED=1`).
+
+When in doubt, start with a coarse span at the boundary of a subsystem, not a fine-grained one at every internal call.