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

docs/broken_links_false_positives.json

@@ -5,3 +5,4 @@
 {"uri": "https://huggingface.co/docs/transformers/en/chat_templating_writing"}
 {"uri": "https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoModelForCausalLM"}
 {"uri": "https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoModelForImageTextToText"}
+{"uri": "https://opentelemetry.io/"}

docs/index.md

@@ -219,6 +219,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
@@ -349,6 +356,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,143 @@
+# Configuration
+
+Telemetry is configured by the `telemetry:` block of your run config. Keep it there: a run's telemetry settings should be recoverable from the file that describes the run, not from whatever happened to be in a shell.
+
+Two things do belong in the environment, because they describe *where* you are running rather than *what* you are measuring: the standard [`OTEL_EXPORTER_OTLP_*`](#standard-otel-sdk-variables) endpoint/protocol/headers, and `OTEL_SERVICE_NAME`.
+
+## 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.
+
+```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)
+  export_sample_rate: 1.0     # for sampled: fraction of worker ranks that export
+  sampler_enabled: false      # drop spans at the SDK level using export_sample_rate
+  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.
+
+`service_name` maps onto the standard `OTEL_SERVICE_NAME` (lens reads it unprefixed), so setting either works.
+
+For the full config model, field semantics, and validation rules, see [lens: configuration](https://github.com/NVIDIA-NeMo/Lens).
+
+### How the settings reach the workers
+
+Ray actors do not inherit the driver's Python objects, so on the driver `init_telemetry_driver` projects the block into `NEMO_RL_OTEL_*` environment variables *before* `init_ray()`; the resulting environment is snapshotted into the Ray `runtime_env` and every worker rebuilds the same config from it.
+
+These variables are a transport, not a second configuration interface. They are listed here so that a `NEMO_RL_OTEL_*` name in a log or a `ps` output is recognisable, and because two of them have no `telemetry:` equivalent:
+
+| Variable | Meaning |
+|---|---|
+| `NEMO_RL_OTEL_RUN_ID` | Correlates the driver and every worker to one run. Generated from `SLURM_JOB_ID` or a random hex string when unset. |
+| `NEMO_RL_OTEL_USER_ID` | Optional user/team label, read by lens. |
+
+The projection uses `os.environ.setdefault`, so a variable already present in the environment wins over the YAML value. That is deliberate for the two above, which a job scheduler supplies. For every other setting, prefer the config: a NeMo-RL toggle set in a shell leaves no trace of who set it or why, and splits a run's configuration between a file and an environment with nothing recording which half came from where. The resolved settings are logged once at init for exactly this reason, and a hydra-style `++telemetry.<field>=<value>` override covers the one-off case without leaving the config record.
+
+## 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` — a deterministic hash of the rank selects `export_sample_rate` of the ranks. The same rank and rate always give the same outcome, so the exporting set is stable across restarts.
+- `first_rank_per_node` — the first local rank on each node exports (reads `LOCAL_RANK`).
+
+`export_sample_rate` applies to `sampled`; it has no effect under the other strategies. `sampler_enabled` is independent of `export_strategy` but asks the same kind of question: it installs lens's rank-aware sampler on the TracerProvider, which hashes the rank against `export_sample_rate` once at startup and then keeps or drops *every* span on that rank. A rank has to clear both filters to emit anything, so leaving the sampler on with a low rate can silence a rank the strategy selected.
+
+The driver is independent of both — it always exports, and its rank sampler is disabled for the same reason (`_unrank` in `nemo_rl/telemetry/setup.py`): a synthetic rank 0 is not a member of the population the filters are selecting from. Singleton actors such as the async trajectory collector are exempt on the same grounds. 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.
+
+`RANK` is **group-local**: the policy group and the generation group each number their workers from zero. So `export_rank: 3` selects rank 3 *of every worker group*, and each group's spans carry an `rl.worker_group` attribute to tell them apart.
+
+## 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 |
+| `rl.worker_group` | worker processes only: the worker group's `name_prefix` (`lm_policy`, `vllm_policy`, ...), from `NRL_WORKER_GROUP` |
+
+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
+
+Each example puts the NeMo-RL settings in the config and only the destination in the environment. The `++` form is a hydra-style CLI override: it is applied to the config and echoed into the run's log, so a one-off stays as traceable as an edit to the YAML.
+
+### Console exporter (no backend)
+
+```bash
+uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml \
+  ++telemetry.enabled=true ++telemetry.exporter=console
+```
+
+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 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 \
+  ++telemetry.enabled=true
+```
+
+See [Observability Stack](observability-stack.md) for the full backend-export setup.
+
+### Per-step granularity
+
+```yaml
+telemetry:
+  enabled: true
+  span_groups: per_step
+```
+
+```bash
+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,175 @@
+# 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.
+
+## The import pattern
+
+Every algorithm instrumentation import should go through
+`nemo_rl.telemetry.instrumentation`, which re-exports the lens primitives with
+`rl.bucket` tagging applied to leaf spans:
+
+```python
+from nemo_rl.telemetry.instrumentation import Bucket, bucket_scope, managed_span, trace_fn
+from nemo_rl.telemetry.setup import get_telemetry_handle
+from nemo_rl.telemetry.span_groups import RLSpanGroup
+```
+
+Never import from `nemo.lens.*` directly in algorithm code — that is how a span
+ends up with no bucket and invisible to the goodput rollup.
+
+### Goodput tagging
+
+`managed_span` / `trace_fn` from `instrumentation` attach ``rl.bucket`` ∈
+``{productive, overhead, idle, wasted}`` for leaf groups (see
+`nemo_rl/telemetry/instrumentation.py`). Umbrella groups (`job`, `step`, `rollout`,
+…) are **not** tagged. Apps do **not** emit rolled-up ``rl.goodput`` /
+``rl.bucket.*`` metrics — the offline monitor SUMs tagged phase / span
+durations by bucket.
+
+To override classification for one site, pass the attribute explicitly:
+
+```python
+with managed_span(
+    RLSpanGroup.GENERATION,
+    "rl.vllm.generate",
+    **{"rl.bucket": "productive"},
+):
+    ...
+```
+
+To reclassify spans opened *below* you — where the callee cannot tell why it was
+called — wrap the region in `bucket_scope` instead. Validation uses this so its
+generation counts as `overhead` rather than goodput:
+
+```python
+from nemo_rl.telemetry.instrumentation import Bucket, bucket_scope
+
+with bucket_scope(Bucket.OVERHEAD):
+    ...  # every leaf span in here is tagged overhead
+```
+
+Umbrellas stay unbucketed inside a scope, and an explicit `rl.bucket=` still
+wins, so a scope cannot make a parent double-count its children. See
+[span groups](span-groups.md).
+
+When adding a **new** span group, update `_DEFAULT_GROUP_BUCKET` /
+`UMBRELLA_GROUPS` in `instrumentation.py` and extend `test_instrumentation.py`.
+
+## Instrumenting inside a Ray actor
+
+Two things that are easy to get wrong, because both fail silently as no-ops
+rather than as errors.
+
+The actor needs its own providers: call `init_telemetry_worker()` in its
+`__init__` (not `post_init`, which some fan-outs run on one rank per group), and
+flush before it dies. An actor reaped with `ray.kill` runs no `atexit` handler,
+so expose a method that calls `shutdown_telemetry()` and have the driver call it
+first — `AsyncTrajectoryCollector.flush_telemetry` is the worked example.
+
+Ray does not propagate OTel context, so the actor's spans form their own trace
+unless you carry the parent across. On the driver, inside the span that should
+be the root, capture `current_trace_carrier()` and pass it to the actor; in the
+actor, wrap the work in `remote_trace_context(carrier)`. Reattach in **every
+thread** the actor spawns — OTel context is a `ContextVar` and threads inherit
+none — and note the carrier is empty (a harmless no-op) whenever the driver's
+enclosing group is disabled. See
+[span groups](span-groups.md#getting-the-collector-into-one-waterfall).
+
+## 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.generation",
+                  **{"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_handle()
+if telemetry is not None:
+    with span_cm("rl.grpo.job", tracer=telemetry.tracer):
+        ...
+```
+
+## Naming conventions
+
+| Kind | Convention | Example |
+|---|---|---|
+| Span name | `rl.<algorithm>.<phase>`, matching the block's `Timer` key (the two umbrella spans excepted — see [span groups](span-groups.md#per-algorithm-span-names)) | `rl.grpo.generation` |
+| 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.efficiency.seconds` |
+
+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. **Leave the fallback stub alone.** The stub `SpanGroup` in that file mirrors only lens's *base* groups and presets, for when nemo-lens is absent; `RLSpanGroup` overrides both `ALL_GROUPS` and `_PRESETS`, so an RL group belongs there and nowhere else. Touch the stub only when lens's own base contract changes.
+3. Document the new group in [Span Groups](span-groups.md), and add it to `EMITTED_GROUPS` in `tests/unit/telemetry/test_span_groups.py` so the preset-reachability test covers it.
+
+Keep the base-class contract — shared with lens and its other consumers — consistent when you do this.
+
+## Adding a metric
+
+NeMo-RL records its `rl.*` metrics from the driver rather than scattering record calls through the algorithm code (see [Metrics](metrics.md)). Two cases:
+
+- **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.
+- **The series is keyed by a NeMo-RL-specific label set** rather than a fixed field — e.g. one value per efficiency category. Lens's `record_rl_metrics` takes fixed keyword fields, so a growing label set doesn't fit it. Define the instrument in `nemo_rl/telemetry/metrics.py` instead (same per-Meter caching pattern) and emit **one dimensioned instrument** with the label as an attribute, not one instrument per label. `rl.efficiency.seconds` is the worked example. This also avoids gating the feature on a lens release.
+
+A training scalar that already flows through `Logger.log_metrics` is a third case with no home yet: mapping NeMo-RL's logger keys onto lens's `record_rl_metrics` fields is still being settled with the lens owners, so the only tee on that hook today is the efficiency one.
+
+Prefer an attribute over a name whenever the label set can grow: `rl.efficiency.seconds{rl.efficiency.category="idle/refit_bubble"}` stays stable as categories come and go, while `rl.efficiency.idle_refit_bubble_seconds` forces an instrument change per category. Keep attribute cardinality bounded — a per-step or per-request value belongs on a span, not a metric label.
+
+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_efficiency_seconds` 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, `telemetry.logs_enabled: true`).
+
+When in doubt, start with a coarse span at the boundary of a subsystem, not a fine-grained one at every internal call.