diff --git a/docker/Dockerfile b/docker/Dockerfile index 1170457ad91..df983326660 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -230,7 +230,7 @@ fi uv sync --link-mode symlink --frozen --extra mcore --no-install-project uv sync --link-mode symlink --frozen --extra automodel --no-install-project uv sync --link-mode symlink --frozen --extra modelopt --no-install-project -uv sync --link-mode symlink --frozen --all-groups --no-install-project +uv sync --link-mode symlink --frozen --all-groups --extra telemetry --no-install-project # Remove the aiohttp in this uv cache dir to fully address CVE GHSA-mqqc-3gqh-h2x8 # The ray install will include the older aiohttp version in its cache @@ -304,7 +304,7 @@ du -sh /root/.cache/uv /root/.cache/trtllm-wheels # Restore the intended default environment in the same layer so the transient # TRT-LLM installation does not add a large intermediate venv layer. -uv sync --link-mode symlink --locked --all-groups --no-install-project +uv sync --link-mode symlink --locked --all-groups --extra telemetry --no-install-project # The final sync can repopulate the shared uv cache, so repeat the security # cleanup performed in the preceding dependency layer. diff --git a/docker/Dockerfile.ngc_pytorch b/docker/Dockerfile.ngc_pytorch index cdd59c9ec9f..1c8b92e162e 100644 --- a/docker/Dockerfile.ngc_pytorch +++ b/docker/Dockerfile.ngc_pytorch @@ -120,7 +120,7 @@ uv sync --link-mode symlink --locked --inexact --extra vllm --no-install-project uv sync --link-mode symlink --locked --inexact --extra mcore --no-install-project $UV_NO_INSTALL_PACKAGES uv sync --link-mode symlink --locked --inexact --extra automodel --no-install-project $UV_NO_INSTALL_PACKAGES uv sync --link-mode symlink --locked --inexact --extra modelopt --no-install-project $UV_NO_INSTALL_PACKAGES -uv sync --link-mode symlink --locked --inexact --all-groups --no-install-project $UV_NO_INSTALL_PACKAGES +uv sync --link-mode symlink --locked --inexact --all-groups --extra telemetry --no-install-project $UV_NO_INSTALL_PACKAGES EOF ENV NEMO_RL_VENV_DIR=/opt/ray_venvs diff --git a/docs/index.md b/docs/index.md index 951f6631716..1d2737ee278 100644 --- a/docs/index.md +++ b/docs/index.md @@ -205,6 +205,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 @@ -339,6 +346,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 00000000000..e07e03d5608 --- /dev/null +++ b/docs/observability/configuration.md @@ -0,0 +1,138 @@ +# Configuration + +Telemetry can be configured two ways, which compose: + +1. A `telemetry:` block in your run config (YAML). +2. `NEMO_RL_OTEL_*` and standard `OTEL_*` environment variables. + +**Raw environment variables always win over the YAML block.** On the driver, the `telemetry:` block is translated into `NEMO_RL_OTEL_*` env vars with `os.environ.setdefault` *before* `init_ray()` — so anything already present in the environment is left untouched, and the resulting environment is snapshotted into the Ray `runtime_env` and inherited by every worker. + +## The `telemetry:` config block + +`telemetry:` is an optional top-level field of every algorithm's `MasterConfig`. It is **documented here, not baked into the exemplar configs** — add it to your own run config, or configure purely via env vars. + +```yaml +telemetry: + enabled: false # master switch; when false, every site is a ~0-cost no-op + service_name: nemo-rl # service.name reported to the backend + span_groups: default # preset (default | per_step | all) or a comma-separated group list + export_strategy: single_rank # single_rank | all_ranks | sampled | first_rank_per_node + export_rank: -1 # for single_rank: which rank exports (-1 = last rank) + traces_enabled: true # emit trace spans + metrics_enabled: true # emit the rl.* metric instruments + logs_enabled: false # bridge Python logging to OTel logs (trace-correlated) + exporter: otlp # otlp | console + vllm_native_tracing: false # opt in to vLLM's own OTLP tracing (gRPC-only — see vllm-tracing.md) +``` + +The defaults above are the field defaults of `TelemetryConfig` (`nemo_rl/telemetry/config.py`). The endpoint, headers, and protocol are **not** in this block — they come from the standard `OTEL_EXPORTER_OTLP_*` env vars (see below). + +The driver always exports (it hosts the training loop and the metrics logger); `export_strategy` / `export_rank` govern the Ray **worker** ranks. + +## NeMo-RL environment variables + +Each `NEMO_RL_OTEL_*` variable maps onto a `NemoLensConfig` field. Lens reads `NEMO_RL_OTEL_` first and falls back to `NEMO_LENS_`, so you can set a shared `NEMO_LENS_*` default and override it per-run with the RL-scoped prefix. + +| Variable | Maps to | Default | +|---|---|---| +| `NEMO_RL_OTEL_ENABLED` | `enabled` | `0` | +| `NEMO_RL_OTEL_SPAN_GROUPS` | `span_groups` | `default` | +| `NEMO_RL_OTEL_EXPORT_STRATEGY` | `export_strategy` | `single_rank` | +| `NEMO_RL_OTEL_EXPORT_RANK` | `export_rank` | `-1` | +| `NEMO_RL_OTEL_TRACES_ENABLED` | `traces_enabled` | `1` | +| `NEMO_RL_OTEL_METRICS_ENABLED` | `metrics_enabled` | `1` | +| `NEMO_RL_OTEL_LOGS_ENABLED` | `logs_enabled` | `0` | +| `NEMO_RL_OTEL_EXPORTER` | `exporter` | `otlp` | +| `NEMO_RL_OTEL_VLLM_NATIVE_TRACING` | `vllm_native_tracing` | `0` | +| `NEMO_RL_OTEL_RUN_ID` | run identifier | (auto) | +| `NEMO_RL_OTEL_USER_ID` | optional user/team label | (empty) | + +`service_name` maps onto the standard `OTEL_SERVICE_NAME` (lens reads it unprefixed). + +For the full config model, field semantics, and validation rules, see [lens: configuration](https://github.com/NVIDIA-NeMo/Lens). + +## Standard OTel SDK variables + +Endpoint, protocol, and headers are honoured by the OTel SDK directly: + +| Variable | Example | +|---|---| +| `OTEL_SERVICE_NAME` | `nemo-rl` | +| `OTEL_EXPORTER_OTLP_ENDPOINT` | `http://localhost:4317` | +| `OTEL_EXPORTER_OTLP_PROTOCOL` | `grpc` or `http/protobuf` | +| `OTEL_EXPORTER_OTLP_HEADERS` | `
=,
=` (e.g. auth headers your backend requires) | + +Pick the protocol to match your backend: a local collector or Jaeger typically speaks gRPC on `:4317`; a direct-to-SaaS OTLP endpoint typically speaks `http/protobuf` on `:443`. See [Observability Stack](observability-stack.md). + +## Export strategy + +`export_strategy` controls which **worker** ranks actually send telemetry: + +- `single_rank` (default) — only the rank named by `export_rank` (`-1` = last rank). +- `all_ranks` — every worker exports. +- `sampled` / `first_rank_per_node` — sample a subset. + +The driver is independent of this — it always exports. Non-exporting ranks get an empty (`frozenset()`) span-group set, so `is_span_group_enabled()` is `False` everywhere and no span objects are created at all. See [lens: sampling](https://github.com/NVIDIA-NeMo/Lens) for the detailed semantics. + +## Run identification + +Every run gets a `run_id` that flows to all backends as a resource attribute and is shared by the driver and every worker. + +**Priority order:** + +1. `NEMO_RL_OTEL_RUN_ID` (explicit, highest priority). +2. `SLURM_JOB_ID` (auto-detected on SLURM clusters). +3. Auto-generated 12-character hex id (fallback). + +The `run_id` is written to the environment on the driver **before** `init_ray()`, so every worker inherits the same value and correlates to the same run. This is also how vLLM's native spans are correlated back to the RL run — see [vLLM Tracing](vllm-tracing.md). + +Filter by `run_id` in your backend to isolate a specific run. + +## Resource attributes + +`init_telemetry_driver` sets stable-for-the-run values on the OTel `Resource`, so they appear on every span/metric as backend "Process" tags: + +| Attribute | Source | +|---|---| +| `rl.algorithm` | the `algorithm=""` passed to `init_telemetry_driver` | +| `rl.model` | `policy.model_name` | +| `nemo.precision` | `policy.precision` | +| `dl.tensor_parallel.size` | `policy.megatron_cfg` / `dtensor_cfg` TP size | +| `dl.pipeline_parallel.size` | `policy.megatron_cfg` PP size | +| `dl.rank`, `dl.world_size` | set automatically by lens | + +Attribute construction is best-effort: a missing config key simply omits that attribute; it never raises. Plus auto-detected host / GPU / SLURM / Kubernetes attributes from lens's resource detection. + +## Typical configurations + +### Console exporter (no backend) + +```bash +export NEMO_RL_OTEL_ENABLED=1 +export NEMO_RL_OTEL_EXPORTER=console +uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml +``` + +Spans and metrics print to stdout — a quick dry run with no backend to stand up. + +### Direct to an OTLP backend (http/protobuf) + +```bash +export NEMO_RL_OTEL_ENABLED=1 +export OTEL_EXPORTER_OTLP_ENDPOINT=https://: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 +``` + +See [Observability Stack](observability-stack.md) for the full backend-export setup. + +### Per-step granularity + +```bash +export NEMO_RL_OTEL_ENABLED=1 +export NEMO_RL_OTEL_SPAN_GROUPS=per_step +export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317 +``` + +`per_step` makes each training step its own root trace (rollout, generation, reward, advantage, policy update). See [Span Groups](span-groups.md). diff --git a/docs/observability/extending.md b/docs/observability/extending.md new file mode 100644 index 00000000000..6e27f021c11 --- /dev/null +++ b/docs/observability/extending.md @@ -0,0 +1,114 @@ +# Extending Instrumentation + +To add new spans or metrics to NeMo-RL code, use the instrumentation primitives from nemo-lens (`managed_span`, `trace_fn`, `span_cm`). The primitives themselves are documented in [lens: instrumentation](https://github.com/NVIDIA-NeMo/Lens); this page covers NeMo-RL conventions. + +```{tip} +If you work in this repo with Claude Code, the `add-span-group` skill (new span group), the `new-instrument` lens skill (new `rl.*` metric), and the `instrumentation-site-helper` agent (new span/metric site) automate the steps below and keep the cross-repo fallback contract in sync. They are optional — everything here can be done by hand. +``` + +## The import / fallback pattern + +Every lens import in NeMo-RL code must go through `nemo_rl.telemetry._fallbacks`, so the code runs unchanged when nemo-lens is not installed: + +```python +from nemo_rl.telemetry._fallbacks import managed_span, trace_fn +from nemo_rl.telemetry.setup import get_telemetry +from nemo_rl.telemetry.span_groups import RLSpanGroup +``` + +`_fallbacks.py` re-exports the real nemo-lens implementations when it is installed, and provides identical no-op stubs when it is not. Never import from `nemo.lens.*` directly in algorithm code. See [lens: optional dependency](https://github.com/NVIDIA-NeMo/Lens). + +## Adding a span + +### Decorator — `trace_fn` + +For a whole function (this is how `rl.vllm.generate` and the `rl..job` spans are done): + +```python +@trace_fn(RLSpanGroup.GENERATION, "rl.vllm.generate") +def generate(self, ...): + ... +``` + +### Group-gated block — `managed_span` + +For a hot path where you want minimal cost when the group is disabled: + +```python +with managed_span(RLSpanGroup.ROLLOUT, "rl.grpo.collect_rollouts", + **{"rl.iteration": iteration}) as span: + result = collect() + if span is not None: + span.set_attribute("rl.num_generations_per_prompt", n) +``` + +`managed_span` yields `None` when the group is disabled; the body still runs, so guard attribute-setting with `if span is not None`. + +### Always-on block — `span_cm` + +`span_cm` always creates a span when telemetry is active (no group gate) — for cold, top-level paths only: + +```python +telemetry = get_telemetry() +if telemetry is not None: + with span_cm("rl.grpo.job", tracer=telemetry.tracer): + ... +``` + +## Naming conventions + +| Kind | Convention | Example | +|---|---|---| +| Span name | `rl..` | `rl.grpo.collect_rollouts` | +| 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.reward.mean` | + +Metric names use the **application scope** (`rl.*`) — never `dl.*`. Attribute names shared across consumers use the constants in `nemo.lens.semconv`; RL-specific short strings are fine hard-coded. + +## Choosing a span group + +Pick from `RLSpanGroup` before inventing a new one: + +- Once per run (setup/whole-job)? → `job` +- Once per training step? → `step` +- Rollout collection? → `rollout`; generation? → `generation` +- Log-probs? → `logprob` (or `reference_policy` for the reference model) +- Reward / advantage / policy update? → `reward` / `advantage` / `policy_update` +- Checkpoint / eval? → `checkpoint` / `evaluate` + +## Adding a new span group + +If nothing fits, add a group to `RLSpanGroup` in `nemo_rl/telemetry/span_groups.py`: + +1. Add the constant, add it to `ALL_GROUPS`, and slot it into the right preset(s) in `_PRESETS`. Decide per preset: `default` is coarse (rarely add here); `per_step` for per-step spans; `all` always includes it. +2. **Update the fallback stub** in the same file — the stub `SpanGroup` used when nemo-lens is absent must keep the same constants and presets in lockstep. +3. Document the new group in [Span Groups](span-groups.md). + +The `add-span-group` skill walks these steps and keeps the base-class contract (shared with lens and the other consumers) consistent. + +## Adding a metric + +The `rl.*` gauges are populated by teeing `Logger.log_metrics` (see [Metrics](metrics.md)) — not by scattering `record_rl_metrics()` calls. So there are two cases: + +- **The scalar already flows through `Logger.log_metrics`** under a `train` prefix but isn't teed. Add a candidate key (or a new field) to `_RL_OTEL_METRIC_MAP` in `nemo_rl/telemetry/metrics.py`, and add the matching field to `record_rl_metrics` in lens's `nemo.lens.instruments.rl`. +- **You need a brand-new instrument** (a new counter/gauge/histogram, or a value that doesn't go through the Logger). Add it to `nemo.lens.instruments.rl` following the per-Meter `WeakKeyDictionary` caching pattern, then record it from the driver via `telemetry.meter`. The `new-instrument` lens skill covers this. + +Keep `rl..` naming and record only non-`None` values. See [lens: metrics](https://github.com/NVIDIA-NeMo/Lens). + +## Testing new instrumentation + +NeMo-RL telemetry tests live under `tests/` and use lens's in-memory exporter fixtures (global OTel state reset per test). When adding a span: + +1. Assert the span is emitted when its group is enabled and absent when disabled. +2. Assert on span name, tags, and parent relationships. + +For a pure metrics-tee change, `map_rl_metrics` in `nemo_rl/telemetry/metrics.py` is a pure function — unit-test the key mapping directly with no OTel setup. + +## When not to add instrumentation + +- Inside a tight inner loop (per-token) — even a gated `managed_span`'s frozenset lookup adds up. +- On high-cardinality attributes (raw prompts, tensor shapes) — cardinality explosion at the backend. +- As a replacement for logging — structured logs belong in logs (correlate via the log bridge, `NEMO_RL_OTEL_LOGS_ENABLED=1`). + +When in doubt, start with a coarse span at the boundary of a subsystem, not a fine-grained one at every internal call. diff --git a/docs/observability/index.md b/docs/observability/index.md new file mode 100644 index 00000000000..59a98c32fd6 --- /dev/null +++ b/docs/observability/index.md @@ -0,0 +1,87 @@ +# 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 reward, loss, KL, throughput, and more. + +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 **entirely optional**. When nemo-lens is not installed or telemetry is disabled, every instrumentation site is a ~0-cost no-op — see `nemo_rl/telemetry/_fallbacks.py`. + +## 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, the `NEMO_RL_OTEL_*` environment variables, 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_OTEL_*` env vars | 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 + +Telemetry needs the nemo-lens SDK (pulled from PyPI by the `telemetry` extra): + +```bash +uv sync --extra telemetry # or, to add just the SDK: uv pip install 'nemo-lens[sdk]' +``` + +## Quick start + +```bash +export NEMO_RL_OTEL_ENABLED=1 +export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317 # your OTLP backend / collector +export NEMO_RL_OTEL_SPAN_GROUPS=default # coarse-grained; safe for production + +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 a steady stream of `rl.*` metrics. Switch to `per_step` for per-step traces (rollout/generation/reward/...), or `all` for everything. + +You do not have to touch the config file: telemetry can be driven purely by env vars, or by adding a `telemetry:` block to your run config. Raw env vars always win over YAML. 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()` at the end of `main()`. + +| Algorithm | Entry point | Representative spans | +|---|---|---| +| GRPO (sync + async) | `examples/run_grpo.py` | `rl.grpo.step`, `rl.grpo.collect_rollouts`, `rl.grpo.compute_rewards`, `rl.grpo.compute_logprobs`, `rl.grpo.compute_advantages`, `rl.grpo.policy_update` | +| PPO | `examples/run_ppo.py` | `rl.ppo.step`, `rl.ppo.collect_rollouts`, `rl.ppo.compute_rewards`, `rl.ppo.compute_advantages`, `rl.ppo.policy_update`, `rl.ppo.value_update` | +| SFT | `examples/run_sft.py` | `rl.sft.step`, `rl.sft.data_processing`, `rl.sft.policy_update` | +| DPO | `examples/run_dpo.py` | `rl.dpo.step`, `rl.dpo.policy_update` | +| RM | `examples/run_rm.py` | `rl.rm.step` | +| Distillation | `examples/run_distillation.py` | `rl.distillation.step`, `rl.distillation.collect_rollouts`, `rl.distillation.teacher_logprobs`, `rl.distillation.policy_update` | +| 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.*` catalog (reward, loss, KL, grad norm, throughput, generation/rollout latency) teed from the driver's metrics logger — see [Metrics](metrics.md). +- **Logs** (optional): via the OTel log bridge when `NEMO_RL_OTEL_LOGS_ENABLED=1` — 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 00000000000..7d096509171 --- /dev/null +++ b/docs/observability/metrics.md @@ -0,0 +1,74 @@ +# Metrics + +NeMo-RL emits two namespaces of metrics: RL training metrics (`rl.*`) 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). + +## RL training metrics (`rl.*`) + +Training has no OTel standard, so NeMo-RL uses a project-specific application scope. These are recorded via `nemo.lens.instruments.rl.record_rl_metrics` from the driver. + +| Metric | Type | Description | +|---|---|---| +| `rl.reward.mean` | Gauge | Mean reward over the batch | +| `rl.kl_divergence` | Gauge | KL divergence from the reference policy | +| `rl.policy_loss` | Gauge | Policy (actor) loss | +| `rl.value_loss` | Gauge | Value (critic) loss — PPO | +| `rl.entropy` | Gauge | Policy entropy | +| `rl.response_length.mean` | Gauge | Mean generated response length (tokens) | +| `rl.grad_norm` | Gauge | Global gradient norm | +| `rl.learning_rate` | Gauge | Current learning rate | +| `rl.tokens_per_sec` | Gauge | Training throughput (tokens/sec) | +| `rl.generation.duration_ms` | Histogram | Generation latency (ms) | +| `rl.rollout.duration_ms` | Histogram | Rollout-collection latency (ms) | + +Loss, reward, KL, grad norm, learning rate, and throughput are **Gauges** (point-in-time value per log step), not Histograms — semantically correct for a value that changes every log interval. + +## How the `rl.*` gauges are populated (the Logger tee) + +NeMo-RL does not sprinkle `record_rl_metrics()` calls through the algorithm code. Instead, `nemo_rl/telemetry/metrics.py` **tees** the scalar metrics that already flow through `nemo_rl.utils.logger.Logger.log_metrics` into nemo-lens: after `log_metrics` fans out to the file / W&B / MLflow backends, it calls `tee_rl_metrics_to_otel(metrics, prefix)`. + +Only the driver's **`train`-prefix** metrics are teed (`prefix in ("train", "")`); other prefixes are skipped. The tee is best-effort — a raw metrics dict is matched against a fixed key map, the first present candidate key wins, and unknown keys or non-scalar values are silently skipped. It is a no-op unless telemetry is actively exporting. + +| Logger metric key (first match wins) | Emitted metric | +|---|---| +| `reward` / `reward_mean` / `mean_reward` | `rl.reward.mean` | +| `kl` / `kl_divergence` / `mean_kl` | `rl.kl_divergence` | +| `loss` / `policy_loss` | `rl.policy_loss` | +| `value_loss` / `critic_loss` | `rl.value_loss` | +| `entropy` | `rl.entropy` | +| `mean_gen_tokens_per_sample` / `response_length_mean` | `rl.response_length.mean` | +| `grad_norm` | `rl.grad_norm` | +| `lr` / `learning_rate` | `rl.learning_rate` | +| `valid_tokens_per_sec_per_gpu` / `tokens_per_sec` | `rl.tokens_per_sec` | + +This means the metrics you already log to W&B are the same series you get in your OTLP backend — no double bookkeeping. If an algorithm logs a scalar under a key not in this map, add a candidate to `_RL_OTEL_METRIC_MAP` (see [Extending](extending.md)). + +## 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 | reward, loss, KL, grad norm, throughput → `rl.*` | +| **Span tag** | categorical per-span context for filtering | `rl.iteration`, `rl.backend`, `rl.num_generations_per_prompt` | +| **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). + +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 00000000000..1018f094079 --- /dev/null +++ b/docs/observability/observability-stack.md @@ -0,0 +1,48 @@ +# 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 + +Set these on the process that runs training (the driver, and via Ray its workers inherit them): + +```bash +NEMO_RL_OTEL_ENABLED=1 +NEMO_RL_OTEL_SPAN_GROUPS=default # start coarse; raise to per_step / all as needed +NEMO_RL_OTEL_METRICS_ENABLED=1 +NEMO_RL_OTEL_LOGS_ENABLED=1 +NEMO_RL_OTEL_VLLM_NATIVE_TRACING=0 # gRPC-only; leave OFF on an http/protobuf path + +# OTLP target — set these for your backend or collector: +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 `NEMO_RL_OTEL_SPAN_GROUPS` to `per_step` (or `all`) for per-step traces. Set an explicit `NEMO_RL_OTEL_RUN_ID` (and optional `NEMO_RL_OTEL_USER_ID`) to name the run instead of taking the auto-generated id. + +## Console / JSON output (no backend) + +To confirm spans and metrics are produced without standing up any backend, use the `console` exporter: + +```bash +NEMO_RL_OTEL_ENABLED=1 NEMO_RL_OTEL_EXPORTER=console \ + uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml +``` + +Each span and metric prints to stdout as **JSON** (`ConsoleSpanExporter` uses `span.to_json()`), so you can capture it to a file: + +```bash +NEMO_RL_OTEL_ENABLED=1 NEMO_RL_OTEL_EXPORTER=console \ + uv run examples/run_grpo.py --config examples/configs/grpo_math_1B.yaml > telemetry.json 2>&1 +``` + +`console` (set via `telemetry.exporter` / `NEMO_RL_OTEL_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 (`NEMO_RL_OTEL_VLLM_NATIVE_TRACING=1`) 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 00000000000..dde0089fc64 --- /dev/null +++ b/docs/observability/span-groups.md @@ -0,0 +1,101 @@ +# Span Groups + +Span granularity in NeMo-RL is controlled by `NEMO_RL_OTEL_SPAN_GROUPS` (or `span_groups` in the `telemetry:` block). 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`, `rollout`, `generation`, `logprob`, `reward`, `advantage`, `policy_update`, `reference_policy`, `data_processing` | 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 | model initialisation spans | +| `load_checkpoint` | base | checkpoint restore spans | +| `step` | base | `rl..step` (one per training step) | +| `forward_backward` | base | forward/backward spans | +| `optimizer` | base | optimizer-step spans | +| `rollout` | RL | `rl..collect_rollouts` | +| `generation` | RL | `rl.` generation + 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 | reference-policy log-prob spans | +| `data_processing` | RL | `rl..data_processing` | + +## Examples + +```bash +# Coarse spans only — default +NEMO_RL_OTEL_SPAN_GROUPS=default + +# Per-step traces (rollout / generation / reward / advantage / policy update) +NEMO_RL_OTEL_SPAN_GROUPS=per_step + +# Coarse job trace + generation spans only +NEMO_RL_OTEL_SPAN_GROUPS=default,generation + +# Everything +NEMO_RL_OTEL_SPAN_GROUPS=all +``` + +## Per-algorithm span names + +Span names follow `rl..`. 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.collect_rollouts`, `rl.grpo.compute_rewards`, `rl.grpo.compute_logprobs`, `rl.grpo.compute_advantages`, `rl.grpo.policy_update`, `rl.grpo.save_checkpoint`, `rl.grpo.evaluate` | +| **PPO** | `rl.ppo.job`, `rl.ppo.step`, `rl.ppo.data_processing`, `rl.ppo.collect_rollouts`, `rl.ppo.compute_rewards`, `rl.ppo.compute_logprobs`, `rl.ppo.compute_advantages`, `rl.ppo.policy_update`, `rl.ppo.value_update`, `rl.ppo.save_checkpoint`, `rl.ppo.evaluate` | +| **SFT** | `rl.sft.job`, `rl.sft.step`, `rl.sft.data_processing`, `rl.sft.policy_update`, `rl.sft.save_checkpoint`, `rl.sft.evaluate` | +| **DPO** | `rl.dpo.job`, `rl.dpo.step`, `rl.dpo.policy_update`, `rl.dpo.save_checkpoint`, `rl.dpo.evaluate` | +| **RM** | `rl.rm.job`, `rl.rm.step`, `rl.rm.save_checkpoint`, `rl.rm.evaluate` | +| **Distillation** | `rl.distillation.job`, `rl.distillation.step`, `rl.distillation.data_processing`, `rl.distillation.collect_rollouts`, `rl.distillation.teacher_logprobs`, `rl.distillation.policy_update`, `rl.distillation.save_checkpoint`, `rl.distillation.evaluate` | +| **vLLM** (driver-side) | `rl.vllm.generate`, `rl.vllm.generate_text` — `generation` group; nested under the active rollout span | + +`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.backend` | generation backend (e.g. `"vllm"`) | + +## 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 (`NEMO_RL_OTEL_ENABLED=0`) | None | Default for smoke tests | +| `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 00000000000..65d2e2a9243 --- /dev/null +++ b/docs/observability/vllm-tracing.md @@ -0,0 +1,49 @@ +# 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: `NEMO_RL_OTEL_VLLM_NATIVE_TRACING=1` | +| 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. + +## Layer 2 — vLLM native OTLP tracing (opt-in) + +vLLM can emit its own OpenTelemetry spans for the engine internals. Enable it with: + +```bash +export NEMO_RL_OTEL_VLLM_NATIVE_TRACING=1 +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 5926a1abaf1..08464fe3b0e 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,6 +68,11 @@ def main() -> None: # Get the next experiment directory with incremented ID config.logger["log_dir"] = get_next_experiment_dir(config.logger["log_dir"]) + # 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") + init_ray() tokenizer = get_tokenizer(config.policy["tokenizer"]) @@ -119,6 +125,9 @@ def main() -> None: master_config, ) + # Flush and shut down telemetry (no-op when telemetry is inactive). + shutdown_telemetry() + if __name__ == "__main__": main() diff --git a/examples/run_dpo.py b/examples/run_dpo.py index df05302528e..4d6212e0294 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,6 +69,11 @@ def main(): f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" ) + # 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") + init_ray() # setup tokenizer @@ -103,6 +109,9 @@ def main(): dpo_save_state, ) + # Flush and shut down telemetry (no-op when telemetry is inactive). + shutdown_telemetry() + if __name__ == "__main__": main() diff --git a/examples/run_grpo.py b/examples/run_grpo.py index 6732ae0115c..e2e6d4022eb 100644 --- a/examples/run_grpo.py +++ b/examples/run_grpo.py @@ -24,6 +24,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, @@ -102,6 +103,11 @@ def main() -> None: f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" ) + # 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") + with rl_init_timer.time("ray_connect"): init_ray() @@ -237,6 +243,9 @@ def _make_policy(**kwargs): master_config, ) + # Flush and shut down telemetry (no-op when telemetry is inactive). + shutdown_telemetry() + if __name__ == "__main__": main() diff --git a/examples/run_ppo.py b/examples/run_ppo.py index 6d148ffd673..e1360cd8a1f 100644 --- a/examples/run_ppo.py +++ b/examples/run_ppo.py @@ -23,6 +23,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, @@ -78,6 +79,11 @@ def main() -> None: f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" ) + # 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") + init_ray() # setup tokenizer @@ -135,6 +141,9 @@ def main() -> None: master_config, ) + # Flush and shut down telemetry (no-op when telemetry is inactive). + shutdown_telemetry() + if __name__ == "__main__": main() diff --git a/examples/run_rm.py b/examples/run_rm.py index 19b6940e016..1ec9c10bef9 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,6 +72,11 @@ def main(): f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" ) + # 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") + init_ray() # setup tokenizer @@ -106,6 +112,9 @@ def main(): rm_save_state, ) + # Flush and shut down telemetry (no-op when telemetry is inactive). + shutdown_telemetry() + if __name__ == "__main__": main() diff --git a/examples/run_sft.py b/examples/run_sft.py index e06011a8ea1..1b7544cb834 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,6 +190,11 @@ def main(is_vlm: bool = False): f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" ) + # 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") + init_ray() # setup tokenizer (or processor) @@ -224,6 +230,9 @@ def main(is_vlm: bool = False): sft_save_state, ) + # Flush and shut down telemetry (no-op when telemetry is inactive). + shutdown_telemetry() + if __name__ == "__main__": main() diff --git a/nemo_rl/algorithms/distillation.py b/nemo_rl/algorithms/distillation.py index 985cb9fd05d..36f66ccbb00 100644 --- a/nemo_rl/algorithms/distillation.py +++ b/nemo_rl/algorithms/distillation.py @@ -91,6 +91,15 @@ ) from nemo_rl.weight_sync.factory import create_weight_synchronizer +try: + from nemo.lens.helpers import managed_span, trace_fn +except ImportError: + from nemo_rl.telemetry._fallbacks import managed_span, trace_fn + +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.setup import get_telemetry +from nemo_rl.telemetry.span_groups import RLSpanGroup + # =============================================================================== # Configuration # =============================================================================== @@ -163,6 +172,7 @@ class MasterConfig(BaseModel, extra="allow"): logger: LoggerConfig # Logger configuration cluster: ClusterConfig # Cluster configuration checkpointing: CheckpointingConfig # Checkpointing configuration + telemetry: Optional[TelemetryConfig] = None # =============================================================================== @@ -676,6 +686,7 @@ def init_nemo_gym(): # =============================================================================== +@trace_fn(RLSpanGroup.JOB, "rl.distillation.job") def distillation_train( student_policy: ColocatablePolicyInterface, teacher_policy: ColocatablePolicyInterface, @@ -693,6 +704,8 @@ def distillation_train( ) -> None: """Run Distillation training algorithm.""" timer = Timer() + _telemetry = get_telemetry() + _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, @@ -775,10 +788,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( @@ -803,7 +831,14 @@ def distillation_train( else: student_generation.prepare_for_generation() - with timer.time("generation"): + with ( + timer.time("generation"), + managed_span( + RLSpanGroup.ROLLOUT, + "rl.distillation.collect_rollouts", + tracer=_tracer, + ), + ): # We cascade NeMo-Gym first since NeMo-Gym requires async rollouts. if use_nemo_gym: generation_config = master_config.policy["generation"] @@ -905,7 +940,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_logprobs", + tracer=_tracer, + ), + ): teacher_topk = teacher_policy.get_topk_logits( train_data, k=master_config.distillation.topk_logits_k, @@ -921,7 +963,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_update", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): train_results = student_policy.train( train_data, loss_fn, @@ -1045,7 +1095,14 @@ def distillation_train( metrics_source[metric_name], ) - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.distillation.save_checkpoint", + tracer=_tracer, + ), + ): print( f"Saving checkpoint for step {total_steps + 1}...", flush=True, @@ -1198,7 +1255,17 @@ def validate( use_nemo_gym = _should_use_nemo_gym(master_config) timer = Timer() - with timer.time("total_validation_time"): + _telemetry = get_telemetry() + _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 de0a60af265..7977a98fb4a 100644 --- a/nemo_rl/algorithms/dpo.py +++ b/nemo_rl/algorithms/dpo.py @@ -43,6 +43,15 @@ from nemo_rl.utils.nsys import maybe_gpu_profile_step from nemo_rl.utils.timer import TimeoutChecker, Timer +try: + from nemo.lens.helpers import managed_span, trace_fn +except ImportError: + from nemo_rl.telemetry._fallbacks import managed_span, trace_fn + +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.setup import get_telemetry +from nemo_rl.telemetry.span_groups import RLSpanGroup + @dataclass class DPOSaveState: @@ -105,6 +114,7 @@ class MasterConfig(BaseModel, extra="allow"): logger: LoggerConfig cluster: ClusterConfig checkpointing: CheckpointingConfig + telemetry: Optional[TelemetryConfig] = None @dataclass @@ -418,8 +428,18 @@ def validate_one_dataset( return timer = Timer() - - with timer.time("total_validation_time"): + _telemetry = get_telemetry() + _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 +543,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 +557,8 @@ def dpo_train( ) -> None: # Run dpo training timer = Timer() + _telemetry = get_telemetry() + _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 +613,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_update", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): train_results = policy.train( batch, loss_fn, @@ -722,7 +761,14 @@ def dpo_train( metrics_source[metric_name], ) - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.dpo.save_checkpoint", + 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 4159e5d5b60..d4f0af63665 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -71,6 +71,7 @@ ) from nemo_rl.data.utils import extract_necessary_env_names, load_dataloader_state from nemo_rl.data_plane.interfaces import DataPlaneConfig +from nemo_rl.telemetry.config import TelemetryConfig from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.distributed.virtual_cluster import ( @@ -132,6 +133,14 @@ ) from nemo_rl.weight_sync.factory import create_weight_synchronizer +try: + from nemo.lens.helpers import managed_span, trace_fn +except ImportError: + from nemo_rl.telemetry._fallbacks import managed_span, trace_fn + +from nemo_rl.telemetry.setup import get_telemetry +from nemo_rl.telemetry.span_groups import RLSpanGroup + # =============================================================================== # Configuration # =============================================================================== @@ -336,6 +345,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 # =============================================================================== @@ -2593,6 +2603,7 @@ def _validation_early_stop_message( ) +@trace_fn(RLSpanGroup.JOB, "rl.grpo.job") def grpo_train( policy: ColocatablePolicyInterface, policy_generation: Optional[GenerationInterface], @@ -2609,6 +2620,8 @@ def grpo_train( ) -> None: """Run GRPO training algorithm.""" timer = Timer(context={"worker": "driver"}) + _telemetry = get_telemetry() + _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, @@ -2737,10 +2750,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, + ), + ): # Repeat batch items repeated_batch: BatchedDataDict[DatumSpec] = ( batch.repeat_interleave( @@ -2812,7 +2840,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.collect_rollouts", + 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() @@ -2907,7 +2945,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.compute_rewards", tracer=_tracer + ), + ): # Extract rewards from final_batch rewards = repeated_batch["total_reward"] @@ -3006,7 +3049,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() @@ -3079,7 +3129,12 @@ 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.compute_logprobs", tracer=_tracer + ), + ): # Custom create this logprob_data so we avoid Ray comm overheads sending unused data to workers. logprob_data = BatchedDataDict[ClippedPGLossDataDict]( { @@ -3150,7 +3205,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.compute_advantages", + tracer=_tracer, + ), + ): print("▶ Computing advantages...", flush=True) # Get token-level mask: token_mask * sample_mask token_mask = train_data["token_mask"] @@ -3195,7 +3257,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_update", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): train_results = policy.train( train_data, loss_fn, @@ -3416,7 +3486,14 @@ def grpo_train( metrics_source[metric_name], ) - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.grpo.save_checkpoint", + tracer=_tracer, + ), + ): # Finalize the previous (possibly async) checkpoint before # starting a new one. No-op with sync save / nothing pending. checkpointer.finalize_pending() @@ -3679,7 +3756,17 @@ def validate( return {}, {} timer = Timer(context={"worker": "validator"}) - with timer.time("total_validation_time"): + _telemetry = get_telemetry() + _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}, + ), + ): print(f"▶ Starting validation at step {step}...", flush=True) total_rewards = [] @@ -3863,6 +3950,7 @@ def aggregate_rollout_metrics( return aggregated +@trace_fn(RLSpanGroup.JOB, "rl.grpo.job") def async_grpo_train( policy: ColocatablePolicyInterface, policy_generation: Optional[GenerationInterface], @@ -3936,6 +4024,8 @@ def async_grpo_train( from nemo_rl.algorithms.async_utils import AsyncTrajectoryCollector, ReplayBuffer timer = Timer(context={"worker": "driver"}) + _telemetry = get_telemetry() + _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"], @@ -4292,7 +4382,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 @@ -4439,7 +4537,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.compute_rewards", 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"]) @@ -4466,7 +4569,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 @@ -4524,7 +4634,12 @@ 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.compute_logprobs", tracer=_tracer + ), + ): if not skip_prev_logprobs: train_data["prev_logprobs"] = policy.get_logprobs( train_data, timer=timer @@ -4573,7 +4688,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.compute_advantages", + tracer=_tracer, + ), + ): print("▶ Computing advantages...", flush=True) # Get token-level mask: token_mask * sample_mask token_mask = train_data["token_mask"] @@ -4634,7 +4756,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_update", + tracer=_tracer, + **{"rl.iteration": step + 1}, + ), + ): train_results = policy.train( train_data, loss_fn, @@ -4869,7 +4999,14 @@ def async_grpo_train( metrics_source[metric_name], ) - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.grpo.save_checkpoint", + tracer=_tracer, + ), + ): # Finalize the previous (possibly async) checkpoint before # starting a new one. No-op with sync save / nothing pending. checkpointer.finalize_pending() diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index f12aaabfb2c..f3441c4beb3 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -89,6 +89,15 @@ from nemo_rl.utils.nsys import maybe_gpu_profile_step from nemo_rl.utils.timer import TimeoutChecker, Timer +try: + from nemo.lens.helpers import managed_span, trace_fn +except ImportError: + from nemo_rl.telemetry._fallbacks import managed_span, trace_fn + +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.setup import get_telemetry +from nemo_rl.telemetry.span_groups import RLSpanGroup + # =============================================================================== # Configuration # =============================================================================== @@ -185,6 +194,7 @@ class MasterConfig(BaseModel, extra="allow"): logger: PPOLoggerConfig cluster: ClusterConfig checkpointing: CheckpointingConfig + telemetry: Optional[TelemetryConfig] = None # =============================================================================== @@ -883,6 +893,7 @@ def _create_advantage_estimator(master_config: MasterConfig): # =============================================================================== +@trace_fn(RLSpanGroup.JOB, "rl.ppo.job") def ppo_train( policy: ColocatablePolicyInterface, policy_generation: Optional[GenerationInterface], @@ -908,6 +919,8 @@ def ppo_train( - Configurable policy training start epoch """ timer = Timer() + _telemetry = get_telemetry() + _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, @@ -997,10 +1010,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"] @@ -1061,7 +1089,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.collect_rollouts", + tracer=_tracer, + ), + ): if policy_generation is not None: policy_generation.clear_logger_metrics() @@ -1133,7 +1168,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.compute_rewards", + tracer=_tracer, + ), + ): rewards = repeated_batch["total_reward"] with timer.time("data_processing"): @@ -1207,7 +1249,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.compute_logprobs", + tracer=_tracer, + ), + ): logprob_data = BatchedDataDict[ClippedPGLossDataDict]( { "input_ids": train_data["input_ids"], @@ -1237,7 +1286,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.compute_advantages", + tracer=_tracer, + ), + ): print("▶ Computing advantages...", flush=True) initial_prompt_message_logs = extract_initial_prompt_messages( repeated_batch["message_log"], @@ -1283,7 +1339,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_update", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): print("▶ Training value...", flush=True) value_results = value_model.train( train_data, @@ -1310,7 +1374,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_update", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): train_results = policy.train( train_data, loss_fn, @@ -1553,7 +1625,14 @@ def ppo_train( metric_name ] - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.ppo.save_checkpoint", + tracer=_tracer, + ), + ): print( f"Saving checkpoint for step {total_steps + 1}...", flush=True, @@ -1737,7 +1816,16 @@ def validate( return {}, {} timer = Timer() - with timer.time("total_validation_time"): + _telemetry = get_telemetry() + _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, + ), + ): 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 0cbe3fc02dc..35cd4b4b383 100644 --- a/nemo_rl/algorithms/rm.py +++ b/nemo_rl/algorithms/rm.py @@ -44,6 +44,15 @@ from nemo_rl.utils.nsys import maybe_gpu_profile_step from nemo_rl.utils.timer import TimeoutChecker, Timer +try: + from nemo.lens.helpers import managed_span, trace_fn +except ImportError: + from nemo_rl.telemetry._fallbacks import managed_span, trace_fn + +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.setup import get_telemetry +from nemo_rl.telemetry.span_groups import RLSpanGroup + @dataclass class RMSaveState: @@ -96,6 +105,7 @@ class MasterConfig(BaseModel, extra="allow"): logger: LoggerConfig cluster: ClusterConfig checkpointing: CheckpointingConfig + telemetry: Optional[TelemetryConfig] = None @dataclass @@ -360,8 +370,17 @@ def validate_one_dataset( return timer = Timer() - - with timer.time("total_validation_time"): + _telemetry = get_telemetry() + _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 +485,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 +499,8 @@ def rm_train( ): # Run basic rm training timer = Timer() + _telemetry = get_telemetry() + _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 +551,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 +687,14 @@ def rm_train( metrics_source[metric_name], ) - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.rm.save_checkpoint", + 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 b6d96778ef0..819e1d199a9 100644 --- a/nemo_rl/algorithms/sft.py +++ b/nemo_rl/algorithms/sft.py @@ -46,6 +46,15 @@ from nemo_rl.utils.nsys import maybe_gpu_profile_step from nemo_rl.utils.timer import TimeoutChecker, Timer +try: + from nemo.lens.helpers import managed_span, trace_fn +except ImportError: + from nemo_rl.telemetry._fallbacks import managed_span, trace_fn + +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.setup import get_telemetry +from nemo_rl.telemetry.span_groups import RLSpanGroup + @dataclass class SFTSaveState: @@ -99,6 +108,7 @@ class MasterConfig(BaseModel, extra="allow"): logger: LoggerConfig cluster: ClusterConfig checkpointing: CheckpointingConfig + telemetry: Optional[TelemetryConfig] = None # ======================================================= @@ -277,8 +287,18 @@ def validate( return {}, {} timer = Timer() - - with timer.time("total_validation_time"): + _telemetry = get_telemetry() + _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 +396,7 @@ def validate( return val_metrics, timing_metrics +@trace_fn(RLSpanGroup.JOB, "rl.sft.job") def sft_train( policy, train_dataloader, @@ -389,6 +410,8 @@ def sft_train( ) -> None: # Run basic sft training timer = Timer() + _telemetry = get_telemetry() + _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 +464,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 +511,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_update", + tracer=_tracer, + **{"rl.iteration": total_steps + 1}, + ), + ): train_results = policy.train( train_data, loss_fn, @@ -579,7 +625,14 @@ def sft_train( metrics_source[metric_name], ) - with timer.time("checkpointing"): + with ( + timer.time("checkpointing"), + managed_span( + RLSpanGroup.CHECKPOINT, + "rl.sft.save_checkpoint", + 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/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 6d0d8d37143..a6702b73883 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -45,9 +45,48 @@ ) from nemo_rl.weight_sync.interfaces import WeightSynchronizer +try: + from nemo.lens.helpers import trace_fn +except ImportError: + from nemo_rl.telemetry._fallbacks import trace_fn + +from nemo_rl.telemetry.setup import get_telemetry +from nemo_rl.telemetry.span_groups import RLSpanGroup + 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() + if telemetry is None or not telemetry.is_exporting: + return + try: + from nemo.lens.instruments.inference import record_inference_metrics + + 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: + logger.debug("nemo-lens: failed to record vLLM metrics", exc_info=True) + + class VllmGeneration(GenerationInterface): @staticmethod def init_cluster_placement_groups( @@ -622,6 +661,7 @@ def init_collective( # this function should co-work with lm_policy, so we should wait for all futures to complete outside return futures + @trace_fn(RLSpanGroup.GENERATION, "rl.vllm.generate") def generate( self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False ) -> BatchedDataDict[GenerationOutputSpec]: @@ -668,8 +708,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]: @@ -714,6 +756,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 ed3f045274e..c4815843837 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -77,6 +77,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. + + Opt-in via ``NEMO_RL_OTEL_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. + """ + if os.environ.get("NEMO_RL_OTEL_VLLM_NATIVE_TRACING", "").strip().lower() not in ( + "1", + "true", + "yes", + "on", + ): + 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 + try: + import inspect + + from vllm.engine.arg_utils import EngineArgs + + supported = set(getattr(EngineArgs, "__dataclass_fields__", {})) | set( + inspect.signature(EngineArgs.__init__).parameters + ) + except Exception: + 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: @@ -596,6 +649,8 @@ def _load_model(self, bundle_indices, seed): if logprobs_mode is not None: llm_kwargs["logprobs_mode"] = logprobs_mode + _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 diff --git a/nemo_rl/telemetry/README.md b/nemo_rl/telemetry/README.md new file mode 100644 index 00000000000..7f686320a9a --- /dev/null +++ b/nemo_rl/telemetry/README.md @@ -0,0 +1,57 @@ +# 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, throughput, generation/rollout latency) 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 / shutdown_telemetry +├── span_groups.py — RLSpanGroup: RL-specific span groups + presets +├── metrics.py — tees Logger.log_metrics scalars into the rl.* instruments +├── _fallbacks.py — no-op shims for when nemo-lens is not installed +└── __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()` at the end of `main()`. `get_telemetry()` returns the process-global `TelemetryHandle`; `init_telemetry_worker()` sets up telemetry inside a Ray actor. + +## Install + +```bash +uv sync --extra telemetry # or, to add just the SDK: uv pip install 'nemo-lens[sdk]' +``` + +## Quick start + +```bash +export NEMO_RL_OTEL_ENABLED=1 +export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317 +export NEMO_RL_OTEL_SPAN_GROUPS=default + +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 00000000000..279130f8ba5 --- /dev/null +++ b/nemo_rl/telemetry/__init__.py @@ -0,0 +1,50 @@ +# 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` / + :func:`~nemo_rl.telemetry.setup.shutdown_telemetry` — lifecycle helpers. + +The instrumentation primitives (``managed_span`` / ``trace_fn`` / ``span_cm`` / +``is_span_group_enabled`` / ``safe_set_span_attributes``) come from +:mod:`nemo_rl.telemetry._fallbacks`, which re-exports the real nemo-lens +implementations when it is installed and no-op stubs when it is not. Importing +this package never requires nemo-lens. +""" + +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.setup import ( + get_telemetry, + init_telemetry_driver, + init_telemetry_worker, + shutdown_telemetry, +) +from nemo_rl.telemetry.span_groups import RLSpanGroup + +__all__ = [ + "TelemetryConfig", + "RLSpanGroup", + "get_telemetry", + "init_telemetry_driver", + "init_telemetry_worker", + "shutdown_telemetry", +] diff --git a/nemo_rl/telemetry/_fallbacks.py b/nemo_rl/telemetry/_fallbacks.py new file mode 100644 index 00000000000..2cd46659810 --- /dev/null +++ b/nemo_rl/telemetry/_fallbacks.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""No-op fallbacks for when nemo-lens is not installed. + +When nemo-lens IS installed, re-exports from nemo.lens.fallbacks for +consistency. When it is NOT installed, provides identical local no-ops. +""" + +try: + from nemo.lens.fallbacks import ( # noqa: F401 + is_span_group_enabled, + managed_span, + safe_set_span_attributes, + span_cm, + trace_fn, + ) +except ImportError: + from contextlib import contextmanager + + def trace_fn(group, name, tracer=None): + """No-op decorator — returns the function unchanged.""" + + def decorator(func): + return func + + return decorator + + @contextmanager + def managed_span(group, name, tracer=None, **attributes): + """No-op context manager — yields None.""" + yield None + + def is_span_group_enabled(group): + """Always returns False when nemo-lens is not installed.""" + return False + + def safe_set_span_attributes(span, attributes, redact_keys=None): + """No-op.""" + pass + + @contextmanager + def span_cm(name, tracer=None, record_exception=True, **attributes): + """No-op context manager — yields None.""" + yield None diff --git a/nemo_rl/telemetry/config.py b/nemo_rl/telemetry/config.py new file mode 100644 index 00000000000..c81f3285c1b --- /dev/null +++ b/nemo_rl/telemetry/config.py @@ -0,0 +1,76 @@ +# 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 pydantic import BaseModel + + +class TelemetryConfig(BaseModel, extra="allow"): + """OpenTelemetry / nemo-lens configuration. + + Telemetry is optional: it activates only when ``enabled`` is true *and* + nemo-lens is installed (``uv sync --extra telemetry``). When either is + absent, every instrumentation site degrades to a ~0-cost no-op. + """ + + 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: str = "single_rank" + """Which ranks export: ``single_rank`` | ``all_ranks`` | ``sampled`` | + ``first_rank_per_node``. The driver always exports (it runs the training + loop and the metrics logger); this governs the Ray worker ranks.""" + + export_rank: int = -1 + """For ``single_rank``: which rank exports (``-1`` = last 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: str = "otlp" + """Exporter backend: ``otlp`` | ``console``. 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/metrics.py b/nemo_rl/telemetry/metrics.py new file mode 100644 index 00000000000..3ad98ce4170 --- /dev/null +++ b/nemo_rl/telemetry/metrics.py @@ -0,0 +1,88 @@ +# 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 scalar metrics into nemo-lens. + +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. +""" + +from __future__ import annotations + +import logging +from typing import Any, Optional + +from nemo_rl.telemetry.setup import get_telemetry + +logger = logging.getLogger(__name__) + +# Map raw Logger metric keys (under the "train"/"" prefix) to +# ``record_rl_metrics`` gauge fields. The first present candidate key wins. +# Best-effort: unmatched keys and non-scalar values are silently skipped. +_RL_OTEL_METRIC_MAP: dict[str, tuple[str, ...]] = { + "reward_mean": ("reward", "reward_mean", "mean_reward"), + "kl_divergence": ("kl", "kl_divergence", "mean_kl"), + "policy_loss": ("loss", "policy_loss"), + "value_loss": ("value_loss", "critic_loss"), + "entropy": ("entropy",), + "response_length_mean": ("mean_gen_tokens_per_sample", "response_length_mean"), + "grad_norm": ("grad_norm",), + "learning_rate": ("lr", "learning_rate"), + "tokens_per_sec": ("valid_tokens_per_sec_per_gpu", "tokens_per_sec"), +} + + +def map_rl_metrics(metrics: dict[str, Any]) -> dict[str, float]: + """Extract the ``record_rl_metrics`` kwargs present in a raw metrics dict. + + Pure function (no OTel side effects) so it is trivially unit-testable. + """ + kwargs: dict[str, float] = {} + for field, candidates in _RL_OTEL_METRIC_MAP.items(): + for key in candidates: + value = metrics.get(key) + # Exclude bools (a subclass of int) and non-numeric values. + if isinstance(value, bool) or not isinstance(value, (int, float)): + continue + kwargs[field] = float(value) + break + return kwargs + + +def tee_rl_metrics_to_otel(metrics: dict[str, Any], prefix: Optional[str]) -> None: + """Mirror standard RL scalar metrics into nemo-lens (no-op unless exporting). + + Only the driver's per-step ``train`` metrics are teed. The OTel instruments + are touched only when telemetry is actively exporting and nemo-lens is + installed; everything else short-circuits to a no-op. + """ + if prefix not in ("train", ""): + return + telemetry = get_telemetry() + if telemetry is None or not telemetry.is_exporting: + return + try: + from nemo.lens.instruments.rl import record_rl_metrics + except ImportError: + return + + kwargs = map_rl_metrics(metrics) + if not kwargs: + return + try: + record_rl_metrics(telemetry.meter, **kwargs) + except Exception: + logger.debug("nemo-lens: failed to tee RL metrics", exc_info=True) diff --git a/nemo_rl/telemetry/setup.py b/nemo_rl/telemetry/setup.py new file mode 100644 index 00000000000..2ea7557e7bd --- /dev/null +++ b/nemo_rl/telemetry/setup.py @@ -0,0 +1,304 @@ +# 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`` 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 ``__init__`` / ``post_init``). It reads the propagated env and sets up + that worker's telemetry. + +Importing this module never requires nemo-lens: every lens import is +function-local and guarded by ``try/except ImportError``. When lens is not +installed (or telemetry is disabled), the init functions return ``None`` and all +instrumentation sites stay no-ops via ``nemo_rl.telemetry._fallbacks``. +""" + +from __future__ import annotations + +import logging +import os +import uuid +from typing import TYPE_CHECKING, Any, Optional + +if TYPE_CHECKING: + from nemo.lens import 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" + +# 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", + "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 _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, + rank: int, + world_size: int, +) -> 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. + """ + 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 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 nemo-lens is not + installed or telemetry is disabled. Idempotent. + """ + global _TELEMETRY_HANDLE, _TELEMETRY_INITIALISED + if _TELEMETRY_INITIALISED: + return _TELEMETRY_HANDLE + _TELEMETRY_INITIALISED = True + + try: + from nemo.lens import NemoLensConfig, setup_telemetry + except ImportError: + return None + + from nemo_rl.telemetry.span_groups import RLSpanGroup + + tel = getattr(master_config, "telemetry", None) + if tel is not None: + _config_to_env(tel) + + config = NemoLensConfig.from_env( + prefix=_OTEL_PREFIX, + fallback_prefix=_OTEL_FALLBACK_PREFIX, + span_group_cls=RLSpanGroup, + ) + if not config.enabled: + return None + + # A friendly default service name if the user set nothing. + if not os.environ.get(_SERVICE_NAME_ENV, "").strip(): + config.service_name = "nemo-rl" + + # 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 + + try: + resource_attrs = _build_resource_attributes( + master_config, algorithm, rank=0, world_size=1 + ) + except Exception: + logger.warning("nemo-lens: failed to build resource attributes", exc_info=True) + resource_attrs = {"rl.algorithm": algorithm} + + handle = setup_telemetry( + config, rank=0, world_size=1, resource_attributes=resource_attrs + ) + _TELEMETRY_HANDLE = handle + + if config.logs_enabled and handle.is_exporting: + try: + from nemo.lens.logging_bridge import setup_logging_bridge + + setup_logging_bridge() + except Exception: + logger.warning("nemo-lens: failed to set up logging bridge", exc_info=True) + + logger.info( + "nemo-lens telemetry initialised (algorithm=%s, exporting=%s, run_id=%s, groups=%s)", + algorithm, + handle.is_exporting, + config.run_id, + config.span_groups, + ) + return handle + + +def init_telemetry_worker( + rank: Optional[int] = None, + world_size: Optional[int] = None, + resource_attributes: Optional[dict] = None, +) -> 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. + + Returns the :class:`TelemetryHandle`, or ``None`` if lens is absent or + telemetry is disabled. Idempotent per process. + """ + global _TELEMETRY_HANDLE, _TELEMETRY_INITIALISED + if _TELEMETRY_INITIALISED: + return _TELEMETRY_HANDLE + _TELEMETRY_INITIALISED = True + + if not ( + _is_env_truthy(f"{_OTEL_PREFIX}_ENABLED") + or _is_env_truthy(f"{_OTEL_FALLBACK_PREFIX}_ENABLED") + ): + return None + + try: + from nemo.lens import NemoLensConfig, setup_telemetry + except ImportError: + return None + + 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=resource_attributes, + ) + _TELEMETRY_HANDLE = handle + return handle + + +def get_telemetry() -> Optional["TelemetryHandle"]: + """Return the process-global telemetry handle (``None`` if uninitialised).""" + return _TELEMETRY_HANDLE + + +def shutdown_telemetry(timeout_ms: int = 5000) -> None: + """Flush and shut down telemetry providers. Call on the driver at job end.""" + 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) diff --git a/nemo_rl/telemetry/span_groups.py b/nemo_rl/telemetry/span_groups.py new file mode 100644 index 00000000000..38f66e02081 --- /dev/null +++ b/nemo_rl/telemetry/span_groups.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +Tries to import the real ``SpanGroup`` from ``nemo.lens``; falls back to a +minimal stub so that NeMo-RL works without nemo-lens installed. +""" + +from typing import ClassVar, Final + +try: + from nemo.lens.groups import SpanGroup +except ImportError: + # TODO(ahmadki): SpanGroups will move from nemo-lens to downstream, + # so this stub will be removed in the future. + class SpanGroup: # type: ignore[no-redef] + """Minimal stub used when nemo-lens is not installed.""" + + JOB = "job" + CHECKPOINT = "checkpoint" + EVALUATE = "evaluate" + MODEL_INIT = "model_init" + LOAD_CHECKPOINT = "load_checkpoint" + STEP = "step" + FORWARD_BACKWARD = "forward_backward" + OPTIMIZER = "optimizer" + + ALL_GROUPS: Final[frozenset] = frozenset( + [ + JOB, + CHECKPOINT, + EVALUATE, + MODEL_INIT, + LOAD_CHECKPOINT, + STEP, + FORWARD_BACKWARD, + OPTIMIZER, + ] + ) + + _PRESETS: ClassVar[dict] = { + "default": frozenset([JOB, CHECKPOINT, EVALUATE]), + "per_step": frozenset( + [ + JOB, + CHECKPOINT, + EVALUATE, + MODEL_INIT, + LOAD_CHECKPOINT, + STEP, + FORWARD_BACKWARD, + OPTIMIZER, + ] + ), + "all": ALL_GROUPS, + } + + @classmethod + def resolve(cls, spec: str) -> frozenset: + raise RuntimeError( + "SpanGroup.resolve() requires nemo-lens. " + "Install it with: uv sync --extra telemetry" + ) + + +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.""" + + # ------------------------------------------------------------------ # + # All groups and presets + # ------------------------------------------------------------------ # + + ALL_GROUPS: Final[frozenset] = SpanGroup.ALL_GROUPS | frozenset( + [ + ROLLOUT, + GENERATION, + LOGPROB, + REWARD, + ADVANTAGE, + POLICY_UPDATE, + REFERENCE_POLICY, + DATA_PROCESSING, + ] + ) + + _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, + SpanGroup.STEP, + ROLLOUT, + GENERATION, + LOGPROB, + REWARD, + ADVANTAGE, + POLICY_UPDATE, + REFERENCE_POLICY, + DATA_PROCESSING, + ] + ), + "all": ALL_GROUPS, + } diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index 75df1c9499a..1d64636dc2e 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -42,6 +42,7 @@ from nemo_rl.data.interfaces import LLMMessageLogType from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.telemetry.metrics import tee_rl_metrics_to_otel # Flag to track if rich logging has been configured _rich_logging_configured = False @@ -1052,6 +1053,8 @@ def log_metrics( for logger in self.loggers: logger.log_metrics(metrics, 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 9958048b578..26b69b0e1ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -110,6 +110,12 @@ dependencies = [ ] [project.optional-dependencies] +# Optional OpenTelemetry instrumentation via nemo-lens (resolved from PyPI). +# When this extra is not installed, every telemetry site degrades to a ~0-cost +# no-op through nemo_rl/telemetry/_fallbacks.py. +telemetry = [ + "nemo-lens[sdk]>=0.1.0", +] fsdp = [ # +cu13 wheels from GitHub match torch cu130; PyPI often resolves to +cu12 (libcudart.so.12). # https://github.com/Dao-AILab/flash-attention/releases/tag/v2.8.1 diff --git a/tests/unit/telemetry/__init__.py b/tests/unit/telemetry/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/telemetry/conftest.py b/tests/unit/telemetry/conftest.py new file mode 100644 index 00000000000..7454307e412 --- /dev/null +++ b/tests/unit/telemetry/conftest.py @@ -0,0 +1,57 @@ +"""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`` env vars before and after each test so nothing leaks. The +nemo-lens resets are guarded so this suite still imports when lens is absent. +""" + +import os + +import pytest + + +def _clear_telemetry_env() -> None: + for key in list(os.environ): + if key.startswith(("NEMO_RL_OTEL", "NEMO_LENS")) or key == "OTEL_SERVICE_NAME": + 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: + 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 + try: + 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()) + except ImportError: + pass + + +@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 00000000000..2f54f2d6ad4 --- /dev/null +++ b/tests/unit/telemetry/test_config.py @@ -0,0 +1,60 @@ +"""Tests for TelemetryConfig and the YAML->env translation.""" + +import os + +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.setup import _config_to_env, _ENV_FIELD_MAP + + +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.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, + 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_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 diff --git a/tests/unit/telemetry/test_fallbacks.py b/tests/unit/telemetry/test_fallbacks.py new file mode 100644 index 00000000000..0c2598e5dd4 --- /dev/null +++ b/tests/unit/telemetry/test_fallbacks.py @@ -0,0 +1,64 @@ +"""Tests that the fallback shims are pure no-ops when nemo-lens is absent. + +Each test forces the ``except ImportError`` branch of ``_fallbacks`` by blocking +``nemo.lens*`` imports and re-importing the module, so the behaviour is verified +regardless of whether nemo-lens happens to be installed in the test env. +""" + +import builtins +import importlib +import sys + +import pytest + + +def _import_fallbacks_without_lens(monkeypatch): + real_import = builtins.__import__ + + def fake_import(name, *args, **kwargs): + if name.startswith("nemo.lens"): + raise ImportError("nemo.lens blocked for test") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fake_import) + sys.modules.pop("nemo_rl.telemetry._fallbacks", None) + return importlib.import_module("nemo_rl.telemetry._fallbacks") + + +@pytest.fixture(autouse=True) +def _restore_real_fallbacks(): + yield + sys.modules.pop("nemo_rl.telemetry._fallbacks", None) + importlib.import_module("nemo_rl.telemetry._fallbacks") + + +def test_managed_span_yields_none(monkeypatch): + fb = _import_fallbacks_without_lens(monkeypatch) + with fb.managed_span("reward", "rl.x", some_attr=1) as span: + assert span is None + + +def test_span_cm_yields_none(monkeypatch): + fb = _import_fallbacks_without_lens(monkeypatch) + with fb.span_cm("rl.x", record_exception=True, attr=2) as span: + assert span is None + + +def test_trace_fn_returns_function_unchanged(monkeypatch): + fb = _import_fallbacks_without_lens(monkeypatch) + + @fb.trace_fn("job", "rl.job") + def add_one(x): + return x + 1 + + assert add_one(41) == 42 + + +def test_is_span_group_enabled_false(monkeypatch): + fb = _import_fallbacks_without_lens(monkeypatch) + assert fb.is_span_group_enabled("reward") is False + + +def test_safe_set_span_attributes_does_not_raise(monkeypatch): + fb = _import_fallbacks_without_lens(monkeypatch) + fb.safe_set_span_attributes(None, {"a": 1}) # must not raise diff --git a/tests/unit/telemetry/test_instrumentation.py b/tests/unit/telemetry/test_instrumentation.py new file mode 100644 index 00000000000..033d8885db4 --- /dev/null +++ b/tests/unit/telemetry/test_instrumentation.py @@ -0,0 +1,76 @@ +"""End-to-end span tests using an in-memory exporter. + +Exercises the same primitives the algorithm loops use (``managed_span`` / +``trace_fn``) and asserts spans are emitted per group, gated off when the group +is disabled, and nest correctly. +""" + +import pytest + +pytest.importorskip("nemo.lens") + +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + +from nemo.lens import NemoLensConfig, setup_telemetry +from nemo.lens.helpers import managed_span, trace_fn +from nemo_rl.telemetry.span_groups import RLSpanGroup + + +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 + + +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" + + +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 + + +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()) + + +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 00000000000..116df074bb5 --- /dev/null +++ b/tests/unit/telemetry/test_metrics.py @@ -0,0 +1,76 @@ +"""Tests for the RL metrics -> nemo-lens tee.""" + +import pytest + +from nemo_rl.telemetry.metrics import map_rl_metrics + + +def test_map_basic_keys(): + mapped = map_rl_metrics( + { + "reward": 0.5, + "loss": 1.0, + "grad_norm": 2.0, + "lr": 1e-4, + "valid_tokens_per_sec_per_gpu": 100.0, + "mean_gen_tokens_per_sample": 64, + } + ) + assert mapped["reward_mean"] == 0.5 + assert mapped["policy_loss"] == 1.0 + assert mapped["grad_norm"] == 2.0 + assert mapped["learning_rate"] == 1e-4 + assert mapped["tokens_per_sec"] == 100.0 + assert mapped["response_length_mean"] == 64.0 + + +def test_map_first_candidate_wins(): + # "loss" is the first candidate for policy_loss. + mapped = map_rl_metrics({"loss": 1.0, "policy_loss": 2.0}) + assert mapped["policy_loss"] == 1.0 + + +def test_map_skips_bool_and_non_numeric(): + mapped = map_rl_metrics({"reward": True, "loss": "nan", "entropy": 0.1}) + assert "reward_mean" not in mapped + assert "policy_loss" not in mapped + assert mapped["entropy"] == 0.1 + + +def test_map_empty_for_unknown_keys(): + assert map_rl_metrics({"unrelated_metric": 1.0}) == {} + + +def test_tee_emits_only_for_train_prefix_when_exporting(): + pytest.importorskip("nemo.lens") + import nemo_rl.telemetry.setup as setup_mod + from nemo.lens import NemoLensConfig, setup_telemetry + from opentelemetry.sdk.metrics.export import InMemoryMetricReader + + from nemo_rl.telemetry.metrics import tee_rl_metrics_to_otel + 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 + ) + + tee_rl_metrics_to_otel({"reward": 0.5, "grad_norm": 1.0}, "train") + tee_rl_metrics_to_otel({"reward": 9.0}, "validation") # wrong prefix -> ignored + + names = { + metric.name + for rm in reader.get_metrics_data().resource_metrics + for sm in rm.scope_metrics + for metric in sm.metrics + } + assert "rl.reward.mean" in names + assert "rl.grad_norm" in names + + +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({"reward": 0.5}, "train") diff --git a/tests/unit/telemetry/test_setup.py b/tests/unit/telemetry/test_setup.py new file mode 100644 index 00000000000..a4e7b4fd222 --- /dev/null +++ b/tests/unit/telemetry/test_setup.py @@ -0,0 +1,73 @@ +"""Tests for the telemetry setup module (driver init, resource attrs, digging).""" + +import pytest + +from nemo_rl.telemetry.config import TelemetryConfig +from nemo_rl.telemetry.setup import ( + _build_resource_attributes, + _dig, + get_telemetry, + init_telemetry_driver, +) + + +class _FakeMasterConfig: + def __init__(self, telemetry=None): + self.telemetry = telemetry + self.policy = { + "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", rank=0, world_size=1 + ) + 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_init_driver_returns_none_when_disabled(): + handle = init_telemetry_driver( + _FakeMasterConfig(TelemetryConfig(enabled=False)), "grpo" + ) + assert handle is None + assert get_telemetry() 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_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() 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 diff --git a/tests/unit/telemetry/test_span_groups.py b/tests/unit/telemetry/test_span_groups.py new file mode 100644 index 00000000000..cc93f0f9b3b --- /dev/null +++ b/tests/unit/telemetry/test_span_groups.py @@ -0,0 +1,61 @@ +"""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", + } +) + + +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_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 940c29c2ed9..a4ba57d51e8 100644 --- a/uv.lock +++ b/uv.lock @@ -4202,6 +4202,25 @@ docs = [ { name = "swagger-plugin-for-sphinx", specifier = ">=6.0.0" }, ] +[[package]] +name = "nemo-lens" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/94/9f/c32809fdb5c375218d8dc4f5762133c0be63a791ced8654a4f2408822513/nemo_lens-0.1.0.tar.gz", hash = "sha256:044621a6d877739e0bc69aced2f2f2a5c565eadc3f9070b99af72bb57abd73aa", size = 46741, upload-time = "2026-07-03T00:49:47.303Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8a/e4/23f7b705bd4102be32f92e651ad5dece122d93e09b8518d3df8c17fd2f6b/nemo_lens-0.1.0-py3-none-any.whl", hash = "sha256:87eca2609dfd41db8f663d943c0c83153e209c1a880bcab816750b6a4f841d58", size = 51451, upload-time = "2026-07-03T00:49:46.109Z" }, +] + +[package.optional-dependencies] +sdk = [ + { name = "opentelemetry-exporter-otlp-proto-grpc", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "opentelemetry-exporter-otlp-proto-http", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "opentelemetry-sdk", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, +] + [[package]] name = "nemo-rl" source = { editable = "." } @@ -4315,6 +4334,9 @@ sglang = [ { name = "sglang-router", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "transformers", version = "5.6.0", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] +telemetry = [ + { name = "nemo-lens", extra = ["sdk"], marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, +] trtllm = [ { name = "tensorrt-llm", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] @@ -4425,6 +4447,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"], marker = "extra == 'telemetry'", specifier = ">=0.1.0" }, { 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" }, @@ -4479,7 +4502,7 @@ requires-dist = [ { name = "wandb", specifier = ">=0.28.0" }, { name = "zstandard" }, ] -provides-extras = ["fsdp", "automodel", "vllm", "sglang", "mcore", "trtllm", "modelopt", "nvrx", "nemo-gym"] +provides-extras = ["telemetry", "fsdp", "automodel", "vllm", "sglang", "mcore", "trtllm", "modelopt", "nvrx", "nemo-gym"] [package.metadata.requires-dev] build = [ @@ -5289,13 +5312,13 @@ name = "opentelemetry-exporter-otlp-proto-grpc" version = "1.44.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "googleapis-common-protos", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "grpcio", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "opentelemetry-api", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "opentelemetry-exporter-otlp-proto-common", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "opentelemetry-proto", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "opentelemetry-sdk", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "typing-extensions", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "googleapis-common-protos", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "grpcio", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "opentelemetry-api", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "opentelemetry-exporter-otlp-proto-common", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "opentelemetry-proto", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "opentelemetry-sdk", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, + { name = "typing-extensions", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/1f/47/80d9e9d468dc5de3af5096f5ccdb065fa4dd1470f74495cc53e59e397f47/opentelemetry_exporter_otlp_proto_grpc-1.44.0.tar.gz", hash = "sha256:40d1ae9e03fcc36de3cbac610cc99f35894938bff9cfd90fc4ec68bd85448463", size = 27225, upload-time = "2026-07-16T15:25:38.308Z" } wheels = [