diff --git a/docs/broken_links_false_positives.json b/docs/broken_links_false_positives.json index ceb2be0a38..b686a625c9 100644 --- a/docs/broken_links_false_positives.json +++ b/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/"} diff --git a/docs/index.md b/docs/index.md index d32943ff39..142cc8be8a 100644 --- a/docs/index.md +++ b/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 diff --git a/docs/observability/configuration.md b/docs/observability/configuration.md new file mode 100644 index 0000000000..998c80869a --- /dev/null +++ b/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.=` 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` | `
=,
=` (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=""` 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://:443 +export OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf +export OTEL_EXPORTER_OTLP_HEADERS="
=" # 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). diff --git a/docs/observability/extending.md b/docs/observability/extending.md new file mode 100644 index 0000000000..e60eb00208 --- /dev/null +++ b/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..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..`, 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.` categorical | `rl.iteration`, `rl.backend` | +| Resource attribute | `rl.` / shared `dl.` | `rl.model`, `dl.tensor_parallel.size` | +| Metric name | `rl..` (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..` 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. diff --git a/docs/observability/index.md b/docs/observability/index.md new file mode 100644 index 0000000000..8b2cd357ea --- /dev/null +++ b/docs/observability/index.md @@ -0,0 +1,91 @@ +# Observability + +NeMo RL is instrumented with [OpenTelemetry](https://opentelemetry.io/) via the [`nemo-lens`](https://github.com/NVIDIA-NeMo/Lens) library. It emits **traces** at RL-algorithm boundaries (rollout, generation, reward, advantage, policy update, ...) and **metrics** for async efficiency accounting and vLLM generation. + +Telemetry exports OTLP and works with any OTLP-compatible backend or an OpenTelemetry Collector (e.g. Jaeger, Grafana Tempo, or an OpenTelemetry Collector that fans out to your backend of choice). + +Telemetry is **off by default**. nemo-lens ships as a base dependency, so it reaches every worker venv, and `telemetry.enabled` is the single switch: while it is false every instrumentation site is a ~0-cost no-op. + +## What's in this section + +```{toctree} +:maxdepth: 1 + +configuration +span-groups +metrics +vllm-tracing +observability-stack +extending +``` + +## Scope + +This documentation covers **NeMo-RL-specific** usage: the `telemetry:` config block, RL span names, `rl.*` metric names, and the two-layer vLLM tracing integration. + +For general concepts — the span-group mechanism, instrumentation primitives, the configuration model, custom exporters, resource detection — see the [lens documentation](https://github.com/NVIDIA-NeMo/Lens). This section links to lens docs when relevant rather than duplicating them. + +| Concern | Owned by | +|---|---| +| `telemetry:` YAML block | NeMo-RL (this section) | +| `RLSpanGroup` groups + presets, `rl.*` span/metric names | NeMo-RL (this section) | +| Driver/worker telemetry lifecycle, vLLM two-layer tracing | NeMo-RL (this section) | +| `managed_span` / `trace_fn` / `span_cm`, config model, exporters, resource detection | [lens](https://github.com/NVIDIA-NeMo/Lens) | + +## Install + +Nothing to install: `nemo-lens[sdk]` is a base dependency, so a normal `uv sync` covers the driver and every worker venv. + +## Quick start + +Add a `telemetry:` block to your run config: + +```yaml +telemetry: + enabled: true + span_groups: default # coarse-grained; safe for production +``` + +Point it at a backend and run: + +```bash +export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317 # your OTLP backend / collector + +uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml +``` + +With `default` span groups, NeMo-RL emits a handful of coarse spans (job, checkpoint, evaluate) plus whatever `rl.*` metrics the driver's logger produces. Switch to `per_step` for per-step traces (rollout/generation/reward/...), or `all` for everything. + +Keeping the settings in the config file is what makes a run's telemetry reproducible from the file alone. The endpoint is the exception: `OTEL_EXPORTER_OTLP_*` are the standard OpenTelemetry variables, and they belong in the environment because they describe where you are running, not what you are measuring. See [Configuration](configuration.md). + +## What gets instrumented + +Each algorithm's `examples/run_.py` calls `init_telemetry_driver(config, algorithm="")` **before** `init_ray()` (so the resolved `NEMO_RL_OTEL_*` settings are snapshotted into the Ray `runtime_env` and inherited by every worker) and `shutdown_telemetry()` from a `finally` block wrapping the whole run, so buffered spans are flushed even when the run fails. + +| Algorithm | Entry point | Representative spans | +|---|---|---| +| GRPO (sync + async) | `examples/run_grpo.py` | `rl.grpo.step`, `rl.grpo.generation`, `rl.grpo.reward_calculation`, `rl.grpo.policy_and_reference_logprobs`, `rl.grpo.advantage_calculation`, `rl.grpo.policy_training` | +| PPO | `examples/run_ppo.py` | `rl.ppo.step`, `rl.ppo.generation`, `rl.ppo.reward_calculation`, `rl.ppo.advantage_calculation`, `rl.ppo.policy_training`, `rl.ppo.value_training` | +| SFT | `examples/run_sft.py` | `rl.sft.step`, `rl.sft.data_processing`, `rl.sft.policy_training` | +| DPO | `examples/run_dpo.py` | `rl.dpo.step`, `rl.dpo.policy_training` | +| RM | `examples/run_rm.py` | `rl.rm.step` | +| Distillation | `examples/run_distillation.py` | `rl.distillation.step`, `rl.distillation.generation`, `rl.distillation.teacher_logprob_inference`, `rl.distillation.policy_training` | +| vLLM generation | `nemo_rl/models/generation/vllm/vllm_generation.py` | `rl.vllm.generate`, `rl.vllm.generate_text` | + +Each span belongs to a **span group** that controls whether it is emitted at runtime. See [Span Groups](span-groups.md) for the full per-algorithm span table. + +## What gets exported + +- **Traces**: any OTLP-compatible backend (Jaeger, Grafana Tempo, an OpenTelemetry Collector, ...) via OTLP. +- **Metrics**: the `rl.efficiency.*` async accounting teed from the driver's metrics logger, plus the vLLM `gen_ai.*` series — see [Metrics](metrics.md). +- **Logs** (optional): via the OTel log bridge when `telemetry.logs_enabled` is true — correlates Python `logging` records with the active span's trace ID. + +By default, only **one rank** exports (`single_rank`, last rank). The driver always exports (it hosts the training loop and the metrics logger). See [Configuration — Export strategy](configuration.md#export-strategy). + +## Related + +- Exporting to an OTLP backend: [Observability Stack](observability-stack.md) +- vLLM tracing (driver spans + native OTLP): [vLLM Tracing](vllm-tracing.md) +- Adding new spans / metrics: [Extending Instrumentation](extending.md) +- Lens configuration model and env vars: [lens: configuration](https://github.com/NVIDIA-NeMo/Lens) +- Instrumentation primitives (`managed_span`, `trace_fn`, `span_cm`): [lens: instrumentation](https://github.com/NVIDIA-NeMo/Lens) diff --git a/docs/observability/metrics.md b/docs/observability/metrics.md new file mode 100644 index 0000000000..00ba9609e1 --- /dev/null +++ b/docs/observability/metrics.md @@ -0,0 +1,97 @@ +# Metrics + +NeMo-RL emits two namespaces of metrics: async efficiency metrics (`rl.efficiency.*`) and vLLM generation metrics (`gen_ai.*`, following the OTel GenAI semantic conventions). + +Metrics are emitted **only when telemetry is exporting** — the driver always exports, so the `rl.*` series come from the driver's metrics logger. For the general instrument pattern (per-Meter caching, `None`-skipping), see [lens: metrics](https://github.com/NVIDIA-NeMo/Lens). + +Training scalars — reward, loss, KL, grad norm, learning rate, throughput — are **not** mirrored to OTel. nemo-lens declares `record_rl_metrics` gauges for most of them, plus `rl.generation.duration_ms` and `rl.rollout.duration_ms` histograms, but NeMo-RL emits none of them: mapping its logger keys onto lens's fixed fields is still being settled with the lens owners. Read those scalars from W&B / TensorBoard, and phase durations from the spans. + +## Async efficiency metrics (`rl.efficiency.*`) + +Async GRPO measures where wall time goes with a `Timer` and logs the result as `efficiency/*` scalars (`print_efficiency_summary` in `nemo_rl/algorithms/utils.py`). Those same values are teed to OTel as one **dimensioned** gauge rather than one instrument per category, so adding a category needs no instrument change. + +The tee lives outside the algorithm code: `nemo_rl/telemetry/metrics.py` hooks `nemo_rl.utils.logger.Logger.log_metrics`, so after `log_metrics` fans a step out to the file / W&B / MLflow backends it calls `tee_rl_metrics_to_otel(metrics, prefix)`. It is best-effort — only the driver's `train` dicts (`prefix in ("train", "")`) carry the efficiency scalars, so other prefixes are skipped, non-scalar values are ignored, and the whole path is a no-op unless telemetry is actively exporting. The efficiency numbers you already see in W&B are therefore the same series you get in your OTLP backend, with no double bookkeeping. + +| Metric | Type | Attributes | Description | +|---|---|---|---| +| `rl.efficiency.seconds` | Gauge (`s`) | `rl.efficiency.category`, `rl.efficiency.measurement`, `rl.efficiency.window`, `rl.bucket` | Time attributed to one efficiency category | +| `rl.efficiency.pct` | Gauge (`%`) | `rl.efficiency.measurement`, `rl.efficiency.window` | Productive share of one step's driver-side wall clock | + +These instruments are defined in `nemo_rl/telemetry/metrics.py` rather than in lens, because they are keyed by NeMo-RL's own efficiency-category labels and there is no fixed lens field for them. + +### Always filter on `rl.efficiency.measurement` + +Some categories are measured on the driver against wall time; others are summed across concurrent collector threads and **can exceed the wall time they happened in**. + +| `rl.efficiency.measurement` | Recorded on | Categories | Safe to sum against elapsed driver time? | +|---|---|---|---| +| `wall_clock` | driver, sequentially | `init/total`, `idle/buffer_starvation`, `idle/refit_bubble`, `idle/validation` | yes | +| `collector_wall_clock` | collector's collection-loop thread, sequentially | `idle/refit_event_wait`, `idle/generation_limit_pause` | no — real durations, but on a timeline that runs concurrently with the driver's | +| `thread_seconds` | collector's batch-worker threads, concurrently | `idle/buffer_full_backoff`, `wasted/failed_trajectory` | no — not durations at all | + +Eight rollout threads each backing off for 10s during the same 10s window produce a `thread_seconds` value of 80, not 10. Summing `rl.efficiency.seconds` by `rl.bucket` without filtering therefore overstates idle time — a wrong answer that looks like a real one. Filter to `rl.efficiency.measurement="wall_clock"` before comparing against elapsed time; read the other two per-phase, `thread_seconds` as a saturation signal. + +The non-`wall_clock` values also carry `rl.bucket` so all three share one vocabulary with the spans; the `measurement` attribute is what keeps a bucket rollup honest. + +Two deliberate metric/span disagreements to know about before comparing a metric rollup against a trace rollup: + +- **The two collector-loop categories** carry `rl.bucket="idle"` as metrics and rely on you filtering by `measurement`, but carry no bucket at all as spans, since a trace has no equivalent filter to rely on. +- **`idle/validation`** is `idle` as a metric and `overhead` on every span covering the same seconds. Both are true of different fleets: the training GPUs are idle, which is what the driver's timer measures, while the generation GPUs are doing necessary non-training work, which is what `bucket_scope(Bucket.OVERHEAD)` in `validate()` tags. Attributing this phase properly needs per-fleet accounting; until then, do not expect the two rollups to agree on a validation step. See [span groups — why `idle/validation` is not a span](span-groups.md#why-idlevalidation-is-not-a-span). + +### `rl.efficiency.window`: what a value covers in time + +`measurement` says whether values may be summed *against each other*; `window` says whether one may be summed *across steps*. They are independent, and the second is the easier one to get silently wrong. + +| `rl.efficiency.window` | Categories | Meaning | +|---|---|---| +| `step` | `idle/buffer_starvation`, `idle/refit_bubble`, `idle/validation` | per-step delta — the driver resets its `Timer` every step, so these sum across steps | +| `run` | `init/total`, and all four collector-side categories | cumulative since the process started — consecutive points already contain each other, so summing across steps multiplies by the step count | + +`init/total` is the driver-side exception: it is measured once, waiting for the first buffer fill before the step loop, then republished unchanged every step so it does not disappear from a dashboard after step 1. Read it as a constant. The collector's `Timer` is never reset, which is why everything from it is `run`. + +`rl.efficiency.pct` is tagged `window="step"` for the same reason its numerator is: the three `step`-window idle categories over that step's wall time. `init/total` is deliberately excluded — it is a run constant, so folding it in would charge the whole startup cost to every step — and so are the collector's categories, which are on another clock. Against the run's elapsed time the ratio would climb toward 100% as the run lengthened no matter what the idle time did, which is why the denominator is one step and not the run. + +## vLLM generation metrics (`gen_ai.*`) + +The driver-side vLLM generation path records token and latency metrics through lens's `record_inference_metrics` with `provider_name="vllm"`, following the [OTel GenAI metrics spec](https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-metrics/). + +| Metric | Type | Description | +|---|---|---| +| `gen_ai.client.token.usage` | Histogram | Tokens per request, split by `gen_ai.token.type` (`input` / `output`) | +| `gen_ai.server.request.duration` | Histogram | End-to-end generation request latency | + +These ride the normal `http/protobuf` OTLP path and reach the same backend as everything else. They are distinct from vLLM's **native** engine metrics (opt-in, gRPC-only) — see [vLLM Tracing](vllm-tracing.md). + +## Metric vs span tag vs resource attribute + +The one rule that trips people up. Classify each value before you emit it: + +| Kind | Use | Example | +|---|---|---| +| **Metric** | numerical value that changes over time | per-category efficiency seconds → `rl.efficiency.seconds` | +| **Span tag** | categorical per-span context for filtering | `rl.iteration`, `rl.bucket`, `rl.num_generations_per_prompt`, `rl.weight_version` | +| **Resource attribute** | stable for the whole run | `rl.algorithm`, `rl.model`, `dl.tensor_parallel.size` | + +Do **not** put a time-series number (loss, reward) on a span attribute — it produces no useful series in your backend and wastes storage. Do **not** put a per-step categorical (iteration number) on a metric label — that is unbounded cardinality. See [lens: metrics — metric vs span attribute vs resource attribute](https://github.com/NVIDIA-NeMo/Lens). + +### Goodput (monitor-derived) + +NeMo-RL does **not** emit `rl.goodput` or `rl.bucket.*` rollup metrics. +Leaf spans carry `rl.bucket` ∈ `{productive, overhead, idle, wasted}`; +umbrella spans (`job` / `step` / `rollout`) omit it. `rl.efficiency.seconds` +carries the same `rl.bucket` tokens, but it is a per-category duration, not a +rollup — and it needs the `rl.efficiency.measurement` filter described above. +Offline monitors (e.g. wandb-monitor) SUM span / phase GPU-time by `rl.bucket` +and compute: + +```text +rl_goodput = productive_gpu_s / (productive + overhead + idle + wasted)_gpu_s +``` + +See [Span groups — goodput buckets](span-groups.md) and `nemo_rl/telemetry/instrumentation.py`. + +Metric names use the **application scope** (`rl.*`); attribute names use the **shared namespace** (`rl.*`, `dl.*`) defined in lens's `semconv.py`. + +## Filtering across runs + +Every `rl.*` data point carries the `run_id` resource attribute. Use it to isolate or compare runs in your backend (Grafana/Prometheus, or any OTLP-compatible backend). See [Configuration — Run identification](configuration.md#run-identification). diff --git a/docs/observability/observability-stack.md b/docs/observability/observability-stack.md new file mode 100644 index 0000000000..fda2edce1b --- /dev/null +++ b/docs/observability/observability-stack.md @@ -0,0 +1,52 @@ +# Exporting to an OTLP backend + +NeMo-RL's telemetry is a standard OpenTelemetry OTLP exporter: enable it, point it at an OTLP endpoint, and run your training. It works with **any OTLP-compatible backend or an OpenTelemetry Collector** — there is nothing NeMo-RL-specific about the backend, and no bundled Jaeger / Prometheus / Grafana. + +Choosing an observability solution — retention, scale, auth, dashboards — is your decision, driven by your organisation's existing stack (e.g. Jaeger, Grafana Tempo, or an OpenTelemetry Collector that fans out to your backend of choice). For backend-specific guidance, see [lens: backends](https://github.com/NVIDIA-NeMo/Lens). + +## Turn it on + +What you measure goes in your run config, so a run's telemetry settings stay recoverable from the file that describes the run: + +```yaml +telemetry: + enabled: true + span_groups: default # start coarse; raise to per_step / all as needed + metrics_enabled: true + logs_enabled: true + vllm_native_tracing: false # gRPC-only; leave off on an http/protobuf path +``` + +Where you send it goes in the environment, since it describes the machine rather than the run. These are the standard OTel SDK variables, read by the SDK directly: + +```bash +OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317 +OTEL_EXPORTER_OTLP_PROTOCOL=grpc # grpc (collector/Jaeger on :4317) or http/protobuf (SaaS OTLP on :443) +# OTEL_EXPORTER_OTLP_HEADERS=
= # optional auth headers, comma-separated +``` + +All three signals (traces, metrics, logs) ship over OTLP to the endpoint you set; on the `http/protobuf` path the SDK appends `/v1/traces`, `/v1/metrics`, `/v1/logs` per signal. 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`. + +Raise `span_groups` to `per_step` (or `all`) for per-step traces. To name the run instead of taking the auto-generated id, set `NEMO_RL_OTEL_RUN_ID` (and optional `NEMO_RL_OTEL_USER_ID`) — the two settings with no `telemetry:` equivalent, since a job scheduler usually supplies them. + +## Console / JSON output (no backend) + +To confirm spans and metrics are produced without standing up any backend, use the `console` exporter. Hydra-style CLI overrides keep a one-off run in the config record — they are echoed into the run's log — so prefer them over exporting a variable: + +```bash +uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml \ + ++telemetry.enabled=true ++telemetry.exporter=console +``` + +Each span and metric prints to stdout as **JSON** (`ConsoleSpanExporter` uses `span.to_json()`), so you can capture it to a file: + +```bash +uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml \ + ++telemetry.enabled=true ++telemetry.exporter=console > telemetry.json 2>&1 +``` + +`console` (set via `telemetry.exporter`) is the only backend-free JSON option nemo-lens exposes. For structured JSON-lines *files*, export OTLP to an OpenTelemetry Collector with a `file` exporter (nemo-lens ships a collector-file config) and point `OTEL_EXPORTER_OTLP_ENDPOINT` at the collector. + +## vLLM native tracing needs a gRPC endpoint + +vLLM's **native** OTLP tracing (`telemetry.vllm_native_tracing: true`) uses a gRPC-only exporter, so it will not ride an `http/protobuf` OTLP endpoint. To capture vLLM's native engine spans, add an OTLP/gRPC receiver (an OTel Collector on `:4317`, or a gRPC-capable backend) that forwards to your backend, and point `OTEL_EXPORTER_OTLP_ENDPOINT` (or `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT`) at it. The driver-side `rl.vllm.*` spans and `gen_ai.*` metrics (Layer 1) reach your backend regardless. See [vLLM Tracing](vllm-tracing.md). diff --git a/docs/observability/span-groups.md b/docs/observability/span-groups.md new file mode 100644 index 0000000000..f6e670d712 --- /dev/null +++ b/docs/observability/span-groups.md @@ -0,0 +1,339 @@ +# Span Groups + +Span granularity in NeMo-RL is controlled by `span_groups` in the [`telemetry:` config block](configuration.md). The spec accepts a preset keyword, a comma-separated list of individual group names, or a mix (e.g. `default,generation,reward`). + +For the general span-group mechanism — how gating works, why a disabled group costs ~nothing — see [lens: span groups](https://github.com/NVIDIA-NeMo/Lens). This page covers NeMo-RL's groups and the per-algorithm span hierarchy. + +## Preset keywords + +| Preset | Groups included | Relative cost | +|---|---|---| +| `default` | `job`, `checkpoint`, `evaluate` | Lowest — safe for production | +| `per_step` | `step`, `checkpoint`, `evaluate`, `model_init`, `rollout`, `generation`, `logprob`, `reward`, `advantage`, `policy_update`, `reference_policy`, `data_processing`, `efficiency` | Moderate | +| `all` | every group (`job` included) | Highest — dev/debug | + +### `per_step` deliberately omits `job` + +`per_step` **excludes** the `job` group on purpose. `job` is the whole-run root span; if it were enabled alongside `step`, every training step would nest under one giant, ever-growing trace. Omitting `job` makes **each training step its own root trace** — bounded in size and easy to search one step at a time. + +`job` lives in `default` (coarse: job + checkpoint + evaluate) and in `all` (one whole-run trace, useful for a short run). Choose `per_step` when you want to inspect individual steps; choose `default`/`all` when you want one trace spanning the run. + +## `RLSpanGroup` + +Defined in `nemo_rl/telemetry/span_groups.py`. Extends lens's base `SpanGroup` with RL-specific groups. + +| Group | Origin | Controls | +|---|---|---| +| `job` | base | the whole-run root span (`rl..job`) | +| `checkpoint` | base | `rl..save_checkpoint` | +| `evaluate` | base | `rl..evaluate` | +| `model_init` | base | `rl.vllm.load_model` (emitted in the generation worker) | +| `load_checkpoint` | base | *reserved — bucketed, but no site emits it yet* | +| `step` | base | `rl..step` (one per training step) | +| `forward_backward` | base | *reserved — bucketed, but no site emits it yet* | +| `optimizer` | base | *reserved — bucketed, but no site emits it yet* | +| `rollout` | RL | `rl..collect_rollouts` | +| `generation` | RL | the driver-side `rl.vllm.generate` / `rl.vllm.generate_text` spans | +| `logprob` | RL | `rl..compute_logprobs` | +| `reward` | RL | `rl..compute_rewards` | +| `advantage` | RL | `rl..compute_advantages` | +| `policy_update` | RL | `rl..policy_update` (and `value_update` for PPO) | +| `reference_policy` | RL | *reserved — bucketed, but no site emits it yet* | +| `data_processing` | RL | `rl..data_processing` | +| `efficiency` | RL | idle phases on async GRPO — driver-side `rl.idle.buffer_starvation`, `rl.idle.refit_bubble`, and collector-side `rl.idle.refit_event_wait`, `rl.idle.generation_limit_pause` | + +## Examples + +```yaml +telemetry: + enabled: true + + # Coarse spans only — default + span_groups: default + + # Per-step traces (rollout / generation / reward / advantage / policy update) + # span_groups: per_step + + # Coarse job trace + generation spans only + # span_groups: default,generation + + # Everything + # span_groups: all +``` + +## Per-algorithm span names + +Span names follow `rl..`, where `` is the `Timer` key the same block records — so a span and the `timing/train/` metric measuring it carry one name rather than two, and correlating a slow span with its timing series needs no mapping. The `Timer` key is the authority: it is pre-existing and already published as a metric name, so a new span takes its name from the timer rather than the reverse. + +Two spans deliberately do not follow it. `rl..step` wraps `total_step_time` and `rl..evaluate` wraps `total_validation_time`: a span's duration is intrinsic, so naming one after a `total_*_time` measurement is tautological, and these two are the umbrella spans a reader meets first in a waterfall. They are named after the operation instead. Every span *inside* them matches its timer key. + +The controlling group is shown for each; a span is only emitted when its group is enabled *and* the rank is exporting. + +| Algorithm | Spans | +|---|---| +| **GRPO** (sync + async) | `rl.grpo.job`, `rl.grpo.step`, `rl.grpo.data_processing`, `rl.grpo.generation`, `rl.grpo.reward_calculation`, `rl.grpo.policy_and_reference_logprobs`, `rl.grpo.advantage_calculation`, `rl.grpo.policy_training`, `rl.grpo.checkpointing`, `rl.grpo.evaluate` | +| **GRPO** (async only) | `rl.idle.buffer_starvation`, `rl.idle.refit_bubble` (driver) and `rl.idle.refit_event_wait`, `rl.idle.generation_limit_pause` (collector actor) — `efficiency` group; named after the `Timer` category, not the algorithm | +| **GRPO** (async only) | `rl.grpo.generation` — `rollout` group, emitted by the collector actor, one span per rollout batch | +| **PPO** | `rl.ppo.job`, `rl.ppo.step`, `rl.ppo.data_processing`, `rl.ppo.generation`, `rl.ppo.reward_calculation`, `rl.ppo.policy_and_reference_logprobs`, `rl.ppo.advantage_calculation`, `rl.ppo.policy_training`, `rl.ppo.value_training`, `rl.ppo.checkpointing`, `rl.ppo.evaluate` | +| **SFT** | `rl.sft.job`, `rl.sft.step`, `rl.sft.data_processing`, `rl.sft.policy_training`, `rl.sft.checkpointing`, `rl.sft.evaluate` | +| **DPO** | `rl.dpo.job`, `rl.dpo.step`, `rl.dpo.policy_training`, `rl.dpo.checkpointing`, `rl.dpo.evaluate` | +| **RM** | `rl.rm.job`, `rl.rm.step`, `rl.rm.checkpointing`, `rl.rm.evaluate` | +| **Distillation** | `rl.distillation.job`, `rl.distillation.step`, `rl.distillation.data_processing`, `rl.distillation.generation`, `rl.distillation.teacher_logprob_inference`, `rl.distillation.policy_training`, `rl.distillation.checkpointing`, `rl.distillation.evaluate` | +| **vLLM** (driver-side) | `rl.vllm.generate`, `rl.vllm.generate_text` — `generation` group; nested under the active rollout span | +| **vLLM** (worker-side) | `rl.vllm.load_model` — `model_init` group; a root span in the generation worker's process, since Ray carries no trace context into `__init__` | + +`rl..job` is a function-level span (via `trace_fn`) wrapping the whole run. Under `per_step` it is suppressed, so each `rl..step` becomes a root trace. + +## Span tags (categorical attributes) + +These are set on spans for filtering — they answer "which one?" / "what kind?", not "how much?". Numerical values that change over time are **metrics**, not span tags (see [Metrics](metrics.md)). + +| Tag | Meaning | +|---|---| +| `rl.iteration` | training iteration index | +| `rl.epoch` | epoch index | +| `rl.step` | step index | +| `rl.num_generations_per_prompt` | GRPO group size | +| `rl.weight_version` / `rl.target_weight_version` | async rollout batch: the weights it generated from, and the training step it targets | +| `rl.num_prompt_groups` | async rollout batch width, so a gap-filling batch is not read as an unexplained speed-up | +| `rl.bucket` | goodput bucket: `productive` / `overhead` / `idle` / `wasted` (omit on umbrellas) | + +### Span group → `rl.bucket` + +Leaf groups are tagged automatically when using +`nemo_rl.telemetry.instrumentation.managed_span` / `trace_fn`. Umbrellas are +timed but **not** tagged so monitors can exclude them from goodput. + +| Group | `rl.bucket` | +|---|---| +| `job`, `step`, `rollout`, `model_init`, `evaluate` | *(none — umbrella)* | +| `generation`, `reward`, `policy_update`, `forward_backward`, `optimizer` | `productive` | +| `data_processing`, `checkpoint`, `load_checkpoint`, `logprob`, `advantage`, `reference_policy` | `overhead` | +| `efficiency` | `idle` for the two driver-side phases; *none* for the two collector-side ones (see below) | + +Rolled-up `rl.goodput` is **monitor-derived**, not emitted by NeMo-RL. + +#### Overriding the bucket for a region: `bucket_scope` + +The table above classifies by *what ran*, but a few phases are productive or not +depending on *why* they ran. `bucket_scope(bucket)` reclassifies every leaf span +opened inside it: + +```python +with bucket_scope(Bucket.OVERHEAD): + ... # generation in here is tagged overhead, not productive +``` + +The one production use is validation, in `grpo.validate` and `ppo.validate`. +Validation generates through the same `generation` group as a training rollout, +but its tokens are scored and discarded, so `productive` would count a +validation pass as goodput. The span is opened by a decorator on +`VllmGeneration.generate` that cannot see its caller, which is why the scope +travels with the execution context (a `ContextVar`) rather than an argument. + +Three properties keep it from creating the double-counting problem it exists to +avoid: umbrellas stay unbucketed inside a scope, an explicit `rl.bucket=` passed +to `managed_span` still wins, and an `efficiency_span` keeps its category's +bucket — that one names the phase it measures, so a caller cannot make +`idle/refit_bubble` productive. It propagates into coroutines started +inside the block — `asyncio.run`, as the rollout entrypoints use — but not into +threads or Ray actors, so a worker-side span is unaffected. + +### The `efficiency` group: idle time on async runs + +Async GRPO measures its stalls with `Timer` under labels like +`idle/buffer_starvation`, on both sides of the run: the driver waiting on the +collector, and the collector waiting on the driver. `efficiency_span` in +`nemo_rl/telemetry/instrumentation.py` emits those as spans, taking the bucket +from `EFFICIENCY_CATEGORY_BUCKET` so `idle/*` lands in `idle` rather than +defaulting to `overhead`. Each span also carries +`rl.efficiency.category` with the raw label, so idle time can be grouped by +cause without parsing the span name. + +Two driver-side phases are wired today, both children of `rl.grpo.step`: + +| Span | Category | Bucket | What the driver is waiting on | +|---|---|---|---| +| `rl.idle.buffer_starvation` | `idle/buffer_starvation` | `idle` | replay buffer is empty — the collector is not keeping up | +| `rl.idle.refit_bubble` | `idle/refit_bubble` | `idle` | collector reaching a safe point, then weight sync | + +With these enabled, a step's child spans account for much more of the step +duration, so a per-step goodput breakdown leaves a smaller unattributed gap. + +#### Why `idle/validation` is not a span + +`idle/validation` is driver-side wall-clock like the two above, but it stays +`Timer`-only, because the window it measures is already accounted as +**`overhead`**: `validate()` wraps its generation in `bucket_scope`, so a second +span calling the same interval `idle` would contradict the label and, wherever +those generate spans exist, be counted twice — a rollup sums durations by +`rl.bucket` with no notion of nesting, so the pass would read as nearly double +its wall time. + +Whether the children exist depends on the rollout path: sync validation +generates through the traced `rl.vllm.generate`, while async validation goes +through `generate_async`, which carries no span today. The `overhead` +attribution is the same either way, which is why this is `Timer`-only in both. + +This is the general rule for `efficiency_span`: **wrap a wait, not a phase that +does instrumented work.** A bucketed span must be a leaf, which is the same +invariant the umbrella groups exist to preserve. + +One gap remains: the phase means different things per fleet. The training GPUs +are idle while the generation GPUs do necessary non-training work, and +`overhead` on the generate span describes the latter only. Attributing the +former needs per-fleet accounting, not a per-phase bucket. Note also that the +`val_at_start` pass has no efficiency timer, so it appears in spans (as +`overhead` generation) but not in `efficiency/*`. + +#### Trace-only: the collector's two loop waits + +`idle/refit_event_wait` and `idle/generation_limit_pause` are emitted as spans +from inside the `AsyncTrajectoryCollector`, but **without** `rl.bucket`: + +| Span | Category | Bucket | What it means | +|---|---|---|---| +| `rl.idle.refit_event_wait` | `idle/refit_event_wait` | *none* | collection loop parked while a refit completes | +| `rl.idle.generation_limit_pause` | `idle/generation_limit_pause` | *none* | every target weight already has enough trajectories | + +Both are `Event.wait()` calls on the single collection-loop thread, so they are +honest wall-clock durations — but the collector's wall clock runs *concurrently* +with the driver's, so summing them against a driver-side denominator would +overcount. Omitting the attribute keeps them out of a bucket rollup by +construction instead of by convention. The membership list is +`COLLECTOR_LOOP_CATEGORIES` in `nemo_rl/telemetry/instrumentation.py`. + +They still carry `rl.efficiency.category`, so they remain identifiable in a +trace and continue to be reported as `efficiency/*` scalars. As metrics they are +labelled `rl.efficiency.measurement="collector_wall_clock"` — sequential and so +real durations, unlike the batch-worker categories, but on the collector's +timeline rather than the driver's. See +[Metrics — always filter on `rl.efficiency.measurement`](metrics.md#always-filter-on-rlefficiencymeasurement). + +#### Still reserved: the rest of the collector-side categories + +`idle/buffer_full_backoff` and `wasted/failed_trajectory` stay `Timer`-only. +Both run in the batch-worker threads, so they are genuinely *thread-seconds* — +several workers accumulate at once and the total can exceed the wall time it +happened in. `idle/buffer_full_backoff` also has no clean block to wrap: it is +recorded as a precomputed duration spanning a retry loop. `wasted/failed_trajectory` +covers the same window as the enclosing `rl.grpo.generation` span, so a span +there would duplicate an existing interval. `init/total` is likewise still +`Timer`-only — it runs before the per-step loop, so it fills no step-level gap. + +So goodput on async runs covers driver idle, but not collector-side idle or +wasted work — use the `efficiency/*` metrics for those. + +### Async rollout spans come from the collector actor + +In an async run, no rollout is generated on the driver. Every trajectory comes +from inside `AsyncTrajectoryCollector`, a separate Ray actor, which calls +`init_telemetry_worker(rank=0, world_size=1, always_export=True)` in its +constructor. Explicit rank because it is a singleton, not a member of a ranked +group, and its `runtime_env` is a copy of the driver's environment, so an +inherited `RANK` must not decide whether it exports. `always_export` goes with +that synthetic rank: an `export_strategy` that selects among a group's ranks has +no meaning applied to a made-up one, and would silently mute the actor — with +`export_rank: 3`, every rollout span in the run would disappear. The driver uses +the same override for the same reason. + +It flushes through `flush_telemetry()`, which the driver calls before `ray.kill`, +since a kill runs no `atexit` handler. That call stops the collection loop and +waits (bounded) for in-flight batch workers first: the shutdown is terminal, so +a still-running thread would keep opening spans against a dead processor. + +Each batch worker opens one `rl.grpo.generation` span — the same name the sync +path uses on the driver, so the two modes read alike — carrying +`rl.weight_version`, `rl.target_weight_version`, `rl.num_generations_per_prompt` +and `rl.num_prompt_groups`. The last one is the batch width: a gap-filling batch +covers a fraction of a full one, so without it a short span looks like an +unexplained speed-up. It is in the `rollout` group, so it is an +umbrella and carries **no** `rl.bucket`: several batch workers run at once, so +their durations sum past wall time and cannot enter a bucket rollup. + +#### Getting the collector into one waterfall + +Ray does not propagate OTel context, so an actor's spans start their own trace +by default. The driver captures its active span as a W3C `traceparent` carrier +with `current_trace_carrier()` — taken inside `rl.grpo.job`, at the point the +collector is constructed — and passes it as the actor's `trace_carrier` +argument. The collector reopens it with `remote_trace_context()` in **both** the +collection-loop thread and every batch-worker thread. Per thread, not once per +process: OTel context is a `ContextVar`, and `threading.Thread` inherits none. + +The result is a single trace per run: + +``` +rl.grpo.job (driver) +├── rl.grpo.step (iteration 1) (driver) +│ ├── rl.idle.buffer_starvation +│ └── rl.grpo.policy_training +├── rl.grpo.generation weight=7 (collector, thread A) +├── rl.idle.generation_limit_pause (collector, loop thread) +├── rl.grpo.generation weight=8 (collector, thread B) +└── rl.grpo.step (iteration 2) (driver) +``` + +**This requires the `job` group to be enabled.** `current_trace_carrier()` +returns an empty dict when no span is recording, and `remote_trace_context({})` +is a no-op, so the collector falls back to root spans. The `default` preset has +`job` but not `rollout`/`efficiency`, so the collector emits nothing; `per_step` +has `rollout`/`efficiency` but deliberately omits `job` so each step is its own +bounded trace. For the unified view, ask for both: + +```yaml +telemetry: + span_groups: per_step,job # or: all +``` + +Be deliberate about it. A run-long root span means one trace accumulating every +step and every rollout batch for the whole job, which is exactly the trace-size +problem `per_step` exists to avoid. Prefer it for debugging a specific run, not +as a standing default on long jobs. + +Two consequences worth internalizing before reading an async trace: + +- **One span per batch, not per sample.** `generate_async` is dispatched one + coroutine per sample, so spanning it would emit thousands of mutually + overlapping spans per step. +- **There is no `productive` generation span in async mode**, and there cannot + usefully be one. `rl.vllm.generate` is only reached through the synchronous + rollout path, and an async run never takes it: `async_grpo_train` requires an + async generation engine, so even validation goes through `generate_async`, + which carries no span today. Generation is a + continuously-batched pipeline overlapping training, so its productive + contribution is a utilization question — answered by fleet metrics — not a + span duration. A span-derived goodput ratio on an async run therefore has no + productive generation term; do not read it as "generation contributed + nothing." + +## Coverage gaps + +A group being enabled does not guarantee spans: something has to emit them. Known +blanks today, so an empty trace is not read as a broken exporter: + +| Area | State | +|---|---| +| SGLang / TRT-LLM / Megatron generation workers | uninstrumented — no `init_telemetry_worker` and no generation spans; only vLLM emits `rl.vllm.*`. (Policy and value workers do initialise telemetry, so their metrics and any future spans are wired.) | +| `VllmGeneration.generate_async` | no span, so async rollouts and async validation have no generate breakdown under `rl.grpo.generation` / `rl.grpo.evaluate` | +| `SyncRolloutActor` | the sync data-plane counterpart of the async collector — uninstrumented, so its rollouts produce no spans | +| Worker flush outside async GRPO | only `async_grpo_train` calls `policy.shutdown()` / `policy_generation.shutdown()`, so on other trainers a worker's last spans depend on the periodic export rather than a flush | +| `load_checkpoint`, `forward_backward`, `optimizer`, `reference_policy` | the groups are defined and bucketed, but no site emits them, so enabling them adds no spans | +| `grpo_sync.py`, `single_controller.py` | no spans; `examples/run_grpo_single_controller.py` also never initialises telemetry, so that entrypoint emits nothing at all | +| `run_vlm_grpo.py`, `run_grpo_sliding_puzzle.py`, `run_xtoken_off_policy_distillation.py`, `run_eval.py` | these call the instrumented loops but never `init_telemetry_driver`, so a `telemetry:` block in their configs parses, the run succeeds, and nothing is emitted — driver or worker | +| Ranked worker spans | separate traces, correlated by `run_id` — only the async collector's context is propagated | + +## Resource attributes (process tags) + +Stable-for-the-run values, set once at init and attached to every span/metric: `rl.algorithm`, `rl.model`, `nemo.precision`, `dl.tensor_parallel.size`, `dl.pipeline_parallel.size`, plus `dl.rank` / `dl.world_size` (set automatically by lens). See [Configuration — Resource attributes](configuration.md#resource-attributes). + +## Granularity guidance + +| Span groups | Relative cost | Recommendation | +|---|---|---| +| Disabled (`telemetry.enabled: false`) | None | The default | +| `default` | Lowest | Safe for all production runs | +| `per_step` | Moderate | Per-step profiling; each step is its own trace | +| `all` | Highest | Development / deep debugging | + +Non-exporting ranks have an empty span-group set — `is_span_group_enabled()` returns `False` everywhere, so no span objects are created at all. The disabled path is a `frozenset` lookup and an immediate return. See [lens: architecture](https://github.com/NVIDIA-NeMo/Lens). diff --git a/docs/observability/vllm-tracing.md b/docs/observability/vllm-tracing.md new file mode 100644 index 0000000000..9b7fbe5223 --- /dev/null +++ b/docs/observability/vllm-tracing.md @@ -0,0 +1,58 @@ +# vLLM Tracing + +Generation is where most of an RL step's wall-clock goes, so NeMo-RL instruments vLLM at **two independent layers**. They answer different questions, ship over different transports, and are enabled independently. + +| | Layer 1 — RL-side spans | Layer 2 — vLLM native OTLP | +|---|---|---| +| What | `rl.vllm.generate` / `rl.vllm.generate_text` spans + token/latency metrics, emitted by NeMo-RL around the vLLM call | vLLM's own internal engine spans (scheduling, prefill/decode, ...) | +| Where | driver, `nemo_rl/models/generation/vllm/vllm_generation.py` | vLLM engine, enabled in `vllm_worker.py` | +| Enabled by | `generation` span group (on by default in `per_step`/`all`) | opt-in: `telemetry.vllm_native_tracing: true` | +| Transport | rides the normal lens OTLP path (`http/protobuf` OK) | **gRPC-only** (needs an OTLP/gRPC endpoint / collector) | +| Correlation | nested under the rollout span (parent-child) | via shared `run_id` / resource attributes (not parent-child) | + +## Layer 1 — RL-side generation spans (default) + +`VllmGeneration.generate` / `generate_text` on the driver are wrapped with `trace_fn(RLSpanGroup.GENERATION, ...)`, emitting `rl.vllm.generate` and `rl.vllm.generate_text` spans. These nest under the active `rl..collect_rollouts` span, so a rollout waterfall shows exactly how long generation took and how it fits inside the step. They also emit the `gen_ai.*` token/latency metrics (see [Metrics](metrics.md)). + +Because these are ordinary lens spans, they travel the same OTLP transport as everything else — including a direct-to-backend `http/protobuf` export path. **Nothing extra is required**: enable the `generation` group (it is in the `per_step` and `all` presets) and they appear. + +This covers the synchronous rollout path only. Async runs drive generation through `generate_async`, which carries no span today, so Layer 1 shows no generate spans there — and the `rl.grpo.generation` span they would nest under comes from the collector actor rather than the driver. See [span groups — coverage gaps](span-groups.md#coverage-gaps). + +## Layer 2 — vLLM native OTLP tracing (opt-in) + +vLLM can emit its own OpenTelemetry spans for the engine internals. Enable it in your run config: + +```yaml +telemetry: + enabled: true + vllm_native_tracing: true +``` + +and point the exporter at a gRPC endpoint: + +```bash +export OTEL_EXPORTER_OTLP_ENDPOINT=http://:4317 # gRPC! +``` + +Under the hood, `_maybe_enable_vllm_native_tracing()` (in `vllm_worker.py`, called from `_load_model`) sets `otlp_traces_endpoint` and `collect_detailed_traces=["all"]` on the vLLM engine args. It reads `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT` if set, otherwise `OTEL_EXPORTER_OTLP_ENDPOINT`. + +### Caveat 1 — vLLM's exporter is gRPC-only + +vLLM's OTLP span exporter speaks **OTLP/gRPC only**. It needs a gRPC OTLP endpoint — a collector on `:4317` or a gRPC-capable backend. It will **not** ride an `http/protobuf` OTLP endpoint, including a direct-to-backend `http/protobuf` path like the one Layer 1 uses. + +So to get vLLM's native spans you need a gRPC OTLP receiver in the picture (e.g. an OTel Collector on `:4317` that forwards to your backend). This is why native tracing is left **off** by default when exporting to an `http/protobuf` endpoint with no collector. See [Observability Stack](observability-stack.md). + +### Caveat 2 — offline generation cannot carry a trace context + +NeMo-RL drives vLLM through the offline `LLM.generate()` API, which does not accept a per-request trace context. So vLLM's native spans **cannot** nest as children of the RL rollout span. Instead they correlate to the RL run through the **shared `run_id` and resource attributes** that every process in the job carries — you line them up by run, not by parent-child edges in one waterfall. + +Practically: Layer 1 gives you generation timing *inside* the RL step tree; Layer 2 gives you vLLM engine internals as a separate set of spans tagged with the same `run_id`. Use both when you need to see why generation was slow at the engine level. + +### Graceful degradation + +If the installed vLLM does not support `otlp_traces_endpoint` (older versions), `_maybe_enable_vllm_native_tracing` logs a warning and skips — it never breaks the run. `collect_detailed_traces` is only set when that engine arg is also supported. If the flag is set but no OTLP endpoint is configured, it logs a warning and does nothing. + +## Which layer do I want? + +- **Just want to see generation cost per rollout?** Layer 1 — enable the `generation` group. Works over any transport, including a direct-to-backend `http/protobuf` path. +- **Debugging vLLM engine internals (scheduling, batching, prefill/decode)?** Add Layer 2 — but stand up a gRPC OTLP collector first, and correlate by `run_id`. diff --git a/examples/run_distillation.py b/examples/run_distillation.py index 5926a1abaf..7b0fb9ad76 100644 --- a/examples/run_distillation.py +++ b/examples/run_distillation.py @@ -22,6 +22,7 @@ from nemo_rl.data.utils import setup_response_data from nemo_rl.distributed.virtual_cluster import init_ray from nemo_rl.models.generation import configure_generation_config +from nemo_rl.telemetry.setup import init_telemetry_driver, shutdown_telemetry from nemo_rl.utils.config import ( load_config, parse_hydra_overrides, @@ -67,57 +68,68 @@ def main() -> None: # Get the next experiment directory with incremented ID config.logger["log_dir"] = get_next_experiment_dir(config.logger["log_dir"]) - init_ray() + # Initialise telemetry on the driver BEFORE init_ray() so the resolved + # NEMO_RL_OTEL_* env is snapshotted into the Ray runtime_env and inherited + # by every worker. No-op unless nemo-lens is installed and telemetry is on. + init_telemetry_driver(config, algorithm="distillation") - tokenizer = get_tokenizer(config.policy["tokenizer"]) + try: + init_ray() - if config.policy["generation"] is not None: - config.policy["generation"] = configure_generation_config( - config.policy["generation"], tokenizer - ) - else: - print(" ⚠️ No generation config found, this may cause issues") - - # setup data - ( - dataset, - val_dataset, - task_to_env, - val_task_to_env, - ) = setup_response_data(tokenizer, config.data, config.env) - - ( - student_policy, - teacher_policy, - student_generation, - _nemo_gym, - dataloader, - val_dataloader, - loss_fn, - logger, - checkpointer, - distillation_state, - master_config, - ) = setup(config, tokenizer, dataset, val_dataset) - - # The checkpointer owns background async-checkpoint finalization threads; - # the context manager guarantees they are flushed (rename + delete) on exit. - with checkpointer: - distillation_train( + tokenizer = get_tokenizer(config.policy["tokenizer"]) + + if config.policy["generation"] is not None: + config.policy["generation"] = configure_generation_config( + config.policy["generation"], tokenizer + ) + else: + print(" ⚠️ No generation config found, this may cause issues") + + # setup data + ( + dataset, + val_dataset, + task_to_env, + val_task_to_env, + ) = setup_response_data(tokenizer, config.data, config.env) + + ( student_policy, teacher_policy, student_generation, + _nemo_gym, dataloader, val_dataloader, - tokenizer, # pass tokenizer parameter loss_fn, - task_to_env, - val_task_to_env, logger, checkpointer, distillation_state, master_config, - ) + ) = setup(config, tokenizer, dataset, val_dataset) + + # The checkpointer owns background async-checkpoint finalization threads; + # the context manager guarantees they are flushed (rename + delete) on exit. + with checkpointer: + distillation_train( + student_policy, + teacher_policy, + student_generation, + dataloader, + val_dataloader, + tokenizer, # pass tokenizer parameter + loss_fn, + task_to_env, + val_task_to_env, + logger, + checkpointer, + distillation_state, + master_config, + ) + finally: + # Flush on the failure paths too, and before cluster teardown: the OTel + # SDK's own atexit hook is registered ahead of Ray's and so runs after + # it. No-op when telemetry is inactive. + shutdown_telemetry() if __name__ == "__main__": diff --git a/examples/run_dpo.py b/examples/run_dpo.py index df05302528..80ded7387d 100644 --- a/examples/run_dpo.py +++ b/examples/run_dpo.py @@ -22,6 +22,7 @@ from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data.utils import setup_preference_data from nemo_rl.distributed.virtual_cluster import init_ray +from nemo_rl.telemetry.setup import init_telemetry_driver, shutdown_telemetry from nemo_rl.utils.config import load_config, parse_hydra_overrides from nemo_rl.utils.logger import get_next_experiment_dir @@ -68,40 +69,51 @@ def main(): f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" ) - init_ray() - - # setup tokenizer - tokenizer = get_tokenizer(config.policy["tokenizer"]) - - # setup data - dataset, val_dataset = setup_preference_data(tokenizer, config.data) - - ( - policy, - cluster, - train_dataloader, - val_dataloader, - loss_fn, - logger, - checkpointer, - dpo_save_state, - master_config, - ) = setup(config, tokenizer, dataset, val_dataset) - - # The checkpointer owns background async-checkpoint finalization threads; - # the context manager guarantees they are flushed (rename + delete) on exit. - with checkpointer: - dpo_train( + # Initialise telemetry on the driver BEFORE init_ray() so the resolved + # NEMO_RL_OTEL_* env is snapshotted into the Ray runtime_env and inherited + # by every worker. No-op unless nemo-lens is installed and telemetry is on. + init_telemetry_driver(config, algorithm="dpo") + + try: + init_ray() + + # setup tokenizer + tokenizer = get_tokenizer(config.policy["tokenizer"]) + + # setup data + dataset, val_dataset = setup_preference_data(tokenizer, config.data) + + ( policy, + cluster, train_dataloader, val_dataloader, - tokenizer, loss_fn, - master_config, logger, checkpointer, dpo_save_state, - ) + master_config, + ) = setup(config, tokenizer, dataset, val_dataset) + + # The checkpointer owns background async-checkpoint finalization threads; + # the context manager guarantees they are flushed (rename + delete) on exit. + with checkpointer: + dpo_train( + policy, + train_dataloader, + val_dataloader, + tokenizer, + loss_fn, + master_config, + logger, + checkpointer, + dpo_save_state, + ) + finally: + # Flush on the failure paths too, and before cluster teardown: the OTel + # SDK's own atexit hook is registered ahead of Ray's and so runs after + # it. No-op when telemetry is inactive. + shutdown_telemetry() if __name__ == "__main__": diff --git a/examples/run_grpo.py b/examples/run_grpo.py index c6f06943b7..b417abe795 100644 --- a/examples/run_grpo.py +++ b/examples/run_grpo.py @@ -30,6 +30,7 @@ from nemo_rl.data_plane.factory import maybe_configure_data_plane_env from nemo_rl.distributed.virtual_cluster import init_ray from nemo_rl.models.generation import configure_generation_config +from nemo_rl.telemetry.setup import init_telemetry_driver, shutdown_telemetry from nemo_rl.utils.config import ( load_config, parse_hydra_overrides, @@ -108,153 +109,164 @@ def main() -> None: f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" ) - with rl_init_timer.time("ray_connect"): - # Must precede init_ray() — see maybe_configure_data_plane_env's docstring. - maybe_configure_data_plane_env(config.data_plane) - init_ray() + # Initialise telemetry on the driver BEFORE init_ray() so the resolved + # NEMO_RL_OTEL_* env is snapshotted into the Ray runtime_env and inherited + # by every worker. No-op unless nemo-lens is installed and telemetry is on. + init_telemetry_driver(config, algorithm="grpo") - # setup tokenizer - with rl_init_timer.time("tokenizer"): - tokenizer = get_tokenizer(config.policy["tokenizer"]) - assert config.policy["generation"] is not None, ( - "A generation config is required for GRPO" - ) - has_refit_draft_weights = bool(config.policy["draft"]["enabled"]) - megatron_cfg = config.policy.get("megatron_cfg") or {} - trains_mtp = bool(megatron_cfg.get("mtp_num_layers")) - config.policy["generation"] = configure_generation_config( - config.policy["generation"], - tokenizer, - has_refit_draft_weights=has_refit_draft_weights, - trains_mtp=trains_mtp, - ) + try: + with rl_init_timer.time("ray_connect"): + # Must precede init_ray() — see maybe_configure_data_plane_env's docstring. + maybe_configure_data_plane_env(config.data_plane) + init_ray() + + # setup tokenizer + with rl_init_timer.time("tokenizer"): + tokenizer = get_tokenizer(config.policy["tokenizer"]) + assert config.policy["generation"] is not None, ( + "A generation config is required for GRPO" + ) + has_refit_draft_weights = bool(config.policy["draft"]["enabled"]) + megatron_cfg = config.policy.get("megatron_cfg") or {} + trains_mtp = bool(megatron_cfg.get("mtp_num_layers")) + config.policy["generation"] = configure_generation_config( + config.policy["generation"], + tokenizer, + has_refit_draft_weights=has_refit_draft_weights, + trains_mtp=trains_mtp, + ) - # setup data - with rl_init_timer.time("data"): - dataset, val_dataset, task_to_env, val_task_to_env = setup_response_data( - tokenizer, config.data, config.env - ) + # setup data + with rl_init_timer.time("data"): + dataset, val_dataset, task_to_env, val_task_to_env = setup_response_data( + tokenizer, config.data, config.env + ) - # Pick the policy factory at the launcher level so the legacy trainer - # stays data-plane-agnostic (architectural invariant — see - # tests/data_plane/unit/test_architecture_invariants.py). - _dp_cfg = config.data_plane or {} - if _dp_cfg.get("enabled", False): - from nemo_rl.models.policy.tq_policy import TQPolicy - - def _make_policy(**kwargs): - return TQPolicy(**kwargs, dp_cfg=_dp_cfg) - - _policy_factory = _make_policy - else: - _policy_factory = None # setup() defaults to plain Policy - - with rl_init_timer.time("setup"): - ( - policy, - policy_generation, - _nemo_gym, - cluster, - dataloader, - val_dataloader, - loss_fn, - logger, - checkpointer, - grpo_state, - master_config, - teacher_worker_groups, - alias_to_group_alias, - ) = setup( - config, - tokenizer, - dataset, - val_dataset, - policy_factory=_policy_factory, - ) + # Pick the policy factory at the launcher level so the legacy trainer + # stays data-plane-agnostic (architectural invariant — see + # tests/data_plane/unit/test_architecture_invariants.py). + _dp_cfg = config.data_plane or {} + if _dp_cfg.get("enabled", False): + from nemo_rl.models.policy.tq_policy import TQPolicy - rl_init_timer.record("total", time.perf_counter() - main_start) + def _make_policy(**kwargs): + return TQPolicy(**kwargs, dp_cfg=_dp_cfg) - rl_init_metrics = rl_init_timer.get_timing_metrics(reduction_op="sum") - print("\n" + "=" * 60) - print(" " * 14 + "RL INIT TIMING BREAKDOWN") - for label, value in sorted(rl_init_metrics.items()): - if isinstance(value, (int, float)): - print(f" {label}: {value:.1f}s") - print("=" * 60 + "\n", flush=True) + _policy_factory = _make_policy + else: + _policy_factory = None # setup() defaults to plain Policy + + with rl_init_timer.time("setup"): + ( + policy, + policy_generation, + _nemo_gym, + cluster, + dataloader, + val_dataloader, + loss_fn, + logger, + checkpointer, + grpo_state, + master_config, + teacher_worker_groups, + alias_to_group_alias, + ) = setup( + config, + tokenizer, + dataset, + val_dataset, + policy_factory=_policy_factory, + ) - try: - # Check if async mode is enabled - if config.grpo.async_grpo.enabled: - # Async GRPO does not support dynamic sampling, reward scaling, or reward shaping (DAPO features) - if config.grpo.use_dynamic_sampling: - raise NotImplementedError( - "use_dynamic_sampling is not supported with async GRPO" - ) - if config.grpo.reward_scaling.enabled: - raise NotImplementedError( - "reward_scaling is not supported with async GRPO" - ) - if config.grpo.reward_shaping.enabled: - raise NotImplementedError( - "reward_shaping is not supported with async GRPO" - ) + rl_init_timer.record("total", time.perf_counter() - main_start) - # Async GRPO does not support multiple dataloaders - if config.data["use_multiple_dataloader"]: - raise NotImplementedError( - "use_multiple_dataloader is not supported with async GRPO" - ) + rl_init_metrics = rl_init_timer.get_timing_metrics(reduction_op="sum") + print("\n" + "=" * 60) + print(" " * 14 + "RL INIT TIMING BREAKDOWN") + for label, value in sorted(rl_init_metrics.items()): + if isinstance(value, (int, float)): + print(f" {label}: {value:.1f}s") + print("=" * 60 + "\n", flush=True) - from nemo_rl.algorithms.grpo import async_grpo_train - - print("🚀 Running async GRPO training") - - # Run async GRPO training - async_grpo_train( - policy=policy, - policy_generation=policy_generation, - dataloader=dataloader, - val_dataloader=val_dataloader, - tokenizer=tokenizer, - loss_fn=loss_fn, - task_to_env=task_to_env, - val_task_to_env=val_task_to_env, - logger=logger, - checkpointer=checkpointer, - grpo_save_state=grpo_state, - master_config=master_config, - max_trajectory_age_steps=config.grpo.async_grpo.max_trajectory_age_steps, - teacher_worker_groups=teacher_worker_groups, - alias_to_group_alias=alias_to_group_alias, - ) - else: - # Two parallel synchronous trainers (verl-style — main_ppo.py vs - # main_ppo_sync.py). data_plane.enabled selects which one runs. - trainer = _select_trainer(master_config) - # grpo_train_sync defers checkpoint finalization to the checkpointer's - # background threads; the context manager guarantees they are flushed on - # exit. (grpo_train also flushes internally; shutdown() is idempotent.) - with checkpointer: - trainer( - policy, - policy_generation, - dataloader, - val_dataloader, - tokenizer, - loss_fn, - task_to_env, - val_task_to_env, - logger, - checkpointer, - grpo_state, - master_config, + try: + # Check if async mode is enabled + if config.grpo.async_grpo.enabled: + # Async GRPO does not support dynamic sampling, reward scaling, or reward shaping (DAPO features) + if config.grpo.use_dynamic_sampling: + raise NotImplementedError( + "use_dynamic_sampling is not supported with async GRPO" + ) + if config.grpo.reward_scaling.enabled: + raise NotImplementedError( + "reward_scaling is not supported with async GRPO" + ) + if config.grpo.reward_shaping.enabled: + raise NotImplementedError( + "reward_shaping is not supported with async GRPO" + ) + + # Async GRPO does not support multiple dataloaders + if config.data["use_multiple_dataloader"]: + raise NotImplementedError( + "use_multiple_dataloader is not supported with async GRPO" + ) + + from nemo_rl.algorithms.grpo import async_grpo_train + + print("🚀 Running async GRPO training") + + # Run async GRPO training + async_grpo_train( + policy=policy, + policy_generation=policy_generation, + dataloader=dataloader, + val_dataloader=val_dataloader, + tokenizer=tokenizer, + loss_fn=loss_fn, + task_to_env=task_to_env, + val_task_to_env=val_task_to_env, + logger=logger, + checkpointer=checkpointer, + grpo_save_state=grpo_state, + master_config=master_config, + max_trajectory_age_steps=config.grpo.async_grpo.max_trajectory_age_steps, + teacher_worker_groups=teacher_worker_groups, + alias_to_group_alias=alias_to_group_alias, ) + else: + # Two parallel synchronous trainers (verl-style — main_ppo.py vs + # main_ppo_sync.py). data_plane.enabled selects which one runs. + trainer = _select_trainer(master_config) + # grpo_train_sync defers checkpoint finalization to the checkpointer's + # background threads; the context manager guarantees they are flushed on + # exit. (grpo_train also flushes internally; shutdown() is idempotent.) + with checkpointer: + trainer( + policy, + policy_generation, + dataloader, + val_dataloader, + tokenizer, + loss_fn, + task_to_env, + val_task_to_env, + logger, + checkpointer, + grpo_state, + master_config, + ) + finally: + shutdown_environments(task_to_env, val_task_to_env) + try: + policy_generation.shutdown() + except Exception as error: + print(f"Error shutting down generation: {error}", flush=True) finally: - shutdown_environments(task_to_env, val_task_to_env) - try: - policy_generation.shutdown() - except Exception as error: - print(f"Error shutting down generation: {error}", flush=True) + # Flush on the failure paths too, and before cluster teardown: the OTel + # SDK's own atexit hook is registered ahead of Ray's and so runs after + # it. No-op when telemetry is inactive. + shutdown_telemetry() if __name__ == "__main__": diff --git a/examples/run_ppo.py b/examples/run_ppo.py index b7e8d35572..bfcc1af7d3 100644 --- a/examples/run_ppo.py +++ b/examples/run_ppo.py @@ -24,6 +24,7 @@ from nemo_rl.distributed.virtual_cluster import init_ray from nemo_rl.models.generation import configure_generation_config from nemo_rl.models.generation.interfaces import GenerationInterface +from nemo_rl.telemetry.setup import init_telemetry_driver, shutdown_telemetry from nemo_rl.utils.config import ( load_config, parse_hydra_overrides, @@ -128,82 +129,93 @@ def main() -> None: f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" ) - init_ray() + # Initialise telemetry on the driver BEFORE init_ray() so the resolved + # NEMO_RL_OTEL_* env is snapshotted into the Ray runtime_env and inherited + # by every worker. No-op unless nemo-lens is installed and telemetry is on. + init_telemetry_driver(config, algorithm="ppo") - # setup tokenizer - tokenizer = get_tokenizer(config.policy["tokenizer"]) - assert config.policy["generation"] is not None, ( - "A generation config is required for PPO" - ) - config.policy["generation"] = configure_generation_config( - config.policy["generation"], tokenizer - ) + try: + init_ray() + + # setup tokenizer + tokenizer = get_tokenizer(config.policy["tokenizer"]) + assert config.policy["generation"] is not None, ( + "A generation config is required for PPO" + ) + config.policy["generation"] = configure_generation_config( + config.policy["generation"], tokenizer + ) - # setup data - ( - dataset, - val_dataset, - task_to_env, - val_task_to_env, - ) = setup_response_data(tokenizer, config.data, config.env) - - ( - policy, - policy_generation, - value_model, - cluster, - dataloader, - val_dataloader, - loss_fn, - value_loss_fn, - logger, - checkpointer, - ppo_state, - master_config, - ) = setup(config, tokenizer, dataset, val_dataset) - - async_config = config.ppo.async_ppo - async_ppo_enabled = async_config.enabled - if async_ppo_enabled: - _validate_async_ppo_config(config, policy_generation) - - with checkpointer: + # setup data + ( + dataset, + val_dataset, + task_to_env, + val_task_to_env, + ) = setup_response_data(tokenizer, config.data, config.env) + + ( + policy, + policy_generation, + value_model, + cluster, + dataloader, + val_dataloader, + loss_fn, + value_loss_fn, + logger, + checkpointer, + ppo_state, + master_config, + ) = setup(config, tokenizer, dataset, val_dataset) + + async_config = config.ppo.async_ppo + async_ppo_enabled = async_config.enabled if async_ppo_enabled: - print("🚀 Running asynchronous PPO training") - async_ppo_train( - policy, - policy_generation, - value_model, - dataloader, - val_dataloader, - tokenizer, - loss_fn, - value_loss_fn, - task_to_env, - val_task_to_env, - logger, - checkpointer, - ppo_state, - master_config, - ) - else: - print("🚀 Running synchronous PPO training") - ppo_train( - policy, - policy_generation, - value_model, - dataloader, - val_dataloader, - tokenizer, - loss_fn, - value_loss_fn, - task_to_env, - val_task_to_env, - logger, - checkpointer, - ppo_state, - master_config, - ) + _validate_async_ppo_config(config, policy_generation) + + with checkpointer: + if async_ppo_enabled: + print("🚀 Running asynchronous PPO training") + async_ppo_train( + policy, + policy_generation, + value_model, + dataloader, + val_dataloader, + tokenizer, + loss_fn, + value_loss_fn, + task_to_env, + val_task_to_env, + logger, + checkpointer, + ppo_state, + master_config, + ) + else: + print("🚀 Running synchronous PPO training") + ppo_train( + policy, + policy_generation, + value_model, + dataloader, + val_dataloader, + tokenizer, + loss_fn, + value_loss_fn, + task_to_env, + val_task_to_env, + logger, + checkpointer, + ppo_state, + master_config, + ) + finally: + # Flush on the failure paths too, and before cluster teardown: the OTel + # SDK's own atexit hook is registered ahead of Ray's and so runs after + # it. No-op when telemetry is inactive. + shutdown_telemetry() if __name__ == "__main__": diff --git a/examples/run_rm.py b/examples/run_rm.py index 19b6940e01..8523a9eaaf 100644 --- a/examples/run_rm.py +++ b/examples/run_rm.py @@ -22,6 +22,7 @@ from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data.utils import setup_preference_data from nemo_rl.distributed.virtual_cluster import init_ray +from nemo_rl.telemetry.setup import init_telemetry_driver, shutdown_telemetry from nemo_rl.utils.config import load_config, parse_hydra_overrides from nemo_rl.utils.logger import get_next_experiment_dir @@ -71,40 +72,51 @@ def main(): f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" ) - init_ray() - - # setup tokenizer - tokenizer = get_tokenizer(config.policy["tokenizer"]) - - # setup data - dataset, val_dataset = setup_preference_data(tokenizer, config.data) - - ( - policy, - cluster, - train_dataloader, - val_dataloader, - loss_fn, - logger, - checkpointer, - rm_save_state, - master_config, - ) = setup(config, tokenizer, dataset, val_dataset) - - # The checkpointer owns background async-checkpoint finalization threads; - # the context manager guarantees they are flushed (rename + delete) on exit. - with checkpointer: - rm_train( + # Initialise telemetry on the driver BEFORE init_ray() so the resolved + # NEMO_RL_OTEL_* env is snapshotted into the Ray runtime_env and inherited + # by every worker. No-op unless nemo-lens is installed and telemetry is on. + init_telemetry_driver(config, algorithm="rm") + + try: + init_ray() + + # setup tokenizer + tokenizer = get_tokenizer(config.policy["tokenizer"]) + + # setup data + dataset, val_dataset = setup_preference_data(tokenizer, config.data) + + ( policy, + cluster, train_dataloader, val_dataloader, - tokenizer, loss_fn, - master_config, logger, checkpointer, rm_save_state, - ) + master_config, + ) = setup(config, tokenizer, dataset, val_dataset) + + # The checkpointer owns background async-checkpoint finalization threads; + # the context manager guarantees they are flushed (rename + delete) on exit. + with checkpointer: + rm_train( + policy, + train_dataloader, + val_dataloader, + tokenizer, + loss_fn, + master_config, + logger, + checkpointer, + rm_save_state, + ) + finally: + # Flush on the failure paths too, and before cluster teardown: the OTel + # SDK's own atexit hook is registered ahead of Ray's and so runs after + # it. No-op when telemetry is inactive. + shutdown_telemetry() if __name__ == "__main__": diff --git a/examples/run_sft.py b/examples/run_sft.py index e06011a8ea..5683a5f664 100644 --- a/examples/run_sft.py +++ b/examples/run_sft.py @@ -30,6 +30,7 @@ update_single_dataset_config, ) from nemo_rl.distributed.virtual_cluster import init_ray +from nemo_rl.telemetry.setup import init_telemetry_driver, shutdown_telemetry from nemo_rl.utils.config import ( load_config, parse_hydra_overrides, @@ -189,40 +190,51 @@ def main(is_vlm: bool = False): f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" ) - init_ray() - - # setup tokenizer (or processor) - tokenizer = get_tokenizer(config.policy["tokenizer"], get_processor=is_vlm) - - # setup data - dataset, val_dataset = setup_data(tokenizer, config.data) - - ( - policy, - cluster, - train_dataloader, - val_dataloader, - loss_fn, - logger, - checkpointer, - sft_save_state, - master_config, - ) = setup(config, tokenizer, dataset, val_dataset) - - # The checkpointer owns background async-checkpoint finalization threads; - # the context manager guarantees they are flushed (rename + delete) on exit. - with checkpointer: - sft_train( + # Initialise telemetry on the driver BEFORE init_ray() so the resolved + # NEMO_RL_OTEL_* env is snapshotted into the Ray runtime_env and inherited + # by every worker. No-op unless nemo-lens is installed and telemetry is on. + init_telemetry_driver(config, algorithm="sft") + + try: + init_ray() + + # setup tokenizer (or processor) + tokenizer = get_tokenizer(config.policy["tokenizer"], get_processor=is_vlm) + + # setup data + dataset, val_dataset = setup_data(tokenizer, config.data) + + ( policy, + cluster, train_dataloader, val_dataloader, - tokenizer, loss_fn, - master_config, logger, checkpointer, sft_save_state, - ) + master_config, + ) = setup(config, tokenizer, dataset, val_dataset) + + # The checkpointer owns background async-checkpoint finalization threads; + # the context manager guarantees they are flushed (rename + delete) on exit. + with checkpointer: + sft_train( + policy, + train_dataloader, + val_dataloader, + tokenizer, + loss_fn, + master_config, + logger, + checkpointer, + sft_save_state, + ) + finally: + # Flush on the failure paths too, and before cluster teardown: the OTel + # SDK's own atexit hook is registered ahead of Ray's and so runs after + # it. No-op when telemetry is inactive. + shutdown_telemetry() if __name__ == "__main__": diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index 19face3d62..3aed00b045 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -64,6 +64,13 @@ GenerationInterface, should_use_async_rollouts, ) +from nemo_rl.telemetry.instrumentation import ( + efficiency_span, + managed_span, + remote_trace_context, +) +from nemo_rl.telemetry.setup import init_telemetry_worker, shutdown_telemetry +from nemo_rl.telemetry.span_groups import RLSpanGroup from nemo_rl.utils.logger import should_log_nemo_gym_full_result_tables from nemo_rl.utils.multimodal_payload_metrics import ( collect_multimodal_payload_metrics, @@ -76,6 +83,10 @@ _MAX_NEMO_GYM_STREAM_RETRIES = 3 _NEMO_GYM_RETRY_DELAY_BASE_SECONDS = 1.0 _REPLAY_BUFFER_MAX_BACKOFF_SECONDS = 0.5 +# How often telemetry teardown re-releases the collection loop's pause events +# while waiting for it to exit. Short enough that the loop is not held for a +# meaningful slice of the quiesce budget, long enough not to spin. +_WAKE_RETRY_INTERVAL_S = 0.1 def _stamped_task_indices(batch: BatchedDataDict[DatumSpec]) -> list[int]: @@ -132,7 +143,23 @@ def __init__( ordinals_frontier_aligned: bool = True, resume_frontier_ordinal: Optional[int] = None, resume_covered_task_indices: Optional[list[int]] = None, + trace_carrier: Optional[dict[str, str]] = None, ) -> None: + # Every rollout in an async run is generated from this process, so + # without this the spans below are no-ops and an async trace has no + # rollout phase at all. rank/world_size are passed explicitly: this + # actor is a singleton rather than a member of a ranked group, and its + # runtime_env is a copy of the driver's environment, so a stray RANK + # there must not decide whether it exports. always_export goes with + # that synthetic rank -- an export_strategy picking among a group's + # ranks would otherwise mute this actor entirely (export_rank: 3 never + # matches rank 0), taking every rollout span with it. + _telemetry = init_telemetry_worker(rank=0, world_size=1, always_export=True) + self._tracer = _telemetry.tracer if _telemetry is not None else None + # The driver's rl.grpo.job span, so this actor's spans land in the run's + # trace rather than as loose roots. Reattached per thread below. + self._trace_carrier = trace_carrier or {} + self.policy_generation = policy_generation self.tokenizer = tokenizer self.task_to_env = task_to_env @@ -214,6 +241,12 @@ def __init__( # Track threads self._inflight_threads: set[_threading.Thread] = set() + # Same threads, different question, so a separate set. A batch worker + # drops out of _inflight_threads as soon as its batch is done -- from + # inside its own finally, which still runs inside the batch's span -- + # whereas telemetry teardown needs to know which threads could still be + # writing spans. Entries here survive until the thread itself is dead. + self._live_threads: set[_threading.Thread] = set() self._threads_lock: _threading.Lock = _threading.Lock() # Simple lock to prevent race conditions when checking/spawning workers @@ -409,7 +442,9 @@ def start_collection( print("Started continuous trajectory collection") - self.collection_thread = _threading.Thread(target=self._collection_loop) + self.collection_thread = _threading.Thread( + target=self._collection_loop, name="collection-loop" + ) self.collection_thread.daemon = True self.collection_thread.start() @@ -449,12 +484,22 @@ def _collection_loop(self): never discarded. The dataloader counts as exhausted only when the iterator drains with no pending prompts remaining. """ - dataloader_exhausted = False if self.dataloader is None: raise RuntimeError( "start_collection must set a dataloader before collection" ) - dataloader_iter = iter(self.dataloader) + # Attached for the life of the loop so the waits below hang off the + # driver's job span. This thread is not the one that captured the + # carrier, so it starts with an empty OTel context. + with remote_trace_context(self._trace_carrier): + self._run_collection_loop(self.dataloader) + + def _run_collection_loop( + self, dataloader: StatefulDataLoader | CyclingDataLoader + ) -> None: + """Body of :meth:`_collection_loop`, inside the driver's trace context.""" + dataloader_exhausted = False + dataloader_iter = iter(dataloader) try: while self.running: # Check if manually paused and wait @@ -464,7 +509,10 @@ def _collection_loop(self): # Check if refit is in progress and wait if not self._refit_pause_cleared.is_set() and self.running: print("⏸️ Pausing collection for refit...") - with self._efficiency_timer.time("idle/refit_event_wait"): + with ( + efficiency_span("idle/refit_event_wait", tracer=self._tracer), + self._efficiency_timer.time("idle/refit_event_wait"), + ): self._refit_pause_cleared.wait() print("▶️ Refit completed, resuming collection") @@ -484,7 +532,12 @@ def _collection_loop(self): self._last_limit_warning_version = self.current_weight_version # Efficiently wait for generation limits to be cleared (no polling!) - with self._efficiency_timer.time("idle/generation_limit_pause"): + with ( + efficiency_span( + "idle/generation_limit_pause", tracer=self._tracer + ), + self._efficiency_timer.time("idle/generation_limit_pause"), + ): self._generation_limit_cleared.wait() # Double-check we're still running after being woken up @@ -828,9 +881,19 @@ def _process_batch( "⏸️ Waiting for refit to complete before starting new " f"generation ({active_threads} threads still active)" ) - with self._efficiency_timer.time("idle/refit_event_wait"): + with ( + efficiency_span("idle/refit_event_wait", tracer=self._tracer), + self._efficiency_timer.time("idle/refit_event_wait"), + ): self._refit_pause_cleared.wait() generation_weight_version = self.current_weight_version + # Teardown clears `running` and then sets this event to wake the + # loop. Without the re-check we would spawn one more batch + # worker, whose spans open against a provider that is on its way + # out. + if not self.running: + self._release_target(reserved_target) + return # Task indices are stamped at yield time in _collection_loop, so # slices and carried-over remainders keep their original ordinals. @@ -857,21 +920,55 @@ def _process_batch( ) def _run_rollout_batch() -> None: - asyncio.run( - self._run_rollout_batch_worker( - repeated_batch=repeated_batch, - generation_weight_version=generation_weight_version, - target_weight_version=reserved_target, - num_generations=num_generations, - use_nemo_gym=use_nemo_gym, - dispatched_task_indices=dispatched_task_indices, + # Reattached here too: this is a fresh thread, which inherits no + # contextvars from the loop thread that spawned it. + with remote_trace_context(self._trace_carrier): + _collect() + + def _collect() -> None: + # The async counterpart of the driver's rl.grpo.generation. + # ROLLOUT is an umbrella group, so this carries no rl.bucket -- + # several batch workers run concurrently, so their durations sum + # past wall time and cannot go into a bucket rollup. It is here + # for the trace: how long a batch took, and at which weight + # version. Deliberately one span per batch, not per sample: + # generate_async is dispatched one coroutine per sample, which + # would be thousands of overlapping spans per step. + with managed_span( + RLSpanGroup.ROLLOUT, + "rl.grpo.generation", + tracer=self._tracer, + **{ + "rl.weight_version": generation_weight_version, + "rl.target_weight_version": reserved_target, + "rl.num_generations_per_prompt": num_generations, + # Batch width, so the duration can be read per rollout. + # A gap-filling batch is a fraction of a full one, and + # without this its span looks like an unexplained + # speed-up. + "rl.num_prompt_groups": num_prompts_to_generate, + }, + ): + asyncio.run( + self._run_rollout_batch_worker( + repeated_batch=repeated_batch, + generation_weight_version=generation_weight_version, + target_weight_version=reserved_target, + num_generations=num_generations, + use_nemo_gym=use_nemo_gym, + dispatched_task_indices=dispatched_task_indices, + ) ) - ) - worker = _threading.Thread(target=_run_rollout_batch, daemon=True) + worker = _threading.Thread( + target=_run_rollout_batch, + daemon=True, + name=f"rollout-batch-target-{reserved_target}", + ) try: with self._threads_lock: self._inflight_threads.add(worker) + self._live_threads.add(worker) if dispatched_task_indices: with self._outstanding_lock: self._outstanding_task_indices.update(dispatched_task_indices) @@ -880,6 +977,7 @@ def _run_rollout_batch() -> None: except Exception: with self._threads_lock: self._inflight_threads.discard(worker) + self._live_threads.discard(worker) if dispatched_task_indices: with self._outstanding_lock: self._outstanding_task_indices.difference_update( @@ -1086,6 +1184,130 @@ def get_efficiency_metrics(self) -> dict[str, float]: self._efficiency_timer.get_timing_metrics(reduction_op="sum"), ) + def flush_telemetry(self, quiesce_timeout_s: float = 5.0) -> None: + """Stop collecting, then export whatever spans are still buffered here. + + The driver reaps this actor with ``ray.kill``, which runs no atexit + handler, so the span processor's pending batch would otherwise be + dropped -- including the last rollout batches of the run. Call this + before the kill. + + Quiesces first because the shutdown is terminal, not a checkpoint: once + the provider is gone, a still-running loop thread or batch worker keeps + opening spans against a dead processor, which drops them and logs a line + per span. Clearing ``running`` and waiting for the in-flight batches is + what makes the flush cover the batches it exists to save. + + Joins ``_live_threads`` rather than ``_inflight_threads``, and joins + rather than polling ``is_alive``: a batch worker leaves the latter from + inside its own ``finally``, which still sits inside the + ``rl.grpo.generation`` span, so neither an empty set nor a set snapshot + means the spans are closed. Thread death does. + + The loop is woken before it is joined, since the refit and + generation-limit waits each hold an open span and each is followed by a + ``running`` re-check, so releasing them makes the loop exit rather than + pick up another batch. It is joined before the batch workers because + while it is alive it can still spawn one. + + The loop gets at most half the budget. The two joins are sequential, so + a loop wedged in a slow batch dispatch would otherwise spend all of it + and leave the batch workers -- whose spans are the ones this exists to + save -- with nothing. + + Bounded rather than reusing :meth:`wait_for_pending_generations`, which + waits indefinitely: that is right mid-run before a refit, but here the + caller is on its way to ``ray.kill``, so a wedged rollout must not be + able to hold the run's teardown open. The caller's RPC timeout has to + leave room for this budget *plus* the export that follows it. + + Args: + quiesce_timeout_s: Total time to wait for the collection loop and + the in-flight batch workers before flushing anyway. + """ + self.running = False + + deadline = time.monotonic() + quiesce_timeout_s + loop = self.collection_thread + if loop is not None: + self._drain_thread( + loop, min(deadline, time.monotonic() + quiesce_timeout_s / 2) + ) + else: + self._wake_waits() + + # Re-read the set each pass rather than snapshotting once: a worker + # spawned just before the loop exited would not be in a snapshot taken + # here. + while time.monotonic() < deadline: + with self._threads_lock: + alive = [t for t in self._live_threads if self._may_still_run(t)] + if not alive: + break + for thread in alive: + self._join_until(thread, deadline) + + with self._threads_lock: + still_running = [ + t.name for t in self._live_threads if self._may_still_run(t) + ] + if loop is not None and self._may_still_run(loop): + still_running.append(loop.name) + if still_running: + print( + f"⚠️ Flushing telemetry with {len(still_running)} thread(s) still " + f"running after {quiesce_timeout_s}s ({', '.join(still_running)}); " + "their spans may be dropped.", + flush=True, + ) + shutdown_telemetry() + + def _wake_waits(self) -> None: + """Release every event the collection loop can be parked on.""" + self._manual_pause_cleared.set() + self._refit_pause_cleared.set() + self._generation_limit_cleared.set() + + def _drain_thread(self, thread: _threading.Thread, deadline: float) -> None: + """Join *thread*, re-arming the pause events until it exits or time is up. + + Re-armed every pass rather than set once, because the loop checks + ``running`` and *then* clears an event: a single set landing in that + window is swallowed by the clear, and the wait that follows it has + nothing left to release it -- the driver is on its way to ``ray.kill``, + so no refit or weight update is coming. + """ + while self._may_still_run(thread): + if time.monotonic() >= deadline: + return + self._wake_waits() + self._join_until( + thread, min(deadline, time.monotonic() + _WAKE_RETRY_INTERVAL_S) + ) + + @staticmethod + def _may_still_run(thread: _threading.Thread) -> bool: + """Whether *thread* could still be opening spans. + + ``is_alive()`` alone is false for a worker that has been registered but + has not reached ``start()`` yet, which would drop it from both the join + set and the warning. + """ + return thread.is_alive() or thread.ident is None + + @staticmethod + def _join_until(thread: _threading.Thread, deadline: float) -> None: + """Join *thread* with whatever is left of the budget, if anything.""" + remaining = deadline - time.monotonic() + if remaining <= 0: + return + try: + thread.join(timeout=remaining) + except RuntimeError: + # Registered before start(), so not joinable yet. Yield rather than + # spin; the caller's drain loop comes back to it. + time.sleep(0.01) + async def drain_payload_metrics(self) -> dict[str, int | float]: """Close one drain-to-drain collector/Gym telemetry interval. @@ -1128,6 +1350,9 @@ def _cleanup_finished_threads(self) -> None: finished = {t for t in self._inflight_threads if not t.is_alive()} for t in finished: self._inflight_threads.remove(t) + # Pruned here (and discarded directly if a thread never started), so + # the set does not grow for the life of the run. + self._live_threads = {t for t in self._live_threads if t.is_alive()} def _release_target(self, target_weight_version: int) -> None: """Release the reservation owned by a completed batch worker.""" diff --git a/nemo_rl/algorithms/distillation.py b/nemo_rl/algorithms/distillation.py index 91e0034657..81e34e4d6a 100644 --- a/nemo_rl/algorithms/distillation.py +++ b/nemo_rl/algorithms/distillation.py @@ -73,6 +73,10 @@ from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface from nemo_rl.models.policy.lm_policy import Policy +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.instrumentation import managed_span, trace_fn +from nemo_rl.telemetry.setup import get_telemetry_handle +from nemo_rl.telemetry.span_groups import RLSpanGroup from nemo_rl.utils.checkpoint import CheckpointingConfig, CheckpointManager from nemo_rl.utils.logger import ( Logger, @@ -160,6 +164,7 @@ class MasterConfig(BaseModel, extra="allow"): logger: LoggerConfig # Logger configuration cluster: ClusterConfig # Cluster configuration checkpointing: CheckpointingConfig # Checkpointing configuration + telemetry: Optional[TelemetryConfig] = None # =============================================================================== @@ -674,6 +679,7 @@ def init_nemo_gym(): # =============================================================================== +@trace_fn(RLSpanGroup.JOB, "rl.distillation.job") def distillation_train( student_policy: ColocatablePolicyInterface, teacher_policy: ColocatablePolicyInterface, @@ -691,6 +697,8 @@ def distillation_train( ) -> None: """Run Distillation training algorithm.""" timer = Timer() + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None timeout = TimeoutChecker( timeout=master_config.checkpointing["checkpoint_must_save_by"], fit_last_save_time=True, @@ -773,10 +781,25 @@ def distillation_train( maybe_gpu_profile_step(student_generation, total_steps + 1) val_metrics, validation_timings = None, None - with timer.time("total_step_time"): + with ( + timer.time("total_step_time"), + managed_span( + RLSpanGroup.STEP, + "rl.distillation.step", + tracer=_tracer, + **{"rl.iteration": total_steps + 1, "rl.epoch": current_epoch + 1}, + ), + ): # Prepare batch print("▶ Preparing batch...", flush=True) - with timer.time("data_processing"): + with ( + timer.time("data_processing"), + managed_span( + RLSpanGroup.DATA_PROCESSING, + "rl.distillation.data_processing", + tracer=_tracer, + ), + ): # Repeat batch items repeated_batch: BatchedDataDict[DatumSpec] = ( batch.repeat_interleave( @@ -801,7 +824,14 @@ def distillation_train( else: student_generation.prepare_for_generation() - with timer.time("generation"): + with ( + timer.time("generation"), + managed_span( + RLSpanGroup.ROLLOUT, + "rl.distillation.generation", + tracer=_tracer, + ), + ): # We cascade NeMo-Gym first since NeMo-Gym requires async rollouts. if use_nemo_gym: generation_config = master_config.policy["generation"] @@ -903,7 +933,14 @@ def distillation_train( teacher_policy.prepare_for_lp_inference() print("▶ Computing teacher logprobs...", flush=True) - with timer.time("teacher_logprob_inference"): + with ( + timer.time("teacher_logprob_inference"), + managed_span( + RLSpanGroup.LOGPROB, + "rl.distillation.teacher_logprob_inference", + tracer=_tracer, + ), + ): teacher_topk = teacher_policy.get_topk_logits( train_data, k=master_config.distillation.topk_logits_k, @@ -919,7 +956,15 @@ def distillation_train( POLICY_GENERATION_STALE = True print("▶ Training policy...", flush=True) - with timer.time("policy_training"): + with ( + timer.time("policy_training"), + managed_span( + RLSpanGroup.POLICY_UPDATE, + "rl.distillation.policy_training", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): train_results = student_policy.train( train_data, loss_fn, @@ -1043,7 +1088,14 @@ def distillation_train( metrics_source[metric_name], ) - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.distillation.checkpointing", + tracer=_tracer, + ), + ): print( f"Saving checkpoint for step {total_steps + 1}...", flush=True, @@ -1196,7 +1248,17 @@ def validate( use_nemo_gym = should_use_nemo_gym(master_config) timer = Timer() - with timer.time("total_validation_time"): + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None + with ( + timer.time("total_validation_time"), + managed_span( + RLSpanGroup.EVALUATE, + "rl.distillation.evaluate", + tracer=_tracer, + **{"rl.step": step}, + ), + ): print(f"▶ Starting validation at step {step}...", flush=True) total_rewards = [] # Can be any metric. Setted to 'accuracy' by default. diff --git a/nemo_rl/algorithms/dpo.py b/nemo_rl/algorithms/dpo.py index de0a60af26..0bbdc3fa6d 100644 --- a/nemo_rl/algorithms/dpo.py +++ b/nemo_rl/algorithms/dpo.py @@ -38,6 +38,10 @@ from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.interfaces import PolicyInterface from nemo_rl.models.policy.lm_policy import Policy +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.instrumentation import managed_span, trace_fn +from nemo_rl.telemetry.setup import get_telemetry_handle +from nemo_rl.telemetry.span_groups import RLSpanGroup from nemo_rl.utils.checkpoint import CheckpointingConfig, CheckpointManager from nemo_rl.utils.logger import Logger, LoggerConfig from nemo_rl.utils.nsys import maybe_gpu_profile_step @@ -105,6 +109,7 @@ class MasterConfig(BaseModel, extra="allow"): logger: LoggerConfig cluster: ClusterConfig checkpointing: CheckpointingConfig + telemetry: Optional[TelemetryConfig] = None @dataclass @@ -418,8 +423,18 @@ def validate_one_dataset( return timer = Timer() - - with timer.time("total_validation_time"): + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None + + with ( + timer.time("total_validation_time"), + managed_span( + RLSpanGroup.EVALUATE, + "rl.dpo.evaluate", + tracer=_tracer, + **{"rl.step": step}, + ), + ): print(f"▶ Starting validation at step {step} for `{dataset_name}` set..") val_metrics = defaultdict(list) @@ -523,6 +538,7 @@ def validate_one_dataset( return val_metrics, timing_metrics +@trace_fn(RLSpanGroup.JOB, "rl.dpo.job") def dpo_train( policy, train_dataloader, @@ -536,6 +552,8 @@ def dpo_train( ) -> None: # Run dpo training timer = Timer() + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None timeout = TimeoutChecker( timeout=master_config.checkpointing["checkpoint_must_save_by"], fit_last_save_time=True, @@ -590,9 +608,25 @@ def dpo_train( maybe_gpu_profile_step(policy, total_steps + 1) val_metrics, validation_timings = None, None - with timer.time("total_step_time"): + with ( + timer.time("total_step_time"), + managed_span( + RLSpanGroup.STEP, + "rl.dpo.step", + tracer=_tracer, + **{"rl.iteration": total_steps + 1, "rl.epoch": current_epoch + 1}, + ), + ): print("▶ Taking a training step...") - with timer.time("policy_training"): + with ( + timer.time("policy_training"), + managed_span( + RLSpanGroup.POLICY_UPDATE, + "rl.dpo.policy_training", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): train_results = policy.train( batch, loss_fn, @@ -722,7 +756,14 @@ def dpo_train( metrics_source[metric_name], ) - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.dpo.checkpointing", + tracer=_tracer, + ), + ): print(f"Saving checkpoint for step {total_steps + 1}...") checkpoint_path = checkpointer.init_tmp_checkpoint( total_steps + 1, vars(dpo_save_state), master_config diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 2b9ab8d5ee..8d36ea3452 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -57,6 +57,7 @@ apply_reward_shaping, ) from nemo_rl.algorithms.utils import ( + WALL_CLOCK_EFFICIENCY_CATEGORIES, calculate_baseline_and_std_per_prompt, get_gdpo_reward_component_keys, log_generation_metrics, @@ -130,6 +131,17 @@ from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface from nemo_rl.models.policy.lm_policy import Policy +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.instrumentation import ( + Bucket, + bucket_scope, + current_trace_carrier, + efficiency_span, + managed_span, + trace_fn, +) +from nemo_rl.telemetry.setup import get_telemetry_handle +from nemo_rl.telemetry.span_groups import RLSpanGroup from nemo_rl.utils.checkpoint import CheckpointingConfig, CheckpointManager from nemo_rl.utils.logger import ( Logger, @@ -423,6 +435,7 @@ class MasterConfig(BaseModel, extra="allow"): reward_penalties: RewardPenaltyConfig = Field(default_factory=RewardPenaltyConfig) data_plane: Optional[DataPlaneConfig] = None on_policy_distillation: Optional[OnPolicyDistillationConfig] = None + telemetry: Optional[TelemetryConfig] = None # =============================================================================== @@ -2845,6 +2858,7 @@ def _validation_early_stop_message( ) +@trace_fn(RLSpanGroup.JOB, "rl.grpo.job") def grpo_train( policy: ColocatablePolicyInterface, policy_generation: Optional[GenerationInterface], @@ -2862,6 +2876,8 @@ def grpo_train( ) -> None: """Run GRPO training algorithm.""" timer = Timer(context={"worker": "driver"}) + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None timeout = TimeoutChecker( timeout=master_config.checkpointing["checkpoint_must_save_by"], fit_last_save_time=True, @@ -2995,10 +3011,25 @@ def grpo_train( maybe_gpu_profile_step(policy_generation, total_steps + 1) val_metrics, validation_timings = None, None - with timer.time("total_step_time"): + with ( + timer.time("total_step_time"), + managed_span( + RLSpanGroup.STEP, + "rl.grpo.step", + tracer=_tracer, + **{"rl.iteration": total_steps + 1, "rl.epoch": current_epoch + 1}, + ), + ): # Prepare batch print("▶ Preparing batch...", flush=True) - with timer.time("data_processing"): + with ( + timer.time("data_processing"), + managed_span( + RLSpanGroup.DATA_PROCESSING, + "rl.grpo.data_processing", + tracer=_tracer, + ), + ): if ( master_config.grpo.deduplicate_multimodal_data and should_use_nemo_gym(master_config) @@ -3092,7 +3123,17 @@ def grpo_train( policy_generation, "snapshot_step_metrics" ): policy_generation.snapshot_step_metrics() - with timer.time("generation"): + with ( + timer.time("generation"), + managed_span( + RLSpanGroup.ROLLOUT, + "rl.grpo.generation", + tracer=_tracer, + **{ + "rl.num_generations_per_prompt": master_config.grpo.num_generations_per_prompt, + }, + ), + ): # Clear logger metrics for each generation step if policy_generation is not None: policy_generation.clear_logger_metrics() @@ -3199,7 +3240,12 @@ def grpo_train( # Calculate rewards & advantages memory_tracker.snapshot_start_of_stage("Processing rewards", dir()) print("▶ Processing rewards...,", flush=True) - with timer.time("reward_calculation"): + with ( + timer.time("reward_calculation"), + managed_span( + RLSpanGroup.REWARD, "rl.grpo.reward_calculation", tracer=_tracer + ), + ): # Extract rewards from final_batch rewards = repeated_batch["total_reward"] @@ -3298,7 +3344,14 @@ def grpo_train( del baseline del std - with timer.time("data_processing"): + with ( + timer.time("data_processing"), + managed_span( + RLSpanGroup.DATA_PROCESSING, + "rl.grpo.data_processing", + tracer=_tracer, + ), + ): use_overlong_filtering = master_config.grpo.overlong_filtering if use_overlong_filtering: loss_multiplier = repeated_batch["loss_multiplier"].clone() @@ -3379,7 +3432,14 @@ def grpo_train( policy.prepare_for_lp_inference() print("▶ Computing logprobs...", flush=True) - with timer.time("policy_and_reference_logprobs"): + with ( + timer.time("policy_and_reference_logprobs"), + managed_span( + RLSpanGroup.LOGPROB, + "rl.grpo.policy_and_reference_logprobs", + tracer=_tracer, + ), + ): # Custom create this logprob_data so we avoid Ray comm overheads sending unused data to workers. logprob_data = BatchedDataDict[ClippedPGLossDataDict]( { @@ -3450,7 +3510,14 @@ def grpo_train( ] = seq_logprob_error_metrics.pop("num_masked_seqs") # Compute advantages with adv_estimator using correct mask and logprobs - with timer.time("advantage_calculation"): + with ( + timer.time("advantage_calculation"), + managed_span( + RLSpanGroup.ADVANTAGE, + "rl.grpo.advantage_calculation", + tracer=_tracer, + ), + ): print("▶ Computing advantages...", flush=True) # Get token-level mask: token_mask * sample_mask token_mask = train_data["token_mask"] @@ -3495,7 +3562,15 @@ def grpo_train( POLICY_GENERATION_STALE = True print("▶ Training policy...", flush=True) - with timer.time("policy_training"): + with ( + timer.time("policy_training"), + managed_span( + RLSpanGroup.POLICY_UPDATE, + "rl.grpo.policy_training", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): train_results = policy.train( train_data, loss_fn, @@ -3732,7 +3807,14 @@ def grpo_train( metrics_source[metric_name], ) - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.grpo.checkpointing", + tracer=_tracer, + ), + ): # Finalize the previous (possibly async) checkpoint before # starting a new one. No-op with sync save / nothing pending. checkpointer.finalize_pending() @@ -3990,7 +4072,26 @@ def validate( return {}, {} timer = Timer(context={"worker": "validator"}) - with timer.time("total_validation_time"): + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None + with ( + timer.time("total_validation_time"), + managed_span( + RLSpanGroup.EVALUATE, + "rl.grpo.evaluate", + tracer=_tracer, + **{"rl.step": step}, + ), + # Validation generates through the same path as training rollouts, but + # its tokens are scored and thrown away — no weights advance. Without + # this the generate spans below land in productive and a validation + # pass reads as goodput. Effective on the sync rollout path, which is + # where those spans exist; async validation goes through + # generate_async, which carries no span yet (see the coverage gaps in + # nemo_rl/telemetry/README.md). The scope is set regardless so it + # applies as soon as that path is instrumented. + bucket_scope(Bucket.OVERHEAD), + ): print(f"▶ Starting validation at step {step}...", flush=True) # >= 1 is validated in setup(). val_num_generations_per_prompt = ( @@ -4226,6 +4327,7 @@ def aggregate_rollout_metrics( return aggregated +@trace_fn(RLSpanGroup.JOB, "rl.grpo.job") def async_grpo_train( policy: ColocatablePolicyInterface, policy_generation: Optional[GenerationInterface], @@ -4307,6 +4409,8 @@ def async_grpo_train( from nemo_rl.algorithms.async_utils import AsyncTrajectoryCollector, ReplayBuffer timer = Timer(context={"worker": "driver"}) + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None training_wall_start = time.perf_counter() timeout = TimeoutChecker( timeout=master_config.checkpointing["checkpoint_must_save_by"], @@ -4499,9 +4603,16 @@ def async_grpo_train( **os.environ, "VIRTUAL_ENV": _tc_py_venv, "UV_PROJECT_ENVIRONMENT": _tc_py_venv, + # Names this actor's spans the way RayWorkerGroup names its groups'. + "NRL_WORKER_GROUP": "trajectory_collector", }, } + # Captured inside rl.grpo.job, so the collector's spans join this run's + # trace instead of starting their own roots. Empty unless the job group is + # enabled (per_step omits it) — see docs/observability/span-groups.md. + _tc_trace_carrier = current_trace_carrier() + # Initialize trajectory collector with synchronized collection trajectory_collector = AsyncTrajectoryCollector.options( runtime_env=_tc_runtime_env @@ -4516,9 +4627,29 @@ def async_grpo_train( alias_to_group_alias=alias_to_group_alias, on_policy_distillation_cfg=opd_module._opd_cfg(master_config), processor=processor, + trace_carrier=_tc_trace_carrier, **collector_start_kwargs, ) + def _flush_collector_telemetry() -> None: + """Export the collector's buffered spans before the actor is reaped. + + ``ray.kill`` runs no atexit handler, so whatever the span processor has + not sent yet goes with the actor -- including the last rollout batches + of the run. Every path that reaps the collector needs this, not only the + normal one, and collection is already running by the time this is + defined. The timeout covers the callee's quiesce budget *plus* its 5s + export; too short and it gives up mid-export, dropping the very spans + it exists to save. + """ + try: + ray.get( + trajectory_collector.flush_telemetry.remote(quiesce_timeout_s=3.0), + timeout=15, + ) + except Exception as e: + print(f"Error flushing trajectory collector telemetry: {e}") + print( f"🚀 Starting async GRPO training with buffer_size={optimal_buffer_size}, " f"max_age={max_trajectory_age_steps} steps, " @@ -4541,6 +4672,7 @@ def async_grpo_train( import traceback traceback.print_exc() + _flush_collector_telemetry() return else: print("🔄 Preparing policy generation for inference...") @@ -4552,6 +4684,7 @@ def async_grpo_train( import traceback traceback.print_exc() + _flush_collector_telemetry() return # Generation must hold the policy's real weights before any backend starts @@ -4622,6 +4755,7 @@ def async_grpo_train( # generation; the remaining actors are reaped when the driver # exits right after this return. checkpointer.shutdown() + _flush_collector_telemetry() try: ray.kill(trajectory_collector) except Exception as e: @@ -4714,7 +4848,9 @@ def async_grpo_train( wait_iterations += 1 time.sleep(1.0) - timer.stop("init/total") + # Retained because the per-step timer.reset() below discards it; the + # efficiency snapshot re-supplies it every step. + init_total_s = timer.stop("init/total") print(f"✅ Buffer ready for step {step}! Starting training loop...") ft_save_period = master_config.checkpointing.get("ft_save_period") @@ -4732,7 +4868,15 @@ def async_grpo_train( if policy != policy_generation: maybe_gpu_profile_step(policy_generation, step + 1) - with timer.time("total_step_time"): + with ( + timer.time("total_step_time"), + managed_span( + RLSpanGroup.STEP, + "rl.grpo.step", + tracer=_tracer, + **{"rl.iteration": step + 1}, + ), + ): num_mask_sample_filtered = 0 # Sample trajectories from replay buffer @@ -4827,7 +4971,10 @@ def async_grpo_train( f"Increase data.train.max_num_epochs or use a larger dataset." ) - with timer.time("idle/buffer_starvation"): + with ( + timer.time("idle/buffer_starvation"), + efficiency_span("idle/buffer_starvation", tracer=_tracer), + ): time.sleep(0.5) continue @@ -4910,7 +5057,12 @@ def async_grpo_train( policy_generation.snapshot_step_metrics() print("▶ Processing rewards...") - with timer.time("reward_calculation"): + with ( + timer.time("reward_calculation"), + managed_span( + RLSpanGroup.REWARD, "rl.grpo.reward_calculation", tracer=_tracer + ), + ): # Must precede prompt extraction: it reuses the same message # dicts, so this also protects the prompt flatten below. backfill_missing_routed_experts(repeated_batch["message_log"]) @@ -4937,7 +5089,14 @@ def async_grpo_train( ) # Prepare training data (same as sync version) - with timer.time("data_processing"): + with ( + timer.time("data_processing"), + managed_span( + RLSpanGroup.DATA_PROCESSING, + "rl.grpo.data_processing", + tracer=_tracer, + ), + ): # Apply overlong filtering - mask out truncated sequences from loss computation with timer.time("overlong_filter"): use_overlong_filtering = master_config.grpo.overlong_filtering @@ -5010,7 +5169,14 @@ def async_grpo_train( policy.prepare_for_lp_inference() print("▶ Computing logprobs...", flush=True) - with timer.time("policy_and_reference_logprobs"): + with ( + timer.time("policy_and_reference_logprobs"), + managed_span( + RLSpanGroup.LOGPROB, + "rl.grpo.policy_and_reference_logprobs", + tracer=_tracer, + ), + ): if not skip_prev_logprobs: train_data["prev_logprobs"] = policy.get_logprobs( train_data, timer=timer @@ -5059,7 +5225,14 @@ def async_grpo_train( ) # Compute advantages with adv_estimator using correct mask and logprobs - with timer.time("advantage_calculation"): + with ( + timer.time("advantage_calculation"), + managed_span( + RLSpanGroup.ADVANTAGE, + "rl.grpo.advantage_calculation", + tracer=_tracer, + ), + ): print("▶ Computing advantages...", flush=True) # Get token-level mask: token_mask * sample_mask token_mask = train_data["token_mask"] @@ -5120,7 +5293,15 @@ def async_grpo_train( POLICY_GENERATION_STALE = True print("▶ Training policy...") - with timer.time("policy_training"): + with ( + timer.time("policy_training"), + managed_span( + RLSpanGroup.POLICY_UPDATE, + "rl.grpo.policy_training", + tracer=_tracer, + **{"rl.iteration": step + 1}, + ), + ): train_results = policy.train( train_data, loss_fn, @@ -5169,41 +5350,45 @@ def async_grpo_train( trajectory_collector.set_weight_version.remote(weight_version) ) else: - timer.start("idle/refit_bubble") - - # Measure pending-generation wait as exposed_generation time - print("🔄 Coordinating with trajectory collector before refit...") - with timer.time("exposed_generation"): - ray.get(trajectory_collector.prepare_for_refit.remote()) - - # Collect generation logger metrics for performance reporting - # inflight batch sizes and num pending samples are collected from each worker - # (colocated collects them before the engine sleeps for training). - if generation_logger_metrics is None: - generation_logger_metrics = ( - policy_generation.get_logger_metrics() - ) - - # Only the actual refit/weight transfer should be counted as weight_sync - print("🔄 Performing policy generation refit...") - with timer.time("weight_sync"): - refit_metrics = refit_policy_generation( - policy, - policy_generation, - colocated_inference, + # A context manager rather than start/stop, so the timer is + # stopped and the span closed even if the refit raises. + with ( + timer.time("idle/refit_bubble"), + efficiency_span("idle/refit_bubble", tracer=_tracer), + ): + # Measure pending-generation wait as exposed_generation time + print( + "🔄 Coordinating with trajectory collector before refit..." ) - POLICY_GENERATION_STALE = False + with timer.time("exposed_generation"): + ray.get(trajectory_collector.prepare_for_refit.remote()) + + # Collect generation logger metrics for performance reporting + # inflight batch sizes and num pending samples are collected from each worker + # (colocated collects them before the engine sleeps for training). + if generation_logger_metrics is None: + generation_logger_metrics = ( + policy_generation.get_logger_metrics() + ) - # Update weight version before resuming trajectory collection so that all trajectories are updated with the new correct weight version - weight_version += 1 - ray.get( - trajectory_collector.set_weight_version.remote( - weight_version + # Only the actual refit/weight transfer should be counted as weight_sync + print("🔄 Performing policy generation refit...") + with timer.time("weight_sync"): + refit_metrics = refit_policy_generation( + policy, + policy_generation, + colocated_inference, ) - ) - ray.get(trajectory_collector.resume_after_refit.remote()) + POLICY_GENERATION_STALE = False - timer.stop("idle/refit_bubble") + # Update weight version before resuming trajectory collection so that all trajectories are updated with the new correct weight version + weight_version += 1 + ray.get( + trajectory_collector.set_weight_version.remote( + weight_version + ) + ) + ray.get(trajectory_collector.resume_after_refit.remote()) # Clear logger metrics after each refit (weight sync), starting a new logging cycle if policy_generation is not None: @@ -5234,6 +5419,13 @@ def async_grpo_train( # Run validation if it's a validation step or last step with val_at_end if should_run_validation: + # Timer only, no efficiency_span: validate() accounts this + # window as overhead (see the bucket_scope in validate), so + # an idle-bucketed span over the same interval would both + # contradict that label and, on the sync rollout path where + # the generate spans below carry it, be double-counted by a + # rollup that sums durations by rl.bucket. The metric has no + # such hierarchy and stays correct. with timer.time("idle/validation"): # No-op on an already-running engine; # wakes the colocated engine when it stayed asleep for a save-bound step. @@ -5411,7 +5603,14 @@ def async_grpo_train( metrics_source[metric_name], ) - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.grpo.checkpointing", + tracer=_tracer, + ), + ): # Finalize the previous (possibly async) checkpoint before # starting a new one. No-op with sync save / nothing pending. checkpointer.finalize_pending() @@ -5603,21 +5802,28 @@ def async_grpo_train( ) driver_efficiency = { cat: timer.reduce(cat, "sum") - for cat in [ - "init/total", - "idle/buffer_starvation", - "idle/refit_bubble", - "idle/validation", - ] + for cat in WALL_CLOCK_EFFICIENCY_CATEGORIES if cat in timer._timers } + # init/total is measured once, before the loop, and the timer.reset() + # at the end of every step drops it -- so re-supply the captured + # value, or the series reports the real startup cost at step 1 and + # zero for the rest of the run. + driver_efficiency["init/total"] = init_total_s merged_efficiency = {**driver_efficiency} for cat, dur in collector_efficiency.items(): merged_efficiency[cat] = merged_efficiency.get(cat, 0.0) + dur total_wall_time = time.perf_counter() - training_wall_start efficiency_loggable = print_efficiency_summary( - merged_efficiency, total_wall_time, step + 1 + merged_efficiency, + total_wall_time, + step + 1, + # The driver's idle categories are per-step (timer.reset() + # below), so the efficiency ratio needs a per-step denominator; + # against the run's elapsed time it would climb toward 100% + # whatever the idle time did. + step_wall_time_s=total_time, ) if master_config.grpo.debug_payload_metrics and not should_run_validation: @@ -5675,6 +5881,7 @@ def async_grpo_train( print(f"Error finalizing pending checkpoint: {e}") print("🛑 Stopping trajectory collection...") + _flush_collector_telemetry() try: ray.kill(trajectory_collector) except Exception as e: diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 8778cd5540..14de4c6797 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -96,6 +96,15 @@ from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.models.value import Value, ValueConfig from nemo_rl.models.value.interfaces import ValueInterface +from nemo_rl.telemetry.config import TelemetryConfig +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 from nemo_rl.utils.checkpoint import ( CheckpointingConfig, CheckpointManager, @@ -275,6 +284,7 @@ class MasterConfig(BaseModel, extra="allow"): logger: PPOLoggerConfig cluster: ClusterConfig checkpointing: CheckpointingConfig + telemetry: Optional[TelemetryConfig] = None # =============================================================================== @@ -1197,6 +1207,7 @@ def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: # =============================================================================== +@trace_fn(RLSpanGroup.JOB, "rl.ppo.job") def ppo_train( policy: ColocatablePolicyInterface, policy_generation: Optional[GenerationInterface], @@ -1222,6 +1233,8 @@ def ppo_train( - Configurable policy training start epoch """ timer = Timer() + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None timeout = TimeoutChecker( timeout=master_config.checkpointing["checkpoint_must_save_by"], fit_last_save_time=True, @@ -1316,10 +1329,25 @@ def ppo_train( maybe_gpu_profile_step(policy_generation, total_steps + 1) val_metrics, validation_timings = None, None - with timer.time("total_step_time"): + with ( + timer.time("total_step_time"), + managed_span( + RLSpanGroup.STEP, + "rl.ppo.step", + tracer=_tracer, + **{"rl.iteration": total_steps + 1, "rl.epoch": current_epoch + 1}, + ), + ): # Prepare batch print("▶ Preparing batch...", flush=True) - with timer.time("data_processing"): + with ( + timer.time("data_processing"), + managed_span( + RLSpanGroup.DATA_PROCESSING, + "rl.ppo.data_processing", + tracer=_tracer, + ), + ): repeated_batch: BatchedDataDict[DatumSpec] = ( batch.repeat_interleave( master_config.ppo.num_generations_per_prompt @@ -1386,7 +1414,14 @@ def ppo_train( policy.offload_to_cpu() policy_generation.prepare_for_generation() - with timer.time("generation"): + with ( + timer.time("generation"), + managed_span( + RLSpanGroup.ROLLOUT, + "rl.ppo.generation", + tracer=_tracer, + ), + ): if policy_generation is not None: policy_generation.clear_logger_metrics() @@ -1458,7 +1493,14 @@ def ppo_train( # Process rewards and build training data memory_tracker.snapshot_start_of_stage("Processing rewards", dir()) print("▶ Processing rewards...", flush=True) - with timer.time("reward_calculation"): + with ( + timer.time("reward_calculation"), + managed_span( + RLSpanGroup.REWARD, + "rl.ppo.reward_calculation", + tracer=_tracer, + ), + ): rewards = repeated_batch["total_reward"] with timer.time("data_processing"): @@ -1532,7 +1574,14 @@ def ppo_train( policy.prepare_for_lp_inference() print("▶ Computing logprobs...", flush=True) - with timer.time("policy_and_reference_logprobs"): + with ( + timer.time("policy_and_reference_logprobs"), + managed_span( + RLSpanGroup.LOGPROB, + "rl.ppo.policy_and_reference_logprobs", + tracer=_tracer, + ), + ): logprob_data = BatchedDataDict[ClippedPGLossDataDict]( { "input_ids": train_data["input_ids"], @@ -1571,7 +1620,14 @@ def ppo_train( # Build prompt IDs for advantage estimation (groups responses from same prompt). # Use the token-length-based extractor so multi-turn prompts containing # assistant messages still resolve to the original prompt only. - with timer.time("advantage_calculation"): + with ( + timer.time("advantage_calculation"), + managed_span( + RLSpanGroup.ADVANTAGE, + "rl.ppo.advantage_calculation", + tracer=_tracer, + ), + ): print("▶ Computing advantages...", flush=True) initial_prompt_message_logs = extract_initial_prompt_messages( repeated_batch["message_log"], @@ -1617,7 +1673,15 @@ def ppo_train( with timer.time("value_training_prep"): value_model.prepare_for_training() - with timer.time("value_training"): + with ( + timer.time("value_training"), + managed_span( + RLSpanGroup.POLICY_UPDATE, + "rl.ppo.value_training", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): print("▶ Training value...", flush=True) value_results = value_model.train( train_data, @@ -1644,7 +1708,15 @@ def ppo_train( POLICY_GENERATION_STALE = True print("▶ Training policy...", flush=True) - with timer.time("policy_training"): + with ( + timer.time("policy_training"), + managed_span( + RLSpanGroup.POLICY_UPDATE, + "rl.ppo.policy_training", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): train_results = policy.train( train_data, loss_fn, @@ -1855,7 +1927,14 @@ def ppo_train( metric_name ] - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.ppo.checkpointing", + tracer=_tracer, + ), + ): print( f"Saving checkpoint for step {total_steps + 1}...", flush=True, @@ -3048,7 +3127,19 @@ def validate( return {}, {} timer = Timer() - with timer.time("total_validation_time"): + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None + with ( + timer.time("total_validation_time"), + managed_span( + RLSpanGroup.EVALUATE, + "rl.ppo.evaluate", + tracer=_tracer, + ), + # Scored-and-discarded generation: overhead, not goodput. See the same + # scope in nemo_rl/algorithms/grpo.py::validate. + bucket_scope(Bucket.OVERHEAD), + ): print(f"▶ Starting validation at step {step}...", flush=True) total_rewards = [] diff --git a/nemo_rl/algorithms/rm.py b/nemo_rl/algorithms/rm.py index 0cbe3fc02d..5b2e7b9216 100644 --- a/nemo_rl/algorithms/rm.py +++ b/nemo_rl/algorithms/rm.py @@ -39,6 +39,10 @@ from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.interfaces import PolicyInterface from nemo_rl.models.policy.lm_policy import Policy +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.instrumentation import managed_span, trace_fn +from nemo_rl.telemetry.setup import get_telemetry_handle +from nemo_rl.telemetry.span_groups import RLSpanGroup from nemo_rl.utils.checkpoint import CheckpointingConfig, CheckpointManager from nemo_rl.utils.logger import Logger, LoggerConfig from nemo_rl.utils.nsys import maybe_gpu_profile_step @@ -96,6 +100,7 @@ class MasterConfig(BaseModel, extra="allow"): logger: LoggerConfig cluster: ClusterConfig checkpointing: CheckpointingConfig + telemetry: Optional[TelemetryConfig] = None @dataclass @@ -360,8 +365,17 @@ def validate_one_dataset( return timer = Timer() - - with timer.time("total_validation_time"): + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None + + with ( + timer.time("total_validation_time"), + managed_span( + RLSpanGroup.EVALUATE, + "rl.rm.evaluate", + tracer=_tracer, + ), + ): print(f"▶ Starting validation at step {step} for `{dataset_name}` set..") # Show a progress indicator for validation @@ -466,6 +480,7 @@ def validate_one_dataset( return val_metrics, timing_metrics +@trace_fn(RLSpanGroup.JOB, "rl.rm.job") def rm_train( policy, train_dataloader, @@ -479,6 +494,8 @@ def rm_train( ): # Run basic rm training timer = Timer() + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None timeout = TimeoutChecker( timeout=master_config.checkpointing["checkpoint_must_save_by"], fit_last_save_time=True, @@ -529,7 +546,15 @@ def rm_train( maybe_gpu_profile_step(policy, total_steps + 1) val_metrics, validation_timings = None, None - with timer.time("total_step_time"): + with ( + timer.time("total_step_time"), + managed_span( + RLSpanGroup.STEP, + "rl.rm.step", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): # Prepare batch and generate responses print("▶ Taking a training step...") @@ -657,7 +682,14 @@ def rm_train( metrics_source[metric_name], ) - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.rm.checkpointing", + tracer=_tracer, + ), + ): print(f"Saving checkpoint for step {total_steps + 1}...") checkpoint_path = checkpointer.init_tmp_checkpoint( total_steps + 1, vars(rm_save_state), master_config diff --git a/nemo_rl/algorithms/sft.py b/nemo_rl/algorithms/sft.py index b6d96778ef..a05f22cd6c 100644 --- a/nemo_rl/algorithms/sft.py +++ b/nemo_rl/algorithms/sft.py @@ -41,6 +41,10 @@ from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.interfaces import PolicyInterface from nemo_rl.models.policy.lm_policy import Policy +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.instrumentation import managed_span, trace_fn +from nemo_rl.telemetry.setup import get_telemetry_handle +from nemo_rl.telemetry.span_groups import RLSpanGroup from nemo_rl.utils.checkpoint import CheckpointingConfig, CheckpointManager from nemo_rl.utils.logger import Logger, LoggerConfig from nemo_rl.utils.nsys import maybe_gpu_profile_step @@ -99,6 +103,7 @@ class MasterConfig(BaseModel, extra="allow"): logger: LoggerConfig cluster: ClusterConfig checkpointing: CheckpointingConfig + telemetry: Optional[TelemetryConfig] = None # ======================================================= @@ -277,8 +282,18 @@ def validate( return {}, {} timer = Timer() - - with timer.time("total_validation_time"): + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None + + with ( + timer.time("total_validation_time"), + managed_span( + RLSpanGroup.EVALUATE, + "rl.sft.evaluate", + tracer=_tracer, + **{"rl.step": step}, + ), + ): print(f"▶ Starting validation at step {step}...") # Show a progress indicator for validation @@ -376,6 +391,7 @@ def validate( return val_metrics, timing_metrics +@trace_fn(RLSpanGroup.JOB, "rl.sft.job") def sft_train( policy, train_dataloader, @@ -389,6 +405,8 @@ def sft_train( ) -> None: # Run basic sft training timer = Timer() + _telemetry = get_telemetry_handle() + _tracer = _telemetry.tracer if _telemetry is not None else None timeout = TimeoutChecker( timeout=master_config.checkpointing["checkpoint_must_save_by"], fit_last_save_time=True, @@ -441,10 +459,25 @@ def sft_train( maybe_gpu_profile_step(policy, total_steps + 1) val_metrics, validation_timings = None, None - with timer.time("total_step_time"): + with ( + timer.time("total_step_time"), + managed_span( + RLSpanGroup.STEP, + "rl.sft.step", + tracer=_tracer, + **{"rl.iteration": total_steps + 1, "rl.epoch": current_epoch + 1}, + ), + ): # Prepare batch and generate responses print("▶ Preparing batch...") - with timer.time("data_processing"): + with ( + timer.time("data_processing"), + managed_span( + RLSpanGroup.DATA_PROCESSING, + "rl.sft.data_processing", + tracer=_tracer, + ), + ): ## add loss mask based on role to every message add_loss_mask_to_message_log( batch["message_log"], @@ -473,7 +506,15 @@ def sft_train( ) print("▶ Taking a training step...") - with timer.time("policy_training"): + with ( + timer.time("policy_training"), + managed_span( + RLSpanGroup.POLICY_UPDATE, + "rl.sft.policy_training", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): train_results = policy.train( train_data, loss_fn, @@ -579,7 +620,14 @@ def sft_train( metrics_source[metric_name], ) - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.sft.checkpointing", + tracer=_tracer, + ), + ): print(f"Saving checkpoint for step {total_steps + 1}...") checkpoint_path = checkpointer.init_tmp_checkpoint( total_steps + 1, vars(sft_save_state), master_config diff --git a/nemo_rl/algorithms/utils.py b/nemo_rl/algorithms/utils.py index 19c9a195ae..f76133443e 100644 --- a/nemo_rl/algorithms/utils.py +++ b/nemo_rl/algorithms/utils.py @@ -967,28 +967,62 @@ def log_generation_metrics( WALL_CLOCK_EFFICIENCY_CATEGORIES + THREAD_ACCUMULATED_EFFICIENCY_CATEGORIES ) +# Wall-clock categories whose value covers the whole run rather than one step. +# The driver's Timer is reset every step, so its idle categories are per-step +# deltas -- but init/total is measured once before the loop and republished +# unchanged afterwards, so it cannot be compared against a single step's wall +# time. Mirrored by _RUN_WINDOW_WALL_CLOCK_CATEGORIES in +# nemo_rl/telemetry/metrics.py, which cannot import this module (torch); a test +# keeps the two in lockstep. +RUN_WINDOW_WALL_CLOCK_CATEGORIES = frozenset({"init/total"}) + +STEP_WINDOW_WALL_CLOCK_CATEGORIES = [ + category + for category in WALL_CLOCK_EFFICIENCY_CATEGORIES + if category not in RUN_WINDOW_WALL_CLOCK_CATEGORIES +] + def print_efficiency_summary( efficiency_metrics: dict[str, float], total_wall_time_s: float, step: int, + step_wall_time_s: Optional[float] = None, ) -> dict[str, float]: """Print a summary table of efficiency metrics and return loggable dict. - Wall-clock categories (driver-side idle) are used for the efficiency - percentage. Collector-side categories are summed across concurrent - threads and reported separately as thread-seconds so they are not - compared directly against a single wall-clock denominator. + Driver-side wall-clock categories drive the efficiency percentage. + Collector-side categories are summed across concurrent threads and reported + separately as thread-seconds so they are not compared directly against a + single wall-clock denominator. + + The efficiency percentage is a *per-step* ratio, because the numerator is: + the driver resets its Timer every step, so its idle categories are per-step + deltas. Dividing them by the cumulative ``total_wall_time_s`` would make the + run look monotonically more efficient the longer it ran, whatever the idle + time actually did. ``init/total`` is excluded from that numerator for the + same reason in reverse -- it is a run-long constant (see + :data:`RUN_WINDOW_WALL_CLOCK_CATEGORIES`) and would otherwise add the whole + startup cost to every single step's waste. Args: efficiency_metrics: Dict mapping category labels to total seconds spent. total_wall_time_s: Total wall-clock time in seconds since training began. + Used for the per-category ``% of Wall`` column. step: Current training step number. + step_wall_time_s: Wall-clock seconds in the step being reported, the + denominator of the efficiency percentage. Defaults to + ``total_wall_time_s`` for callers with no per-step measurement. Returns: Dict of metrics suitable for logging to WandB/TensorBoard, including per-category seconds and percentages, total waste, and efficiency_pct. """ + # Captured before the fallback below, because it decides which window the + # efficiency percentage actually covers. + pct_is_per_step = step_wall_time_s is not None + if step_wall_time_s is None: + step_wall_time_s = total_wall_time_s print(f"\n📊 Efficiency Summary (Step {step}):") print(f" {'Category':<35} {'Time (s)':>10} {'% of Wall':>10}") print(f" {'─' * 57}") @@ -1010,17 +1044,17 @@ def print_efficiency_summary( loggable[f"efficiency/{category}_s"] = duration wall_waste = sum( - efficiency_metrics.get(cat, 0.0) for cat in WALL_CLOCK_EFFICIENCY_CATEGORIES + efficiency_metrics.get(cat, 0.0) for cat in STEP_WINDOW_WALL_CLOCK_CATEGORIES ) - if total_wall_time_s > 0 and wall_waste > total_wall_time_s: - wall_waste = total_wall_time_s + if step_wall_time_s > 0 and wall_waste > step_wall_time_s: + wall_waste = step_wall_time_s - productive = max(0.0, total_wall_time_s - wall_waste) + productive = max(0.0, step_wall_time_s - wall_waste) efficiency_pct = ( - (productive / total_wall_time_s * 100) if total_wall_time_s > 0 else 100.0 + (productive / step_wall_time_s * 100) if step_wall_time_s > 0 else 100.0 ) efficiency_pct = min(100.0, max(0.0, efficiency_pct)) - waste_pct = (wall_waste / total_wall_time_s * 100) if total_wall_time_s > 0 else 0.0 + waste_pct = (wall_waste / step_wall_time_s * 100) if step_wall_time_s > 0 else 0.0 waste_pct = min(100.0, max(0.0, waste_pct)) print(f" {'─' * 57}") @@ -1028,14 +1062,26 @@ def print_efficiency_summary( f" {'Collector thread-seconds (info)':<35} " f"{thread_seconds_total:>10.2f} {'n/a':>10}" ) - print(f" {'Wall-clock waste':<35} {wall_waste:>10.2f} {waste_pct:>9.2f}%") - print(f" {'Productive time':<35} {productive:>10.2f} {efficiency_pct:>9.2f}%") - print(f" {'Efficiency':<35} {'':>10} {efficiency_pct:>9.2f}%") + # Labelled "this step" because the denominator is the step's wall time, not + # the run's -- the column header above it is a share of the run. + print( + f" {'Wall-clock waste (this step)':<35} {wall_waste:>10.2f} {waste_pct:>9.2f}%" + ) + print( + f" {'Productive time (this step)':<35} {productive:>10.2f} {efficiency_pct:>9.2f}%" + ) + print(f" {'Efficiency (this step)':<35} {'':>10} {efficiency_pct:>9.2f}%") loggable["efficiency/thread_seconds_total_s"] = thread_seconds_total loggable["efficiency/total_waste_s"] = wall_waste loggable["efficiency/productive_time_s"] = productive loggable["efficiency/efficiency_pct"] = efficiency_pct loggable["efficiency/total_wall_time_s"] = total_wall_time_s + # Carries the window of the ratio above to the OTel tee, which tags it. A + # caller that supplies no per-step denominator gets a run-cumulative ratio, + # and publishing that as per-step would state the opposite of what it is. + # A float, not the string it represents, because everything in this dict is + # also logged to WandB/TensorBoard as a scalar. + loggable["efficiency/efficiency_pct_is_per_step"] = float(pct_is_per_step) return loggable diff --git a/nemo_rl/distributed/worker_groups.py b/nemo_rl/distributed/worker_groups.py index 905854e0ea..7d88794f54 100644 --- a/nemo_rl/distributed/worker_groups.py +++ b/nemo_rl/distributed/worker_groups.py @@ -550,6 +550,10 @@ def _create_workers_from_bundle_indices( "NODE_RANK": str(pg_idx), "AVAILABLE_ADDR_LIST": str(available_addresses), "AVAILABLE_PORT_LIST": str(available_ports), + # RANK is group-local, so it alone cannot tell a policy + # worker from a generation worker. Observability consumers + # need the group to disambiguate them. + "NRL_WORKER_GROUP": self.name_prefix, } ) # Remove Ray-specific environment variables, let the worker itself set them. diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index bef0f00f20..f15c2bf081 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -47,6 +47,10 @@ compute_spec_decode_metrics, resolve_generation_worker_cls, ) +from nemo_rl.telemetry.instrumentation import trace_fn +from nemo_rl.telemetry.metrics import warn_once +from nemo_rl.telemetry.setup import get_telemetry_handle +from nemo_rl.telemetry.span_groups import RLSpanGroup from nemo_rl.utils.multimodal_payload_metrics import ( collect_multimodal_payload_metrics, collect_sharded_multimodal_payload_metrics, @@ -58,6 +62,40 @@ logger = logging.getLogger(__name__) +def _record_vllm_generation_metrics( + model_name: str | None, + data: BatchedDataDict, + combined: BatchedDataDict, +) -> None: + """Record vLLM token-usage metrics to nemo-lens (no-op unless exporting).""" + telemetry = get_telemetry_handle() + if telemetry is None or not telemetry.is_exporting: + return + from nemo.lens.instruments.inference import record_inference_metrics + + # Guards only the recording: this runs per generation call, so it must not + # break generation, but a permanently dead metric should still be visible + # once at default verbosity rather than only under debug. + try: + input_tokens = ( + int(data["input_lengths"].sum()) if "input_lengths" in data else None + ) + output_tokens = ( + int(combined["generation_lengths"].sum()) + if "generation_lengths" in combined + else None + ) + record_inference_metrics( + telemetry.meter, + model=model_name or "", + input_tokens=input_tokens, + output_tokens=output_tokens, + provider_name="vllm", + ) + except Exception: + warn_once("vllm_inference_metrics", "nemo-lens: failed to record vLLM metrics") + + class VllmGeneration(GenerationInterface): @staticmethod def init_cluster_placement_groups( @@ -720,6 +758,7 @@ def rebuild_collective( ) return futures + @trace_fn(RLSpanGroup.GENERATION, "rl.vllm.generate") def generate( self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False ) -> BatchedDataDict[GenerationOutputSpec]: @@ -773,8 +812,10 @@ def generate( f"Missing required keys for GenerationOutputSpec: {missing_keys}" ) + _record_vllm_generation_metrics(self.cfg.get("model_name"), data, combined) return combined + @trace_fn(RLSpanGroup.GENERATION, "rl.vllm.generate_text") def generate_text( self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False ) -> BatchedDataDict[GenerationOutputSpec]: @@ -826,6 +867,7 @@ def generate_text( f"Missing required keys for GenerationOutputSpec: {missing_keys}" ) + _record_vllm_generation_metrics(self.cfg.get("model_name"), data, combined) return combined async def _async_generate_base( diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index 1a0b9a6e80..ba9ad5fb83 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -14,6 +14,7 @@ import copy import gc +import inspect import logging import os import sys @@ -60,6 +61,14 @@ ) from nemo_rl.models.huggingface.common import ModelFlag from nemo_rl.models.policy.utils import is_vllm_v1_engine_enabled +from nemo_rl.telemetry.instrumentation import trace_fn +from nemo_rl.telemetry.setup import ( + init_telemetry_worker, + shutdown_telemetry, + telemetry_enabled_in_env, + vllm_native_tracing_requested, +) +from nemo_rl.telemetry.span_groups import RLSpanGroup from nemo_rl.utils.nsys import wrap_with_nvtx_name from nemo_rl.utils.nvml import log_gpu_memory_diagnostics from nemo_rl.weight_sync.checkpoint_engine_config import ( @@ -82,6 +91,59 @@ def _context_capped_max_new_tokens( return min(configured_max_new_tokens, remaining_context) +def _maybe_enable_vllm_native_tracing(llm_kwargs: dict[str, Any]) -> None: + """Optionally enable vLLM's native OpenTelemetry tracing on the engine. + + Requires both ``telemetry.enabled`` and ``telemetry.vllm_native_tracing`` + (plus an OTLP endpoint). vLLM's OTLP span exporter is gRPC-only, so the + endpoint must speak OTLP/gRPC (e.g. a collector on ``:4317`` or a + gRPC-capable backend) — it will not reach an ``http/protobuf`` OTLP + endpoint. Degrades to a no-op if the installed vLLM lacks these engine args. + """ + # The master switch is re-checked here because _config_to_env() exports + # every field before init_telemetry_driver's `enabled` check, so + # vllm_native_tracing would otherwise survive enabled=false and turn on + # per-request tracing on exactly the runs that disabled telemetry. + if not (telemetry_enabled_in_env() and vllm_native_tracing_requested()): + return + endpoint = ( + os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", "").strip() + or os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "").strip() + ) + if not endpoint: + logger.warning( + "nemo-lens: NEMO_RL_OTEL_VLLM_NATIVE_TRACING is set but no OTLP " + "endpoint is configured; skipping native vLLM tracing." + ) + return + # Narrow: this introspects a third-party surface that moves across vLLM + # versions, which justifies tolerating a missing attribute or a moved + # module -- but not swallowing every failure inside a vLLM worker. + try: + from vllm.engine.arg_utils import EngineArgs + + supported = set(getattr(EngineArgs, "__dataclass_fields__", {})) | set( + inspect.signature(EngineArgs.__init__).parameters + ) + except (ImportError, AttributeError, ValueError, TypeError): + logger.warning( + "nemo-lens: could not introspect vLLM EngineArgs; skipping native " + "vLLM tracing.", + exc_info=True, + ) + return + if "otlp_traces_endpoint" not in supported: + logger.warning( + "nemo-lens: installed vLLM does not support 'otlp_traces_endpoint'; " + "skipping native vLLM tracing." + ) + return + llm_kwargs.setdefault("otlp_traces_endpoint", endpoint) + if "collect_detailed_traces" in supported: + llm_kwargs.setdefault("collect_detailed_traces", ["all"]) + logger.info("nemo-lens: enabled vLLM native OTLP tracing -> %s", endpoint) + + def _resolve_enable_prefix_caching(vllm_cfg: dict[str, Any]) -> bool: enable_prefix_caching = vllm_cfg.get("enable_prefix_caching", None) if enable_prefix_caching is None: @@ -324,6 +386,10 @@ def __init__( if bundle_indices is not None and len(bundle_indices) == 1: bind_to_gpu_numa(int(ray.get_gpu_ids()[0])) + # OTel providers are process-global, so the driver's setup does not + # reach this actor. No-op unless telemetry is enabled. + init_telemetry_worker() + self._init_config( config, bundle_indices, fraction_of_gpus, seed, extra_env_vars ) @@ -389,6 +455,7 @@ def _init_config( self.rank = 0 self.world_size = 1 + @trace_fn(RLSpanGroup.MODEL_INIT, "rl.vllm.load_model") def _load_model(self, bundle_indices, seed): """Perform the heavy model loading and engine creation. @@ -650,6 +717,9 @@ def _load_model(self, bundle_indices, seed): sampling_style=video_config.sampling_style, temporal_patch_size=video_config.temporal_patch_size, ) + + _maybe_enable_vllm_native_tracing(llm_kwargs) + self._create_engine(llm_kwargs) log_gpu_memory_diagnostics( label="after_engine_create", worker_type="VllmGenerationWorker", device_id=0 @@ -1428,6 +1498,9 @@ def shutdown(self) -> bool: except Exception as e: print(f"Error during vLLM shutdown: {e}") return False + finally: + # Flush buffered spans/metrics before the actor goes away. + shutdown_telemetry() @ray.remote( diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index bc5c978898..01076d1948 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -54,6 +54,7 @@ from nemo_rl.models.generation.openai_server_utils import ( replace_prefix_tokens, ) +from nemo_rl.telemetry.setup import shutdown_telemetry LOGGER = logging.getLogger(__name__) @@ -1711,6 +1712,13 @@ async def shutdown(self) -> bool: except Exception as e: print(f"Error during vLLM shutdown: {e}") return False + finally: + # Flush buffered spans/metrics before the actor goes away. Off the + # event loop: the flush blocks on a network export with a 5s + # timeout, and this is an async actor whose other coroutines -- + # including in-flight generate requests -- share this loop. Same + # reason the sparse-refit shutdown above is offloaded. + await asyncio.to_thread(shutdown_telemetry) @ray.remote( diff --git a/nemo_rl/models/policy/workers/base_policy_worker.py b/nemo_rl/models/policy/workers/base_policy_worker.py index 9dea9ecd9a..693cfe84e0 100644 --- a/nemo_rl/models/policy/workers/base_policy_worker.py +++ b/nemo_rl/models/policy/workers/base_policy_worker.py @@ -19,6 +19,7 @@ from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.policy.interfaces import ReferenceLogprobOutputSpec +from nemo_rl.telemetry.setup import shutdown_telemetry from nemo_rl.utils.nsys import wrap_with_nvtx_name @@ -292,6 +293,12 @@ def shutdown(self) -> bool: return True except Exception: return False + finally: + # Flush buffered spans/metrics before the actor goes away. Only the + # async GRPO trainer calls shutdown() today; elsewhere Ray reaps the + # actor and this never runs, so worker telemetry relies on the batch + # processor's periodic export. + shutdown_telemetry() def start_gpu_profiling(self) -> None: """Start GPU profiling.""" diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker.py b/nemo_rl/models/policy/workers/dtensor_policy_worker.py index a71a4e415d..d460b2c650 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker.py @@ -93,6 +93,7 @@ PolicyCheckpointEngineMixin, maybe_preinit_nixl_checkpoint_engine, ) +from nemo_rl.telemetry.setup import init_telemetry_worker from nemo_rl.utils.grad_norm import warn_if_inf_grad_norm from nemo_rl.utils.native_checkpoint import ( load_checkpoint, @@ -229,6 +230,10 @@ def __init__( # affinity file, and reading it does not initialize CUDA. bind_to_gpu_numa(int(ray.get_gpu_ids()[0])) + # OTel providers are process-global, so the driver's setup does not + # reach this actor. No-op unless telemetry is enabled. + init_telemetry_worker() + self.tokenizer = tokenizer self.processor = processor self.is_vlm = processor is not None diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py index 135b6cc4f3..058d620898 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py @@ -80,6 +80,7 @@ from nemo_rl.models.policy.workers.patches import ( apply_transformer_engine_patch, ) +from nemo_rl.telemetry.setup import init_telemetry_worker from nemo_rl.utils.checkpoint import CheckpointingConfig from nemo_rl.utils.grad_norm import warn_if_inf_grad_norm from nemo_rl.utils.nsys import wrap_with_nvtx_name @@ -263,6 +264,10 @@ def __init__( # file, and reading it does not initialize CUDA. bind_to_gpu_numa(int(ray.get_gpu_ids()[0])) + # OTel providers are process-global, so the driver's setup does not + # reach this actor. No-op unless telemetry is enabled. + init_telemetry_worker() + # Store configuration self.cfg = config diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 99035d6664..904a1359b5 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -117,6 +117,7 @@ maybe_preinit_nixl_checkpoint_engine, ) from nemo_rl.models.policy.workers.patches import apply_transformer_engine_patch +from nemo_rl.telemetry.setup import init_telemetry_worker from nemo_rl.utils.grad_norm import warn_if_inf_grad_norm from nemo_rl.utils.nsys import wrap_with_nvtx_name from nemo_rl.utils.nvml import log_gpu_memory_diagnostics @@ -460,6 +461,10 @@ def __init__( # set by configure_worker), so it can't identify this worker's GPU. bind_to_gpu_numa(local_rank) + # OTel providers are process-global, so the driver's setup does not + # reach this actor. No-op unless telemetry is enabled. + init_telemetry_worker() + self.cfg = config self._router_replay_enabled = router_replay_enabled(config) self._nixl_preinit_agent = maybe_preinit_nixl_checkpoint_engine(config) diff --git a/nemo_rl/models/value/workers/dtensor_value_worker_v2.py b/nemo_rl/models/value/workers/dtensor_value_worker_v2.py index 0ad1b4fedf..f68e29b951 100644 --- a/nemo_rl/models/value/workers/dtensor_value_worker_v2.py +++ b/nemo_rl/models/value/workers/dtensor_value_worker_v2.py @@ -57,6 +57,7 @@ from nemo_rl.models.policy.workers.patches import apply_transformer_engine_patch from nemo_rl.models.value.config import ValueConfig from nemo_rl.models.value.interfaces import ValueOutputSpec +from nemo_rl.telemetry.setup import init_telemetry_worker from nemo_rl.utils.checkpoint import CheckpointingConfig from nemo_rl.utils.nsys import wrap_with_nvtx_name @@ -156,6 +157,10 @@ def __init__( # file, and reading it does not initialize CUDA. bind_to_gpu_numa(int(ray.get_gpu_ids()[0])) + # OTel providers are process-global, so the driver's setup does not + # reach this actor. No-op unless telemetry is enabled. + init_telemetry_worker() + # Store configuration and tokenizer self.cfg = config self.tokenizer = tokenizer diff --git a/nemo_rl/models/value/workers/megatron_value_worker.py b/nemo_rl/models/value/workers/megatron_value_worker.py index f3063888ca..9c61c422ab 100644 --- a/nemo_rl/models/value/workers/megatron_value_worker.py +++ b/nemo_rl/models/value/workers/megatron_value_worker.py @@ -84,6 +84,7 @@ from nemo_rl.models.policy.workers.patches import apply_transformer_engine_patch from nemo_rl.models.value.config import ValueConfig from nemo_rl.models.value.interfaces import ValueOutputSpec +from nemo_rl.telemetry.setup import init_telemetry_worker from nemo_rl.utils.nsys import wrap_with_nvtx_name TokenizerType = TypeVar("TokenizerType", bound=PreTrainedTokenizerBase) @@ -353,6 +354,10 @@ def __init__( # (RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1). bind_to_gpu_numa(local_rank) + # OTel providers are process-global, so the driver's setup does not + # reach this actor. No-op unless telemetry is enabled. + init_telemetry_worker() + self.cfg = config self.rank = get_rank_safe() diff --git a/nemo_rl/telemetry/README.md b/nemo_rl/telemetry/README.md new file mode 100644 index 0000000000..268ac19444 --- /dev/null +++ b/nemo_rl/telemetry/README.md @@ -0,0 +1,78 @@ +# NeMo-RL OpenTelemetry Instrumentation + +This module contains NeMo-RL's OpenTelemetry integration, built on top of [`nemo-lens`](https://github.com/NVIDIA-NeMo/Lens). + +It emits **traces** at RL-algorithm boundaries (rollout, generation, reward, advantage, policy update, checkpoint, evaluate) and **metrics** (`rl.*`: reward, loss, KL, grad norm, learning rate, throughput) that export to any OTLP-compatible backend. + +Telemetry is **optional**: it activates only when `enabled` is true *and* nemo-lens is installed. When either is absent, every instrumentation site degrades to a ~0-cost no-op. + +## Contents + +``` +nemo_rl/telemetry/ +├── config.py — TelemetryConfig: the telemetry: config block +├── setup.py — init_telemetry_driver / init_telemetry_worker / get_telemetry_handle / shutdown_telemetry +├── span_groups.py — RLSpanGroup: RL-specific span groups + presets +├── instrumentation.py — managed_span/trace_fn wrappers + phase/group → rl.bucket map (monitor derives goodput) +├── metrics.py — tees Logger.log_metrics scalars into the rl.* instruments +└── __init__.py +``` + +Metric instruments, resource detection, and the instrumentation primitives themselves live in `nemo-lens`. This module is a thin integration layer. + +## Wiring + +Each `examples/run_.py` calls `init_telemetry_driver(config, algorithm="")` **before** `init_ray()` (so `NEMO_RL_OTEL_*` is snapshotted into the Ray `runtime_env` and inherited by workers) and `shutdown_telemetry()` from a `finally` block wrapping the whole run, so buffered spans are flushed on the failure path too. `get_telemetry_handle()` returns the process-global `TelemetryHandle`. + +OTel providers are process-global, so each Ray actor sets up its own: the policy, value and vLLM generation workers call `init_telemetry_worker()` from `__init__` and `shutdown_telemetry()` from `shutdown` (the latter matters — span/metric processors buffer in the background, and an actor that exits without flushing drops whatever it had not exported). Worker ranks come from the `RANK` / `WORLD_SIZE` env vars, and `RayWorkerGroup` also exports `NRL_WORKER_GROUP` so a worker's spans carry `rl.worker_group` — `RANK` is group-local, so it alone cannot distinguish `lm_policy` rank 3 from `vllm_policy` rank 3. + +The async trajectory collector is a singleton actor rather than a group member, so it passes `rank=0, world_size=1, always_export=True`: an `export_strategy` selecting among a group's ranks cannot meaningfully be applied to a synthetic rank, and would otherwise mute every rollout span in the run. It flushes on demand via `flush_telemetry()` because the driver reaps it with `ray.kill`, which runs no `atexit` handler. + +Trace context does not cross the Ray call boundary on its own, so a worker's spans are roots of their own traces, correlated to the driver by `run_id` rather than parented to it. The one exception is the async trajectory collector: the driver hands it a W3C carrier at construction and it reattaches that context per thread, so its rollout spans nest under `rl.grpo.job`. See [span groups — getting the collector into one waterfall](../../docs/observability/span-groups.md#getting-the-collector-into-one-waterfall). + +Not yet wired: + +| Gap | Effect | +|---|---| +| SGLang, TRT-LLM and Megatron generation workers | no `init_telemetry_worker`, no generation spans — only vLLM is instrumented | +| `grpo_sync.py`, `single_controller.py` | no spans at all; `examples/run_grpo_single_controller.py` never calls `init_telemetry_driver`, so that entrypoint emits no telemetry | +| `run_vlm_grpo.py`, `run_grpo_sliding_puzzle.py`, `run_xtoken_off_policy_distillation.py`, `run_eval.py` | no `init_telemetry_driver`, so a `telemetry:` block in those configs parses and the run succeeds while emitting nothing, driver *and* worker | +| `VllmGeneration.generate_async` | no `rl.vllm.generate` span, so async rollouts and async validation show `rl.grpo.generation` / `rl.grpo.evaluate` with no generate breakdown inside | +| `SyncRolloutActor` | the sync data-plane counterpart of the collector; no `init_telemetry_worker`, no rollout spans | +| Worker `shutdown()` on non-async trainers | only `async_grpo_train` calls `policy.shutdown()` / `policy_generation.shutdown()`; elsewhere Ray reaps the actors, so the worker's final flush never runs and its telemetry depends on the periodic export | +| Trace context into ranked workers | policy/value/vLLM worker spans stay separate traces (see above) | + +## Install + +Nothing to install. `nemo-lens[sdk]` is a base dependency, which is what gets it into the *worker* interpreters: Ray actors run under the `PY_EXECUTABLES` entries in `nemo_rl/distributed/virtual_cluster.py` (`uv run --locked --extra vllm`, `--extra mcore`, ...), and those resolve the base dependencies plus one backend extra. + +## Quick start + +```yaml +# in your run config +telemetry: + enabled: true + span_groups: default +``` + +```bash +export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317 + +uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml +``` + +## Full documentation + +See `docs/observability/` in this repository: + +| Topic | Doc | +|---|---| +| Overview | [docs/observability/index.md](../../docs/observability/index.md) | +| Configuration (`telemetry:` block, env vars) | [docs/observability/configuration.md](../../docs/observability/configuration.md) | +| Span groups and per-algorithm span names | [docs/observability/span-groups.md](../../docs/observability/span-groups.md) | +| `rl.*` metrics and the Logger tee | [docs/observability/metrics.md](../../docs/observability/metrics.md) | +| vLLM tracing (driver spans + native OTLP) | [docs/observability/vllm-tracing.md](../../docs/observability/vllm-tracing.md) | +| Exporting to an OTLP backend | [docs/observability/observability-stack.md](../../docs/observability/observability-stack.md) | +| Adding new instrumentation | [docs/observability/extending.md](../../docs/observability/extending.md) | + +For the generic `nemo-lens` documentation (configuration model, instrumentation primitives, custom exporters, design decisions), see the lens docs at . diff --git a/nemo_rl/telemetry/__init__.py b/nemo_rl/telemetry/__init__.py new file mode 100644 index 0000000000..1750b29262 --- /dev/null +++ b/nemo_rl/telemetry/__init__.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""NeMo-RL telemetry: optional OpenTelemetry instrumentation via nemo-lens. + +Public surface: + +* :class:`~nemo_rl.telemetry.config.TelemetryConfig` — the ``telemetry:`` config + block. +* :class:`~nemo_rl.telemetry.span_groups.RLSpanGroup` — RL span-group presets. +* :func:`~nemo_rl.telemetry.setup.init_telemetry_driver` / + :func:`~nemo_rl.telemetry.setup.init_telemetry_worker` / + :func:`~nemo_rl.telemetry.setup.get_telemetry_handle` / + :func:`~nemo_rl.telemetry.setup.shutdown_telemetry` — lifecycle helpers. + +The instrumentation primitives (``managed_span`` / ``trace_fn`` / ``span_cm`` / +``is_span_group_enabled``) come from :mod:`nemo_rl.telemetry.instrumentation`, +which re-exports them from nemo-lens with ``rl.bucket`` efficiency tagging +applied. +""" + +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.instrumentation import Bucket, bucket_for_span_group +from nemo_rl.telemetry.setup import ( + get_telemetry_handle, + init_telemetry_driver, + init_telemetry_worker, + shutdown_telemetry, +) +from nemo_rl.telemetry.span_groups import RLSpanGroup + +__all__ = [ + "TelemetryConfig", + "RLSpanGroup", + "Bucket", + "bucket_for_span_group", + "get_telemetry_handle", + "init_telemetry_driver", + "init_telemetry_worker", + "shutdown_telemetry", +] diff --git a/nemo_rl/telemetry/config.py b/nemo_rl/telemetry/config.py new file mode 100644 index 0000000000..3820b07226 --- /dev/null +++ b/nemo_rl/telemetry/config.py @@ -0,0 +1,97 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Telemetry configuration schema for NeMo-RL. + +The ``telemetry:`` block of a run config. :mod:`nemo_rl.telemetry.setup` +translates it into ``NEMO_RL_OTEL_*`` environment variables on the driver +*before* ``init_ray()``, so every Ray worker inherits the same settings via the +Ray ``runtime_env``. Raw ``NEMO_RL_OTEL_*`` / ``OTEL_EXPORTER_OTLP_*`` env vars +always win over these YAML values (they are applied with ``setdefault``). + +This module imports only ``pydantic`` — it never requires nemo-lens, so it is +safe to import unconditionally from the algorithm ``MasterConfig`` classes. +""" + +from __future__ import annotations + +from typing import Annotated, Literal + +from pydantic import BaseModel, Field + + +class TelemetryConfig(BaseModel, extra="allow"): + """OpenTelemetry / nemo-lens configuration. + + Telemetry activates only when ``enabled`` is true; otherwise every + instrumentation site degrades to a ~0-cost no-op. + + Fields with a fixed set of valid values are typed so that a typo is + rejected when the YAML is parsed, in every process, rather than surfacing + later on a GPU node -- or not at all, when ``enabled`` is false and the + driver returns before it validates anything. + """ + + enabled: bool = False + """Master switch. When false, all instrumentation is a ~0-cost no-op.""" + + service_name: str = "nemo-rl" + """``service.name`` reported to the OTLP backend.""" + + span_groups: str = "default" + """Span-group spec: a preset (``default`` | ``per_step`` | ``all``) or a + comma-separated list of individual group names (e.g. + ``"default,generation,reward"``). See ``RLSpanGroup``.""" + + export_strategy: Literal[ + "single_rank", "all_ranks", "sampled", "first_rank_per_node" + ] = "single_rank" + """Which ranks export. The driver always exports (it runs the training loop + and the metrics logger); this governs the Ray worker ranks. nemo-lens owns + the strategy registry, so the driver re-checks this name against it.""" + + export_rank: Annotated[int, Field(ge=-1)] = -1 + """For ``single_rank``: which rank exports (``-1`` = last rank).""" + + export_sample_rate: Annotated[float, Field(ge=0.0, le=1.0)] = 1.0 + """For ``sampled``: fraction of worker ranks that export, in ``[0.0, 1.0]``. + Also the sampling rate used by the span sampler when ``sampler_enabled`` is + true. ``1.0`` means every rank considered by the strategy exports.""" + + sampler_enabled: bool = False + """Enable lens's rank-aware span sampler on the TracerProvider. It drops + spans at the SDK level — cheaper than exporting and filtering downstream — + and decides all-or-nothing per rank, from a hash of the rank against + ``export_sample_rate``. This is a *second*, independent filter: a rank has + to pass both it and ``export_strategy`` to emit anything. The driver and + singleton actors are exempt from both, having no real rank.""" + + traces_enabled: bool = True + """Emit trace spans.""" + + metrics_enabled: bool = True + """Emit metric instruments (the ``rl.*`` gauges/histograms).""" + + logs_enabled: bool = False + """Bridge Python logging to OTel logs (exported with trace correlation).""" + + exporter: Literal["otlp", "console"] = "otlp" + """Exporter backend. The OTLP endpoint / headers / protocol come from the + standard ``OTEL_EXPORTER_OTLP_*`` env vars, so any OTLP-compatible backend + or an OpenTelemetry Collector works.""" + + vllm_native_tracing: bool = False + """Enable vLLM's own OTLP tracing inside generation workers (opt-in). vLLM's + exporter is gRPC-only, so this needs a gRPC OTLP endpoint / collector — it + does not ride an ``http/protobuf`` OTLP endpoint used by lens.""" diff --git a/nemo_rl/telemetry/instrumentation.py b/nemo_rl/telemetry/instrumentation.py new file mode 100644 index 0000000000..3035e079ca --- /dev/null +++ b/nemo_rl/telemetry/instrumentation.py @@ -0,0 +1,354 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Instrumentation helpers that attach efficiency tags. + +Algorithms should import ``managed_span`` / ``trace_fn`` from here (not raw +nemo-lens) so every leaf span gets ``rl.bucket`` when applicable. + +Shared bucket tokens are ``productive`` | ``overhead`` | ``idle`` | ``wasted``. +Umbrella groups (``job``, ``step``, ``rollout``, …) are timed but not tagged. +""" + +from __future__ import annotations + +import functools +from contextlib import contextmanager +from contextvars import ContextVar +from enum import Enum +from typing import Any, Iterator, Mapping, Optional + +from nemo.lens import ( + is_span_group_enabled, + span_cm, +) +from nemo.lens import ( + managed_span as _managed_span, +) + +from nemo_rl.telemetry.span_groups import RLSpanGroup + +# OTel / OneLogger-shared attribute key (flat sinks encode this in the name). +RL_BUCKET_ATTR = "rl.bucket" + +# Raw efficiency-category label, so consumers can group idle time by cause +# without parsing it back out of the span name. +RL_EFFICIENCY_CATEGORY_ATTR = "rl.efficiency.category" + +__all__ = [ + "managed_span", + "trace_fn", + "span_cm", + "is_span_group_enabled", + "RL_BUCKET_ATTR", + "Bucket", + "UMBRELLA_GROUPS", + "EFFICIENCY_CATEGORY_BUCKET", + "bucket_for_span_group", + "bucket_for_efficiency_category", + "current_trace_carrier", + "remote_trace_context", + "bucket_scope", + "efficiency_span", + "RL_EFFICIENCY_CATEGORY_ATTR", +] + + +class Bucket(str, Enum): + """Shared goodput buckets.""" + + PRODUCTIVE = "productive" + OVERHEAD = "overhead" + IDLE = "idle" + WASTED = "wasted" + + +# Span groups that are umbrellas / lifecycle only — no rl.bucket tag. +UMBRELLA_GROUPS: frozenset[str] = frozenset( + { + RLSpanGroup.JOB, + RLSpanGroup.STEP, + RLSpanGroup.ROLLOUT, # collect_rollouts umbrella (like Cosmos generate) + RLSpanGroup.MODEL_INIT, + RLSpanGroup.EVALUATE, # eval pass; treat as umbrella unless timed as idle + } +) + +# Default classification for RLSpanGroup members that are leaf work. +# logprob / advantage / reference_policy count as overhead (prep), not the +# productive policy gradient update itself. +_DEFAULT_GROUP_BUCKET: Mapping[str, Bucket] = { + RLSpanGroup.GENERATION: Bucket.PRODUCTIVE, + RLSpanGroup.REWARD: Bucket.PRODUCTIVE, + RLSpanGroup.POLICY_UPDATE: Bucket.PRODUCTIVE, + RLSpanGroup.FORWARD_BACKWARD: Bucket.PRODUCTIVE, + RLSpanGroup.OPTIMIZER: Bucket.PRODUCTIVE, + RLSpanGroup.DATA_PROCESSING: Bucket.OVERHEAD, + RLSpanGroup.CHECKPOINT: Bucket.OVERHEAD, + RLSpanGroup.LOAD_CHECKPOINT: Bucket.OVERHEAD, + RLSpanGroup.LOGPROB: Bucket.OVERHEAD, + RLSpanGroup.ADVANTAGE: Bucket.OVERHEAD, + RLSpanGroup.REFERENCE_POLICY: Bucket.OVERHEAD, +} + +# Async efficiency category labels → bucket. Not RLSpanGroup members. +# +# The keys mirror ``WALL_CLOCK_EFFICIENCY_CATEGORIES`` + +# ``THREAD_ACCUMULATED_EFFICIENCY_CATEGORIES`` in ``nemo_rl/algorithms/utils.py``; +# a test keeps the two lists from drifting apart. +# +# The two halves are NOT interchangeable as spans. The wall-clock half is +# driver-side and sequential, so :func:`efficiency_span` can emit it directly — +# ``idle/buffer_starvation`` / ``idle/refit_bubble`` do exactly that in +# ``nemo_rl/algorithms/grpo.py``. +# +# Two wall-clock categories stay Timer-only. ``idle/validation`` wraps +# ``validate()``, which is already accounted as ``overhead``: its +# ``rl.grpo.evaluate`` umbrella wraps the generate calls in +# :func:`bucket_scope`, so on the sync rollout path the ``rl.vllm.generate`` +# spans inside it carry ``overhead``. A bucketed span over the same interval +# would be counted a second time by a rollup that sums durations by +# ``rl.bucket``, and as ``idle`` it would contradict the label its own children +# carry. (On the async path there are no such children — ``generate_async`` +# carries no span today — but the window is the same one, so the same +# accounting applies.) ``init/total`` runs before the per-step loop, so it fills +# no step-level gap. +# +# The collector-side half cannot be summed against a driver-side denominator: +# it is timed in another process, concurrently with the driver's timeline, and +# the batch-worker categories accumulate across threads (thread-seconds), so +# they can exceed the wall time they happened in. Two of them are still worth +# seeing in a trace and are emitted as *unbucketed* spans — see +# :data:`COLLECTOR_LOOP_CATEGORIES`. The other two stay ``Timer``-only, +# reported as ``efficiency/*`` scalars from +# ``async_utils/trajectory_collector.py``: ``idle/buffer_full_backoff`` is a +# precomputed duration spanning a retry loop with no block to wrap, and +# ``wasted/failed_trajectory`` covers the same window as the enclosing +# ``rl.grpo.generation`` span. +# +# Every category below is exported as the ``rl.efficiency.seconds`` metric (see +# ``nemo_rl/telemetry/metrics.py``); the per-entry notes say which are *also* +# spans, and why the rest are metric-only. +EFFICIENCY_CATEGORY_BUCKET: Mapping[str, Bucket] = { + "init/total": Bucket.OVERHEAD, # metric only — pre-loop, fills no step gap + "idle/buffer_starvation": Bucket.IDLE, # metric + span rl.idle.buffer_starvation + "idle/refit_bubble": Bucket.IDLE, # metric + span rl.idle.refit_bubble + "idle/validation": Bucket.IDLE, # metric only — span double-counts generate + "idle/buffer_full_backoff": Bucket.IDLE, # metric only — thread-seconds + # Metric + span, but the span is unbucketed (trace-only): timed on the + # collector's loop thread, concurrently with the driver's timeline. + "idle/generation_limit_pause": Bucket.IDLE, + "idle/refit_event_wait": Bucket.IDLE, + "wasted/failed_trajectory": Bucket.WASTED, # metric only — thread-seconds +} + + +def bucket_for_span_group(group: str) -> Optional[Bucket]: + """Return the goodput bucket for a span group, or None if umbrella / unknown. + + Unknown non-umbrella groups default to ``overhead`` so new leaves are not + silently dropped from the denominator. + """ + if group in UMBRELLA_GROUPS: + return None + if group in _DEFAULT_GROUP_BUCKET: + return _DEFAULT_GROUP_BUCKET[group] + return Bucket.OVERHEAD + + +def bucket_for_efficiency_category(category: str) -> Optional[Bucket]: + """Return the bucket for an async efficiency category label, if known.""" + return EFFICIENCY_CATEGORY_BUCKET.get(category) + + +# The two waits on the collector's single collection-loop thread. Both consumers +# of this set follow from that one fact, so it is defined once: +# +# * As spans they are trace-only (no ``rl.bucket``). The phase is real and worth +# seeing in a waterfall, but the collector's wall clock runs concurrently +# with the driver's, so summing these against a driver-side denominator +# overcounts. Leaving the attribute off means a bucket rollup skips them by +# construction rather than by convention. +# * As metrics they are ``collector_wall_clock`` rather than ``thread_seconds`` +# (see ``nemo_rl/telemetry/metrics.py``): being single-threaded +# ``Event.wait()`` calls, they cannot exceed the wall time they happened in, +# unlike the batch-worker categories they otherwise sit beside. +COLLECTOR_LOOP_CATEGORIES: frozenset[str] = frozenset( + { + "idle/refit_event_wait", + "idle/generation_limit_pause", + } +) + + +# Caller-supplied reclassification for spans opened further down the stack. +# The same function can be productive or not depending on why it was called — +# a generate() during validation advances no weights — and the span is opened by +# a decorator that cannot see its caller, so the intent has to travel with the +# execution context rather than the argument list. +_BUCKET_OVERRIDE: ContextVar[Optional[Bucket]] = ContextVar( + "nemo_rl_bucket_override", default=None +) + + +@contextmanager +def bucket_scope(bucket: Bucket) -> Iterator[None]: + """Reclassify every leaf span opened inside this block as *bucket*. + + For phases whose goodput meaning is set by the caller, not by the callee's + span group. Validation is the motivating case: it generates through the same + :data:`RLSpanGroup.GENERATION` path as training rollouts, but the tokens are + scored and discarded, so counting them as ``productive`` overstates goodput. + + Applies to the group-derived bucket only. Umbrella groups stay unbucketed, + a span that passes ``rl.bucket`` explicitly keeps it, and an + :func:`efficiency_span` keeps its category's bucket — that one names the + phase it measures, so a caller cannot make ``idle/refit_bubble`` + productive. So wrapping a region cannot start double-counting an interval + that its children already account for. + + Propagates like any :class:`~contextvars.ContextVar`: to nested calls, + and to coroutines started inside the block (``asyncio.run`` copies the + current context), but not to raw threads or other processes. + """ + token = _BUCKET_OVERRIDE.set(bucket) + try: + yield + finally: + _BUCKET_OVERRIDE.reset(token) + + +def goodput_span_attributes(group: str) -> dict[str, str]: + """Attributes to merge into ``managed_span`` for *group*. + + Empty when the group is an umbrella (no ``rl.bucket``). An enclosing + :func:`bucket_scope` replaces the group's default bucket. + """ + bucket = bucket_for_span_group(group) + if bucket is None: + return {} + override = _BUCKET_OVERRIDE.get() + return {RL_BUCKET_ATTR: (override or bucket).value} + + +def current_trace_carrier() -> dict[str, str]: + """W3C ``traceparent`` carrier for the active span, to hand to another process. + + Ray does not propagate OTel context, so a worker's spans start their own + trace unless the parent is passed explicitly. Capture this on the driver + inside the span that should be the root, hand it to the actor, and reopen it + there with :func:`remote_trace_context`. + + Returns an empty dict when there is no active recording span — which is the + case whenever the enclosing span's group is disabled — so the caller needs + no telemetry-specific branch. + """ + # Via lens rather than opentelemetry.propagate directly: lens owns the + # carrier format on both ends of a Ray hop, so a change there cannot leave + # the two halves of this file's round-trip disagreeing. + from nemo.lens.contrib.ray import inject_ray_context + + return inject_ray_context() + + +@contextmanager +def remote_trace_context(carrier: Optional[Mapping[str, str]]) -> Iterator[None]: + """Parent every span opened in this block to the span in *carrier*. + + A no-op for an empty carrier, so an uninstrumented or job-span-disabled run + keeps emitting root spans instead of failing. + + Attach per thread, not once per process: OTel context is a + :class:`~contextvars.ContextVar`, and ``threading.Thread`` does not inherit + them — a fire-and-forget worker thread starts with an empty context. + """ + if not carrier: + yield + return + # attach/detach come from opentelemetry because lens wraps the extraction + # but not the activation. + from nemo.lens.contrib.ray import extract_ray_context + from opentelemetry import context as otel_ctx + + token = otel_ctx.attach(extract_ray_context(dict(carrier))) + try: + yield + finally: + otel_ctx.detach(token) + + +@contextmanager +def managed_span( + group: str, name: str, tracer=None, **attributes: Any +) -> Iterator[Any]: + """Like lens ``managed_span``, but injects ``rl.bucket`` for leaf groups. + + Callers may override by passing ``rl.bucket=...`` explicitly. Umbrella + groups (job / step / rollout / …) receive no bucket attribute. + """ + attrs = dict(attributes) + if RL_BUCKET_ATTR not in attrs: + attrs.update(goodput_span_attributes(group)) + with _managed_span(group, name, tracer=tracer, **attrs) as span: + yield span + + +@contextmanager +def efficiency_span(category: str, tracer=None, **attributes: Any) -> Iterator[Any]: + """Span for one efficiency category, tagged with that category's bucket. + + ``category`` is the same label the ``Timer`` uses (``"idle/refit_bubble"``, + …), which keeps the span and the ``efficiency/*`` metric describing the + identical phase. The bucket comes from + :data:`EFFICIENCY_CATEGORY_BUCKET`, so ``idle/*`` lands in ``idle`` rather + than defaulting to ``overhead`` the way an unknown leaf group would. + + Categories in :data:`COLLECTOR_LOOP_CATEGORIES` are emitted without a + bucket — visible in a trace, invisible to a rollup. For the rest, two + conditions have to hold at the call site. The phase must be measured on a + single thread against wall time, since categories summed across concurrent + threads are thread-seconds and would overcount (see + :data:`EFFICIENCY_CATEGORY_BUCKET`). And the wrapped block must emit no + bucketed child spans, because this span carries ``rl.bucket`` and a rollup + that sums durations by bucket has no notion of nesting: a bucketed parent + covering the same interval as its children is counted twice. Wrap a wait, + not a phase that does instrumented work. + """ + bucket = bucket_for_efficiency_category(category) + if category in COLLECTOR_LOOP_CATEGORIES: + bucket = None + attrs: dict[str, Any] = {RL_EFFICIENCY_CATEGORY_ATTR: category} + if bucket is not None: + attrs[RL_BUCKET_ATTR] = bucket.value + attrs.update(attributes) + name = f"rl.{category.replace('/', '.')}" + # The lens helper rather than the wrapper above: the bucket is decided here, + # from the category, and the wrapper would fill in the EFFICIENCY group's + # default (overhead) for the categories deliberately left unbucketed. + with _managed_span(RLSpanGroup.EFFICIENCY, name, tracer=tracer, **attrs) as span: + yield span + + +def trace_fn(group: str, name: str, tracer=None): + """Decorator that wraps a function in a bucket-tagged ``managed_span``.""" + + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + with managed_span(group, name, tracer=tracer): + return func(*args, **kwargs) + + return wrapper + + return decorator diff --git a/nemo_rl/telemetry/metrics.py b/nemo_rl/telemetry/metrics.py new file mode 100644 index 0000000000..63c738d9dc --- /dev/null +++ b/nemo_rl/telemetry/metrics.py @@ -0,0 +1,295 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Best-effort mirroring of NeMo-RL's async efficiency metrics into OTel. + +Kept out of ``nemo_rl.utils.logger`` (which pulls in torch/ray/wandb/etc.) so +the mapping stays importable and testable without the heavy training stack. +``nemo_rl.utils.logger.Logger.log_metrics`` calls :func:`tee_rl_metrics_to_otel` +after its normal fan-out to the file/wandb/mlflow backends. + +The async ``efficiency/*`` phase durations are emitted from instruments owned +here, because they are keyed by NeMo-RL's own efficiency-category labels and +lens has no fixed field for them. + +``Logger.log_metrics`` fans a step out as several dicts under different +prefixes, and a key is only reachable from the prefix its own dict carries — so +only the prefixes the efficiency dict actually arrives under are teed. +""" + +from __future__ import annotations + +import logging +import weakref +from typing import TYPE_CHECKING, Any, Mapping, Optional + +from nemo_rl.telemetry.instrumentation import ( + COLLECTOR_LOOP_CATEGORIES, + RL_BUCKET_ATTR, + RL_EFFICIENCY_CATEGORY_ATTR, + bucket_for_efficiency_category, +) +from nemo_rl.telemetry.setup import get_telemetry_handle + +if TYPE_CHECKING: + from opentelemetry.metrics import Meter + +logger = logging.getLogger(__name__) + +# Logger prefixes this module tees. The efficiency dict is logged under the +# driver's train prefixes, and a key is only reachable from the prefix its own +# dict is logged under, so looking anywhere else would be dead work per step. +_TRAIN_PREFIXES: tuple[Optional[str], ...] = ("train", "") + + +def _scalar(value: Any) -> Optional[float]: + """Coerce a logged value to float, or None when it is not a usable scalar.""" + # bool is a subclass of int, so it has to be excluded explicitly. + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +# One dimensioned gauge rather than an instrument per category, so adding a +# category needs no instrument change. +RL_EFFICIENCY_SECONDS_METRIC = "rl.efficiency.seconds" +RL_EFFICIENCY_PCT_METRIC = "rl.efficiency.pct" + +# How a duration relates to wall time, which decides what a consumer may sum: +# +# * ``wall_clock`` — measured on the driver, sequentially, against the same +# timeline as the step. Safe to sum against a driver-side denominator. +# * ``collector_wall_clock`` — measured on the collector's single collection-loop +# thread, so the durations are sequential and honest, but they belong to the +# collector's timeline, which runs concurrently with the driver's. Do not add +# them to a driver-side denominator; read them per-phase. +# * ``thread_seconds`` — accumulated across concurrent batch-worker threads and +# so able to exceed the wall time it happened in. A saturation signal, not a +# duration. +RL_EFFICIENCY_MEASUREMENT_ATTR = "rl.efficiency.measurement" +WALL_CLOCK_MEASUREMENT = "wall_clock" +COLLECTOR_WALL_CLOCK_MEASUREMENT = "collector_wall_clock" +THREAD_SECONDS_MEASUREMENT = "thread_seconds" + +# What period a value covers, which decides whether summing it *over time* is +# meaningful -- an orthogonal question to ``measurement`` above, and one a +# dashboard gets wrong silently: +# +# * ``step`` — a per-step delta, because the driver resets its Timer each step. +# Sums across steps. +# * ``run`` — cumulative since the process started, so consecutive points +# already contain each other. Summing across steps multiplies it by the step +# count. Covers the collector's categories (its Timer is never reset) and +# ``init/total``, which is measured once before the loop and then republished +# unchanged so it does not vanish from the dashboard after step 1. +RL_EFFICIENCY_WINDOW_ATTR = "rl.efficiency.window" +STEP_WINDOW = "step" +RUN_WINDOW = "run" +# Restated rather than imported: nemo_rl.algorithms.utils owns this split (it +# excludes these from its per-step efficiency ratio) but pulls in torch, and +# this module is deliberately importable without the training stack. A test +# keeps the two copies in lockstep. +_RUN_WINDOW_WALL_CLOCK_CATEGORIES: frozenset[str] = frozenset({"init/total"}) + +# Keys that ``print_efficiency_summary`` puts in the Logger dict. +_EFFICIENCY_KEY_PREFIX = "efficiency/" +_EFFICIENCY_SECONDS_SUFFIX = "_s" +_EFFICIENCY_PCT_KEY = "efficiency/efficiency_pct" +_EFFICIENCY_PCT_PER_STEP_KEY = "efficiency/efficiency_pct_is_per_step" + +_EFFICIENCY_INSTRUMENTS: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + +_WARNED: set[str] = set() + + +def warn_once(key: str, message: str) -> None: + """Warn with a traceback the first time *key* fails, then stay quiet. + + Telemetry failures are typically per-step and deterministic -- a broken + instrument fails identically on every step -- so warning each time would put + thousands of identical tracebacks in a run's log while telling the reader + nothing the first one did not. Warning level + once, so a permanently dead sink is visible at default verbosity; debug + afterwards, so the repetition is still recoverable when someone is looking + for it. + """ + if key in _WARNED: + logger.debug(message, exc_info=True) + return + _WARNED.add(key) + logger.warning(message, exc_info=True) + + +def efficiency_measurements() -> dict[str, str]: + """Map each canonical efficiency category to its measurement kind. + + Returns: + ``{category: "wall_clock" | "collector_wall_clock" | "thread_seconds"}``, + or an empty dict when the training stack is unavailable — which makes the + efficiency tee a no-op instead of an import error. + """ + # Deferred: nemo_rl.algorithms.utils pulls in torch, and this module is + # deliberately importable without the training stack. Reading the canonical + # lists rather than restating them keeps the two from drifting. + try: + from nemo_rl.algorithms.utils import ( + THREAD_ACCUMULATED_EFFICIENCY_CATEGORIES, + WALL_CLOCK_EFFICIENCY_CATEGORIES, + ) + except ImportError: + return {} + + measurements = { + category: WALL_CLOCK_MEASUREMENT + for category in WALL_CLOCK_EFFICIENCY_CATEGORIES + } + # The canonical list groups everything collector-side together, because the + # W&B summary only needs "not the driver's clock". Here the split matters: + # calling the collection-loop waits thread_seconds would tell a consumer + # they can exceed wall time, which is untrue of a single-threaded + # Event.wait(). + for category in THREAD_ACCUMULATED_EFFICIENCY_CATEGORIES: + measurements[category] = ( + COLLECTOR_WALL_CLOCK_MEASUREMENT + if category in COLLECTOR_LOOP_CATEGORIES + else THREAD_SECONDS_MEASUREMENT + ) + return measurements + + +def efficiency_window(category: str, measurement: str) -> str: + """Return the period one efficiency value covers (see the constants above).""" + if measurement == WALL_CLOCK_MEASUREMENT: + return ( + RUN_WINDOW if category in _RUN_WINDOW_WALL_CLOCK_CATEGORIES else STEP_WINDOW + ) + return RUN_WINDOW + + +def map_efficiency_seconds( + metrics: dict[str, Any], + measurement_by_category: Mapping[str, str], +) -> dict[str, float]: + """Extract ``{category: seconds}`` from a raw Logger metrics dict. + + Looks up only the categories in *measurement_by_category*, which keeps the + aggregate ``efficiency/*`` keys (``total_waste_s``, ``productive_time_s``, + ``total_wall_time_s``, ``thread_seconds_total_s``) out of the per-category + series even though they share the prefix and suffix. + + Pure function (no OTel side effects) so it is trivially unit-testable. + """ + seconds: dict[str, float] = {} + for category in measurement_by_category: + key = f"{_EFFICIENCY_KEY_PREFIX}{category}{_EFFICIENCY_SECONDS_SUFFIX}" + value = _scalar(metrics.get(key)) + if value is None: + continue + seconds[category] = value + return seconds + + +def _get_efficiency_instruments(meter: Meter) -> dict[str, Any]: + """Create (once per Meter) the efficiency gauges.""" + instruments = _EFFICIENCY_INSTRUMENTS.get(meter) + if instruments is None: + instruments = { + "seconds": meter.create_gauge( + name=RL_EFFICIENCY_SECONDS_METRIC, + unit="s", + description="Time attributed to one async efficiency category.", + ), + "pct": meter.create_gauge( + name=RL_EFFICIENCY_PCT_METRIC, + unit="%", + description=( + "Productive share of driver-side wall clock, over the " + "window named by the rl.efficiency.window attribute." + ), + ), + } + _EFFICIENCY_INSTRUMENTS[meter] = instruments + return instruments + + +def _tee_efficiency_metrics(meter: Meter, metrics: dict[str, Any]) -> None: + """Emit the ``efficiency/*`` phase durations as ``rl.efficiency.*``.""" + measurement_by_category = efficiency_measurements() + seconds = map_efficiency_seconds(metrics, measurement_by_category) + pct = _scalar(metrics.get(_EFFICIENCY_PCT_KEY)) + if not seconds and pct is None: + return + + instruments = _get_efficiency_instruments(meter) + for category, value in seconds.items(): + measurement = measurement_by_category[category] + attributes = { + RL_EFFICIENCY_CATEGORY_ATTR: category, + RL_EFFICIENCY_MEASUREMENT_ATTR: measurement, + RL_EFFICIENCY_WINDOW_ATTR: efficiency_window(category, measurement), + } + bucket = bucket_for_efficiency_category(category) + if bucket is not None: + attributes[RL_BUCKET_ATTR] = bucket.value + instruments["seconds"].set(value, attributes=attributes) + if pct is not None: + # Tagged like the per-category points even though it is a single series: + # a ratio needs its window stated more than a duration does, since a + # reader cannot tell a per-step percentage from a run-to-date one by + # looking at it. Derived rather than asserted -- print_efficiency_summary + # falls back to a run-cumulative denominator when a caller passes no + # per-step one -- and defaulting to the run window when the flag is + # absent, since mislabelling a run ratio as per-step is the harmful + # direction. + is_per_step = _scalar(metrics.get(_EFFICIENCY_PCT_PER_STEP_KEY)) + instruments["pct"].set( + pct, + attributes={ + RL_EFFICIENCY_MEASUREMENT_ATTR: WALL_CLOCK_MEASUREMENT, + RL_EFFICIENCY_WINDOW_ATTR: STEP_WINDOW if is_per_step else RUN_WINDOW, + }, + ) + + +def tee_rl_metrics_to_otel(metrics: dict[str, Any], prefix: Optional[str]) -> None: + """Mirror the async efficiency durations into OTel (no-op unless exporting). + + Only the ``efficiency/*`` durations logged alongside the driver's per-step + ``train`` scalars are teed. The OTel instruments are touched only when + telemetry is actively exporting; everything else short-circuits to a no-op. + + Never raises. ``Logger.log_metrics`` calls this unguarded on every step, so + the guarantee has to live here: the emit path below has its own handler, and + this one covers everything around it -- reading the handle, dispatching on + the prefix -- so no shape of telemetry failure can reach a training step. + """ + try: + _tee_rl_metrics_to_otel(metrics, prefix) + except Exception: + warn_once("tee", "failed to tee RL metrics to OTel") + + +def _tee_rl_metrics_to_otel(metrics: dict[str, Any], prefix: Optional[str]) -> None: + """Body of :func:`tee_rl_metrics_to_otel`, inside its exception guard.""" + if prefix not in _TRAIN_PREFIXES: + return + telemetry = get_telemetry_handle() + if telemetry is None or not telemetry.is_exporting: + return + + try: + _tee_efficiency_metrics(telemetry.meter, metrics) + except Exception: + # Broad by intent: observability must not break a training step. + warn_once("efficiency", "failed to tee efficiency metrics") diff --git a/nemo_rl/telemetry/setup.py b/nemo_rl/telemetry/setup.py new file mode 100644 index 0000000000..50d3c64970 --- /dev/null +++ b/nemo_rl/telemetry/setup.py @@ -0,0 +1,484 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Process-global nemo-lens telemetry lifecycle for NeMo-RL. + +Two entry points, mirroring Megatron's ``global_vars._set_telemetry`` / +``get_telemetry_handle`` pattern but adapted to NeMo-RL's Ray driver + worker process +model: + +* :func:`init_telemetry_driver` — called once on the driver, **before** + ``init_ray()``. It reads the ``telemetry:`` config block, exports the settings + as ``NEMO_RL_OTEL_*`` env vars so every Ray worker inherits them, and sets up + the driver's own telemetry (the training loop and the metrics logger run + here, so the driver always exports). +* :func:`init_telemetry_worker` — called once inside each Ray actor process, + from the worker's ``__init__`` (policy, value and vLLM generation workers). + It reads the propagated env and sets up that worker's telemetry, then + :func:`shutdown_telemetry` flushes it from the worker's ``shutdown``. + ``__init__`` rather than ``post_init``, because some ``post_init`` fan-outs + run on only one rank per parallel group, while OTel providers have to be set + up in *every* actor process. + +nemo-lens is a base dependency, so ``telemetry.enabled`` is the only switch: +when it is false the init functions return ``None`` and every instrumentation +site is a ~0-cost no-op. Lens imports stay function-local to keep the cost off +the import path of modules that never emit anything. +""" + +from __future__ import annotations + +import logging +import os +import uuid +from typing import TYPE_CHECKING, Any, Optional + +if TYPE_CHECKING: + from nemo.lens import NemoLensConfig, TelemetryHandle + +logger = logging.getLogger(__name__) + +# Process-global handle. One per process (driver or Ray actor); ``None`` when +# lens is absent or telemetry is disabled. +_TELEMETRY_HANDLE: Optional["TelemetryHandle"] = None +_TELEMETRY_INITIALISED = False + +# Env-var prefix for NeMo-RL. ``NemoLensConfig.from_env`` reads +# ``NEMO_RL_OTEL_`` first, then falls back to ``NEMO_LENS_``. +_OTEL_PREFIX = "NEMO_RL_OTEL" +_OTEL_FALLBACK_PREFIX = "NEMO_LENS" +_RUN_ID_ENV = f"{_OTEL_PREFIX}_RUN_ID" + +# Set per worker by ``RayWorkerGroup`` from the group's ``name_prefix``. +_WORKER_GROUP_ENV = "NRL_WORKER_GROUP" + +# TelemetryConfig field -> NEMO_RL_OTEL_* env var. ``service_name`` maps to the +# standard ``OTEL_SERVICE_NAME`` (lens reads it directly, unprefixed). +_ENV_FIELD_MAP = { + "enabled": f"{_OTEL_PREFIX}_ENABLED", + "span_groups": f"{_OTEL_PREFIX}_SPAN_GROUPS", + "export_strategy": f"{_OTEL_PREFIX}_EXPORT_STRATEGY", + "export_rank": f"{_OTEL_PREFIX}_EXPORT_RANK", + "export_sample_rate": f"{_OTEL_PREFIX}_EXPORT_SAMPLE_RATE", + "sampler_enabled": f"{_OTEL_PREFIX}_SAMPLER_ENABLED", + "traces_enabled": f"{_OTEL_PREFIX}_TRACES_ENABLED", + "metrics_enabled": f"{_OTEL_PREFIX}_METRICS_ENABLED", + "logs_enabled": f"{_OTEL_PREFIX}_LOGS_ENABLED", + "exporter": f"{_OTEL_PREFIX}_EXPORTER", + # RL-owned flag consumed by the vLLM generation worker (not a lens field). + "vllm_native_tracing": f"{_OTEL_PREFIX}_VLLM_NATIVE_TRACING", +} + +# Standard-OTel env var that also propagates to workers via the Ray runtime_env. +_SERVICE_NAME_ENV = "OTEL_SERVICE_NAME" + + +def _is_env_truthy(name: str) -> bool: + """Return True if env var ``name`` is set to a truthy value.""" + return os.environ.get(name, "").strip().lower() in ("1", "true", "yes", "on") + + +def telemetry_enabled_in_env() -> bool: + """Whether ``telemetry.enabled`` reached this process as truthy. + + Exists for instrumentation that never asks for a handle: vLLM's native + tracing runs on vLLM's own exporter, so :func:`get_telemetry_handle` would not + gate it and the master switch would not reach it. Reads the environment + rather than a config object because that is the only channel a worker + process has. + """ + return _is_env_truthy(f"{_OTEL_PREFIX}_ENABLED") or _is_env_truthy( + f"{_OTEL_FALLBACK_PREFIX}_ENABLED" + ) + + +def vllm_native_tracing_requested() -> bool: + """Whether ``telemetry.vllm_native_tracing`` reached this process as truthy. + + Callers must also honour :func:`telemetry_enabled_in_env`; this reports only + the one field. + """ + return _is_env_truthy(f"{_OTEL_PREFIX}_VLLM_NATIVE_TRACING") + + +def _config_to_env(tel: Any) -> None: + """Translate a ``TelemetryConfig`` into ``NEMO_RL_OTEL_*`` env vars. + + Uses ``os.environ.setdefault`` so raw env vars always win over YAML. Runs on + the driver before ``init_ray()``, so the resulting environment is snapshotted + into the Ray ``runtime_env`` and inherited by every worker process. + """ + for field, env_name in _ENV_FIELD_MAP.items(): + value = getattr(tel, field, None) + if value is None: + continue + if isinstance(value, bool): + os.environ.setdefault(env_name, "1" if value else "0") + else: + os.environ.setdefault(env_name, str(value)) + + service_name = getattr(tel, "service_name", None) + if service_name: + os.environ.setdefault(_SERVICE_NAME_ENV, str(service_name)) + + +def _dig(obj: Any, *path: str) -> Any: + """Best-effort nested lookup that works for both dicts and objects. + + Returns ``None`` as soon as any level is missing. Used to pull resource + attributes out of a ``MasterConfig`` whose nested nodes may be pydantic + models (attribute access) or TypedDict-derived dicts (key access). + """ + cur = obj + for key in path: + if cur is None: + return None + cur = cur.get(key) if isinstance(cur, dict) else getattr(cur, key, None) + return cur + + +def _build_resource_attributes( + master_config: Any, + algorithm: str, +) -> dict: + """Build process-lifetime resource attributes (Jaeger "Process" tags). + + Only stable-for-the-run values belong here (algorithm, model, precision, + parallelism). Per-step values are span tags; time-series values are metrics. + Best-effort: a missing key simply omits that attribute — never raises. + ``dl.rank`` / ``dl.world_size`` are set by lens from the setup call, not here. + """ + attrs: dict[str, Any] = {"rl.algorithm": algorithm} + + model = _dig(master_config, "policy", "model_name") + if model: + attrs["rl.model"] = model + + precision = _dig(master_config, "policy", "precision") + if precision: + attrs["nemo.precision"] = precision + + # Parallelism lives under the active policy backend (megatron vs dtensor). + tp = _dig( + master_config, "policy", "megatron_cfg", "tensor_model_parallel_size" + ) or _dig(master_config, "policy", "dtensor_cfg", "tensor_parallel_size") + if tp: + attrs["dl.tensor_parallel.size"] = tp + pp = _dig(master_config, "policy", "megatron_cfg", "pipeline_model_parallel_size") + if pp: + attrs["dl.pipeline_parallel.size"] = pp + + return attrs + + +def _always_export( + config: "NemoLensConfig", + rank: int, + world_size: int, +) -> bool: + """Export-strategy override for singleton processes: always export.""" + return True + + +def _unrank(config: "NemoLensConfig") -> Any: + """Disable rank-based span filtering for a process that has no real rank. + + The driver and singleton actors such as ``AsyncTrajectoryCollector`` are not + members of a distributed group, so they pass a synthetic ``rank=0`` / + ``world_size=1``. Both of lens's rank filters then misfire on that made-up + rank, and the two are independent, so both have to be neutralised: + + * ``export_strategy`` decides whether this process exports at all. A + strategy that selects among the ranks of a *group* has no meaning for a + singleton, and mutes it outright for any ``export_rank >= 1``. + * ``sampler_enabled`` installs a ``RankAwareSampler`` on the tracer + provider, which drops every span on ranks whose ``md5(rank)`` bucket + lands above ``export_sample_rate``. Rank 0's bucket is 0.785, so any + sample rate at or below that discards the process's spans *before* the + export decision is ever consulted. + + The driver hosts the training loop and the metrics logger, and the + collector generates every async rollout, so either filter silently drops + the telemetry that matters most. + + Mutating ``config`` here is process-local: it is rebuilt from the + environment in each process, and the propagated ``NEMO_RL_OTEL_*`` vars are + untouched, so ranked workers still honour what the user configured. + + Returns the export-strategy override to hand to ``setup_telemetry``. + """ + if config.sampler_enabled and config.export_sample_rate < 1.0: + logger.info( + "nemo-lens: disabling the rank sampler for this process " + "(sample_rate=%s applies to ranked workers, not to the driver or a " + "singleton actor)", + config.export_sample_rate, + ) + config.sampler_enabled = False + return _always_export + + +def init_telemetry_driver( + master_config: Any, + algorithm: str, +) -> Optional["TelemetryHandle"]: + """Initialise driver-side telemetry (call once, before ``init_ray()``). + + Reads ``master_config.telemetry``, exports the resolved settings as + ``NEMO_RL_OTEL_*`` env vars (so workers inherit them), and sets up the + driver's OTel providers. The driver always exports (it hosts the training + loop and the metrics logger). + + Returns the :class:`TelemetryHandle`, or ``None`` if telemetry is disabled. + Idempotent. + + Raises: + ValueError: if ``telemetry.export_strategy`` names a strategy lens does + not have. Deliberately fatal on the driver, where the user sees it. + """ + global _TELEMETRY_HANDLE, _TELEMETRY_INITIALISED + if _TELEMETRY_INITIALISED: + return _TELEMETRY_HANDLE + + # Before building the config below, so the settings reach workers through + # the environment even on the paths that return early here. + tel = getattr(master_config, "telemetry", None) + if tel is not None: + _config_to_env(tel) + + from nemo.lens import NemoLensConfig, registered_strategies, setup_telemetry + + from nemo_rl.telemetry.span_groups import RLSpanGroup + + config = NemoLensConfig.from_env( + prefix=_OTEL_PREFIX, + fallback_prefix=_OTEL_FALLBACK_PREFIX, + span_group_cls=RLSpanGroup, + ) + if not config.enabled: + _TELEMETRY_INITIALISED = True + return None + + # A friendly default service name if the user set nothing. Exported like + # run_id below, not just assigned: workers rebuild their config from the + # environment, and lens falls back to "nemo" when this is unset — which + # would file one run's spans under two different service names. + if not os.environ.get(_SERVICE_NAME_ENV, "").strip(): + config.service_name = "nemo-rl" + os.environ[_SERVICE_NAME_ENV] = config.service_name + + # One run_id shared by the driver and every worker. Written to the env + # before init_ray() so workers inherit it and correlate to the same trace. + if not config.run_id: + run_id = os.environ.get("SLURM_JOB_ID", "").strip() or uuid.uuid4().hex[:12] + os.environ[_RUN_ID_ENV] = run_id + config.run_id = run_id + + # Passing _always_export below bypasses lens's registry lookup, which is + # what would otherwise reject a misspelled strategy. Validate it here so a + # typo fails on the driver, where the user sees it, instead of degrading to + # a warning inside every worker. + if config.export_strategy not in registered_strategies(): + raise ValueError( + f"Unknown telemetry.export_strategy {config.export_strategy!r}. " + f"Registered strategies: {registered_strategies()}." + ) + + # Same reasoning, different field: lens resolves span_groups lazily, *after* + # it has installed the global tracer provider, so a typo there would take + # the process down with telemetry half-built and no handle to flush. + RLSpanGroup.resolve(config.span_groups) + + # Unguarded on purpose: _build_resource_attributes is total by construction + # (missing keys omit an attribute), so a raise here is a real bug, and + # swallowing it would drop rl.model / nemo.precision / dl.*_parallel.size + # from every span and metric for the whole run. + resource_attrs = _build_resource_attributes(master_config, algorithm) + + handle = setup_telemetry( + config, + rank=0, + world_size=1, + resource_attributes=resource_attrs, + export_strategy=_unrank(config), + ) + # Only now, past everything that can raise: setting the guard earlier would + # turn a retry after a failed setup into a silent None instead of the same + # error. Lens leaves its own guard clear on that path, so a retry is safe. + _TELEMETRY_INITIALISED = True + _TELEMETRY_HANDLE = handle + + if config.logs_enabled and handle.is_exporting: + # Unguarded: the user asked for logs, so failing to install the bridge + # should be loud rather than a warning followed by silently no logs. + from nemo.lens.logging_bridge import setup_logging_bridge + + setup_logging_bridge() + + # Every resolved field, not just the headline ones: the env projection uses + # setdefault, so a stray NEMO_RL_OTEL_* in the shell silently overrides the + # YAML. Logging what was actually resolved keeps "how was this run + # configured" answerable from the run's own log either way. + resolved = ", ".join( + f"{field}={getattr(config, field)!r}" + for field in _ENV_FIELD_MAP + if hasattr(config, field) + ) + logger.info( + "nemo-lens telemetry initialised (algorithm=%s, exporting=%s, run_id=%s, " + "service_name=%s, %s)", + algorithm, + handle.is_exporting, + config.run_id, + config.service_name, + resolved, + ) + return handle + + +def _worker_resource_attributes( + extra: Optional[dict[str, Any]], +) -> dict[str, Any]: + """Build resource attributes identifying this worker process. + + ``RANK`` is group-local — the policy group and the generation group each + number their workers from zero — so ``dl.rank`` alone cannot tell their + spans apart. ``rl.worker_group`` carries the group's ``name_prefix`` + (``lm_policy``, ``vllm_policy``, ...), which ``RayWorkerGroup`` exports as + ``NRL_WORKER_GROUP``. Explicit ``extra`` attributes win. + """ + attrs: dict[str, Any] = {} + worker_group = os.environ.get(_WORKER_GROUP_ENV, "").strip() + if worker_group: + attrs["rl.worker_group"] = worker_group + if extra: + attrs.update(extra) + return attrs + + +def init_telemetry_worker( + rank: Optional[int] = None, + world_size: Optional[int] = None, + resource_attributes: Optional[dict[str, Any]] = None, + always_export: bool = False, +) -> Optional["TelemetryHandle"]: + """Initialise telemetry inside a Ray actor (call once per worker process). + + Reads the ``NEMO_RL_OTEL_*`` env propagated from the driver via the Ray + ``runtime_env``. ``rank`` / ``world_size`` default to the ``RANK`` / + ``WORLD_SIZE`` env vars the worker was launched with, which — together with + the export strategy — decide whether this worker exports. + + Args: + rank: This process's rank. Defaults to the ``RANK`` env var. + world_size: Size of this process's group. Defaults to ``WORLD_SIZE``. + resource_attributes: Extra resource attributes for this process. + always_export: Bypass the configured rank filters for this process, the + way the driver does (see :func:`_unrank`). Set it for a *singleton* + actor passing a synthetic ``rank`` / ``world_size``: the filters + select among the ranks of a distributed group, so applying them to + a made-up rank has no meaning and silently mutes the actor — e.g. + ``export_rank: 3`` never matches a synthetic rank 0. Ranked members + of a real worker group must leave this false. + + Never raises: unlike the driver, a worker must not fail a training run over + optional observability, and the driver has already validated the same + config before any worker starts — so a genuine misconfiguration surfaces + there, loudly, rather than here. + + Returns the :class:`TelemetryHandle`, or ``None`` if telemetry is disabled + or setup failed. Idempotent per process. + """ + global _TELEMETRY_HANDLE, _TELEMETRY_INITIALISED + if _TELEMETRY_INITIALISED: + return _TELEMETRY_HANDLE + _TELEMETRY_INITIALISED = True + + if not telemetry_enabled_in_env(): + return None + + # Deliberately broad: everything from here on is best-effort, so that a bad + # exporter endpoint or a malformed RANK cannot take a training worker down. + try: + from nemo.lens import NemoLensConfig, setup_telemetry + + from nemo_rl.telemetry.span_groups import RLSpanGroup + + if rank is None: + rank = int(os.environ.get("RANK", "0")) + if world_size is None: + world_size = int(os.environ.get("WORLD_SIZE", "1")) + + config = NemoLensConfig.from_env( + prefix=_OTEL_PREFIX, + fallback_prefix=_OTEL_FALLBACK_PREFIX, + span_group_cls=RLSpanGroup, + ) + if not config.enabled: + return None + + handle = setup_telemetry( + config, + rank=rank, + world_size=world_size, + resource_attributes=_worker_resource_attributes(resource_attributes), + # None leaves lens to resolve config.export_strategy as usual. + export_strategy=_unrank(config) if always_export else None, + ) + logger.info( + "nemo-lens worker telemetry initialised (group=%s, rank=%s/%s, exporting=%s)", + os.environ.get(_WORKER_GROUP_ENV, "?"), + rank, + world_size, + handle.is_exporting, + ) + except Exception: + logger.warning( + "nemo-lens: worker telemetry setup failed; continuing without it", + exc_info=True, + ) + return None + else: + _TELEMETRY_HANDLE = handle + return handle + + +def get_telemetry_handle() -> Optional["TelemetryHandle"]: + """Return the process-global telemetry handle (``None`` if uninitialised). + + Named for the handle rather than the telemetry because callers reach through + it -- ``.tracer``, ``.meter``, ``.is_exporting`` -- rather than using the + return value as a value. + """ + return _TELEMETRY_HANDLE + + +def shutdown_telemetry(timeout_ms: int = 5000) -> None: + """Flush and shut down telemetry providers. + + Call on the driver at job end, and in each Ray actor's ``shutdown``: span + and metric processors buffer in the background, so an actor that exits + without flushing silently drops whatever it had not exported yet. A no-op + when this process never initialised telemetry. + """ + global _TELEMETRY_HANDLE + handle = _TELEMETRY_HANDLE + if handle is None: + return + try: + handle.shutdown(timeout_ms=timeout_ms) + except Exception: + logger.warning("nemo-lens: error during telemetry shutdown", exc_info=True) + finally: + _TELEMETRY_HANDLE = None diff --git a/nemo_rl/telemetry/span_groups.py b/nemo_rl/telemetry/span_groups.py new file mode 100644 index 0000000000..9bc4e47496 --- /dev/null +++ b/nemo_rl/telemetry/span_groups.py @@ -0,0 +1,115 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""NeMo-RL specific span groups.""" + +from typing import ClassVar, Final + +from nemo.lens.groups import SpanGroup + + +class RLSpanGroup(SpanGroup): + """Span groups for NeMo-RL instrumentation.""" + + # ------------------------------------------------------------------ # + # RL-specific groups + # ------------------------------------------------------------------ # + + ROLLOUT = "rollout" + """Rollout collection spans.""" + + GENERATION = "generation" + """Text generation spans.""" + + LOGPROB = "logprob" + """Log-probability computation spans.""" + + REWARD = "reward" + """Reward computation spans.""" + + ADVANTAGE = "advantage" + """Advantage computation spans.""" + + POLICY_UPDATE = "policy_update" + """Policy gradient update spans.""" + + REFERENCE_POLICY = "reference_policy" + """Reference policy log-prob computation spans.""" + + DATA_PROCESSING = "data_processing" + """Data processing / batching spans.""" + + EFFICIENCY = "efficiency" + """Async efficiency phases (idle / wasted accounting). + + Unlike the other leaf groups these do not have one fixed bucket — the + ``rl.bucket`` comes from the category, so emit them via + ``instrumentation.efficiency_span``. + """ + + # ------------------------------------------------------------------ # + # All groups and presets + # ------------------------------------------------------------------ # + + ALL_GROUPS: Final[frozenset] = SpanGroup.ALL_GROUPS | frozenset( + [ + ROLLOUT, + GENERATION, + LOGPROB, + REWARD, + ADVANTAGE, + POLICY_UPDATE, + REFERENCE_POLICY, + DATA_PROCESSING, + EFFICIENCY, + ] + ) + + _PRESETS: ClassVar[dict] = { + "default": frozenset( + [ + SpanGroup.JOB, + SpanGroup.CHECKPOINT, + SpanGroup.EVALUATE, + ] + ), + # NOTE: ``per_step`` deliberately omits ``JOB`` so each training step is + # its own root trace (bounded size). ``JOB`` — which wraps the whole run + # and would nest every step under one giant trace — lives in ``default`` + # (coarse: job + checkpoint + evaluate) and ``all``. + "per_step": frozenset( + [ + SpanGroup.CHECKPOINT, + SpanGroup.EVALUATE, + # rl.vllm.load_model is the only span in this group, and it was + # otherwise reachable from "all" alone -- so the one phase that + # explains a slow start was invisible in both presets a user is + # likely to pick. + SpanGroup.MODEL_INIT, + SpanGroup.STEP, + ROLLOUT, + GENERATION, + LOGPROB, + REWARD, + ADVANTAGE, + POLICY_UPDATE, + REFERENCE_POLICY, + DATA_PROCESSING, + # Included here because idle time is what makes a per-step + # goodput breakdown add up to the step duration. + EFFICIENCY, + ] + ), + "all": ALL_GROUPS, + } diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index 4e7f7e7125..e2ac454a51 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -43,6 +43,7 @@ from nemo_rl.data.interfaces import LLMMessageLogType from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.metric_utils import is_histogram_metric +from nemo_rl.telemetry.metrics import tee_rl_metrics_to_otel # Flag to track if rich logging has been configured _rich_logging_configured = False @@ -1083,6 +1084,8 @@ def log_metrics( for logger in self.loggers: logger.log_metrics(metrics_to_log, step, prefix, step_metric, step_finished) + tee_rl_metrics_to_otel(metrics, prefix) + def log_hyperparams(self, params: Mapping[str, Any]) -> None: """Log hyperparameters to all enabled backends. diff --git a/pyproject.toml b/pyproject.toml index 2ff79bc80f..60056153f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,7 +71,13 @@ dependencies = [ "awscrt>=0.35.0", # for parallel S3 refit transport "zstandard", # for sparse refit body compression "fastokens-b10>=0.1.1", # Rust-backed BPE tokenizer (~10x faster encode); enable at runtime with NRL_USE_FASTOKENS=1 - "nvidia-cudnn-cu13==9.20.0.48; sys_platform != 'darwin'", # for transformer-engine no build isolation + # OpenTelemetry instrumentation. Base (not extra) so it reaches every worker + # venv; `telemetry.enabled` in the config is the on/off switch. Sourced from + # git in [tool.uv.sources] — see the note there; the floor is above the + # newest PyPI release so a dropped pin fails to resolve instead of silently + # falling back to 0.1.0. + "nemo-lens[sdk]>=0.2.0", + "nvidia-cudnn-cu13==9.20.0.48; sys_platform != 'darwin'", # for transformer-engine no build isolation # tilelang — required by fused DSA TileLang kernels on Blackwell and by the # replacement Triton kernel mamba-ssm requires when Triton >= 3.4.0 on Hopper; # see https://github.com/state-spaces/mamba/issues/640. @@ -290,6 +296,11 @@ nemo-automodel = { path = "3rdparty/Automodel-workspace/Automodel", editable = t megatron-bridge = { path = "3rdparty/Megatron-Bridge-workspace/Megatron-Bridge", editable = true } nemo_gym = { workspace = true } nemo_run = { git = "https://github.com/NVIDIA-NeMo/Run", rev = "414f0077c648fde2c71bb1186e97ccbf96d6844c" } +# Must stay on the rev megatron-core pins for its dev/otel extras: the workspace +# resolves a single nemo-lens, and PyPI's newest release (0.1.0) is older than +# that rev, so declaring this dependency without the pin downgrades lens for +# megatron-core too. +nemo-lens = { git = "https://github.com/NVIDIA-NeMo/Lens.git", rev = "b85578fc2b736a1804705e537001b5f45e9c715d" } sglang = { git = "https://github.com/sgl-project/sglang.git", rev = "3003d70f680d41c59d1b7acbf65cb47795dfd19e", subdirectory = "python" } # torch/torchvision/triton all come from the torch index in order to pick up aarch64 wheels torch = [ diff --git a/pyrefly.toml b/pyrefly.toml index a412196f88..726e5e5ec2 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -235,6 +235,12 @@ project-includes = [ "nemo_rl/models/value/config.py", "nemo_rl/models/value/tq_value.py", "nemo_rl/models/value/workers/__init__.py", + "nemo_rl/telemetry/__init__.py", + "nemo_rl/telemetry/config.py", + "nemo_rl/telemetry/instrumentation.py", + "nemo_rl/telemetry/metrics.py", + "nemo_rl/telemetry/setup.py", + "nemo_rl/telemetry/span_groups.py", "nemo_rl/utils/__init__.py", "nemo_rl/utils/checkpoint.py", "nemo_rl/utils/checkpoint_engines/__init__.py", diff --git a/tests/unit/algorithms/test_async_utils.py b/tests/unit/algorithms/test_async_utils.py index 2783f223d2..b5a2ebb27c 100644 --- a/tests/unit/algorithms/test_async_utils.py +++ b/tests/unit/algorithms/test_async_utils.py @@ -2180,6 +2180,240 @@ def test_rollouts_state_roundtrips_pending_batch(self, tmp_path): assert status["data_exhausted"] is False assert status["errored"] is False + def test_flush_telemetry_waits_for_inflight_batches(self, monkeypatch): + """The flush must cover the batches it exists to save. + + Shutting the provider down is terminal, so a batch worker still running + when it happens loses its span -- exactly the last rollouts of the run + this call is meant to rescue. + """ + collector = self.create_local_collector() + collector.running = True + alive_at_shutdown = [] + monkeypatch.setattr( + "nemo_rl.algorithms.async_utils.trajectory_collector.shutdown_telemetry", + lambda: alive_at_shutdown.append(worker.is_alive()), + ) + + release = threading.Event() + worker = threading.Thread(target=lambda: release.wait(timeout=10), daemon=True) + collector._inflight_threads.add(worker) + collector._live_threads.add(worker) + worker.start() + threading.Timer(0.2, release.set).start() + + collector.flush_telemetry(quiesce_timeout_s=10.0) + + assert collector.running is False + assert alive_at_shutdown == [False] + + def test_flush_telemetry_gives_up_on_a_wedged_batch(self, monkeypatch): + # The caller is on its way to ray.kill, so the wait is bounded: losing a + # wedged worker's span beats hanging the run's teardown. + collector = self.create_local_collector() + collector.running = True + shutdown_calls = [] + monkeypatch.setattr( + "nemo_rl.algorithms.async_utils.trajectory_collector.shutdown_telemetry", + lambda: shutdown_calls.append(True), + ) + + release = threading.Event() + worker = threading.Thread(target=lambda: release.wait(timeout=30), daemon=True) + collector._inflight_threads.add(worker) + collector._live_threads.add(worker) + worker.start() + try: + started = time.monotonic() + collector.flush_telemetry(quiesce_timeout_s=0.3) + elapsed = time.monotonic() - started + finally: + release.set() + + assert shutdown_calls == [True] + # Bounded by the budget, not by the wedged worker's 30s: it must have + # waited (so the budget is respected) and then given up (so teardown + # cannot be held open). + assert 0.3 <= elapsed < 3.0 + + def test_flush_telemetry_waits_past_the_inflight_bookkeeping(self, monkeypatch): + """Leaving ``_inflight_threads`` does not mean the span is closed. + + A batch worker discards itself from the set inside its own ``finally``, + which still runs inside the ``rl.grpo.generation`` span. Watching the set + would let the flush start while that span is still open; joining the + thread waits for the span to end, because thread death follows it. + """ + collector = self.create_local_collector() + collector.running = True + span_closed = threading.Event() + span_closed_at_shutdown = [] + monkeypatch.setattr( + "nemo_rl.algorithms.async_utils.trajectory_collector.shutdown_telemetry", + lambda: span_closed_at_shutdown.append(span_closed.is_set()), + ) + + def _batch_worker(): + # Mirrors the real ordering: bookkeeping first, span close after. + with collector._threads_lock: + collector._inflight_threads.discard(threading.current_thread()) + time.sleep(0.3) + span_closed.set() + + worker = threading.Thread(target=_batch_worker, daemon=True) + collector._inflight_threads.add(worker) + collector._live_threads.add(worker) + worker.start() + + collector.flush_telemetry(quiesce_timeout_s=10.0) + + assert span_closed_at_shutdown == [True] + + def test_flush_telemetry_wakes_a_parked_collection_loop(self, monkeypatch): + """The loop's own waits hold open spans, so it has to be woken. + + Both collection-loop waits are ``Event.wait()`` calls inside an + ``efficiency_span``. Joining a loop parked in one of them would burn the + whole budget and then flush with that span still open. + """ + collector = self.create_local_collector() + collector.running = True + monkeypatch.setattr( + "nemo_rl.algorithms.async_utils.trajectory_collector.shutdown_telemetry", + lambda: None, + ) + + collector._refit_pause_cleared.clear() + loop = threading.Thread( + target=collector._refit_pause_cleared.wait, kwargs={"timeout": 30} + ) + collector.collection_thread = loop + loop.start() + + started = time.monotonic() + collector.flush_telemetry(quiesce_timeout_s=10.0) + elapsed = time.monotonic() - started + + assert not loop.is_alive() + assert elapsed < 3.0 + + def test_flush_telemetry_keeps_budget_for_the_batch_workers(self, monkeypatch): + """A wedged loop must not spend the whole budget. + + The loop and the batch workers are joined in sequence, and the workers + hold the rollout spans this flush exists to save -- so a loop that never + exits has to be capped rather than allowed to starve them. + """ + collector = self.create_local_collector() + collector.running = True + monkeypatch.setattr( + "nemo_rl.algorithms.async_utils.trajectory_collector.shutdown_telemetry", + lambda: None, + ) + + # Ignores every wake, so it holds the loop join for its whole cap. + wedged_loop = threading.Thread(target=lambda: time.sleep(30), daemon=True) + collector.collection_thread = wedged_loop + wedged_loop.start() + + joined = threading.Event() + worker = threading.Thread(target=joined.set, daemon=True) + collector._live_threads.add(worker) + worker.start() + + started = time.monotonic() + collector.flush_telemetry(quiesce_timeout_s=1.0) + elapsed = time.monotonic() - started + + # The worker was reached, so the drain loop still had budget left. + assert joined.is_set() + assert not worker.is_alive() + # Capped at half the budget on the loop, and the whole call stays inside + # it despite the loop never exiting. + assert 0.5 <= elapsed < 2.0 + + def test_flush_telemetry_survives_a_clear_that_races_the_wake(self, monkeypatch): + """The loop clears a pause event just after testing ``running``. + + A single wake landing in that window is swallowed by the clear, and + nothing else will ever set the event -- the driver is on its way to + ``ray.kill``, so no refit or weight update is coming. Re-arming on each + pass is what keeps that from wedging teardown for the full budget. + """ + collector = self.create_local_collector() + collector.running = True + monkeypatch.setattr( + "nemo_rl.algorithms.async_utils.trajectory_collector.shutdown_telemetry", + lambda: None, + ) + + entered = threading.Event() + + def _loop_body(): + # The real ordering: check running, then clear, then wait. The sleep + # widens the window so the flush's first wake lands inside it. + entered.set() + if collector.running: + time.sleep(0.2) + collector._generation_limit_cleared.clear() + collector._generation_limit_cleared.wait(timeout=30) + + loop = threading.Thread(target=_loop_body, daemon=True) + collector.collection_thread = loop + loop.start() + entered.wait(timeout=5) + + started = time.monotonic() + collector.flush_telemetry(quiesce_timeout_s=10.0) + elapsed = time.monotonic() - started + + assert not loop.is_alive() + assert elapsed < 3.0 + + def test_flush_telemetry_waits_for_a_worker_that_has_not_started(self, monkeypatch): + """A registered thread with no ``ident`` yet is about to open a span. + + Spawning registers the worker under the lock and starts it just after, + so ``is_alive()`` is briefly false for a thread that is about to run. + Filtering on it alone would neither wait for that worker nor warn about + it. + """ + collector = self.create_local_collector() + collector.running = True + shutdown_calls = [] + monkeypatch.setattr( + "nemo_rl.algorithms.async_utils.trajectory_collector.shutdown_telemetry", + lambda: shutdown_calls.append(True), + ) + + unstarted = threading.Thread(target=lambda: None, daemon=True) + collector._live_threads.add(unstarted) + assert not unstarted.is_alive() + + threading.Timer(0.2, unstarted.start).start() + collector.flush_telemetry(quiesce_timeout_s=10.0) + + # Waited for the late start and then for the thread itself, rather than + # treating "not alive yet" as "already done". + assert unstarted.ident is not None + assert not unstarted.is_alive() + assert shutdown_calls == [True] + + def test_cleanup_prunes_dead_live_threads(self): + # _live_threads outlives the batch accounting on purpose, so cleanup is + # the only thing keeping a long run from holding every Thread it spawned. + collector = self.create_local_collector() + dead = threading.Thread(target=lambda: None) + dead.start() + dead.join() + collector._live_threads.add(dead) + collector._inflight_threads.add(dead) + + collector._cleanup_finished_threads() + + assert collector._live_threads == set() + assert collector._inflight_threads == set() + def create_mock_config(self) -> MasterConfig: """Create a mock master config for testing.""" return MasterConfig.model_construct( @@ -2682,9 +2916,12 @@ def __init__(self): started = [] class RecordingThread: - def __init__(self, *, target, daemon): + def __init__(self, *, target, daemon, name): assert daemon + # Named so teardown can say which thread it is still waiting on. + assert name self.target = target + self.name = name def start(self): started.append(self) @@ -2736,9 +2973,12 @@ class FakeReplayBuffer: started_threads = [] class RecordingThread: - def __init__(self, *, target, daemon): + def __init__(self, *, target, daemon, name): assert daemon + # Named so teardown can say which thread it is still waiting on. + assert name self.target = target + self.name = name def start(self): started_threads.append(self) @@ -2802,9 +3042,12 @@ class FakeReplayBuffer: started_threads = [] class RecordingThread: - def __init__(self, *, target, daemon): + def __init__(self, *, target, daemon, name): assert daemon + # Named so teardown can say which thread it is still waiting on. + assert name self.target = target + self.name = name def start(self): started_threads.append(self) diff --git a/tests/unit/algorithms/test_utils.py b/tests/unit/algorithms/test_utils.py index bc30730a92..559d2f5c13 100755 --- a/tests/unit/algorithms/test_utils.py +++ b/tests/unit/algorithms/test_utils.py @@ -24,6 +24,7 @@ from nemo_rl.algorithms.ppo import PPOConfig from nemo_rl.algorithms.utils import ( EFFICIENCY_CATEGORIES, + STEP_WINDOW_WALL_CLOCK_CATEGORIES, WALL_CLOCK_EFFICIENCY_CATEGORIES, calculate_baseline_and_std_per_prompt, get_tokenizer, @@ -798,9 +799,11 @@ def test_basic_efficiency_calculation(self, capsys): result = print_efficiency_summary(metrics, total_wall, step=1) - assert result["efficiency/total_waste_s"] == 20.0 - assert result["efficiency/productive_time_s"] == 80.0 - assert result["efficiency/efficiency_pct"] == pytest.approx(80.0) + # init/total is not waste for this step: it is a run-long constant, so + # only the three per-step idle categories count (5 + 3 + 2). + assert result["efficiency/total_waste_s"] == 10.0 + assert result["efficiency/productive_time_s"] == 90.0 + assert result["efficiency/efficiency_pct"] == pytest.approx(90.0) assert result["efficiency/total_wall_time_s"] == 100.0 assert result["efficiency/init/total_s"] == 10.0 @@ -812,7 +815,51 @@ def test_basic_efficiency_calculation(self, capsys): captured = capsys.readouterr() assert "Efficiency Summary (Step 1)" in captured.out - assert "80.00%" in captured.out + assert "90.00%" in captured.out + + def test_startup_cost_is_not_charged_to_every_step(self): + """init/total is republished every step, so counting it would recur. + + The driver measures it once before the loop and re-supplies the same + value on every step so the series does not zero out. Folding it into the + waste aggregate would subtract the whole startup cost from every step. + """ + idle_only = {"idle/refit_bubble": 4.0} + with_startup = {**idle_only, "init/total": 300.0} + + assert print_efficiency_summary( + with_startup, total_wall_time_s=1000.0, step=9, step_wall_time_s=20.0 + )["efficiency/efficiency_pct"] == pytest.approx( + print_efficiency_summary( + idle_only, total_wall_time_s=1000.0, step=9, step_wall_time_s=20.0 + )["efficiency/efficiency_pct"] + ) + + def test_efficiency_is_a_per_step_ratio(self): + """The numerator is per-step, so the denominator has to be too. + + Against the run's elapsed time, a fixed per-step idle cost would look + like it was shrinking: the same 5s of idle in a 20s step is 25% waste at + step 2 and 25% at step 500, but 5/1000 at step 500 if divided by the run. + """ + metrics = {"idle/buffer_starvation": 5.0} + + result = print_efficiency_summary( + metrics, total_wall_time_s=1000.0, step=50, step_wall_time_s=20.0 + ) + + assert result["efficiency/efficiency_pct"] == pytest.approx(75.0) + # The per-category column stays a share of the run, which is what makes + # a cumulative denominator the right one there. + assert result["efficiency/idle/buffer_starvation_pct"] == pytest.approx(0.5) + + def test_step_wall_time_defaults_to_the_run(self): + """Callers with no per-step measurement keep the old denominator.""" + metrics = {"idle/refit_bubble": 25.0} + + result = print_efficiency_summary(metrics, total_wall_time_s=100.0, step=1) + + assert result["efficiency/efficiency_pct"] == pytest.approx(75.0) def test_zero_wall_time(self): """Test that zero wall time produces 100% efficiency.""" @@ -831,10 +878,11 @@ def test_all_categories_present(self): if cat in WALL_CLOCK_EFFICIENCY_CATEGORIES: assert f"efficiency/{cat}_pct" in result - # total_waste_s reflects wall-clock categories only (4 x 1.0); collector - # thread-seconds are reported separately, not folded into wall waste. - assert result["efficiency/total_waste_s"] == 4.0 - assert result["efficiency/efficiency_pct"] == pytest.approx(96.0) + # total_waste_s reflects the per-step wall-clock categories only + # (3 x 1.0): collector thread-seconds are reported separately rather + # than folded into wall waste, and init/total is a run constant. + assert result["efficiency/total_waste_s"] == 3.0 + assert result["efficiency/efficiency_pct"] == pytest.approx(97.0) def test_efficiency_with_collector_metrics_merge(self): """Test merging driver and collector metrics before computing efficiency.""" @@ -847,13 +895,16 @@ def test_efficiency_with_collector_metrics_merge(self): result = print_efficiency_summary(merged, total_wall_time_s=50.0, step=3) - assert result["efficiency/total_waste_s"] == 8.0 + # Only the driver's per-step idle counts: refit_bubble (3.0). The + # collector's two categories are thread-seconds, init/total is a run + # constant. + assert result["efficiency/total_waste_s"] == 3.0 assert result["efficiency/thread_seconds_total_s"] == 3.0 - assert result["efficiency/efficiency_pct"] == pytest.approx(84.0) + assert result["efficiency/efficiency_pct"] == pytest.approx(94.0) def test_wall_waste_clamped_to_wall_time(self): """Wall-clock waste is capped so efficiency stays in [0, 100].""" - metrics = {cat: 30.0 for cat in WALL_CLOCK_EFFICIENCY_CATEGORIES} + metrics = {cat: 30.0 for cat in STEP_WINDOW_WALL_CLOCK_CATEGORIES} result = print_efficiency_summary(metrics, total_wall_time_s=60.0, step=2) assert result["efficiency/total_waste_s"] == 60.0 diff --git a/tests/unit/reference_configs/distillation_math.yaml b/tests/unit/reference_configs/distillation_math.yaml index 3eec82229b..801b553e2f 100644 --- a/tests/unit/reference_configs/distillation_math.yaml +++ b/tests/unit/reference_configs/distillation_math.yaml @@ -288,3 +288,7 @@ cluster: master_port_range_low: 1400 master_port_range_high: 1999 segment_size: null + +# OpenTelemetry instrumentation; null = disabled (default). Mirrors the optional +# field on the MasterConfig, which stays None unless a run opts in. +telemetry: null diff --git a/tests/unit/reference_configs/dpo.yaml b/tests/unit/reference_configs/dpo.yaml index 3f9bd65b83..c329b8c361 100755 --- a/tests/unit/reference_configs/dpo.yaml +++ b/tests/unit/reference_configs/dpo.yaml @@ -309,3 +309,7 @@ cluster: gpus_per_node: 1 num_nodes: 1 segment_size: null + +# OpenTelemetry instrumentation; null = disabled (default). Mirrors the optional +# field on the MasterConfig, which stays None unless a run opts in. +telemetry: null diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index c51186f9ae..968745b675 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -539,3 +539,7 @@ data_plane: # Multi-Teacher On-Policy Distillation (MOPD); null = disabled (default). Mirrors # the field on the GRPO MasterConfig added for MOPD support. on_policy_distillation: null + +# OpenTelemetry instrumentation; null = disabled (default). Mirrors the optional +# field on the MasterConfig, which stays None unless a run opts in. +telemetry: null diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 71b3bff39a..706fde4c3d 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -454,3 +454,7 @@ cluster: master_port_range_low: 1400 master_port_range_high: 1999 segment_size: null + +# OpenTelemetry instrumentation; null = disabled (default). Mirrors the optional +# field on the MasterConfig, which stays None unless a run opts in. +telemetry: null diff --git a/tests/unit/reference_configs/rm.yaml b/tests/unit/reference_configs/rm.yaml index 4b931c8f7d..fad467ad90 100644 --- a/tests/unit/reference_configs/rm.yaml +++ b/tests/unit/reference_configs/rm.yaml @@ -226,3 +226,7 @@ cluster: gpus_per_node: 1 num_nodes: 1 segment_size: null + +# OpenTelemetry instrumentation; null = disabled (default). Mirrors the optional +# field on the MasterConfig, which stays None unless a run opts in. +telemetry: null diff --git a/tests/unit/reference_configs/sft.yaml b/tests/unit/reference_configs/sft.yaml index 3924de619e..ec02bdc4ea 100644 --- a/tests/unit/reference_configs/sft.yaml +++ b/tests/unit/reference_configs/sft.yaml @@ -289,3 +289,7 @@ cluster: gpus_per_node: 1 num_nodes: 1 segment_size: null + +# OpenTelemetry instrumentation; null = disabled (default). Mirrors the optional +# field on the MasterConfig, which stays None unless a run opts in. +telemetry: null diff --git a/tests/unit/telemetry/__init__.py b/tests/unit/telemetry/__init__.py new file mode 100644 index 0000000000..52a7a9daf0 --- /dev/null +++ b/tests/unit/telemetry/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unit/telemetry/conftest.py b/tests/unit/telemetry/conftest.py new file mode 100644 index 0000000000..d3fc9cefc7 --- /dev/null +++ b/tests/unit/telemetry/conftest.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fixtures for NeMo-RL telemetry unit tests. + +Resets global OpenTelemetry providers, the nemo-lens init guard, span-group +state, the process-global telemetry handle, and the ``NEMO_RL_OTEL_*`` / +``OTEL_SERVICE_NAME`` / ``NRL_WORKER_GROUP`` env vars before and after each test +so nothing leaks. +""" + +import ast +import os +from pathlib import Path + +import pytest + + +def string_constants(node: ast.AST | None) -> set[str]: + """The string literals in a list/set/tuple display, ignoring anything else.""" + return { + elt.value + for elt in getattr(node, "elts", []) + if isinstance(elt, ast.Constant) and isinstance(elt.value, str) + } + + +def algorithms_utils_categories(*names: str) -> dict[str, set[str]]: + """Read named category collections out of ``nemo_rl/algorithms/utils.py``. + + That module owns the canonical efficiency-category lists but imports torch, + which is far too heavy for this suite — and several telemetry constants + deliberately restate those lists so they stay importable without the + training stack. Parsing the source is what keeps the copies honest. + + Handles list, set, and ``frozenset({...})`` literals; string elements only. + """ + source = Path(__file__).resolve().parents[3] / "nemo_rl" / "algorithms" / "utils.py" + wanted = set(names) + found: dict[str, set[str]] = {} + for node in ast.parse(source.read_text()).body: + if not isinstance(node, ast.AnnAssign | ast.Assign): + continue + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + for target in targets: + if not isinstance(target, ast.Name) or target.id not in wanted: + continue + value = node.value + # frozenset({...}) / set([...]) wrap the literal in a call. + if isinstance(value, ast.Call) and value.args: + value = value.args[0] + found[target.id] = string_constants(value) + assert wanted == set(found), f"could not locate {wanted - set(found)} in {source}" + return found + + +def _clear_telemetry_env() -> None: + for key in list(os.environ): + if key.startswith(("NEMO_RL_OTEL", "NEMO_LENS")) or key in ( + "OTEL_SERVICE_NAME", + "NRL_WORKER_GROUP", + ): + del os.environ[key] + + +def _reset_rl_telemetry() -> None: + import nemo_rl.telemetry.setup as setup_mod + + setup_mod._TELEMETRY_HANDLE = None + setup_mod._TELEMETRY_INITIALISED = False + + +def _reset_otel_and_lens() -> None: + # Before the providers are dropped: several tests build real console + # providers, and simply nulling the globals would leave their exporter + # threads and periodic metric readers running for the rest of the session. + try: + from nemo_rl.telemetry.setup import shutdown_telemetry + + shutdown_telemetry() + except Exception: + pass + try: + import opentelemetry.metrics._internal as _metrics_mod + import opentelemetry.trace as _trace_mod + from opentelemetry.util._once import Once + + _trace_mod._TRACER_PROVIDER = None + _trace_mod._TRACER_PROVIDER_SET_ONCE = Once() + _metrics_mod._METER_PROVIDER = None + _metrics_mod._METER_PROVIDER_SET_ONCE = Once() + except Exception: + pass + import nemo.lens.handle as _handle_mod + from nemo.lens.state import set_enabled_span_groups + + _handle_mod._INITIALIZED = False + set_enabled_span_groups(frozenset()) + + +@pytest.fixture(autouse=True) +def _reset_telemetry_state(): + _clear_telemetry_env() + _reset_otel_and_lens() + _reset_rl_telemetry() + yield + _reset_otel_and_lens() + _reset_rl_telemetry() + _clear_telemetry_env() diff --git a/tests/unit/telemetry/test_config.py b/tests/unit/telemetry/test_config.py new file mode 100644 index 0000000000..af15a5db08 --- /dev/null +++ b/tests/unit/telemetry/test_config.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for TelemetryConfig and the YAML->env translation.""" + +import os + +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.setup import _ENV_FIELD_MAP, _config_to_env + + +def test_defaults(): + cfg = TelemetryConfig() + assert cfg.enabled is False + assert cfg.service_name == "nemo-rl" + assert cfg.span_groups == "default" + assert cfg.export_strategy == "single_rank" + assert cfg.export_rank == -1 + # Must match the nemo-lens defaults, since propagating them is unconditional. + assert cfg.export_sample_rate == 1.0 + assert cfg.sampler_enabled is False + assert cfg.traces_enabled is True + assert cfg.metrics_enabled is True + assert cfg.logs_enabled is False + assert cfg.exporter == "otlp" + assert cfg.vllm_native_tracing is False + + +def test_extra_keys_allowed(): + cfg = TelemetryConfig(enabled=True, future_unknown_key="x") + assert cfg.enabled is True + + +def test_config_to_env_translation(): + tel = TelemetryConfig( + enabled=True, + span_groups="per_step", + export_rank=3, + export_strategy="sampled", + export_sample_rate=0.25, + sampler_enabled=True, + vllm_native_tracing=True, + exporter="console", + service_name="my-rl", + ) + _config_to_env(tel) + assert os.environ["NEMO_RL_OTEL_ENABLED"] == "1" + assert os.environ["NEMO_RL_OTEL_SPAN_GROUPS"] == "per_step" + assert os.environ["NEMO_RL_OTEL_EXPORT_RANK"] == "3" + assert os.environ["NEMO_RL_OTEL_EXPORT_STRATEGY"] == "sampled" + assert os.environ["NEMO_RL_OTEL_EXPORT_SAMPLE_RATE"] == "0.25" + assert os.environ["NEMO_RL_OTEL_SAMPLER_ENABLED"] == "1" + assert os.environ["NEMO_RL_OTEL_VLLM_NATIVE_TRACING"] == "1" + assert os.environ["NEMO_RL_OTEL_EXPORTER"] == "console" + assert os.environ["OTEL_SERVICE_NAME"] == "my-rl" + + +def test_disabled_translates_to_zero(): + _config_to_env(TelemetryConfig(enabled=False)) + assert os.environ["NEMO_RL_OTEL_ENABLED"] == "0" + + +def test_env_wins_over_yaml(): + os.environ["NEMO_RL_OTEL_SPAN_GROUPS"] = "all" + _config_to_env(TelemetryConfig(enabled=True, span_groups="per_step")) + # setdefault must not overwrite a pre-existing env var. + assert os.environ["NEMO_RL_OTEL_SPAN_GROUPS"] == "all" + + +def test_env_field_map_fields_exist_on_config(): + cfg = TelemetryConfig() + for field in _ENV_FIELD_MAP: + assert hasattr(cfg, field), field + + +def test_every_config_field_propagates_to_workers(): + # A field that exists on the config but is missing from _ENV_FIELD_MAP is + # settable in YAML yet never reaches a worker. service_name is the one + # exception: it maps onto the unprefixed OTEL_SERVICE_NAME. + declared = set(TelemetryConfig.model_fields) - {"service_name"} + assert declared == set(_ENV_FIELD_MAP) diff --git a/tests/unit/telemetry/test_instrumentation.py b/tests/unit/telemetry/test_instrumentation.py new file mode 100644 index 0000000000..bbcb6dea7c --- /dev/null +++ b/tests/unit/telemetry/test_instrumentation.py @@ -0,0 +1,423 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for ``nemo_rl.telemetry.instrumentation``. + +Two layers: + +* Goodput bucket classification (``Bucket`` / ``bucket_for_span_group`` / + ``goodput_span_attributes`` / efficiency-category mapping) — pure functions, + no nemo-lens required. +* End-to-end span emission via the ``managed_span`` / ``trace_fn`` primitives + the algorithm loops use, asserting spans emit per group, gate off when the + group is disabled, carry ``rl.bucket`` on leaf groups, and nest correctly — + requires nemo-lens. +""" + +import pytest + +from nemo_rl.telemetry.instrumentation import ( + EFFICIENCY_CATEGORY_BUCKET, + RL_BUCKET_ATTR, + RL_EFFICIENCY_CATEGORY_ATTR, + UMBRELLA_GROUPS, + Bucket, + bucket_for_efficiency_category, + bucket_for_span_group, + bucket_scope, + current_trace_carrier, + efficiency_span, + goodput_span_attributes, + managed_span, + remote_trace_context, + trace_fn, +) +from nemo_rl.telemetry.span_groups import RLSpanGroup + +try: + from nemo.lens import NemoLensConfig, setup_telemetry + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + _HAS_LENS = True +except ImportError: + _HAS_LENS = False + +requires_lens = pytest.mark.skipif( + not _HAS_LENS, reason="nemo-lens (+ opentelemetry sdk) not installed" +) + + +# --------------------------------------------------------------------------- # +# Goodput bucket classification (pure functions) # +# --------------------------------------------------------------------------- # +def test_shared_bucket_tokens(): + assert {b.value for b in Bucket} == { + "productive", + "overhead", + "idle", + "wasted", + } + + +def test_umbrellas_have_no_bucket(): + for group in ( + RLSpanGroup.JOB, + RLSpanGroup.STEP, + RLSpanGroup.ROLLOUT, + RLSpanGroup.MODEL_INIT, + RLSpanGroup.EVALUATE, + ): + assert group in UMBRELLA_GROUPS + assert bucket_for_span_group(group) is None + assert goodput_span_attributes(group) == {} + + +def test_leaf_groups_map_to_expected_buckets(): + assert bucket_for_span_group(RLSpanGroup.GENERATION) is Bucket.PRODUCTIVE + assert bucket_for_span_group(RLSpanGroup.REWARD) is Bucket.PRODUCTIVE + assert bucket_for_span_group(RLSpanGroup.POLICY_UPDATE) is Bucket.PRODUCTIVE + assert bucket_for_span_group(RLSpanGroup.DATA_PROCESSING) is Bucket.OVERHEAD + assert bucket_for_span_group(RLSpanGroup.CHECKPOINT) is Bucket.OVERHEAD + assert bucket_for_span_group(RLSpanGroup.LOGPROB) is Bucket.OVERHEAD + assert bucket_for_span_group(RLSpanGroup.ADVANTAGE) is Bucket.OVERHEAD + assert bucket_for_span_group(RLSpanGroup.REFERENCE_POLICY) is Bucket.OVERHEAD + + +def test_goodput_span_attributes_shape(): + attrs = goodput_span_attributes(RLSpanGroup.GENERATION) + assert attrs == {RL_BUCKET_ATTR: "productive"} + + +def test_unknown_non_umbrella_defaults_to_overhead(): + assert bucket_for_span_group("brand_new_leaf") is Bucket.OVERHEAD + assert goodput_span_attributes("brand_new_leaf")[RL_BUCKET_ATTR] == "overhead" + + +def test_efficiency_categories_mapped(): + assert bucket_for_efficiency_category("idle/buffer_starvation") is Bucket.IDLE + assert bucket_for_efficiency_category("wasted/failed_trajectory") is Bucket.WASTED + assert bucket_for_efficiency_category("init/total") is Bucket.OVERHEAD + assert set(EFFICIENCY_CATEGORY_BUCKET) # non-empty + + +def _efficiency_categories_from_algorithms_utils() -> set[str]: + """Every category async GRPO records, read from the canonical source.""" + from tests.unit.telemetry.conftest import algorithms_utils_categories + + found = algorithms_utils_categories( + "WALL_CLOCK_EFFICIENCY_CATEGORIES", + "THREAD_ACCUMULATED_EFFICIENCY_CATEGORIES", + ) + return set().union(*found.values()) + + +def test_efficiency_category_bucket_matches_production_categories(): + """EFFICIENCY_CATEGORY_BUCKET duplicates the category strings that async + GRPO actually records, so a new ``idle/*`` timer must not silently land + without a bucket. + """ + assert set(EFFICIENCY_CATEGORY_BUCKET) == ( + _efficiency_categories_from_algorithms_utils() + ) + + +@requires_lens +def test_efficiency_span_carries_bucket_and_category(): + handle, exporter = _setup("all") + with efficiency_span("idle/refit_bubble", tracer=handle.tracer) as span: + assert span is not None + handle.shutdown() + + (emitted,) = exporter.get_finished_spans() + assert emitted.name == "rl.idle.refit_bubble" + assert emitted.attributes[RL_BUCKET_ATTR] == "idle" + assert emitted.attributes[RL_EFFICIENCY_CATEGORY_ATTR] == "idle/refit_bubble" + + +@requires_lens +def test_efficiency_span_wasted_category_is_not_tagged_idle(): + handle, exporter = _setup("all") + with efficiency_span("wasted/failed_trajectory", tracer=handle.tracer): + pass + handle.shutdown() + + (emitted,) = exporter.get_finished_spans() + assert emitted.attributes[RL_BUCKET_ATTR] == "wasted" + + +@requires_lens +def test_trace_carrier_reparents_spans_in_another_context(): + """A carrier moves the trace across a process boundary Ray does not. + + Simulates the collector: capture inside the driver's job span, then reopen + it where no span is active and check the child joins the same trace. + """ + handle, exporter = _setup("all") + with managed_span(RLSpanGroup.JOB, "rl.grpo.job", tracer=handle.tracer) as parent: + carrier = current_trace_carrier() + parent_ctx = parent.get_span_context() + + # Outside the parent block: nothing is active, so this would be a root. + with remote_trace_context(carrier): + with managed_span( + RLSpanGroup.ROLLOUT, "rl.grpo.generation", tracer=handle.tracer + ): + pass + handle.shutdown() + + child = next( + span + for span in exporter.get_finished_spans() + if span.name == "rl.grpo.generation" + ) + assert child.parent is not None + assert child.parent.span_id == parent_ctx.span_id + assert child.context.trace_id == parent_ctx.trace_id + + +@requires_lens +def test_span_is_a_root_without_a_carrier(): + """No job span (the per_step case) must degrade to a root, not an error.""" + handle, exporter = _setup("all") + carrier = current_trace_carrier() + assert carrier == {} + with remote_trace_context(carrier): + with managed_span( + RLSpanGroup.ROLLOUT, "rl.grpo.generation", tracer=handle.tracer + ): + pass + handle.shutdown() + + (emitted,) = exporter.get_finished_spans() + assert emitted.parent is None + + +@requires_lens +def test_collector_efficiency_spans_are_trace_only(): + """Collector-thread waits are visible in a trace but absent from rollups. + + They are timed concurrently with the driver's timeline, so a bucket would + be summed against a wall-clock denominator it does not belong to. Omitting + the attribute excludes them by construction instead of by convention. + """ + handle, exporter = _setup("all") + with efficiency_span("idle/generation_limit_pause", tracer=handle.tracer): + pass + with efficiency_span("idle/refit_event_wait", tracer=handle.tracer): + pass + handle.shutdown() + + emitted = exporter.get_finished_spans() + assert [span.name for span in emitted] == [ + "rl.idle.generation_limit_pause", + "rl.idle.refit_event_wait", + ] + for span in emitted: + assert RL_BUCKET_ATTR not in span.attributes + # The category still identifies the phase, and still says "idle/…", + # so the span is greppable without being summable. + assert span.attributes[RL_EFFICIENCY_CATEGORY_ATTR].startswith("idle/") + + +@requires_lens +def test_efficiency_span_is_gated_by_span_group(): + # The efficiency group is absent from the coarse "default" preset, so idle + # spans must not appear there. + handle, exporter = _setup("default") + with efficiency_span("idle/buffer_starvation", tracer=handle.tracer): + pass + handle.shutdown() + assert exporter.get_finished_spans() == () + + +def test_efficiency_group_is_in_per_step_preset(): + # Per-step goodput only adds up if idle is included alongside the phases. + assert RLSpanGroup.EFFICIENCY in RLSpanGroup._PRESETS["per_step"] + assert RLSpanGroup.EFFICIENCY in RLSpanGroup.ALL_GROUPS + assert RLSpanGroup.EFFICIENCY not in RLSpanGroup._PRESETS["default"] + + +def test_bucket_scope_replaces_leaf_bucket_and_restores_it(): + assert goodput_span_attributes(RLSpanGroup.GENERATION) == { + RL_BUCKET_ATTR: "productive" + } + with bucket_scope(Bucket.OVERHEAD): + assert goodput_span_attributes(RLSpanGroup.GENERATION) == { + RL_BUCKET_ATTR: "overhead" + } + assert goodput_span_attributes(RLSpanGroup.GENERATION) == { + RL_BUCKET_ATTR: "productive" + } + + +def test_bucket_scope_leaves_umbrellas_unbucketed(): + """An override must not start bucketing umbrellas. + + ``rl..evaluate`` encloses the generate spans it reclassifies, so + tagging it too would count the same interval twice. + """ + with bucket_scope(Bucket.OVERHEAD): + assert goodput_span_attributes(RLSpanGroup.EVALUATE) == {} + + +def test_every_rl_span_group_is_classified(): + """Every known RLSpanGroup is either umbrella or has an explicit/default bucket.""" + for group in RLSpanGroup.ALL_GROUPS: + bucket = bucket_for_span_group(group) + if group in UMBRELLA_GROUPS: + assert bucket is None, group + else: + assert bucket in Bucket, group + + +# --------------------------------------------------------------------------- # +# Span emission via managed_span / trace_fn (in-memory exporter) # +# --------------------------------------------------------------------------- # +def _setup(groups): + exporter = InMemorySpanExporter() + cfg = NemoLensConfig(enabled=True, span_groups=groups, _span_group_cls=RLSpanGroup) + handle = setup_telemetry(cfg, rank=0, world_size=1, span_exporter=exporter) + return handle, exporter + + +@requires_lens +def test_managed_span_emits_when_group_enabled(): + handle, exporter = _setup("generation") + with managed_span( + RLSpanGroup.GENERATION, + "rl.vllm.generate", + tracer=handle.tracer, + **{"rl.backend": "vllm"}, + ) as span: + assert span is not None + handle.shutdown() + spans = exporter.get_finished_spans() + assert [s.name for s in spans] == ["rl.vllm.generate"] + assert spans[0].attributes["rl.backend"] == "vllm" + # Leaf groups carry rl.bucket for offline goodput rollup. + assert spans[0].attributes[RL_BUCKET_ATTR] == "productive" + + +@requires_lens +def test_umbrella_span_has_no_bucket(): + handle, exporter = _setup("all") + with managed_span(RLSpanGroup.STEP, "rl.grpo.step", tracer=handle.tracer) as span: + assert span is not None + handle.shutdown() + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert RL_BUCKET_ATTR not in spans[0].attributes + + +@requires_lens +def test_managed_span_noop_when_group_disabled(): + # "generation" is not part of the "default" preset. + handle, exporter = _setup("default") + with managed_span( + RLSpanGroup.GENERATION, "rl.vllm.generate", tracer=handle.tracer + ) as span: + assert span is None + handle.shutdown() + assert len(exporter.get_finished_spans()) == 0 + + +@requires_lens +def test_trace_fn_job_span(): + handle, exporter = _setup("all") + + @trace_fn(RLSpanGroup.JOB, "rl.grpo.job") + def train(): + return 42 + + assert train() == 42 + handle.shutdown() + assert any(s.name == "rl.grpo.job" for s in exporter.get_finished_spans()) + + +@requires_lens +def test_validation_generation_is_overhead_not_productive(): + """Mirror of ``validate()``: evaluate umbrella + a generate span inside it. + + The generate span is opened by a decorator on ``VllmGeneration.generate`` + that cannot see whether the caller is a training rollout or validation, so + the reclassification has to come from the enclosing scope. + """ + handle, exporter = _setup("all") + + @trace_fn(RLSpanGroup.GENERATION, "rl.vllm.generate", tracer=handle.tracer) + def generate(): + return "tokens" + + with ( + managed_span(RLSpanGroup.EVALUATE, "rl.grpo.evaluate", tracer=handle.tracer), + bucket_scope(Bucket.OVERHEAD), + ): + generate() + generate() # a training rollout, outside the scope + handle.shutdown() + + spans = exporter.get_finished_spans() + evaluate = next(s for s in spans if s.name == "rl.grpo.evaluate") + validation_gen, train_gen = (s for s in spans if s.name == "rl.vllm.generate") + assert RL_BUCKET_ATTR not in evaluate.attributes + assert validation_gen.attributes[RL_BUCKET_ATTR] == "overhead" + assert train_gen.attributes[RL_BUCKET_ATTR] == "productive" + + +@requires_lens +def test_bucket_scope_reaches_generation_under_asyncio_run(): + """``run_multi_turn_rollout`` drives generation through ``asyncio.run``. + + A ``ContextVar`` survives that (the task copies the current context), which + is what makes the scope usable from the synchronous ``validate()``. + """ + import asyncio + + handle, exporter = _setup("all") + + @trace_fn(RLSpanGroup.GENERATION, "rl.vllm.generate", tracer=handle.tracer) + def generate(): + return "tokens" + + async def rollout(): + generate() + + with bucket_scope(Bucket.OVERHEAD): + asyncio.run(rollout()) + handle.shutdown() + + (emitted,) = exporter.get_finished_spans() + assert emitted.attributes[RL_BUCKET_ATTR] == "overhead" + + +@requires_lens +def test_explicit_bucket_wins_over_scope(): + handle, exporter = _setup("all") + with bucket_scope(Bucket.OVERHEAD): + with managed_span( + RLSpanGroup.GENERATION, + "rl.vllm.generate", + tracer=handle.tracer, + **{RL_BUCKET_ATTR: "wasted"}, + ): + pass + handle.shutdown() + + (emitted,) = exporter.get_finished_spans() + assert emitted.attributes[RL_BUCKET_ATTR] == "wasted" + + +@requires_lens +def test_step_nests_under_job(): + handle, exporter = _setup("all") + with managed_span(RLSpanGroup.JOB, "rl.grpo.job", tracer=handle.tracer): + with managed_span(RLSpanGroup.STEP, "rl.grpo.step", tracer=handle.tracer): + pass + handle.shutdown() + spans = {s.name: s for s in exporter.get_finished_spans()} + assert "rl.grpo.job" in spans and "rl.grpo.step" in spans + step, job = spans["rl.grpo.step"], spans["rl.grpo.job"] + assert step.parent is not None + assert step.parent.span_id == job.context.span_id diff --git a/tests/unit/telemetry/test_metrics.py b/tests/unit/telemetry/test_metrics.py new file mode 100644 index 0000000000..e60efc3f82 --- /dev/null +++ b/tests/unit/telemetry/test_metrics.py @@ -0,0 +1,465 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the async efficiency metrics -> OTel tee.""" + +import logging + +import pytest + +from nemo_rl.telemetry.metrics import map_efficiency_seconds + + +def test_tee_noop_when_not_exporting(): + # No telemetry handle set -> must be a silent no-op (no exception). + from nemo_rl.telemetry.metrics import tee_rl_metrics_to_otel + + tee_rl_metrics_to_otel({"efficiency/idle/refit_bubble_s": 12.0}, "train") + + +_MEASUREMENTS = { + "idle/refit_bubble": "wall_clock", + "idle/buffer_full_backoff": "thread_seconds", +} + + +def test_map_efficiency_seconds_reads_per_category_keys(): + seconds = map_efficiency_seconds( + { + "efficiency/idle/refit_bubble_s": 12.0, + "efficiency/idle/buffer_full_backoff_s": 80, + }, + _MEASUREMENTS, + ) + assert seconds == {"idle/refit_bubble": 12.0, "idle/buffer_full_backoff": 80.0} + + +def test_map_efficiency_seconds_ignores_aggregate_keys(): + # These share the efficiency/ prefix and _s suffix but are not categories, + # so a prefix/suffix parse would invent bogus category series from them. + seconds = map_efficiency_seconds( + { + "efficiency/total_waste_s": 30.0, + "efficiency/productive_time_s": 70.0, + "efficiency/total_wall_time_s": 100.0, + "efficiency/thread_seconds_total_s": 400.0, + }, + _MEASUREMENTS, + ) + assert seconds == {} + + +def test_map_efficiency_seconds_skips_bool_and_non_numeric(): + seconds = map_efficiency_seconds( + { + "efficiency/idle/refit_bubble_s": True, + "efficiency/idle/buffer_full_backoff_s": "nan", + }, + _MEASUREMENTS, + ) + assert seconds == {} + + +def test_efficiency_measurements_classifies_every_category(): + # Drift guard: a category added to algorithms/utils.py without a + # classification would silently drop out of the OTel series. + pytest.importorskip("nemo_rl.algorithms.utils") + from nemo_rl.algorithms.utils import EFFICIENCY_CATEGORIES + from nemo_rl.telemetry.metrics import ( + COLLECTOR_WALL_CLOCK_MEASUREMENT, + THREAD_SECONDS_MEASUREMENT, + WALL_CLOCK_MEASUREMENT, + efficiency_measurements, + ) + + measurements = efficiency_measurements() + assert set(measurements) == set(EFFICIENCY_CATEGORIES) + assert set(measurements.values()) <= { + WALL_CLOCK_MEASUREMENT, + COLLECTOR_WALL_CLOCK_MEASUREMENT, + THREAD_SECONDS_MEASUREMENT, + } + + +def test_collector_loop_waits_are_not_labelled_thread_seconds(): + """The two loop waits are sequential, so they are real durations. + + They sit in ``THREAD_ACCUMULATED_EFFICIENCY_CATEGORIES`` because the W&B + summary only splits driver from collector, but labelling them + ``thread_seconds`` would tell a consumer they can exceed wall time -- which + is untrue of a single-threaded ``Event.wait()``. + """ + pytest.importorskip("nemo_rl.algorithms.utils") + from nemo_rl.telemetry.metrics import ( + COLLECTOR_WALL_CLOCK_MEASUREMENT, + THREAD_SECONDS_MEASUREMENT, + efficiency_measurements, + ) + + measurements = efficiency_measurements() + assert measurements["idle/refit_event_wait"] == COLLECTOR_WALL_CLOCK_MEASUREMENT + assert ( + measurements["idle/generation_limit_pause"] == COLLECTOR_WALL_CLOCK_MEASUREMENT + ) + # The batch-worker categories keep the label that warns about summing. + assert measurements["idle/buffer_full_backoff"] == THREAD_SECONDS_MEASUREMENT + assert measurements["wasted/failed_trajectory"] == THREAD_SECONDS_MEASUREMENT + + +def _gauge_points(data, name): + """Data points for gauge *name* in one already-collected metrics batch. + + ``InMemoryMetricReader.get_metrics_data()`` drains what it collects, so + callers must collect once and filter the result rather than calling the + reader per metric name. + """ + if data is None: + return [] + return [ + point + for rm in data.resource_metrics + for sm in rm.scope_metrics + for metric in sm.metrics + if metric.name == name + for point in metric.data.data_points + ] + + +def _start_exporting_telemetry(): + """Install an exporting telemetry handle; returns its metric reader.""" + from nemo.lens import NemoLensConfig, setup_telemetry + from opentelemetry.sdk.metrics.export import InMemoryMetricReader + + import nemo_rl.telemetry.setup as setup_mod + from nemo_rl.telemetry.span_groups import RLSpanGroup + + reader = InMemoryMetricReader() + cfg = NemoLensConfig(enabled=True, _span_group_cls=RLSpanGroup) + setup_mod._TELEMETRY_HANDLE = setup_telemetry( + cfg, rank=0, world_size=1, metric_reader=reader + ) + return reader + + +def test_tee_emits_efficiency_seconds_tagged_by_measurement(monkeypatch): + pytest.importorskip("nemo.lens") + import nemo_rl.telemetry.metrics as metrics_mod + from nemo_rl.telemetry.instrumentation import ( + RL_BUCKET_ATTR, + RL_EFFICIENCY_CATEGORY_ATTR, + ) + from nemo_rl.telemetry.metrics import ( + RL_EFFICIENCY_MEASUREMENT_ATTR, + RL_EFFICIENCY_SECONDS_METRIC, + tee_rl_metrics_to_otel, + ) + + reader = _start_exporting_telemetry() + # Pin the classification so this test does not move when the canonical + # category lists change. + monkeypatch.setattr(metrics_mod, "efficiency_measurements", lambda: _MEASUREMENTS) + + tee_rl_metrics_to_otel( + { + "efficiency/idle/refit_bubble_s": 12.0, + "efficiency/idle/buffer_full_backoff_s": 80.0, + "efficiency/total_waste_s": 999.0, + }, + "", + ) + + points = _gauge_points(reader.get_metrics_data(), RL_EFFICIENCY_SECONDS_METRIC) + by_category = { + point.attributes[RL_EFFICIENCY_CATEGORY_ATTR]: point for point in points + } + assert set(by_category) == set(_MEASUREMENTS) + + wall = by_category["idle/refit_bubble"] + assert wall.value == 12.0 + assert wall.attributes[RL_EFFICIENCY_MEASUREMENT_ATTR] == "wall_clock" + assert wall.attributes[RL_BUCKET_ATTR] == "idle" + + thread = by_category["idle/buffer_full_backoff"] + assert thread.value == 80.0 + assert thread.attributes[RL_EFFICIENCY_MEASUREMENT_ATTR] == "thread_seconds" + # Same bucket as the wall-clock point, which is exactly why the measurement + # attribute has to be present for a bucket rollup to stay honest. + assert thread.attributes[RL_BUCKET_ATTR] == "idle" + + +def test_efficiency_window_separates_per_step_from_cumulative(): + """Summing over time needs its own attribute, not the measurement kind. + + The driver resets its ``Timer`` each step and the collector never does, and + ``init/total`` is a run constant republished every step -- so a dashboard + summing ``wall_clock`` across steps would count startup once per step. + """ + from nemo_rl.telemetry.metrics import ( + COLLECTOR_WALL_CLOCK_MEASUREMENT, + RUN_WINDOW, + STEP_WINDOW, + THREAD_SECONDS_MEASUREMENT, + WALL_CLOCK_MEASUREMENT, + efficiency_window, + ) + + assert efficiency_window("idle/refit_bubble", WALL_CLOCK_MEASUREMENT) == STEP_WINDOW + assert efficiency_window("init/total", WALL_CLOCK_MEASUREMENT) == RUN_WINDOW + assert ( + efficiency_window("idle/refit_event_wait", COLLECTOR_WALL_CLOCK_MEASUREMENT) + == RUN_WINDOW + ) + assert ( + efficiency_window("wasted/failed_trajectory", THREAD_SECONDS_MEASUREMENT) + == RUN_WINDOW + ) + + +def test_run_window_categories_match_the_efficiency_summary(): + """The window split is restated here; the summary excludes the same set. + + ``print_efficiency_summary`` keeps run-window categories out of its per-step + waste ratio using its own copy of this set. If the two drifted, a category + could be tagged ``step`` here while being excluded from the step ratio + there, or charged to every step while advertised as a run constant. + """ + from nemo_rl.telemetry.metrics import _RUN_WINDOW_WALL_CLOCK_CATEGORIES + from tests.unit.telemetry.conftest import algorithms_utils_categories + + canonical = algorithms_utils_categories("RUN_WINDOW_WALL_CLOCK_CATEGORIES") + + assert ( + set(_RUN_WINDOW_WALL_CLOCK_CATEGORIES) + == canonical["RUN_WINDOW_WALL_CLOCK_CATEGORIES"] + ) + + +def test_tee_tags_the_pct_as_a_per_step_wall_clock_ratio(monkeypatch): + """The one aggregate point needs its window stated most. + + A percentage carries no unit to hint at what it covers, so an untagged + point invites a reader to treat a per-step ratio as run-to-date. + """ + pytest.importorskip("nemo.lens") + import nemo_rl.telemetry.metrics as metrics_mod + from nemo_rl.telemetry.metrics import ( + RL_EFFICIENCY_MEASUREMENT_ATTR, + RL_EFFICIENCY_PCT_METRIC, + RL_EFFICIENCY_WINDOW_ATTR, + STEP_WINDOW, + WALL_CLOCK_MEASUREMENT, + tee_rl_metrics_to_otel, + ) + + reader = _start_exporting_telemetry() + monkeypatch.setattr(metrics_mod, "efficiency_measurements", lambda: _MEASUREMENTS) + + tee_rl_metrics_to_otel( + { + "efficiency/efficiency_pct": 91.5, + "efficiency/efficiency_pct_is_per_step": 1.0, + }, + "", + ) + + (point,) = _gauge_points(reader.get_metrics_data(), RL_EFFICIENCY_PCT_METRIC) + assert point.value == 91.5 + assert point.attributes[RL_EFFICIENCY_WINDOW_ATTR] == STEP_WINDOW + assert point.attributes[RL_EFFICIENCY_MEASUREMENT_ATTR] == WALL_CLOCK_MEASUREMENT + + +def test_a_run_cumulative_pct_is_not_published_as_per_step(monkeypatch): + """print_efficiency_summary falls back to a run-to-date denominator. + + Callers that pass no per-step wall time get a ratio over the whole run, so + the window has to be derived from what the summary did rather than asserted + here -- and an absent flag has to mean ``run``, since mislabelling a + cumulative ratio as per-step is the direction that misleads. + """ + pytest.importorskip("nemo.lens") + import nemo_rl.telemetry.metrics as metrics_mod + from nemo_rl.telemetry.metrics import ( + RL_EFFICIENCY_PCT_METRIC, + RL_EFFICIENCY_WINDOW_ATTR, + RUN_WINDOW, + tee_rl_metrics_to_otel, + ) + + reader = _start_exporting_telemetry() + monkeypatch.setattr(metrics_mod, "efficiency_measurements", lambda: _MEASUREMENTS) + + tee_rl_metrics_to_otel( + { + "efficiency/efficiency_pct": 91.5, + "efficiency/efficiency_pct_is_per_step": 0.0, + }, + "", + ) + (point,) = _gauge_points(reader.get_metrics_data(), RL_EFFICIENCY_PCT_METRIC) + assert point.attributes[RL_EFFICIENCY_WINDOW_ATTR] == RUN_WINDOW + + +def test_a_pct_with_no_window_flag_defaults_to_the_run_window(monkeypatch): + """An absent flag must not be read as per-step. + + Any caller predating the flag, or one that builds the dict by hand, would + otherwise have its run-to-date ratio published as though it covered a + single step. + """ + pytest.importorskip("nemo.lens") + import nemo_rl.telemetry.metrics as metrics_mod + from nemo_rl.telemetry.metrics import ( + RL_EFFICIENCY_PCT_METRIC, + RL_EFFICIENCY_WINDOW_ATTR, + RUN_WINDOW, + tee_rl_metrics_to_otel, + ) + + reader = _start_exporting_telemetry() + monkeypatch.setattr(metrics_mod, "efficiency_measurements", lambda: _MEASUREMENTS) + + tee_rl_metrics_to_otel({"efficiency/efficiency_pct": 91.5}, "") + + (point,) = _gauge_points(reader.get_metrics_data(), RL_EFFICIENCY_PCT_METRIC) + assert point.attributes[RL_EFFICIENCY_WINDOW_ATTR] == RUN_WINDOW + + +def test_tee_tags_efficiency_seconds_with_the_window(monkeypatch): + pytest.importorskip("nemo.lens") + import nemo_rl.telemetry.metrics as metrics_mod + from nemo_rl.telemetry.instrumentation import RL_EFFICIENCY_CATEGORY_ATTR + from nemo_rl.telemetry.metrics import ( + RL_EFFICIENCY_SECONDS_METRIC, + RL_EFFICIENCY_WINDOW_ATTR, + tee_rl_metrics_to_otel, + ) + + reader = _start_exporting_telemetry() + monkeypatch.setattr( + metrics_mod, + "efficiency_measurements", + lambda: {**_MEASUREMENTS, "init/total": "wall_clock"}, + ) + + tee_rl_metrics_to_otel( + { + "efficiency/idle/refit_bubble_s": 12.0, + "efficiency/init/total_s": 300.0, + }, + "", + ) + + points = _gauge_points(reader.get_metrics_data(), RL_EFFICIENCY_SECONDS_METRIC) + windows = { + point.attributes[RL_EFFICIENCY_CATEGORY_ATTR]: point.attributes[ + RL_EFFICIENCY_WINDOW_ATTR + ] + for point in points + } + assert windows["idle/refit_bubble"] == "step" + assert windows["init/total"] == "run" + + +def test_tee_emits_efficiency_pct(monkeypatch): + pytest.importorskip("nemo.lens") + import nemo_rl.telemetry.metrics as metrics_mod + from nemo_rl.telemetry.metrics import ( + RL_EFFICIENCY_PCT_METRIC, + tee_rl_metrics_to_otel, + ) + + reader = _start_exporting_telemetry() + monkeypatch.setattr(metrics_mod, "efficiency_measurements", lambda: _MEASUREMENTS) + + tee_rl_metrics_to_otel({"efficiency/efficiency_pct": 87.5}, "") + + points = _gauge_points(reader.get_metrics_data(), RL_EFFICIENCY_PCT_METRIC) + assert [point.value for point in points] == [87.5] + + +def test_tee_skips_prefixes_other_than_train(): + # The efficiency dict is logged under "train"/"" -- looking for it under the + # other prefixes log_metrics fans a step out to would be dead work on every + # step, and admitting one would attribute another dict's values to a + # category series. + pytest.importorskip("nemo.lens") + from nemo_rl.telemetry.metrics import ( + RL_EFFICIENCY_SECONDS_METRIC, + tee_rl_metrics_to_otel, + ) + + reader = _start_exporting_telemetry() + tee_rl_metrics_to_otel({"efficiency/idle/refit_bubble_s": 12.0}, "performance") + tee_rl_metrics_to_otel({"efficiency/idle/refit_bubble_s": 12.0}, "validation") + + assert _gauge_points(reader.get_metrics_data(), RL_EFFICIENCY_SECONDS_METRIC) == [] + + +def test_tee_skips_efficiency_instruments_without_efficiency_keys(monkeypatch): + # A plain train-metrics dict must not create the efficiency series at all. + pytest.importorskip("nemo.lens") + import nemo_rl.telemetry.metrics as metrics_mod + from nemo_rl.telemetry.metrics import ( + RL_EFFICIENCY_PCT_METRIC, + RL_EFFICIENCY_SECONDS_METRIC, + tee_rl_metrics_to_otel, + ) + + reader = _start_exporting_telemetry() + monkeypatch.setattr(metrics_mod, "efficiency_measurements", lambda: _MEASUREMENTS) + + tee_rl_metrics_to_otel({"reward": 0.5}, "train") + + data = reader.get_metrics_data() + assert _gauge_points(data, RL_EFFICIENCY_SECONDS_METRIC) == [] + assert _gauge_points(data, RL_EFFICIENCY_PCT_METRIC) == [] + + +def test_tee_never_raises_into_the_training_step(monkeypatch, caplog): + """Logger.log_metrics calls this unguarded, so it has to swallow everything. + + Exercised through a failure *outside* the inner handler -- reading the + telemetry handle -- since that one is already covered separately. + """ + import nemo_rl.telemetry.metrics as metrics_mod + from nemo_rl.telemetry.metrics import tee_rl_metrics_to_otel + + def _boom(): + raise RuntimeError("handle is wedged") + + monkeypatch.setattr(metrics_mod, "get_telemetry_handle", _boom) + monkeypatch.setattr(metrics_mod, "_WARNED", set()) + + with caplog.at_level(logging.WARNING, logger=metrics_mod.__name__): + tee_rl_metrics_to_otel({"reward": 1.0}, "train") + + assert [r.levelno for r in caplog.records] == [logging.WARNING] + + +def test_a_broken_tee_warns_once_not_once_per_step(monkeypatch, caplog): + """A deterministic failure repeats every step, so it must not log every step.""" + pytest.importorskip("nemo.lens") + import nemo_rl.telemetry.metrics as metrics_mod + from nemo_rl.telemetry.metrics import tee_rl_metrics_to_otel + + _start_exporting_telemetry() + monkeypatch.setattr(metrics_mod, "efficiency_measurements", lambda: _MEASUREMENTS) + monkeypatch.setattr(metrics_mod, "_WARNED", set()) + + def _boom(*args, **kwargs): + raise RuntimeError("instrument is gone") + + monkeypatch.setattr(metrics_mod, "map_efficiency_seconds", _boom) + + with caplog.at_level(logging.DEBUG, logger=metrics_mod.__name__): + for _ in range(3): + tee_rl_metrics_to_otel({"efficiency/idle/refit_bubble_s": 1.0}, "train") + + records = [r for r in caplog.records if "efficiency" in r.message] + assert [r.levelno for r in records] == [ + logging.WARNING, + logging.DEBUG, + logging.DEBUG, + ] + # The traceback survives the demotion, so the repetition stays diagnosable. + assert all(r.exc_info is not None for r in records) diff --git a/tests/unit/telemetry/test_setup.py b/tests/unit/telemetry/test_setup.py new file mode 100644 index 0000000000..ffb6fb096c --- /dev/null +++ b/tests/unit/telemetry/test_setup.py @@ -0,0 +1,514 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the telemetry setup module (driver init, resource attrs, digging).""" + +import logging +import os +from types import SimpleNamespace + +import pytest +from pydantic import ValidationError + +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.setup import ( + _build_resource_attributes, + _dig, + _worker_resource_attributes, + get_telemetry_handle, + init_telemetry_driver, + init_telemetry_worker, + shutdown_telemetry, + telemetry_enabled_in_env, + vllm_native_tracing_requested, +) + + +class _FakeMasterConfig: + def __init__(self, telemetry=None, policy=None): + self.telemetry = telemetry + self.policy = policy or { + "model_name": "org/Model-1B", + "precision": "bfloat16", + "megatron_cfg": { + "tensor_model_parallel_size": 2, + "pipeline_model_parallel_size": 1, + }, + } + + +def test_dig_handles_dicts_objects_and_missing(): + assert _dig({"a": {"b": 7}}, "a", "b") == 7 + assert _dig({"a": {}}, "a", "missing") is None + assert _dig(None, "a") is None + + class Node: + x = {"y": 9} + + assert _dig(Node(), "x", "y") == 9 + + +def test_build_resource_attributes(): + attrs = _build_resource_attributes(_FakeMasterConfig(), "grpo") + assert attrs["rl.algorithm"] == "grpo" + assert attrs["rl.model"] == "org/Model-1B" + assert attrs["nemo.precision"] == "bfloat16" + assert attrs["dl.tensor_parallel.size"] == 2 + assert attrs["dl.pipeline_parallel.size"] == 1 + + +def test_build_resource_attributes_dtensor_tp(): + cfg = _FakeMasterConfig( + policy={ + "model_name": "org/Model-1B", + "precision": "bfloat16", + "dtensor_cfg": {"tensor_parallel_size": 4}, + } + ) + attrs = _build_resource_attributes(cfg, "grpo") + assert attrs["dl.tensor_parallel.size"] == 4 + assert "dl.pipeline_parallel.size" not in attrs + + +def test_init_driver_returns_none_when_disabled(): + handle = init_telemetry_driver( + _FakeMasterConfig(TelemetryConfig(enabled=False)), "grpo" + ) + assert handle is None + assert get_telemetry_handle() is None + + +def test_init_driver_returns_none_when_no_telemetry_block(): + handle = init_telemetry_driver(_FakeMasterConfig(telemetry=None), "grpo") + assert handle is None + + +def test_init_worker_returns_none_when_disabled(): + # No NEMO_RL_OTEL_ENABLED / NEMO_LENS_ENABLED — every actor takes this path. + handle = init_telemetry_worker() + assert handle is None + assert get_telemetry_handle() is None + + +def test_config_rejects_an_unknown_export_strategy_at_parse_time(): + # The Literal is what a user actually hits: it fires wherever the YAML is + # read, in every process, and regardless of `enabled` -- so a typo cannot + # sit dormant in a disabled block until someone switches telemetry on. + with pytest.raises(ValidationError): + TelemetryConfig(enabled=True, export_strategy="single_ranks") + + +def test_init_driver_rejects_a_strategy_lens_does_not_register(): + # Second line of defence, for drift between the Literal above and lens's + # registry rather than for user typos. It matters because the driver + # overrides the strategy with _always_export, which bypasses the registry + # lookup that would otherwise reject an unknown name. model_construct skips + # validation, standing in for a Literal that has gained a value lens + # dropped. + pytest.importorskip("nemo.lens") + tel = TelemetryConfig.model_construct(enabled=True, export_strategy="single_ranks") + config = SimpleNamespace(telemetry=tel) + with pytest.raises(ValueError, match="Unknown telemetry.export_strategy"): + init_telemetry_driver(config, algorithm="grpo") + # Again, because the init guard must not be set by a path that raised: + # otherwise the second call reports success-with-nothing instead of the bug. + with pytest.raises(ValueError, match="Unknown telemetry.export_strategy"): + init_telemetry_driver(config, algorithm="grpo") + + +def test_config_rejects_out_of_range_export_bounds(): + # export_rank had no check anywhere before, and setup.py documents the + # failure it causes: a rank that never matches silently mutes the actor. + with pytest.raises(ValidationError): + TelemetryConfig(export_rank=-2) + with pytest.raises(ValidationError): + TelemetryConfig(export_sample_rate=1.5) + with pytest.raises(ValidationError): + TelemetryConfig(exporter="consoel") + + +def test_init_driver_rejects_unknown_span_group(): + # Lens resolves span_groups only after installing the global tracer + # provider, so without this check a typo kills the run mid-setup. + pytest.importorskip("nemo.lens") + config = SimpleNamespace( + telemetry=TelemetryConfig(enabled=True, span_groups="per_stp") + ) + with pytest.raises(ValueError): + init_telemetry_driver(config, algorithm="grpo") + + +def test_documented_export_strategies_are_registered(): + pytest.importorskip("nemo.lens") + from nemo.lens import registered_strategies + + documented = {"single_rank", "all_ranks", "sampled", "first_rank_per_node"} + assert documented <= set(registered_strategies()), ( + "TelemetryConfig.export_strategy docstring lists a strategy lens does " + "not register" + ) + + +@pytest.mark.parametrize( + "strategy", + ["single_rank", "all_ranks", "sampled", "first_rank_per_node"], +) +def test_init_driver_accepts_every_documented_export_strategy(strategy): + # Exercises init end to end rather than just the registry: a documented + # value has to survive the validation branch and lens's own setup. (Whether + # the strategy *selects* this process is not in play here -- the driver + # overrides it -- see the worker tests for that.) + pytest.importorskip("nemo.lens") + cfg = TelemetryConfig(enabled=True, exporter="console", export_strategy=strategy) + handle = init_telemetry_driver(_FakeMasterConfig(cfg), "grpo") + assert handle is not None + assert handle.is_exporting + + +def test_init_driver_publishes_default_service_name_to_env(): + # Workers rebuild their config from the environment, so a service name only + # assigned on the config object would leave them defaulting to lens's + # "nemo" and split one run across two services. + pytest.importorskip("nemo.lens") + # Empty rather than the field default, so _config_to_env skips it and only + # the branch under test can populate the env var. + cfg = TelemetryConfig(enabled=True, exporter="console", service_name="") + handle = init_telemetry_driver(_FakeMasterConfig(cfg), "grpo") + assert handle is not None + assert os.environ["OTEL_SERVICE_NAME"] == "nemo-rl" + + +def test_init_driver_keeps_user_service_name(monkeypatch): + pytest.importorskip("nemo.lens") + monkeypatch.setenv("OTEL_SERVICE_NAME", "my-service") + handle = init_telemetry_driver( + _FakeMasterConfig(TelemetryConfig(enabled=True, exporter="console")), "grpo" + ) + assert handle is not None + assert os.environ["OTEL_SERVICE_NAME"] == "my-service" + + +def test_init_driver_publishes_env_even_when_telemetry_is_disabled(): + """``_config_to_env`` runs before the ``enabled`` early return. + + Deliberate -- the env is the only channel to a Ray worker, so it must not + depend on which branch the driver took. It is also why anything reading + these variables has to check the master switch itself; see + :func:`telemetry_enabled_in_env`. + """ + cfg = TelemetryConfig(enabled=False, span_groups="per_step") + + assert init_telemetry_driver(_FakeMasterConfig(cfg), "grpo") is None + + assert os.environ["NEMO_RL_OTEL_SPAN_GROUPS"] == "per_step" + assert not telemetry_enabled_in_env() + + +def test_vllm_native_tracing_needs_the_master_switch(monkeypatch): + # The field is exported regardless of `enabled`, so reading it alone would + # leave per-request vLLM tracing on for a run that turned telemetry off. + monkeypatch.setenv("NEMO_RL_OTEL_VLLM_NATIVE_TRACING", "1") + monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "0") + + assert vllm_native_tracing_requested() + assert not telemetry_enabled_in_env() + + monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "1") + assert telemetry_enabled_in_env() + + +def test_init_driver_disables_the_rank_sampler(monkeypatch): + # The driver is not one of the ranks the user asked to sample; leaving the + # sampler on drops the training loop's own spans. + pytest.importorskip("nemo.lens") + captured = {} + + def _fake_setup_telemetry(config, **kwargs): + captured["sampler_enabled"] = config.sampler_enabled + return SimpleNamespace(is_exporting=True, tracer=None) + + monkeypatch.setattr("nemo.lens.setup_telemetry", _fake_setup_telemetry) + cfg = TelemetryConfig( + enabled=True, + exporter="console", + sampler_enabled=True, + export_sample_rate=0.1, + ) + + assert init_telemetry_driver(_FakeMasterConfig(cfg), "grpo") is not None + assert captured["sampler_enabled"] is False + # Process-local: what ranked workers inherit must still say the user asked + # for sampling. + assert os.environ["NEMO_RL_OTEL_SAMPLER_ENABLED"] == "1" + + +def test_worker_resource_attributes_carries_worker_group(monkeypatch): + monkeypatch.setenv("NRL_WORKER_GROUP", "vllm_policy") + assert _worker_resource_attributes(None) == {"rl.worker_group": "vllm_policy"} + + +def test_worker_resource_attributes_without_group_env(): + # Workers not created by RayWorkerGroup simply omit the attribute. + assert _worker_resource_attributes(None) == {} + + +def test_worker_resource_attributes_explicit_extra_wins(monkeypatch): + monkeypatch.setenv("NRL_WORKER_GROUP", "lm_policy") + attrs = _worker_resource_attributes({"rl.worker_group": "override", "k": 1}) + assert attrs == {"rl.worker_group": "override", "k": 1} + + +def test_init_worker_sets_worker_group_attribute(monkeypatch): + pytest.importorskip("nemo.lens") + captured = {} + + def _fake_setup_telemetry(config, **kwargs): + captured.update(kwargs) + return SimpleNamespace(is_exporting=True) + + monkeypatch.setattr("nemo.lens.setup_telemetry", _fake_setup_telemetry) + monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "1") + monkeypatch.setenv("NRL_WORKER_GROUP", "vllm_policy") + monkeypatch.setenv("RANK", "3") + monkeypatch.setenv("WORLD_SIZE", "8") + + handle = init_telemetry_worker() + assert handle is not None + assert captured["rank"] == 3 + assert captured["world_size"] == 8 + assert captured["resource_attributes"] == {"rl.worker_group": "vllm_policy"} + + +def test_init_worker_explicit_rank_overrides_env(monkeypatch): + """Singleton actors pass their own rank instead of reading the env. + + The trajectory collector's runtime_env is a copy of the driver's + environment, so a ``RANK`` inherited from the launcher would otherwise + decide whether the collector exports. + """ + pytest.importorskip("nemo.lens") + captured = {} + + def _fake_setup_telemetry(config, **kwargs): + captured.update(kwargs) + return SimpleNamespace(is_exporting=True) + + monkeypatch.setattr("nemo.lens.setup_telemetry", _fake_setup_telemetry) + monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "1") + monkeypatch.setenv("RANK", "5") + monkeypatch.setenv("WORLD_SIZE", "8") + + assert init_telemetry_worker(rank=0, world_size=1) is not None + assert captured["rank"] == 0 + assert captured["world_size"] == 1 + + +def test_init_worker_honours_export_strategy_by_default(monkeypatch): + # The baseline for the always_export test below: a ranked worker must obey + # whatever the user configured. + pytest.importorskip("nemo.lens") + monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "1") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORTER", "console") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORT_STRATEGY", "single_rank") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORT_RANK", "3") + + handle = init_telemetry_worker(rank=0, world_size=1) + assert handle is not None + assert not handle.is_exporting + + +def test_init_worker_always_export_overrides_export_rank(monkeypatch): + """A singleton actor's synthetic rank must not be subject to the strategy. + + The trajectory collector reports ``rank=0, world_size=1`` because it is not + a member of a ranked group. Applying ``export_rank: 3`` to that made-up rank + would silently mute the actor -- taking every async rollout span with it. + """ + pytest.importorskip("nemo.lens") + monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "1") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORTER", "console") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORT_STRATEGY", "single_rank") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORT_RANK", "3") + + handle = init_telemetry_worker(rank=0, world_size=1, always_export=True) + assert handle is not None + assert handle.is_exporting + + +def test_init_worker_always_export_overrides_sample_rate(monkeypatch): + # rank 0 hashes into the 0.785 bucket, so a low sample rate excludes it. + pytest.importorskip("nemo.lens") + monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "1") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORTER", "console") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORT_STRATEGY", "sampled") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORT_SAMPLE_RATE", "0.1") + + handle = init_telemetry_worker(rank=0, world_size=1, always_export=True) + assert handle is not None + assert handle.is_exporting + + +def test_init_worker_always_export_disables_the_rank_sampler(monkeypatch): + """Exporting is not enough: the sampler drops spans before that decision. + + ``sampler_enabled`` installs lens's ``RankAwareSampler`` on the tracer + provider, which filters on the same rank hash independently of + ``export_strategy``. Left on, it discards every span of a synthetic rank 0 + while ``is_exporting`` still reports True -- telemetry that looks wired and + produces nothing. + """ + pytest.importorskip("nemo.lens") + captured = {} + + def _fake_setup_telemetry(config, **kwargs): + captured["sampler_enabled"] = config.sampler_enabled + return SimpleNamespace(is_exporting=True) + + monkeypatch.setattr("nemo.lens.setup_telemetry", _fake_setup_telemetry) + monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "1") + monkeypatch.setenv("NEMO_RL_OTEL_SAMPLER_ENABLED", "1") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORT_SAMPLE_RATE", "0.1") + + assert init_telemetry_worker(rank=0, world_size=1, always_export=True) is not None + assert captured["sampler_enabled"] is False + + +def test_init_worker_leaves_the_rank_sampler_alone_for_ranked_workers(monkeypatch): + # The counterpart: a real group member is part of the population the user + # asked to sample, so the sampler must stay on. + pytest.importorskip("nemo.lens") + captured = {} + + def _fake_setup_telemetry(config, **kwargs): + captured["sampler_enabled"] = config.sampler_enabled + return SimpleNamespace(is_exporting=True) + + monkeypatch.setattr("nemo.lens.setup_telemetry", _fake_setup_telemetry) + monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "1") + monkeypatch.setenv("NEMO_RL_OTEL_SAMPLER_ENABLED", "1") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORT_SAMPLE_RATE", "0.1") + + assert init_telemetry_worker(rank=2, world_size=8) is not None + assert captured["sampler_enabled"] is True + + +def test_init_worker_survives_a_failing_setup(monkeypatch, caplog): + # A worker must not take a training run down over optional observability, + # but the reason has to reach the log or the run is silently half-traced. + pytest.importorskip("nemo.lens") + + def _boom(config, **kwargs): + raise ImportError("OpenTelemetry SDK is required for telemetry export") + + monkeypatch.setattr("nemo.lens.setup_telemetry", _boom) + monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "1") + + with caplog.at_level(logging.WARNING): + assert init_telemetry_worker() is None + assert "OpenTelemetry SDK is required" in caplog.text + + +def test_init_worker_stays_quiet_when_telemetry_is_disabled(monkeypatch, caplog): + # Nothing to warn about: the user did not ask for telemetry. + pytest.importorskip("nemo.lens") + monkeypatch.delenv("NEMO_RL_OTEL_ENABLED", raising=False) + monkeypatch.delenv("NEMO_LENS_ENABLED", raising=False) + + with caplog.at_level(logging.WARNING): + assert init_telemetry_worker() is None + assert caplog.text == "" + + +def test_init_worker_never_raises_on_setup_failure(monkeypatch): + # A worker must not fail a training run over optional observability. + pytest.importorskip("nemo.lens") + + def _boom(config, **kwargs): + raise ValueError("unknown export_strategy 'typo'") + + monkeypatch.setattr("nemo.lens.setup_telemetry", _boom) + monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "1") + + assert init_telemetry_worker() is None + assert get_telemetry_handle() is None + + +def test_init_driver_enabled_is_idempotent(): + pytest.importorskip("nemo.lens") + cfg = TelemetryConfig(enabled=True, span_groups="default", exporter="console") + handle1 = init_telemetry_driver(_FakeMasterConfig(cfg), "grpo") + assert handle1 is not None + assert handle1.is_exporting + assert get_telemetry_handle() is handle1 + # Second call must not re-init or raise; returns the same handle. + handle2 = init_telemetry_driver( + _FakeMasterConfig(TelemetryConfig(enabled=True)), "grpo" + ) + assert handle2 is handle1 + + +def test_init_driver_exports_despite_nonzero_export_rank(): + # export_rank selects among the Ray worker ranks; it must not switch off the + # driver, whose rank=0/world_size=1 are synthetic. + pytest.importorskip("nemo.lens") + cfg = TelemetryConfig( + enabled=True, + exporter="console", + export_strategy="single_rank", + export_rank=3, + ) + handle = init_telemetry_driver(_FakeMasterConfig(cfg), "grpo") + assert handle is not None + assert handle.is_exporting + + +def test_init_driver_exports_under_sampled_strategy(monkeypatch): + # rank 0 hashes into the 0.785 bucket, so a lower sample rate would exclude + # the driver without the export-strategy override. + pytest.importorskip("nemo.lens") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORT_SAMPLE_RATE", "0.1") + cfg = TelemetryConfig(enabled=True, exporter="console", export_strategy="sampled") + handle = init_telemetry_driver(_FakeMasterConfig(cfg), "grpo") + assert handle is not None + assert handle.is_exporting + + +def test_shutdown_is_a_noop_without_telemetry(): + # The common case by far: every run with no `telemetry:` block reaches the + # driver's `finally` and each actor's shutdown hook with no handle. + import nemo_rl.telemetry.setup as setup_mod + + setup_mod._TELEMETRY_HANDLE = None + shutdown_telemetry() + + +def test_shutdown_clears_the_handle_so_a_second_call_cannot_reach_a_dead_provider(): + # Both the driver's `finally` and the worker shutdown hooks call this, and + # flushing an already-shut-down provider is what logs "Already shutdown". + import nemo_rl.telemetry.setup as setup_mod + + calls = [] + setup_mod._TELEMETRY_HANDLE = SimpleNamespace( + shutdown=lambda timeout_ms: calls.append(timeout_ms) + ) + + shutdown_telemetry(timeout_ms=1234) + shutdown_telemetry(timeout_ms=1234) + + assert calls == [1234] + assert get_telemetry_handle() is None + + +def test_shutdown_swallows_a_failing_flush(): + # It runs in a `finally`, so raising here would replace whatever exception + # actually ended the run -- and still has to clear the handle. + import nemo_rl.telemetry.setup as setup_mod + + def _boom(timeout_ms): + raise RuntimeError("exporter is gone") + + setup_mod._TELEMETRY_HANDLE = SimpleNamespace(shutdown=_boom) + + shutdown_telemetry() + + assert get_telemetry_handle() is None diff --git a/tests/unit/telemetry/test_source_drift.py b/tests/unit/telemetry/test_source_drift.py new file mode 100644 index 0000000000..1c358d5944 --- /dev/null +++ b/tests/unit/telemetry/test_source_drift.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Guards against declarations drifting from the call sites that use them. + +The other suites in this directory compare one declaration to another, which +catches a constant updated in one place and not the other. These tests close the +remaining direction -- a *call site* naming something no declaration knows about +-- which is the direction that fails silently. + +Source is parsed rather than imported: the algorithm modules pull in torch, and +none of this needs it (or a GPU, or nemo-lens). +""" + +import ast +import re +from pathlib import Path + +from tests.unit.telemetry.conftest import algorithms_utils_categories + +_REPO = Path(__file__).resolve().parents[3] +_ALGORITHMS = _REPO / "nemo_rl" / "algorithms" + +# Efficiency timers are driver-side only, but spans are not: rl.vllm.* is +# emitted from the generation worker and belongs in the doc tables too. +_SPAN_EMITTING_DIRS = (_ALGORITHMS, _REPO / "nemo_rl" / "models" / "generation") + +# Timer methods that take a category label as their first argument. +_TIMER_METHODS = frozenset({"time", "reduce", "record", "start", "stop"}) + +# Prefixes owned by the efficiency accounting in nemo_rl/algorithms/utils.py. +_EFFICIENCY_PREFIXES = ("idle/", "wasted/") + + +def _python_sources(root: Path) -> list[Path]: + return sorted(root.rglob("*.py")) + + +def _string_arg(node: ast.Call, index: int = 0) -> str | None: + """The *index*-th positional argument of *node*, if it is a string literal.""" + if len(node.args) <= index: + return None + arg = node.args[index] + return ( + arg.value + if isinstance(arg, ast.Constant) and isinstance(arg.value, str) + else None + ) + + +def test_every_efficiency_timer_at_a_call_site_is_declared(): + """An undeclared idle/wasted timer is counted as *productive*, silently. + + ``print_efficiency_summary`` derives waste by iterating the declared lists, + not by reading what the Timer recorded, and productive time is + ``step_wall_time - waste``. So a category nothing declares does not go + missing from the report -- it inverts, and efficiency reads higher than + reality. Nothing warns, and ``bucket_for_efficiency_category`` returns None + for it, leaving any matching span unbucketed too. + """ + declared: set[str] = set().union( + *algorithms_utils_categories( + "WALL_CLOCK_EFFICIENCY_CATEGORIES", + "THREAD_ACCUMULATED_EFFICIENCY_CATEGORIES", + ).values() + ) + + used: dict[str, str] = {} + for source in _python_sources(_ALGORITHMS): + for node in ast.walk(ast.parse(source.read_text())): + if not isinstance(node, ast.Call): + continue + called = getattr(node.func, "attr", None) or getattr(node.func, "id", None) + if called not in _TIMER_METHODS | {"efficiency_span"}: + continue + label = _string_arg(node) + if label and label.startswith(_EFFICIENCY_PREFIXES): + used[label] = source.relative_to(_REPO).as_posix() + + undeclared = { + label: where for label, where in used.items() if label not in declared + } + assert not undeclared, ( + "efficiency categories measured but not declared in " + f"nemo_rl/algorithms/utils.py, so their time is charged to productive: " + f"{undeclared}" + ) + # Sanity: the walk found something, so a refactor that moves these calls + # cannot quietly turn this test into a tautology. + assert used, "found no idle/* or wasted/* timers -- has the matcher gone stale?" + + +def test_every_emitted_span_name_is_documented(): + """Direction is ``emitted <= documented``. + + The docs may list a span that is not wired yet (``forward_backward``, + ``optimizer``), but a name that is emitted and undocumented leaves someone + filtering a Tempo/Jaeger query on a name that does not exist -- zero + results, nothing to explain it. A typo at an emit site fails the same way, + and produces a real span under the wrong name, since the goodput rollup keys + on ``rl.bucket`` rather than the name. + """ + documented = set( + _span_names_in( + (_REPO / "docs" / "observability" / "span-groups.md").read_text() + ) + ) + + emitted: dict[str, str] = {} + for directory in _SPAN_EMITTING_DIRS: + for source in _python_sources(directory): + for node in ast.walk(ast.parse(source.read_text())): + if not isinstance(node, ast.Call): + continue + called = getattr(node.func, "attr", None) or getattr( + node.func, "id", None + ) + if called not in ("managed_span", "trace_fn"): + continue + # Span name is the second positional arg, after the span group. + name = _string_arg(node, index=1) + if name: + emitted[name] = source.relative_to(_REPO).as_posix() + + undocumented = { + name: where for name, where in emitted.items() if name not in documented + } + assert not undocumented, ( + "spans emitted but absent from docs/observability/span-groups.md: " + f"{undocumented}" + ) + assert emitted, "found no span names -- has the matcher gone stale?" + + +def _span_names_in(markdown: str) -> set[str]: + """Every ``rl.*`` name in backticks, which is how the doc tables list them.""" + return set(re.findall(r"`(rl\.[a-z0-9_.]+)`", markdown)) diff --git a/tests/unit/telemetry/test_span_groups.py b/tests/unit/telemetry/test_span_groups.py new file mode 100644 index 0000000000..6a55ed1b99 --- /dev/null +++ b/tests/unit/telemetry/test_span_groups.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for RLSpanGroup presets and resolution.""" + +import pytest + +# ``resolve()`` requires the real nemo-lens SpanGroup base class. +pytest.importorskip("nemo.lens") + +from nemo_rl.telemetry.span_groups import RLSpanGroup + +RL_GROUPS = frozenset( + { + "rollout", + "generation", + "logprob", + "reward", + "advantage", + "policy_update", + "reference_policy", + "data_processing", + "efficiency", + } +) + +# Every group NeMo-RL emits a span in, RL-specific and inherited alike. Keep in +# sync when instrumenting a new group -- that is the point of +# ``test_every_emitted_group_is_reachable_from_a_shipped_preset``. +# +# A superset on purpose: it also holds the groups that are defined and bucketed +# but have no call site yet (``reference_policy``; see the coverage gaps in +# docs/observability/span-groups.md), so the preset wiring is already correct +# when one of them is instrumented rather than needing a second edit here. +EMITTED_GROUPS = RL_GROUPS | frozenset( + {"job", "step", "checkpoint", "evaluate", "model_init"} +) + + +def test_all_groups_includes_base_and_rl(): + assert RL_GROUPS <= RLSpanGroup.ALL_GROUPS + assert {"job", "checkpoint", "evaluate", "step"} <= RLSpanGroup.ALL_GROUPS + + +def test_default_preset_is_coarse(): + assert RLSpanGroup.resolve("default") == frozenset( + {"job", "checkpoint", "evaluate"} + ) + + +def test_per_step_has_step_and_phases_but_not_job(): + per_step = RLSpanGroup.resolve("per_step") + assert "step" in per_step + assert RL_GROUPS <= per_step + # per_step deliberately omits JOB so each step is its own root trace. + assert "job" not in per_step + + +def test_every_emitted_group_is_reachable_from_a_shipped_preset(): + """A group only in ``all`` is invisible to both presets users pick. + + ``model_init`` was in exactly that position: its one span, + ``rl.vllm.load_model``, could not appear under ``default`` or ``per_step``, + so the phase that explains a slow start was unobservable in practice. + """ + reachable = RLSpanGroup.resolve("default") | RLSpanGroup.resolve("per_step") + assert EMITTED_GROUPS <= reachable, ( + f"only reachable from 'all': {sorted(EMITTED_GROUPS - reachable)}" + ) + + +def test_all_preset_matches_all_groups(): + resolved = RLSpanGroup.resolve("all") + assert "job" in resolved + assert resolved == RLSpanGroup.ALL_GROUPS + + +def test_resolve_comma_list(): + assert RLSpanGroup.resolve("reward,generation") == frozenset( + {"reward", "generation"} + ) + + +def test_resolve_is_case_insensitive(): + assert RLSpanGroup.resolve("DEFAULT") == RLSpanGroup.resolve("default") + + +def test_resolve_unknown_raises(): + with pytest.raises(ValueError): + RLSpanGroup.resolve("nonexistent_group") diff --git a/uv.lock b/uv.lock index 6a5914ccb3..1af18d185f 100644 --- a/uv.lock +++ b/uv.lock @@ -4312,6 +4312,13 @@ dependencies = [ { name = "opentelemetry-api" }, ] +[package.optional-dependencies] +sdk = [ + { name = "opentelemetry-exporter-otlp-proto-grpc" }, + { name = "opentelemetry-exporter-otlp-proto-http" }, + { name = "opentelemetry-sdk" }, +] + [[package]] name = "nemo-rl" source = { editable = "." } @@ -4331,6 +4338,7 @@ dependencies = [ { name = "mlflow" }, { name = "mooncake-transfer-engine-cuda13" }, { name = "nccl4py" }, + { name = "nemo-lens", extra = ["sdk"] }, { name = "ninja" }, { name = "nixl" }, { name = "num2words" }, @@ -4550,6 +4558,7 @@ requires-dist = [ { name = "nccl4py", marker = "sys_platform != 'darwin'" }, { name = "nemo-automodel", extras = ["moe"], marker = "extra == 'automodel'", editable = "3rdparty/Automodel-workspace/Automodel" }, { name = "nemo-gym", marker = "extra == 'nemo-gym'", editable = "3rdparty/Gym-workspace/Gym" }, + { name = "nemo-lens", extras = ["sdk"], git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b85578fc2b736a1804705e537001b5f45e9c715d" }, { name = "ninja" }, { name = "nixl", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", specifier = "==1.3.0" }, { name = "num2words", specifier = ">=0.5.14" },