From eeb02f1a9adc1023662c0def712332a61a5c247d Mon Sep 17 00:00:00 2001 From: BruceLoveDecimal Date: Wed, 22 Jul 2026 14:42:27 +0800 Subject: [PATCH] feat(observability): add NVTX tracing and Prometheus metrics Signed-off-by: BruceLoveDecimal --- csrc/ops.cpp | 103 +++- csrc/utils/nvtx_utils.h | 76 +++ docs/.nav.yml | 1 + docs/getting_started/observability.md | 158 ++++++ envs.py | 1 + examples/grafana-dashboard.json | 226 +++++++++ examples/grpo_single_gpu.py | 96 ++-- pyproject.toml | 5 +- rl_engine/executors/deepspeed_trainer.py | 73 ++- rl_engine/executors/paged_kv_baseline.py | 9 + rl_engine/executors/rollout.py | 24 +- rl_engine/executors/stateless_executor.py | 87 +++- rl_engine/executors/vllm_sampler.py | 59 ++- rl_engine/kernels/registry.py | 34 +- rl_engine/observability/__init__.py | 21 + rl_engine/observability/metrics.py | 473 ++++++++++++++++++ rl_engine/observability/nvtx.py | 62 +++ rl_engine/observability/server.py | 83 ++++ setup.py | 4 + tests/test_observability.py | 560 ++++++++++++++++++++++ 20 files changed, 2048 insertions(+), 107 deletions(-) create mode 100644 csrc/utils/nvtx_utils.h create mode 100644 docs/getting_started/observability.md create mode 100644 examples/grafana-dashboard.json create mode 100644 rl_engine/observability/__init__.py create mode 100644 rl_engine/observability/metrics.py create mode 100644 rl_engine/observability/nvtx.py create mode 100644 rl_engine/observability/server.py create mode 100644 tests/test_observability.py diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 61ba4a3b..bd50b566 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -4,6 +4,8 @@ #include #include +#include "utils/nvtx_utils.h" + // Fused LogP Declarations torch::Tensor fused_logp_forward(torch::Tensor logits, torch::Tensor token_ids); @@ -125,49 +127,98 @@ at::Tensor prefix_shared_attention( PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "RL-Kernel High-Performance Operator Extension Library"; - m.def("fused_logp", &fused_logp_forward, "Fused logp forward fallback"); + m.def("fused_logp", rlk::nvtx_wrap("rlk::fused_logp", &fused_logp_forward), + "Fused logp forward fallback"); #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_SM90) - m.def("fused_logp_sm90", &fused_logp_sm90_forward, "TMA-accelerated Online Softmax Fused LogP"); - m.def("fused_linear_logp_sm90", &fused_linear_logp_sm90_forward, + m.def("fused_logp_sm90", rlk::nvtx_wrap("rlk::fused_logp_sm90", &fused_logp_sm90_forward), + "TMA-accelerated Online Softmax Fused LogP"); + m.def("fused_linear_logp_sm90", + rlk::nvtx_wrap("rlk::fused_linear_logp_sm90", &fused_linear_logp_sm90_forward), "TMA+WGMMA fused linear log-prob (hidden @ W^T -> selected-token logp), SM90"); - m.def("fused_linear_logp_sm90_global_target", &fused_linear_logp_sm90_global_target_forward, + m.def("fused_linear_logp_sm90_global_target", + rlk::nvtx_wrap("rlk::fused_linear_logp_sm90_global_target", + &fused_linear_logp_sm90_global_target_forward), "TMA+WGMMA local-shard target-logit/lse for vocab-parallel linear log-prob, SM90"); - m.def("fused_linear_logp_sm90_backward", &fused_linear_logp_sm90_backward, + m.def("fused_linear_logp_sm90_backward", + rlk::nvtx_wrap("rlk::fused_linear_logp_sm90_backward", + &fused_linear_logp_sm90_backward), "CUDA fused backward for linear log-prob, SM90 backend"); - m.def("linear_logp_probs_bf16_forward", &linear_logp_probs_bf16_forward, + m.def("linear_logp_probs_bf16_forward", + rlk::nvtx_wrap("rlk::linear_logp_probs_bf16_forward", &linear_logp_probs_bf16_forward), "Build bf16 softmax probabilities and selected log-prob from bf16 logits"); - m.def("linear_logp_bf16_forward", &linear_logp_bf16_forward, + m.def("linear_logp_bf16_forward", + rlk::nvtx_wrap("rlk::linear_logp_bf16_forward", &linear_logp_bf16_forward), "Build selected log-prob and lse from bf16 logits without saving probabilities"); - m.def("linear_logp_local_probs_bf16_forward", &linear_logp_local_probs_bf16_forward, + m.def("linear_logp_local_probs_bf16_forward", + rlk::nvtx_wrap("rlk::linear_logp_local_probs_bf16_forward", + &linear_logp_local_probs_bf16_forward), "Build local bf16 softmax probabilities, target logits, and lse from bf16 logits"); - m.def("linear_logp_local_bf16_forward", &linear_logp_local_bf16_forward, + m.def("linear_logp_local_bf16_forward", + rlk::nvtx_wrap("rlk::linear_logp_local_bf16_forward", &linear_logp_local_bf16_forward), "Build local target logits and lse from bf16 logits without saving probabilities"); - m.def("linear_logp_probs_bf16_to_dlogits_", &linear_logp_probs_bf16_to_dlogits_, + m.def("linear_logp_probs_bf16_to_dlogits_", + rlk::nvtx_wrap("rlk::linear_logp_probs_bf16_to_dlogits_", + &linear_logp_probs_bf16_to_dlogits_), "In-place bf16 probs -> dlogits for selected log-prob backward"); m.def("linear_logp_local_probs_bf16_to_dlogits_", - &linear_logp_local_probs_bf16_to_dlogits_, + rlk::nvtx_wrap("rlk::linear_logp_local_probs_bf16_to_dlogits_", + &linear_logp_local_probs_bf16_to_dlogits_), "In-place local bf16 probs -> TP dlogits for selected log-prob backward"); - m.def("linear_logp_logits_bf16_to_dlogits", &linear_logp_logits_bf16_to_dlogits, + m.def("linear_logp_logits_bf16_to_dlogits", + rlk::nvtx_wrap("rlk::linear_logp_logits_bf16_to_dlogits", + &linear_logp_logits_bf16_to_dlogits), "Build bf16 dlogits from bf16 logits and fp32 lse"); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) - m.def("fused_logp_forward_out", &fused_logp_forward_out, "Fused logp out"); - m.def("fused_logp_forward_fp32", &fused_logp_forward_fp32, "Fused logp fp32"); - m.def("fused_logp_forward_indexed_out", &fused_logp_forward_indexed_out, "Fused logp indexed out"); - m.def("fused_logp_forward_indexed_fp32", &fused_logp_forward_indexed_fp32, "Fused logp indexed fp32"); - m.def("fused_logp_forward_online_out", &fused_logp_forward_online_out, "Fused logp online out"); - m.def("fused_logp_forward_online_fp32", &fused_logp_forward_online_fp32, "Fused logp online fp32"); - m.def("fused_logp_forward_online_indexed_out", &fused_logp_forward_online_indexed_out, "Fused logp online indexed out"); - m.def("fused_logp_forward_online_indexed_fp32", &fused_logp_forward_online_indexed_fp32, "Fused logp online indexed fp32"); - m.def("deterministic_logp", &deterministic_logp_forward, "Batch-invariant deterministic logp"); - m.def("deterministic_logp_forward_out", &deterministic_logp_forward_out, "Batch-invariant deterministic logp out"); - m.def("deterministic_logp_forward_fp32", &deterministic_logp_forward_fp32, "Batch-invariant deterministic logp fp32"); - m.def("deterministic_logp_forward_indexed_out", &deterministic_logp_forward_indexed_out, "Batch-invariant deterministic logp indexed out"); - m.def("deterministic_logp_forward_indexed_fp32", &deterministic_logp_forward_indexed_fp32, "Batch-invariant deterministic logp indexed fp32"); + m.def("fused_logp_forward_out", + rlk::nvtx_wrap("rlk::fused_logp_forward_out", &fused_logp_forward_out), + "Fused logp out"); + m.def("fused_logp_forward_fp32", + rlk::nvtx_wrap("rlk::fused_logp_forward_fp32", &fused_logp_forward_fp32), + "Fused logp fp32"); + m.def("fused_logp_forward_indexed_out", + rlk::nvtx_wrap("rlk::fused_logp_forward_indexed_out", &fused_logp_forward_indexed_out), + "Fused logp indexed out"); + m.def("fused_logp_forward_indexed_fp32", + rlk::nvtx_wrap("rlk::fused_logp_forward_indexed_fp32", &fused_logp_forward_indexed_fp32), + "Fused logp indexed fp32"); + m.def("fused_logp_forward_online_out", + rlk::nvtx_wrap("rlk::fused_logp_forward_online_out", &fused_logp_forward_online_out), + "Fused logp online out"); + m.def("fused_logp_forward_online_fp32", + rlk::nvtx_wrap("rlk::fused_logp_forward_online_fp32", &fused_logp_forward_online_fp32), + "Fused logp online fp32"); + m.def("fused_logp_forward_online_indexed_out", + rlk::nvtx_wrap("rlk::fused_logp_forward_online_indexed_out", + &fused_logp_forward_online_indexed_out), + "Fused logp online indexed out"); + m.def("fused_logp_forward_online_indexed_fp32", + rlk::nvtx_wrap("rlk::fused_logp_forward_online_indexed_fp32", + &fused_logp_forward_online_indexed_fp32), + "Fused logp online indexed fp32"); + m.def("deterministic_logp", + rlk::nvtx_wrap("rlk::deterministic_logp", &deterministic_logp_forward), + "Batch-invariant deterministic logp"); + m.def("deterministic_logp_forward_out", + rlk::nvtx_wrap("rlk::deterministic_logp_forward_out", &deterministic_logp_forward_out), + "Batch-invariant deterministic logp out"); + m.def("deterministic_logp_forward_fp32", + rlk::nvtx_wrap("rlk::deterministic_logp_forward_fp32", &deterministic_logp_forward_fp32), + "Batch-invariant deterministic logp fp32"); + m.def("deterministic_logp_forward_indexed_out", + rlk::nvtx_wrap("rlk::deterministic_logp_forward_indexed_out", + &deterministic_logp_forward_indexed_out), + "Batch-invariant deterministic logp indexed out"); + m.def("deterministic_logp_forward_indexed_fp32", + rlk::nvtx_wrap("rlk::deterministic_logp_forward_indexed_fp32", + &deterministic_logp_forward_indexed_fp32), + "Batch-invariant deterministic logp indexed fp32"); // registry Prefix-Shared Attention - m.def("prefix_shared_attention", &prefix_shared_attention, "Prefix-Shared Fused Attention for GRPO"); + m.def("prefix_shared_attention", + rlk::nvtx_wrap("rlk::prefix_shared_attention", &prefix_shared_attention), + "Prefix-Shared Fused Attention for GRPO"); #endif } diff --git a/csrc/utils/nvtx_utils.h b/csrc/utils/nvtx_utils.h new file mode 100644 index 00000000..204a37d6 --- /dev/null +++ b/csrc/utils/nvtx_utils.h @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors + +// NVTX range helpers for the pybind operator bindings. +// +// NVTX v3 is header-only and ships with CUDA >= 10, so no linker change is +// needed. ROCm builds and builds with KERNEL_ALIGN_DISABLE_NVTX take the +// identity pass-through branch, which generates no code. + +#pragma once + +#include +#include +#include + +#if defined(KERNEL_ALIGN_WITH_CUDA) && !defined(KERNEL_ALIGN_DISABLE_NVTX) +#include + +namespace rlk { + +inline bool nvtx_enabled() { + static const bool enabled = [] { + const char* value = std::getenv("RL_KERNEL_NVTX"); + return value != nullptr && + (std::strcmp(value, "1") == 0 || std::strcmp(value, "true") == 0 || + std::strcmp(value, "yes") == 0 || std::strcmp(value, "on") == 0); + }(); + return enabled; +} + +class NvtxRange { + public: + explicit NvtxRange(const char* name) : active_(nvtx_enabled()) { + if (active_) { + nvtxRangePushA(name); + } + } + ~NvtxRange() { + if (active_) { + nvtxRangePop(); + } + } + NvtxRange(const NvtxRange&) = delete; + NvtxRange& operator=(const NvtxRange&) = delete; + + private: + bool active_; +}; + +// Wraps a free function so the pybind-bound call is enclosed in an NVTX range. +template +auto nvtx_wrap(const char* name, Ret (*fn)(Args...)) { + return [name, fn](Args... args) -> Ret { + NvtxRange guard{name}; + return fn(std::forward(args)...); + }; +} + +} // namespace rlk + +#define RLK_NVTX_RANGE(name) ::rlk::NvtxRange _rlk_nvtx_guard{(name)} + +#else // ROCm / NVTX-disabled builds: identity pass-through, zero code generated. + +namespace rlk { + +template +auto nvtx_wrap(const char*, F fn) { + return fn; +} + +} // namespace rlk + +#define RLK_NVTX_RANGE(name) + +#endif diff --git a/docs/.nav.yml b/docs/.nav.yml index 8794c731..dc3a3674 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -12,6 +12,7 @@ nav: - getting_started/installation.md - getting_started/faq.md - Hardware Profiling Guide: getting_started/hardware-profiling.md + - Observability: getting_started/observability.md - Operators: - operators/README.md - operators/activation.md diff --git a/docs/getting_started/observability.md b/docs/getting_started/observability.md new file mode 100644 index 00000000..b0db55d7 --- /dev/null +++ b/docs/getting_started/observability.md @@ -0,0 +1,158 @@ +# Observability + +RL-Kernel ships a two-layer observability stack: + +- **NVTX tracing (micro):** opt-in per-operator `rlk::` ranges around every C++ binding plus + Python stage ranges (`rlk::rollout.generate`, `rlk::score`, `rlk::train_step`, ...), + visible on an [Nsight Systems](https://developer.nvidia.com/nsight-systems) timeline. + Set `RL_KERNEL_NVTX=1` only for profiling runs; normal runs retain just a cached branch. +- **Prometheus metrics (macro):** a per-process `/metrics` endpoint with token throughput, + rollout QPS, stage latencies, hardware fallback state, KV-cache fragmentation, peak GPU + memory, and training loss — plus a ready-made Grafana dashboard. + +## Installation + +Prometheus support is an optional extra; the engine runs unchanged without it: + +```bash +pip install -e ".[observability]" # adds prometheus-client +``` + +Without `prometheus-client` installed, every metrics call degrades to a silent no-op. + +## Prometheus metrics + +### Starting the exporter + +The exporter never starts as an import side effect. Either call it explicitly: + +```python +from rl_engine.observability import start_metrics_server + +port = start_metrics_server() # default 0.0.0.0:8000; port=0 -> ephemeral +``` + +or opt in via environment variables when running the examples: + +```bash +RL_KERNEL_METRICS=1 python examples/grpo_single_gpu.py --steps 3 --device cpu & +curl -s localhost:8000/metrics | grep rl_kernel_ +``` + +| Variable | Default | Meaning | +| --- | --- | --- | +| `RL_KERNEL_METRICS` | unset | `1` → examples/entrypoints start the exporter | +| `RL_KERNEL_METRICS_PORT` | `8000` | Exporter port; `0` = ephemeral | +| `RL_KERNEL_METRICS_ADDR` | `0.0.0.0` | Bind address override | + +For multi-process (e.g. Ray) deployments each worker exposes its own endpoint; use +`RL_KERNEL_METRICS_PORT=0` to avoid port collisions and let your scrape topology discover +the workers. + +### Metric schema + +All metrics use the `rl_kernel_` prefix and Prometheus base units (seconds, bytes). + +| Metric | Type | Labels | +| --- | --- | --- | +| `rl_kernel_build_info` | Gauge (=1) | `version`, `device_type`, `backend_version`, `ext_available` | +| `rl_kernel_stage_duration_seconds` | Histogram | `stage` | +| `rl_kernel_requests_total` | Counter | `stage`, `status` (`ok`/`error`) | +| `rl_kernel_output_tokens_total` | Counter | `stage` | +| `rl_kernel_selected_backend_info` | Gauge (=1) | `op_type`, `backend`, `is_fallback` | +| `rl_kernel_hardware_fallback_total` | Counter | `op_type`, `requested_backend`, `selected_backend`, `reason` | +| `rl_kernel_kv_cache_blocks` | Gauge | `kind` (`required`/`reserved`) | +| `rl_kernel_kv_cache_fragmentation_ratio` | Gauge | — | +| `rl_kernel_gpu_peak_memory_bytes` | Gauge | `stage`, `kind` (`allocated`/`reserved`) | +| `rl_kernel_training_loss` | Gauge | — | +| `rl_kernel_weight_version` | Gauge | `role` (`published`/`consumed`) | + +Fallbacks are reported at two layers: `selected_backend_info` captures the *state* chosen +once per process at kernel dispatch time (`is_fallback="true"` means a lower-priority +backend was selected), while `hardware_fallback_total` counts *events* — backend +import/instantiation failures and per-call attention-backend fallbacks +(flash-attention → eager). + +### Useful queries + +```promql +rate(rl_kernel_output_tokens_total[1m]) # token throughput +rate(rl_kernel_requests_total{stage="rollout_generate"}[1m]) # rollout QPS +histogram_quantile(0.95, rate(rl_kernel_stage_duration_seconds_bucket[5m])) +sum(rl_kernel_selected_backend_info{is_fallback="true"}) / sum(rl_kernel_selected_backend_info) +``` + +### Prometheus scrape config + +```yaml +scrape_configs: + - job_name: rl-kernel + scrape_interval: 5s + static_configs: + - targets: ["localhost:8000"] +``` + +### Grafana dashboard + +Import [`examples/grafana-dashboard.json`](https://github.com/RL-Align/RL-Kernel/blob/main/examples/grafana-dashboard.json) +(Grafana ≥ 10, *Dashboards → Import*) and select your Prometheus datasource when prompted — +the dashboard binds it through the `${DS_PROMETHEUS}` input, no hardcoded UID. + +### Using the facade in your own code + +```python +from rl_engine.observability import metrics + +with metrics.stage_timer("my_stage"): # duration histogram + ok/error counter + run_stage() +metrics.add_output_tokens("my_stage", n_tokens) +``` + +Rules the built-in instrumentation follows (and yours should too): + +- per step / per request granularity — never per token in hot loops; +- label values only from finite enums in code, never user data; +- exceptions propagate unchanged (recorded with `status="error"`); +- never re-query something the stage already computed — see below. + +### Overhead + +Metrics use exact 16-call batched updates for the module-level registry: counters and +histogram buckets remain exact, gauges retain the latest value, and pending updates flush +before each scrape. This moves Prometheus locks and repeated fixed-label lookups off the +hot path. On an RTX 5090 with torch 2.8/CUDA 12.8, the complete default score path measured +below 1% overhead at the benchmark's toy shape (paired, randomly interleaved calls). + +| Workload | Baseline stage time | With metrics | Overhead | +| --- | --- | --- | --- | +| `benchmark_stateless_executor` defaults (toy model, hidden 256 / vocab 4096) | 1.038 ms | 1.047 ms | +0.90 % | + +With `RL_KERNEL_NVTX=1`, Python and C++ ranges intentionally add profiler annotation cost; +do not use that mode for the normal metrics-overhead acceptance run. + +If you add instrumentation, be careful what you *read* to populate it. Querying +`torch.cuda.max_memory_allocated()` / `max_memory_reserved()` costs ~85 µs for the pair — +each one rebuilds the full `torch.cuda.memory_stats()` dict. Prefer values the stage has +already computed, or read `memory_stats()` once and pull both peaks out of it. + +## NVTX tracing with Nsight Systems + +Capture a timeline on a CUDA machine: + +```bash +RL_KERNEL_NVTX=1 nsys profile -t nvtx,cuda -o rlk_trace \ + python examples/grpo_single_gpu.py --steps 3 +nsys stats --report nvtx_sum rlk_trace.nsys-rep # lists rlk::* ranges +``` + +On the timeline the Python stage ranges (`rlk::train_step` with `rlk::train.forward` / +`rlk::train.backward` / `rlk::train.optim_step` inside) contain the per-operator +`rlk::` ranges emitted by the C++ extension, correlated with kernel +execution on the CUDA streams. + +Notes: + +- NVTX v3 is header-only and ships with the CUDA Toolkit; no extra link dependency. +- ROCm builds compile the tracing macros to no-ops. +- To compile the extension without NVTX, build with `KERNEL_ALIGN_DISABLE_NVTX=1`. +- Python stage ranges silently no-op when CUDA is unavailable. diff --git a/envs.py b/envs.py index 34aded7c..9e1b0690 100644 --- a/envs.py +++ b/envs.py @@ -26,3 +26,4 @@ def env_flag(name: str, default: bool = False) -> bool: KERNEL_ALIGN_NCU_LINEINFO = "KERNEL_ALIGN_NCU_LINEINFO" KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC = "KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC" KERNEL_ALIGN_FORCE_SM90 = "KERNEL_ALIGN_FORCE_SM90" +KERNEL_ALIGN_DISABLE_NVTX = "KERNEL_ALIGN_DISABLE_NVTX" diff --git a/examples/grafana-dashboard.json b/examples/grafana-dashboard.json new file mode 100644 index 00000000..1487c78d --- /dev/null +++ b/examples/grafana-dashboard.json @@ -0,0 +1,226 @@ +{ + "__inputs": [ + { + "name": "DS_PROMETHEUS", + "label": "Prometheus", + "description": "Prometheus datasource scraping RL-Kernel /metrics endpoints", + "type": "datasource", + "pluginId": "prometheus", + "pluginName": "Prometheus" + } + ], + "__requires": [ + { "type": "grafana", "id": "grafana", "name": "Grafana", "version": "10.0.0" }, + { "type": "datasource", "id": "prometheus", "name": "Prometheus", "version": "1.0.0" } + ], + "title": "RL-Kernel Observability", + "uid": "rl-kernel-observability", + "tags": ["rl-kernel"], + "schemaVersion": 39, + "version": 1, + "refresh": "10s", + "time": { "from": "now-30m", "to": "now" }, + "templating": { "list": [] }, + "annotations": { "list": [] }, + "editable": true, + "panels": [ + { + "id": 1, + "type": "timeseries", + "title": "Output token throughput (tokens/s)", + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 0 }, + "fieldConfig": { "defaults": { "unit": "ops" }, "overrides": [] }, + "targets": [ + { + "refId": "A", + "expr": "rate(rl_kernel_output_tokens_total[1m])", + "legendFormat": "{{stage}}" + } + ] + }, + { + "id": 2, + "type": "timeseries", + "title": "Rollout QPS", + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 0 }, + "fieldConfig": { "defaults": { "unit": "reqps" }, "overrides": [] }, + "targets": [ + { + "refId": "A", + "expr": "rate(rl_kernel_requests_total{stage=\"rollout_generate\",status=\"ok\"}[1m])", + "legendFormat": "rollout ok" + }, + { + "refId": "B", + "expr": "rate(rl_kernel_requests_total{status=\"error\"}[1m])", + "legendFormat": "{{stage}} errors" + } + ] + }, + { + "id": 3, + "type": "timeseries", + "title": "Stage latency p95 (s)", + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 8 }, + "fieldConfig": { "defaults": { "unit": "s" }, "overrides": [] }, + "targets": [ + { + "refId": "A", + "expr": "histogram_quantile(0.95, sum by (le, stage) (rate(rl_kernel_stage_duration_seconds_bucket[5m])))", + "legendFormat": "{{stage}} p95" + }, + { + "refId": "B", + "expr": "histogram_quantile(0.50, sum by (le, stage) (rate(rl_kernel_stage_duration_seconds_bucket[5m])))", + "legendFormat": "{{stage}} p50" + } + ] + }, + { + "id": 4, + "type": "timeseries", + "title": "Training loss", + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 8 }, + "fieldConfig": { "defaults": { "unit": "none" }, "overrides": [] }, + "targets": [ + { + "refId": "A", + "expr": "rl_kernel_training_loss", + "legendFormat": "loss" + } + ] + }, + { + "id": 5, + "type": "table", + "title": "Selected kernel backends (dispatch state)", + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 16 }, + "fieldConfig": { "defaults": {}, "overrides": [] }, + "targets": [ + { + "refId": "A", + "expr": "rl_kernel_selected_backend_info", + "format": "table", + "instant": true + } + ] + }, + { + "id": 6, + "type": "stat", + "title": "Dispatch fallback share", + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "gridPos": { "h": 8, "w": 6, "x": 12, "y": 16 }, + "fieldConfig": { + "defaults": { "unit": "percentunit", "min": 0, "max": 1 }, + "overrides": [] + }, + "targets": [ + { + "refId": "A", + "expr": "sum(rl_kernel_selected_backend_info{is_fallback=\"true\"}) / sum(rl_kernel_selected_backend_info)", + "legendFormat": "fallback share" + } + ] + }, + { + "id": 7, + "type": "timeseries", + "title": "Hardware fallback events (events/s)", + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "gridPos": { "h": 8, "w": 6, "x": 18, "y": 16 }, + "fieldConfig": { "defaults": { "unit": "ops" }, "overrides": [] }, + "targets": [ + { + "refId": "A", + "expr": "rate(rl_kernel_hardware_fallback_total[5m])", + "legendFormat": "{{op_type}}: {{requested_backend}} -> {{selected_backend}} ({{reason}})" + } + ] + }, + { + "id": 8, + "type": "timeseries", + "title": "Paged KV-cache blocks", + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "gridPos": { "h": 8, "w": 8, "x": 0, "y": 24 }, + "fieldConfig": { "defaults": { "unit": "none" }, "overrides": [] }, + "targets": [ + { + "refId": "A", + "expr": "rl_kernel_kv_cache_blocks", + "legendFormat": "{{kind}}" + } + ] + }, + { + "id": 9, + "type": "timeseries", + "title": "KV-cache fragmentation ratio", + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "gridPos": { "h": 8, "w": 8, "x": 8, "y": 24 }, + "fieldConfig": { + "defaults": { "unit": "percentunit", "min": 0, "max": 1 }, + "overrides": [] + }, + "targets": [ + { + "refId": "A", + "expr": "rl_kernel_kv_cache_fragmentation_ratio", + "legendFormat": "fragmentation" + } + ] + }, + { + "id": 10, + "type": "timeseries", + "title": "Peak GPU memory (bytes)", + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "gridPos": { "h": 8, "w": 8, "x": 16, "y": 24 }, + "fieldConfig": { "defaults": { "unit": "bytes" }, "overrides": [] }, + "targets": [ + { + "refId": "A", + "expr": "rl_kernel_gpu_peak_memory_bytes", + "legendFormat": "{{stage}} {{kind}}" + } + ] + }, + { + "id": 11, + "type": "timeseries", + "title": "Weight version (published vs consumed)", + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 32 }, + "fieldConfig": { "defaults": { "unit": "none" }, "overrides": [] }, + "targets": [ + { + "refId": "A", + "expr": "rl_kernel_weight_version", + "legendFormat": "{{role}}" + } + ] + }, + { + "id": 12, + "type": "table", + "title": "Build info", + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 32 }, + "fieldConfig": { "defaults": {}, "overrides": [] }, + "targets": [ + { + "refId": "A", + "expr": "rl_kernel_build_info", + "format": "table", + "instant": true + } + ] + } + ] +} diff --git a/examples/grpo_single_gpu.py b/examples/grpo_single_gpu.py index 0929478a..2d6fd136 100644 --- a/examples/grpo_single_gpu.py +++ b/examples/grpo_single_gpu.py @@ -22,6 +22,9 @@ sys.path.insert(0, str(REPO_ROOT)) from rl_engine.kernels.registry import resolve_logp_op_type # noqa: E402 +from rl_engine.observability import maybe_start_metrics_server_from_env # noqa: E402 +from rl_engine.observability import metrics as obs_metrics # noqa: E402 +from rl_engine.observability import nvtx_range # noqa: E402 from rl_engine.testing import ( # noqa: E402 active_token_count, compute_policy_ratio, @@ -186,6 +189,9 @@ def run_training(args: argparse.Namespace) -> list[StepMetrics]: if args.steps <= 0: raise ValueError("--steps must be greater than zero") + # Opt-in Prometheus exporter: only starts when RL_KERNEL_METRICS=1. + maybe_start_metrics_server_from_env() + torch.manual_seed(args.seed) device = select_device(args.device) logp_op = resolve_logp_op( @@ -247,48 +253,58 @@ def run_training(args: argparse.Namespace) -> list[StepMetrics]: ) for step in range(args.steps): - optimizer.zero_grad(set_to_none=True) - logits = policy(batch.token_ids) - reference_logps = selected_logprobs_reference( - logits, - batch.token_ids, - mask=batch.completion_mask, - output_dtype=torch.float32, - ) - kernel_logps = selected_logps_with_op( - logp_op, - logits, - batch.token_ids, - batch.completion_mask, - ) - drift = summarize_kernel_drift( - kernel_logps.detach(), - reference_logps.detach(), - batch.completion_mask, + with obs_metrics.stage_timer("train_step"), nvtx_range("rlk::train_step"): + optimizer.zero_grad(set_to_none=True) + with nvtx_range("rlk::train.forward"): + logits = policy(batch.token_ids) + reference_logps = selected_logprobs_reference( + logits, + batch.token_ids, + mask=batch.completion_mask, + output_dtype=torch.float32, + ) + kernel_logps = selected_logps_with_op( + logp_op, + logits, + batch.token_ids, + batch.completion_mask, + ) + drift = summarize_kernel_drift( + kernel_logps.detach(), + reference_logps.detach(), + batch.completion_mask, + ) + + if kernel_logps.requires_grad: + train_logps = kernel_logps + train_source = backend_name + else: + train_logps = reference_logps + train_source = "autograd_reference" + + loss, policy_loss, kl = grpo_loss( + train_logps, + old_logps, + ref_logps, + advantages, + batch.completion_mask, + args.clip_eps, + args.beta, + ) + if not torch.isfinite(loss): + raise RuntimeError(f"non-finite loss at step {step}: {loss.item()}") + + with nvtx_range("rlk::train.backward"): + loss.backward() + with nvtx_range("rlk::train.optim_step"): + optimizer.step() + + obs_metrics.set_training_loss(float(loss.detach().cpu().item())) + obs_metrics.add_output_tokens( + "train_step", + int(active_token_count(batch.completion_mask).item()), ) - if kernel_logps.requires_grad: - train_logps = kernel_logps - train_source = backend_name - else: - train_logps = reference_logps - train_source = "autograd_reference" - - loss, policy_loss, kl = grpo_loss( - train_logps, - old_logps, - ref_logps, - advantages, - batch.completion_mask, - args.clip_eps, - args.beta, - ) - if not torch.isfinite(loss): - raise RuntimeError(f"non-finite loss at step {step}: {loss.item()}") - - loss.backward() - optimizer.step() - step_metrics = StepMetrics( step=step, loss=float(loss.detach().cpu().item()), diff --git a/pyproject.toml b/pyproject.toml index b0b80a1a..a31bcd67 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,10 @@ dependencies = [ cuda = ["flashinfer-python>=0.1.6", "nvidia-ml-py"] rocm = ["aiter"] vllm = ["vllm>=0.6.0"] -dev = ["pytest", "black", "isort", "ruff", "mypy", "pre-commit"] +observability = ["prometheus-client>=0.20"] +# prometheus-client is repeated in dev so the CPU CI job (installs ".[dev]") +# exercises the real exporter path, not only the no-op degradation path. +dev = ["pytest", "black", "isort", "ruff", "mypy", "pre-commit", "prometheus-client>=0.20"] [tool.setuptools.packages.find] where = ["."] diff --git a/rl_engine/executors/deepspeed_trainer.py b/rl_engine/executors/deepspeed_trainer.py index 9eb276c8..14483a95 100644 --- a/rl_engine/executors/deepspeed_trainer.py +++ b/rl_engine/executors/deepspeed_trainer.py @@ -29,7 +29,10 @@ ) from rl_engine.kernels.ops.pytorch.loss.linear_logp import NativeLinearLogpOp from rl_engine.kernels.registry import kernel_registry +from rl_engine.observability.metrics import metrics as obs_metrics +from rl_engine.observability.nvtx import nvtx_range from rl_engine.testing import compute_policy_ratio, compute_reference_kl, masked_mean +from rl_engine.utils.logger import logger _TDestination = TypeVar("_TDestination", bound=dict[str, Any]) @@ -127,7 +130,16 @@ def __init__( self._linear_logp = _linear_logp_op_for_device(self.device) def train(self, rollout: RolloutStageResult) -> TrainingStageResult: + with obs_metrics.stage_timer("train_step"), nvtx_range("rlk::train_step"): + result = self._train_impl(rollout) + _record_training_observability(result.metrics, device=self.device) + return result + + def _train_impl(self, rollout: RolloutStageResult) -> TrainingStageResult: started_at = time.perf_counter() + cuda_tracking = self.device.type == "cuda" and torch.cuda.is_available() + if cuda_tracking: + torch.cuda.reset_peak_memory_stats(self.device) batch, payload_metrics = self._batch_from_rollout_or_synthetic(rollout) training_model = _unwrap_training_model(self.engine, self.model) training_embedding = _embedding_layer(training_model) @@ -150,14 +162,15 @@ def train(self, rollout: RolloutStageResult) -> TrainingStageResult: zero_stage=self._deepspeed_zero_stage, world_size=self._engine_world_size(), ): - current_logps = _extract_logps( - self.engine(batch.token_ids.long()), - training_model, - batch.token_ids, - batch.completion_mask, - self._linear_logp, - output_dtype=torch.float32, - ) + with nvtx_range("rlk::train.forward"): + current_logps = _extract_logps( + self.engine(batch.token_ids.long()), + training_model, + batch.token_ids, + batch.completion_mask, + self._linear_logp, + output_dtype=torch.float32, + ) old_logps = current_logps.detach() - 0.01 ref_logps = objective_reference_logps(current_logps, batch) ratio = compute_policy_ratio(current_logps, old_logps, batch.completion_mask) @@ -166,8 +179,10 @@ def train(self, rollout: RolloutStageResult) -> TrainingStageResult: policy_loss = -torch.minimum(unclipped, clipped) kl = compute_reference_kl(current_logps, ref_logps, batch.completion_mask) loss = masked_mean(policy_loss + 0.01 * kl, batch.completion_mask) - self.engine.backward(loss) - self.engine.step() + with nvtx_range("rlk::train.backward"): + self.engine.backward(loss) + with nvtx_range("rlk::train.optim_step"): + self.engine.step() finished_at = time.perf_counter() published = self._next_published_weight_version(rollout.weight_version) @@ -216,6 +231,20 @@ def publish_weights( runtime cannot provide a safe full-state view, the worker fails explicitly instead of publishing a shard. """ + with obs_metrics.stage_timer("publish_weights"), nvtx_range("rlk::publish_weights"): + manifest = self._publish_weights_impl( + weight_version=weight_version, + metadata=metadata, + ) + obs_metrics.set_weight_version("published", int(weight_version)) + return manifest + + def _publish_weights_impl( + self, + *, + weight_version: int, + metadata: Optional[Mapping[str, Any]] = None, + ) -> WeightUpdateManifest: manifest_metadata = dict(metadata or {}) layout = { "kind": "full-state", @@ -312,6 +341,30 @@ def _resolved_deepspeed_config(self) -> dict[str, Any]: return _deep_merge(base, dict(self.config.deepspeed_config)) +def _record_training_observability( + result_metrics: Mapping[str, Any], + *, + device: torch.device, +) -> None: + """Export training metrics to the Prometheus facade; never raises.""" + try: + loss = result_metrics.get("loss") + if loss is not None: + obs_metrics.set_training_loss(float(loss)) + obs_metrics.add_output_tokens("train_step", int(result_metrics.get("active_tokens", 0))) + if device.type == "cuda" and torch.cuda.is_available(): + # One memory_stats() read instead of two max_memory_* calls, which + # would each rebuild the same ~40us allocator-statistics dict. + stats = torch.cuda.memory_stats(device) + obs_metrics.set_gpu_peak_memory( + "train_step", + allocated_bytes=int(stats.get("allocated_bytes.all.peak", 0)), + reserved_bytes=int(stats.get("reserved_bytes.all.peak", 0)), + ) + except Exception as exc: + logger.warn_once(f"Failed to record training metrics: {exc}") + + def _load_deepspeed() -> Any: _configure_cuda_home_from_python_packages() try: diff --git a/rl_engine/executors/paged_kv_baseline.py b/rl_engine/executors/paged_kv_baseline.py index 0c465720..5e18c5d8 100644 --- a/rl_engine/executors/paged_kv_baseline.py +++ b/rl_engine/executors/paged_kv_baseline.py @@ -23,6 +23,8 @@ score_rewards, summarize_tensor_tree, ) +from rl_engine.observability.metrics import metrics as obs_metrics +from rl_engine.utils.logger import logger @dataclass(frozen=True) @@ -270,6 +272,13 @@ def collect_paged_kv_metrics( input_ids = inputs.input_ids active_tokens = int(_bool_mask(inputs.completion_mask, device=input_ids.device).sum().item()) total_kv_cache_bytes = reservation.reserved_bytes + model_kv_cache_summary.total_bytes + try: + obs_metrics.set_kv_cache_blocks( + required=int(reservation.required_blocks), + reserved=int(reservation.reserved_blocks), + ) + except Exception as exc: + logger.warn_once(f"Failed to record paged KV-cache metrics: {exc}") metrics: dict[str, float | int | str | bool] = { "baseline_kind": "generation_engine_paged_kv_reservation", "baseline_includes_model_kv_cache": model_kv_cache_summary.tensor_count > 0, diff --git a/rl_engine/executors/rollout.py b/rl_engine/executors/rollout.py index 1832ae14..4cf77f80 100644 --- a/rl_engine/executors/rollout.py +++ b/rl_engine/executors/rollout.py @@ -16,6 +16,8 @@ ) from rl_engine.executors.vllm_sampler import VLLMSamplerConfig, VLLMSharedPrefixSampler from rl_engine.kernels.registry import kernel_registry, resolve_logp_op_type +from rl_engine.observability.metrics import metrics as obs_metrics +from rl_engine.observability.nvtx import nvtx_range from rl_engine.utils.logger import logger @@ -67,17 +69,19 @@ def update_weights(self, manifest: WeightUpdateManifest) -> Mapping[str, torch.T manifest.transport, ) previous_update_id = self.active_weight_update_id - try: - imported = dict(self.bridge.import_update(manifest)) - if self.weight_install_adapter is not None: - self.weight_install_adapter.install(manifest, imported) - self.bridge.acknowledge(manifest.update_id) - except Exception as exc: + with obs_metrics.stage_timer("update_weights"), nvtx_range("rlk::update_weights"): try: - self.bridge.reject(manifest.update_id, f"rollout weight install failed: {exc}") - except Exception: - logger.exception("Failed to reject weight update %s", manifest.update_id) - raise + imported = dict(self.bridge.import_update(manifest)) + if self.weight_install_adapter is not None: + self.weight_install_adapter.install(manifest, imported) + self.bridge.acknowledge(manifest.update_id) + except Exception as exc: + try: + self.bridge.reject(manifest.update_id, f"rollout weight install failed: {exc}") + except Exception: + logger.exception("Failed to reject weight update %s", manifest.update_id) + raise + obs_metrics.set_weight_version("consumed", int(manifest.weight_version)) self.shared_weights = imported self.active_weight_version = manifest.weight_version self.active_weight_update_id = manifest.update_id diff --git a/rl_engine/executors/stateless_executor.py b/rl_engine/executors/stateless_executor.py index 2047218f..d302866d 100644 --- a/rl_engine/executors/stateless_executor.py +++ b/rl_engine/executors/stateless_executor.py @@ -11,7 +11,10 @@ import torch +from rl_engine.observability.metrics import metrics as obs_metrics +from rl_engine.observability.nvtx import nvtx_range from rl_engine.testing.reference_ops import selected_logprobs_reference +from rl_engine.utils.logger import logger StatelessForwardMode = Literal["reference", "reward", "both"] StatelessAttentionBackend = Literal["flash_attention_2", "sdpa", "eager", "model_default"] @@ -111,6 +114,14 @@ def __init__( self.reward_adapter = reward_adapter or default_reward_adapter def score(self, inputs: StatelessForwardInputs) -> StatelessForwardResult: + # The error recorder must wrap input validation too: a shape/device + # mismatch is the most common failure and has to be counted as an + # errored stage, not slip past uninstrumented. + stage_started = time.perf_counter() + with _StageErrorRecorder("score", stage_started): + return self._score_impl(inputs) + + def _score_impl(self, inputs: StatelessForwardInputs) -> StatelessForwardResult: _validate_inputs(inputs, self.config) device = inputs.input_ids.device @@ -120,7 +131,10 @@ def score(self, inputs: StatelessForwardInputs) -> StatelessForwardResult: torch.cuda.synchronize(device) started_at = time.perf_counter() - with _temporarily_configure_stateless_model(self.model, self.config) as no_cache_policy: + with ( + nvtx_range("rlk::score"), + _temporarily_configure_stateless_model(self.model, self.config) as no_cache_policy, + ): with torch.no_grad(): try: raw_outputs, use_cache_passed = _run_no_cache_forward(self.model, inputs) @@ -201,6 +215,7 @@ def score(self, inputs: StatelessForwardInputs) -> StatelessForwardResult: no_cache_policy=no_cache_policy, kv_cache_summary=kv_cache_summary, ) + _record_score_observability(metrics) return StatelessForwardResult( reference_logps=reference_logps, rewards=rewards, @@ -209,6 +224,76 @@ def score(self, inputs: StatelessForwardInputs) -> StatelessForwardResult: ) +def _record_score_observability( + result_metrics: Mapping[str, float | int | str | bool], +) -> None: + """Export scoring metrics to the Prometheus facade; never raises. + + Peak-memory values are read back from the metrics dict rather than queried + again: each ``torch.cuda.max_memory_*`` call goes through ``memory_stats()``, + which costs ~40us, and ``collect_stateless_metrics`` has already paid for it. + """ + try: + allocated_bytes = ( + int(float(result_metrics["peak_allocated_mb"]) * 1_048_576) + if "peak_allocated_mb" in result_metrics + else None + ) + reserved_bytes = ( + int(float(result_metrics["peak_reserved_mb"]) * 1_048_576) + if "peak_reserved_mb" in result_metrics + else None + ) + obs_metrics.record_stage( + "score", + float(result_metrics["elapsed_ms"]) / 1000.0, + status="ok", + output_tokens=int(result_metrics.get("active_completion_tokens", 0)), + allocated_bytes=allocated_bytes, + reserved_bytes=reserved_bytes, + ) + if result_metrics.get("attention_backend_fallback"): + reason = str(result_metrics.get("attention_backend_fallback_reason", "")) + obs_metrics.record_hardware_fallback( + "attention_backend", + requested_backend=str(result_metrics.get("attention_backend_requested", "")), + selected_backend="eager", + # Only the exception type is kept: the free-form message would + # be an unbounded Prometheus label value. + reason=reason.split(":", 1)[0] or "unknown", + ) + except Exception as exc: + logger.warn_once(f"Failed to record stateless scoring metrics: {exc}") + + +class _StageErrorRecorder: + """Record an errored stage exit. + + The success path records ``status="ok"`` itself (with the executor's own + forward-region timing), so this only fires when the wrapped block raises, + and reports the full-method wall time from ``score`` down to the failure. + """ + + def __init__(self, stage: str, started_at: float) -> None: + self._stage = stage + self._started_at = started_at + + def __enter__(self) -> None: + return None + + def __exit__(self, exc_type: Any, _exc: Any, _traceback: Any) -> None: + if exc_type is None: + return + try: + obs_metrics.record_stage( + self._stage, + time.perf_counter() - self._started_at, + status="error", + ) + except Exception as metrics_exc: + logger.warn_once(f"Failed to record stage error metrics: {metrics_exc}") + + def score_reference_logprobs( logits: torch.Tensor, inputs: StatelessForwardInputs, diff --git a/rl_engine/executors/vllm_sampler.py b/rl_engine/executors/vllm_sampler.py index 0f10fb75..fcf834b2 100644 --- a/rl_engine/executors/vllm_sampler.py +++ b/rl_engine/executors/vllm_sampler.py @@ -7,6 +7,10 @@ from dataclasses import asdict, dataclass, field from typing import Any, Mapping, Optional, Sequence +from rl_engine.observability.metrics import metrics as obs_metrics +from rl_engine.observability.nvtx import nvtx_range +from rl_engine.utils.logger import logger + @dataclass(frozen=True) class VLLMSamplerConfig: @@ -104,24 +108,28 @@ def generate( sampling_params: Optional[Mapping[str, Any]] = None, ) -> dict[str, Any]: """Generate grouped candidates while keeping each prompt prefix byte-identical.""" - prompt_list = _normalize_prompts(prompts) - generations = num_generations or self.config.num_generations - if generations < 1: - raise ValueError("num_generations must be >= 1") - - expanded_prompts = _expand_prompts(prompt_list, generations) - params = self._build_sampling_params(sampling_params) - outputs = self.engine.generate(expanded_prompts, params) - grouped_outputs = _group_outputs(outputs, len(prompt_list), generations) - - return { - "backend": self.config.backend, - "prefix_cache_enabled": self.config.enable_prefix_caching, - "num_prompts": len(prompt_list), - "num_generations": generations, - "outputs": grouped_outputs, - "normalized_outputs": normalize_grouped_outputs(grouped_outputs), - } + with obs_metrics.stage_timer("rollout_generate"), nvtx_range("rlk::rollout.generate"): + prompt_list = _normalize_prompts(prompts) + generations = num_generations or self.config.num_generations + if generations < 1: + raise ValueError("num_generations must be >= 1") + + expanded_prompts = _expand_prompts(prompt_list, generations) + params = self._build_sampling_params(sampling_params) + outputs = self.engine.generate(expanded_prompts, params) + grouped_outputs = _group_outputs(outputs, len(prompt_list), generations) + + normalized_outputs = normalize_grouped_outputs(grouped_outputs) + result = { + "backend": self.config.backend, + "prefix_cache_enabled": self.config.enable_prefix_caching, + "num_prompts": len(prompt_list), + "num_generations": generations, + "outputs": grouped_outputs, + "normalized_outputs": normalized_outputs, + } + _record_rollout_observability(normalized_outputs) + return result def _build_engine(self) -> Any: if not self.config.model: @@ -157,6 +165,21 @@ def _load_vllm_classes(self) -> tuple[type, type]: return self._llm_cls, self._sampling_params_cls +def _record_rollout_observability( + normalized_outputs: Sequence[Sequence[NormalizedRolloutCandidate]], +) -> None: + """Export generated-token counts to the Prometheus facade; never raises.""" + try: + generated_tokens = sum( + len(candidate.token_ids or []) + for candidate_group in normalized_outputs + for candidate in candidate_group + ) + obs_metrics.add_output_tokens("rollout_generate", generated_tokens) + except Exception as exc: + logger.warn_once(f"Failed to record rollout generation metrics: {exc}") + + def _normalize_prompts( prompts: str | Mapping[str, Any] | Sequence[str | Mapping[str, Any]], ) -> list[str | Mapping[str, Any]]: diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 041ed3e1..3f7170cb 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -6,6 +6,7 @@ from enum import Enum, EnumMeta from typing import Any, Dict, Optional, Set, Type +from rl_engine.observability.metrics import metrics from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -293,8 +294,10 @@ def get_op(self, op_type: str) -> Any: platform = "cpu" candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) - for backend in candidates: + failed_this_call: list[tuple[str, str]] = [] + for index, backend in enumerate(candidates): if backend.name in self._instance_cache: + self._record_dispatch(op_type, backend.name, index, failed_this_call) return self._instance_cache[backend.name] if backend.name in self._failed_backends: @@ -305,15 +308,44 @@ def get_op(self, op_type: str) -> Any: try: op_instance = op_class() self._instance_cache[backend.name] = op_instance + self._record_dispatch(op_type, backend.name, index, failed_this_call) return op_instance except Exception as e: logger.error(f"Failed to instantiate {backend.name}: {e}") self._failed_backends.add(backend.name) + failed_this_call.append((backend.name, "instantiation_error")) else: self._failed_backends.add(backend.name) + failed_this_call.append((backend.name, "import_error")) + self._record_dispatch(op_type, "none", -1, failed_this_call) raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + def _record_dispatch( + self, + op_type: str, + selected_backend: str, + candidate_index: int, + failed_candidates: list[tuple[str, str]], + ) -> None: + """Export dispatch-time state and per-call fallback events; never raises.""" + try: + for failed_backend, reason in failed_candidates: + metrics.record_hardware_fallback( + op_type, + requested_backend=failed_backend, + selected_backend=selected_backend, + reason=reason, + ) + if candidate_index >= 0: + metrics.set_selected_backend( + op_type, + selected_backend, + is_fallback=candidate_index > 0, + ) + except Exception as e: # observability must never break dispatch + logger.warn_once(f"Failed to record kernel dispatch metrics: {e}") + def _load_backend(self, backend: OpBackend) -> Optional[Type]: """Dynamic loading technique: Import modules only when needed and check environment dependencies. diff --git a/rl_engine/observability/__init__.py b/rl_engine/observability/__init__.py new file mode 100644 index 00000000..3f1ef80c --- /dev/null +++ b/rl_engine/observability/__init__.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RL-Kernel observability: Prometheus metrics facade + NVTX stage ranges. + +Importing this package has no side effects: nothing is registered globally and +no HTTP server starts until ``start_metrics_server`` is called explicitly. +""" + +from rl_engine.observability.metrics import SCHEMA_METRIC_NAMES, MetricsRegistry, metrics +from rl_engine.observability.nvtx import nvtx_range +from rl_engine.observability.server import maybe_start_metrics_server_from_env, start_metrics_server + +__all__ = [ + "SCHEMA_METRIC_NAMES", + "MetricsRegistry", + "metrics", + "nvtx_range", + "maybe_start_metrics_server_from_env", + "start_metrics_server", +] diff --git a/rl_engine/observability/metrics.py b/rl_engine/observability/metrics.py new file mode 100644 index 00000000..2d1aa465 --- /dev/null +++ b/rl_engine/observability/metrics.py @@ -0,0 +1,473 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Prometheus metrics facade for RL-Kernel. + +The facade exposes domain-level recording methods instead of raw prometheus +objects. ``prometheus_client`` is an optional dependency: when it cannot be +imported every method degrades to a silent no-op, so engine code may call the +module-level ``metrics`` singleton unconditionally. +""" + +from __future__ import annotations + +import bisect +import threading +import time +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Any, Iterator, Optional + +from rl_engine.utils.logger import logger + +# Metric families exported by this facade. tests/test_observability.py uses this +# schema to guard examples/grafana-dashboard.json against drift. +SCHEMA_METRIC_NAMES: frozenset[str] = frozenset( + { + "rl_kernel_build_info", + "rl_kernel_stage_duration_seconds", + "rl_kernel_requests_total", + "rl_kernel_output_tokens_total", + "rl_kernel_selected_backend_info", + "rl_kernel_hardware_fallback_total", + "rl_kernel_kv_cache_blocks", + "rl_kernel_kv_cache_fragmentation_ratio", + "rl_kernel_gpu_peak_memory_bytes", + "rl_kernel_training_loss", + "rl_kernel_weight_version", + } +) + +# Stage latencies span sub-millisecond kernel paths to multi-minute rollouts. +_STAGE_DURATION_BUCKETS = ( + 0.005, + 0.01, + 0.025, + 0.05, + 0.1, + 0.25, + 0.5, + 1.0, + 2.5, + 5.0, + 10.0, + 30.0, + 60.0, + 120.0, + 300.0, + 600.0, +) + + +@dataclass +class _PendingStageUpdate: + """Bounded aggregate for amortized Prometheus child updates.""" + + count: int = 0 + output_tokens: int = 0 + duration_sum: float = 0.0 + durations: list[float] = field(default_factory=list) + bucket_counts: list[int] = field( + default_factory=lambda: [0] * (len(_STAGE_DURATION_BUCKETS) + 1) + ) + allocated_bytes: Optional[int] = None + reserved_bytes: Optional[int] = None + + +class MetricsRegistry: + """Prometheus facade. Every method is a safe no-op when prometheus_client is absent. + + Each instance owns a dedicated ``CollectorRegistry`` (never the process-global + default) so instances cannot collide and tests stay isolated. + """ + + def __init__(self, *, stage_batch_size: int = 1) -> None: + if stage_batch_size < 1: + raise ValueError("stage_batch_size must be at least 1") + self._initialized = False + self._enabled = False + self._registry: Any = None + self._stage_batch_size = stage_batch_size + self._pending_stage_updates: dict[tuple[str, str], _PendingStageUpdate] = {} + self._pending_stage_lock = threading.Lock() + # Resolving a labelled child costs a dict lookup plus a lock on every + # call, so cache them: instrumented sites run per step and label values + # come from finite enums, which bounds this cache by construction. + self._children: dict[tuple[str, tuple[str, ...]], Any] = {} + self._stage_duration: Any = None + self._requests: Any = None + self._output_tokens: Any = None + self._selected_backend: Any = None + self._hardware_fallback: Any = None + self._kv_cache_blocks: Any = None + self._kv_cache_fragmentation: Any = None + self._gpu_peak_memory: Any = None + self._training_loss: Any = None + self._weight_version: Any = None + + @property + def enabled(self) -> bool: + """Whether prometheus_client is importable and collectors are registered.""" + return self._ensure_initialized() + + @property + def collector_registry(self) -> Optional[Any]: + """The dedicated CollectorRegistry, or None when disabled.""" + return self._registry if self._ensure_initialized() else None + + # ------------------------------------------------------------------ + # Domain-level recording methods + # ------------------------------------------------------------------ + + @contextmanager + def stage_timer(self, stage: str) -> Iterator[None]: + """Time a stage; records duration and a request with ok/error status. + + Exceptions propagate unchanged after being recorded with status="error". + """ + start = time.perf_counter() + try: + yield + except BaseException: + self.record_stage(stage, time.perf_counter() - start, status="error") + raise + self.record_stage(stage, time.perf_counter() - start, status="ok") + + def record_stage( + self, + stage: str, + duration_seconds: float, + *, + status: str, + output_tokens: int = 0, + allocated_bytes: Optional[int] = None, + reserved_bytes: Optional[int] = None, + ) -> None: + """Record one completed stage through a single facade call. + + Executors that already measure their own stage duration should use this + method instead of wrapping the same work in ``stage_timer``. Besides + keeping the exported duration aligned with their result contract, the + combined update avoids repeated initialization checks and labelled-child + cache lookups on latency-sensitive paths. + """ + if not self._ensure_initialized(): + return + if self._stage_batch_size == 1: + self._record_stage_now( + stage, + duration_seconds, + status=status, + output_tokens=output_tokens, + allocated_bytes=allocated_bytes, + reserved_bytes=reserved_bytes, + ) + return + + key = (stage, status) + with self._pending_stage_lock: + pending = self._pending_stage_updates.get(key) + if pending is None: + pending = _PendingStageUpdate() + self._pending_stage_updates[key] = pending + pending.count += 1 + pending.output_tokens += max(0, output_tokens) + pending.duration_sum += duration_seconds + pending.durations.append(duration_seconds) + bucket_index = bisect.bisect_left(_STAGE_DURATION_BUCKETS, duration_seconds) + pending.bucket_counts[bucket_index] += 1 + if allocated_bytes is not None: + pending.allocated_bytes = allocated_bytes + if reserved_bytes is not None: + pending.reserved_bytes = reserved_bytes + if pending.count < self._stage_batch_size: + return + del self._pending_stage_updates[key] + self._flush_stage_update(stage, status, pending) + + def flush(self) -> None: + """Flush pending batched updates before collection or shutdown.""" + if not self._ensure_initialized(): + return + with self._pending_stage_lock: + pending_updates = self._pending_stage_updates + self._pending_stage_updates = {} + for (stage, status), pending in pending_updates.items(): + self._flush_stage_update(stage, status, pending) + + def add_output_tokens(self, stage: str, count: int) -> None: + if not self._ensure_initialized() or count <= 0: + return + self._child(self._output_tokens, "tokens", (stage,)).inc(count) + + def set_selected_backend(self, op_type: str, backend: str, *, is_fallback: bool) -> None: + if not self._ensure_initialized(): + return + labels = (op_type, backend, str(is_fallback).lower()) + self._child(self._selected_backend, "backend", labels).set(1) + + def record_hardware_fallback( + self, + op_type: str, + *, + requested_backend: str, + selected_backend: str, + reason: str, + ) -> None: + if not self._ensure_initialized(): + return + labels = (op_type, requested_backend, selected_backend, reason) + self._child(self._hardware_fallback, "fallback", labels).inc() + + def set_kv_cache_blocks(self, *, required: int, reserved: int) -> None: + if not self._ensure_initialized(): + return + self._child(self._kv_cache_blocks, "kv", ("required",)).set(required) + self._child(self._kv_cache_blocks, "kv", ("reserved",)).set(reserved) + fragmentation = 0.0 if reserved <= 0 else 1.0 - (required / reserved) + self._kv_cache_fragmentation.set(fragmentation) + + def set_gpu_peak_memory( + self, + stage: str, + *, + allocated_bytes: int, + reserved_bytes: int, + ) -> None: + if not self._ensure_initialized(): + return + self._child(self._gpu_peak_memory, "mem", (stage, "allocated")).set(allocated_bytes) + self._child(self._gpu_peak_memory, "mem", (stage, "reserved")).set(reserved_bytes) + + def set_training_loss(self, loss: float) -> None: + if not self._ensure_initialized(): + return + self._training_loss.set(loss) + + def set_weight_version(self, role: str, version: int) -> None: + if not self._ensure_initialized(): + return + self._child(self._weight_version, "weights", (role,)).set(version) + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + def _record_stage_now( + self, + stage: str, + duration_seconds: float, + *, + status: str, + output_tokens: int, + allocated_bytes: Optional[int], + reserved_bytes: Optional[int], + ) -> None: + self._child(self._stage_duration, "duration", (stage,)).observe(duration_seconds) + self._child(self._requests, "requests", (stage, status)).inc() + if output_tokens > 0: + self._child(self._output_tokens, "tokens", (stage,)).inc(output_tokens) + if allocated_bytes is not None: + self._child(self._gpu_peak_memory, "mem", (stage, "allocated")).set( + allocated_bytes + ) + if reserved_bytes is not None: + self._child(self._gpu_peak_memory, "mem", (stage, "reserved")).set(reserved_bytes) + + def _flush_stage_update( + self, + stage: str, + status: str, + pending: _PendingStageUpdate, + ) -> None: + duration = self._child(self._stage_duration, "duration", (stage,)) + # prometheus_client has no public batch-observe API. Its Histogram child + # stores the non-cumulative bucket values and sum as thread-safe Value + # objects; incrementing those once per batch preserves exactly the same + # samples as calling observe() for every stage, while moving locks and + # label lookups off the hot path. + histogram_sum = getattr(duration, "_sum", None) + histogram_buckets = getattr(duration, "_buckets", None) + if ( + histogram_sum is not None + and histogram_buckets is not None + and len(histogram_buckets) == len(pending.bucket_counts) + ): + histogram_sum.inc(pending.duration_sum) + for bucket, count in zip( + histogram_buckets, + pending.bucket_counts, + strict=True, + ): + if count: + bucket.inc(count) + else: # preserve correctness if prometheus_client changes its internals + for duration_seconds in pending.durations: + duration.observe(duration_seconds) + self._child(self._requests, "requests", (stage, status)).inc(pending.count) + if pending.output_tokens > 0: + self._child(self._output_tokens, "tokens", (stage,)).inc(pending.output_tokens) + if pending.allocated_bytes is not None: + self._child(self._gpu_peak_memory, "mem", (stage, "allocated")).set( + pending.allocated_bytes + ) + if pending.reserved_bytes is not None: + self._child(self._gpu_peak_memory, "mem", (stage, "reserved")).set( + pending.reserved_bytes + ) + + def _child(self, metric: Any, key: str, labels: tuple[str, ...]) -> Any: + """Return a labelled child, caching it across calls.""" + cache_key = (key, labels) + child = self._children.get(cache_key) + if child is None: + child = metric.labels(*labels) + self._children[cache_key] = child + return child + + def _ensure_initialized(self) -> bool: + if self._initialized: + return self._enabled + self._initialized = True + try: + import prometheus_client + except ImportError: + logger.info_once( + "prometheus_client is not installed; RL-Kernel metrics are disabled. " + "Install with `pip install rl-kernel[observability]` to enable them." + ) + self._enabled = False + return False + try: + self._build_collectors(prometheus_client) + self._enabled = True + except Exception as exc: # exporter failures must never break engine code + logger.warning(f"Failed to initialize RL-Kernel metrics collectors: {exc}") + self._enabled = False + return self._enabled + + def _build_collectors(self, prometheus_client: Any) -> None: + owner = self + + class _FlushingCollectorRegistry(prometheus_client.CollectorRegistry): + def collect(registry_self) -> Iterator[Any]: + owner.flush() + yield from super().collect() + + registry = _FlushingCollectorRegistry() + self._registry = registry + + for collector_name in ("ProcessCollector", "PlatformCollector"): + collector_cls = getattr(prometheus_client, collector_name, None) + if collector_cls is None: + continue + try: + collector_cls(registry=registry) + except Exception: + # Platform/process collectors are best-effort (e.g. non-Linux). + pass + + self._stage_duration = prometheus_client.Histogram( + "rl_kernel_stage_duration_seconds", + "Wall-clock duration of an instrumented RL-Kernel stage.", + labelnames=("stage",), + buckets=_STAGE_DURATION_BUCKETS, + registry=registry, + ) + self._requests = prometheus_client.Counter( + "rl_kernel_requests", + "Completed stage executions by status.", + labelnames=("stage", "status"), + registry=registry, + ) + self._output_tokens = prometheus_client.Counter( + "rl_kernel_output_tokens", + "Tokens processed per stage (generated for rollout, active for train/score).", + labelnames=("stage",), + registry=registry, + ) + self._selected_backend = prometheus_client.Gauge( + "rl_kernel_selected_backend_info", + "Kernel backend selected at dispatch time (value is always 1).", + labelnames=("op_type", "backend", "is_fallback"), + registry=registry, + ) + self._hardware_fallback = prometheus_client.Counter( + "rl_kernel_hardware_fallback", + "Runtime fallback events from a requested backend to a selected backend.", + labelnames=("op_type", "requested_backend", "selected_backend", "reason"), + registry=registry, + ) + self._kv_cache_blocks = prometheus_client.Gauge( + "rl_kernel_kv_cache_blocks", + "Paged KV-cache blocks by kind.", + labelnames=("kind",), + registry=registry, + ) + self._kv_cache_fragmentation = prometheus_client.Gauge( + "rl_kernel_kv_cache_fragmentation_ratio", + "Paged KV-cache fragmentation: 1 - required/reserved (0 when reserved is 0).", + registry=registry, + ) + self._gpu_peak_memory = prometheus_client.Gauge( + "rl_kernel_gpu_peak_memory_bytes", + "Peak GPU memory observed during a stage.", + labelnames=("stage", "kind"), + registry=registry, + ) + self._training_loss = prometheus_client.Gauge( + "rl_kernel_training_loss", + "Most recent training loss.", + registry=registry, + ) + self._weight_version = prometheus_client.Gauge( + "rl_kernel_weight_version", + "Latest published/consumed weight version.", + labelnames=("role",), + registry=registry, + ) + self._set_build_info(prometheus_client, registry) + + def _set_build_info(self, prometheus_client: Any, registry: Any) -> None: + build_info = prometheus_client.Gauge( + "rl_kernel_build_info", + "RL-Kernel build/platform information (value is always 1).", + labelnames=("version", "device_type", "backend_version", "ext_available"), + registry=registry, + ) + build_info.labels( + version=_package_version(), + device_type=_device_label("device_type"), + backend_version=_device_label("backend_version"), + ext_available=str(_extension_available()).lower(), + ).set(1) + + +def _package_version() -> str: + try: + from importlib.metadata import version + + return version("RL-Kernel") + except Exception: + return "unknown" + + +def _device_label(attribute: str) -> str: + try: + from rl_engine.platforms.device import device_ctx + + return str(getattr(device_ctx, attribute)) + except Exception: + return "unknown" + + +def _extension_available() -> bool: + try: + from rl_engine.kernels.ops.base import _EXT_AVAILABLE + + return bool(_EXT_AVAILABLE) + except Exception: + return False + + +metrics = MetricsRegistry(stage_batch_size=16) # module-level singleton, mirrors utils/logger.py diff --git a/rl_engine/observability/nvtx.py b/rl_engine/observability/nvtx.py new file mode 100644 index 00000000..39eefba1 --- /dev/null +++ b/rl_engine/observability/nvtx.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Python-side NVTX stage ranges. + +Stage ranges (``rlk::rollout.generate``, ``rlk::train_step``, ...) group the +per-operator ``rlk::`` ranges emitted by the C++ bindings on an Nsight +Systems timeline. Uses ``torch.cuda.nvtx`` (no extra dependency) and silently +no-ops when CUDA is unavailable or a profiler API is missing. +""" + +from __future__ import annotations + +import os +from contextlib import AbstractContextManager, nullcontext +from functools import lru_cache +from typing import Any, Optional + +RL_KERNEL_NVTX = "RL_KERNEL_NVTX" + + +@lru_cache(maxsize=1) +def _nvtx_enabled() -> bool: + """Return whether runtime NVTX emission was explicitly requested.""" + return os.environ.get(RL_KERNEL_NVTX, "").strip().lower() in {"1", "true", "yes", "on"} + + +@lru_cache(maxsize=1) +def _nvtx_module() -> Optional[Any]: + if not _nvtx_enabled(): + return None + try: + import torch + + if not torch.cuda.is_available(): + return None + nvtx = torch.cuda.nvtx + if not hasattr(nvtx, "range_push") or not hasattr(nvtx, "range_pop"): + return None + return nvtx + except Exception: + return None + + +class _NvtxRange(AbstractContextManager[None]): + def __init__(self, nvtx: Any, name: str) -> None: + self._nvtx = nvtx + self._name = name + + def __enter__(self) -> None: + self._nvtx.range_push(self._name) + + def __exit__(self, *_exc_info: Any) -> None: + self._nvtx.range_pop() + + +def nvtx_range(name: str) -> AbstractContextManager[None]: + """Return an opted-in NVTX range, or a cheap null context by default.""" + nvtx = _nvtx_module() + if nvtx is None: + return nullcontext() + return _NvtxRange(nvtx, name) diff --git a/rl_engine/observability/server.py b/rl_engine/observability/server.py new file mode 100644 index 00000000..52dd036c --- /dev/null +++ b/rl_engine/observability/server.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Prometheus /metrics HTTP exporter for RL-Kernel. + +The server is never started as an import side effect. Application code (or an +example honoring ``RL_KERNEL_METRICS=1``) must call ``start_metrics_server`` +explicitly; library executors only record into the facade. +""" + +from __future__ import annotations + +import os +import threading +from typing import Optional + +from rl_engine.observability.metrics import metrics +from rl_engine.utils.logger import logger + +RL_KERNEL_METRICS = "RL_KERNEL_METRICS" +RL_KERNEL_METRICS_PORT = "RL_KERNEL_METRICS_PORT" +RL_KERNEL_METRICS_ADDR = "RL_KERNEL_METRICS_ADDR" + +_DEFAULT_PORT = 8000 +_DEFAULT_ADDR = "0.0.0.0" + +_lock = threading.Lock() +_bound_port: Optional[int] = None + + +def start_metrics_server(port: Optional[int] = None, addr: Optional[str] = None) -> int: + """Start the /metrics HTTP server for this process and return the bound port. + + Idempotent per process: subsequent calls return the port of the already + running server. ``port=0`` binds an ephemeral port. Defaults come from + ``RL_KERNEL_METRICS_PORT`` / ``RL_KERNEL_METRICS_ADDR``. + + Raises RuntimeError when prometheus_client is not installed; use + ``maybe_start_metrics_server_from_env`` for a best-effort opt-in start. + """ + global _bound_port + + with _lock: + if _bound_port is not None: + logger.warn_once( + "start_metrics_server called more than once; reusing the running exporter." + ) + return _bound_port + + registry = metrics.collector_registry + if registry is None: + raise RuntimeError( + "prometheus_client is required for the metrics exporter. " + "Install it with `pip install rl-kernel[observability]`." + ) + from prometheus_client import start_http_server + + resolved_port = ( + port + if port is not None + else int(os.environ.get(RL_KERNEL_METRICS_PORT, str(_DEFAULT_PORT))) + ) + resolved_addr = addr or os.environ.get(RL_KERNEL_METRICS_ADDR, _DEFAULT_ADDR) + + server, _thread = start_http_server(resolved_port, addr=resolved_addr, registry=registry) + _bound_port = int(server.server_port) + logger.info(f"RL-Kernel metrics exporter listening on {resolved_addr}:{_bound_port}") + return _bound_port + + +def maybe_start_metrics_server_from_env() -> Optional[int]: + """Start the exporter iff ``RL_KERNEL_METRICS=1``; never raises. + + Returns the bound port, or None when disabled or unavailable. Intended for + examples and entrypoints, not library code. + """ + if os.environ.get(RL_KERNEL_METRICS, "").strip() != "1": + return None + try: + return start_metrics_server() + except Exception as exc: + logger.warning(f"RL_KERNEL_METRICS=1 but the metrics exporter could not start: {exc}") + return None diff --git a/setup.py b/setup.py index 6f94e040..27191981 100644 --- a/setup.py +++ b/setup.py @@ -134,6 +134,9 @@ def get_extensions(): nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] + if envs.env_flag(envs.KERNEL_ALIGN_DISABLE_NVTX): + cxx_flags.append("-DKERNEL_ALIGN_DISABLE_NVTX") + nvcc_flags.append("-DKERNEL_ALIGN_DISABLE_NVTX") extra_link_args = list(torch_rpath) sm90_srcs = [ @@ -187,6 +190,7 @@ def get_cmdclass(): "cuda": ["flashinfer"], "rocm": ["aiter"], "vllm": ["vllm>=0.6.0"], + "observability": ["prometheus-client>=0.20"], }, python_requires=">=3.10", include_package_data=True, diff --git a/tests/test_observability.py b/tests/test_observability.py new file mode 100644 index 00000000..227aa14b --- /dev/null +++ b/tests/test_observability.py @@ -0,0 +1,560 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import importlib +import json +import pathlib +import re +import sys +import urllib.request +from typing import Any, Optional + +import pytest +import torch + +from rl_engine.executors.paged_kv_baseline import PagedKVScoringConfig, reserve_paged_kv_cache +from rl_engine.executors.stateless_executor import ( + StatelessForwardConfig, + StatelessForwardExecutor, + StatelessForwardInputs, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend +from rl_engine.observability import SCHEMA_METRIC_NAMES, MetricsRegistry, metrics, nvtx_range +from rl_engine.platforms.device import device_ctx + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[1] +DASHBOARD_PATH = REPO_ROOT / "examples" / "grafana-dashboard.json" + + +def _sample_value( + registry: MetricsRegistry, + name: str, + labels: Optional[dict[str, str]] = None, +) -> Optional[float]: + collector_registry = registry.collector_registry + assert collector_registry is not None + return collector_registry.get_sample_value(name, labels or {}) + + +def _make_score_inputs(batch_size: int = 2, seq_len: int = 6) -> StatelessForwardInputs: + torch.manual_seed(0) + input_ids = torch.randint(0, 16, (batch_size, seq_len)) + attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long) + completion_mask = torch.ones(batch_size, seq_len, dtype=torch.bool) + completion_mask[:, 0] = False + return StatelessForwardInputs( + input_ids=input_ids, + attention_mask=attention_mask, + completion_mask=completion_mask, + ) + + +class _TinyLogitsModel(torch.nn.Module): + def __init__(self, vocab_size: int = 16, hidden_dim: int = 8): + super().__init__() + self.embedding = torch.nn.Embedding(vocab_size, hidden_dim) + self.proj = torch.nn.Linear(hidden_dim, vocab_size) + + def forward(self, input_ids: torch.Tensor, attention_mask=None, use_cache=None): + del attention_mask, use_cache + return {"logits": self.proj(self.embedding(input_ids))} + + +class _FlashFailingModel(_TinyLogitsModel): + """Fails the first forward like a missing flash-attention install.""" + + def __init__(self): + super().__init__() + self.calls = 0 + + def forward(self, input_ids: torch.Tensor, attention_mask=None, use_cache=None): + self.calls += 1 + if self.calls == 1: + raise ImportError("flash attention kernels are unavailable") + return super().forward(input_ids, attention_mask, use_cache) + + +class _AlwaysFailingModel(_TinyLogitsModel): + def forward(self, input_ids: torch.Tensor, attention_mask=None, use_cache=None): + raise RuntimeError("score exploded") + + +@pytest.fixture() +def blocked_prometheus(monkeypatch): + """Make `import prometheus_client` fail inside the fixture scope.""" + monkeypatch.setitem(sys.modules, "prometheus_client", None) + + +@pytest.fixture() +def fresh_metrics(): + pytest.importorskip("prometheus_client") + return MetricsRegistry() + + +# --------------------------------------------------------------------------- +# 1. No-op degradation without prometheus_client +# --------------------------------------------------------------------------- + + +def test_facade_is_noop_without_prometheus(blocked_prometheus): + registry = MetricsRegistry() + + assert registry.enabled is False + assert registry.collector_registry is None + + with registry.stage_timer("score"): + pass + registry.add_output_tokens("score", 3) + registry.set_selected_backend("logp", "PYTORCH_NATIVE", is_fallback=False) + registry.record_hardware_fallback( + "logp", + requested_backend="a", + selected_backend="b", + reason="import_error", + ) + registry.set_kv_cache_blocks(required=1, reserved=2) + registry.set_gpu_peak_memory("score", allocated_bytes=1, reserved_bytes=2) + registry.set_training_loss(0.5) + registry.set_weight_version("published", 1) + + +def test_stage_timer_reraises_without_prometheus(blocked_prometheus): + registry = MetricsRegistry() + with pytest.raises(ValueError, match="boom"): + with registry.stage_timer("score"): + raise ValueError("boom") + + +def test_executor_runs_unchanged_without_prometheus(blocked_prometheus, monkeypatch): + disabled = MetricsRegistry() + monkeypatch.setattr("rl_engine.executors.stateless_executor.obs_metrics", disabled) + + executor = StatelessForwardExecutor( + _TinyLogitsModel(), + StatelessForwardConfig(mode="reference", attention_backend="model_default"), + ) + result = executor.score(_make_score_inputs()) + + assert result.reference_logps is not None + assert result.metrics["active_completion_tokens"] == 10 + + +# --------------------------------------------------------------------------- +# 2. /metrics endpoint +# --------------------------------------------------------------------------- + + +def test_metrics_endpoint_serves_schema_families(): + pytest.importorskip("prometheus_client") + from rl_engine.observability.server import start_metrics_server + + # Touch a couple of facade methods so families beyond build_info have data. + metrics.set_training_loss(1.25) + with metrics.stage_timer("score"): + pass + + port = start_metrics_server(port=0, addr="127.0.0.1") + assert port > 0 + # Idempotent second call reuses the running exporter. + assert start_metrics_server(port=0, addr="127.0.0.1") == port + + body = urllib.request.urlopen(f"http://127.0.0.1:{port}/metrics", timeout=5).read().decode() + assert "rl_kernel_build_info" in body + assert "rl_kernel_training_loss" in body + assert "rl_kernel_stage_duration_seconds" in body + assert "rl_kernel_requests_total" in body + + +def test_env_helper_disabled_by_default(monkeypatch): + from rl_engine.observability.server import maybe_start_metrics_server_from_env + + monkeypatch.delenv("RL_KERNEL_METRICS", raising=False) + assert maybe_start_metrics_server_from_env() is None + + +# --------------------------------------------------------------------------- +# 3. Kernel registry dispatch + fallback instrumentation +# --------------------------------------------------------------------------- + + +def _active_platform() -> str: + """Priority-map key `KernelRegistry.get_op` resolves to on this host.""" + if device_ctx.is_rocm: + return "rocm" + return "cuda" if device_ctx.device_type == "cuda" else "cpu" + + +def test_registry_records_selected_backend(fresh_metrics, monkeypatch): + monkeypatch.setattr("rl_engine.kernels.registry.metrics", fresh_metrics) + registry = KernelRegistry() + # Pin a single known backend so the assertion holds on CPU and GPU hosts alike. + monkeypatch.setitem( + registry._priority_map[_active_platform()], + "logp", + [OpBackend.PYTORCH_NATIVE], + ) + + op = registry.get_op("logp") + + assert op is not None + assert ( + _sample_value( + fresh_metrics, + "rl_kernel_selected_backend_info", + {"op_type": "logp", "backend": "PYTORCH_NATIVE", "is_fallback": "false"}, + ) + == 1 + ) + + +def test_registry_records_hardware_fallback(fresh_metrics, monkeypatch): + monkeypatch.setattr("rl_engine.kernels.registry.metrics", fresh_metrics) + registry = KernelRegistry() + monkeypatch.setitem( + registry._priority_map[_active_platform()], + "logp", + [OpBackend.FLASH_ATTN, OpBackend.PYTORCH_NATIVE], + ) + + real_import_module = importlib.import_module + + def failing_import(path: str, *args: Any, **kwargs: Any): + if "flash_attn" in path: + raise ModuleNotFoundError("No module named 'flash_attn'", name="flash_attn") + return real_import_module(path, *args, **kwargs) + + monkeypatch.setattr(importlib, "import_module", failing_import) + + op = registry.get_op("logp") + + assert op is not None + assert ( + _sample_value( + fresh_metrics, + "rl_kernel_hardware_fallback_total", + { + "op_type": "logp", + "requested_backend": "FLASH_ATTN", + "selected_backend": "PYTORCH_NATIVE", + "reason": "import_error", + }, + ) + == 1 + ) + assert ( + _sample_value( + fresh_metrics, + "rl_kernel_selected_backend_info", + {"op_type": "logp", "backend": "PYTORCH_NATIVE", "is_fallback": "true"}, + ) + == 1 + ) + + +# --------------------------------------------------------------------------- +# 4. Timer semantics +# --------------------------------------------------------------------------- + + +def test_stage_timer_records_ok_and_error(fresh_metrics): + with fresh_metrics.stage_timer("train_step"): + pass + with pytest.raises(RuntimeError, match="exploded"): + with fresh_metrics.stage_timer("train_step"): + raise RuntimeError("exploded") + + ok = _sample_value( + fresh_metrics, + "rl_kernel_requests_total", + {"stage": "train_step", "status": "ok"}, + ) + error = _sample_value( + fresh_metrics, + "rl_kernel_requests_total", + {"stage": "train_step", "status": "error"}, + ) + duration_count = _sample_value( + fresh_metrics, + "rl_kernel_stage_duration_seconds_count", + {"stage": "train_step"}, + ) + assert ok == 1 + assert error == 1 + assert duration_count == 2 + + +def test_batched_stage_updates_flush_exact_values_on_collection(): + registry = MetricsRegistry(stage_batch_size=8) + registry.record_stage( + "score", + 0.002, + status="ok", + output_tokens=3, + allocated_bytes=10, + reserved_bytes=20, + ) + registry.record_stage( + "score", + 0.020, + status="ok", + output_tokens=4, + allocated_bytes=30, + reserved_bytes=40, + ) + + assert ( + _sample_value( + registry, + "rl_kernel_requests_total", + {"stage": "score", "status": "ok"}, + ) + == 2 + ) + assert ( + _sample_value( + registry, + "rl_kernel_stage_duration_seconds_count", + {"stage": "score"}, + ) + == 2 + ) + assert _sample_value( + registry, + "rl_kernel_stage_duration_seconds_sum", + {"stage": "score"}, + ) == pytest.approx(0.022) + assert ( + _sample_value( + registry, + "rl_kernel_output_tokens_total", + {"stage": "score"}, + ) + == 7 + ) + assert ( + _sample_value( + registry, + "rl_kernel_gpu_peak_memory_bytes", + {"stage": "score", "kind": "allocated"}, + ) + == 30 + ) + + +def test_metrics_registry_rejects_invalid_stage_batch_size(): + with pytest.raises(ValueError, match="stage_batch_size"): + MetricsRegistry(stage_batch_size=0) + + +# --------------------------------------------------------------------------- +# 5. Executor instrumentation +# --------------------------------------------------------------------------- + + +def test_score_records_stage_and_tokens(fresh_metrics, monkeypatch): + monkeypatch.setattr("rl_engine.executors.stateless_executor.obs_metrics", fresh_metrics) + + executor = StatelessForwardExecutor( + _TinyLogitsModel(), + StatelessForwardConfig(mode="reference", attention_backend="model_default"), + ) + result = executor.score(_make_score_inputs()) + + assert result.metrics["active_completion_tokens"] == 10 + assert ( + _sample_value( + fresh_metrics, + "rl_kernel_requests_total", + {"stage": "score", "status": "ok"}, + ) + == 1 + ) + assert _sample_value(fresh_metrics, "rl_kernel_output_tokens_total", {"stage": "score"}) == 10 + + +def test_score_records_attention_fallback_event(fresh_metrics, monkeypatch): + monkeypatch.setattr("rl_engine.executors.stateless_executor.obs_metrics", fresh_metrics) + + executor = StatelessForwardExecutor( + _FlashFailingModel(), + StatelessForwardConfig(mode="reference", attention_backend="flash_attention_2"), + ) + result = executor.score(_make_score_inputs()) + + assert result.metrics["attention_backend_fallback"] is True + assert ( + _sample_value( + fresh_metrics, + "rl_kernel_hardware_fallback_total", + { + "op_type": "attention_backend", + "requested_backend": "flash_attention_2", + "selected_backend": "eager", + "reason": "ImportError", + }, + ) + == 1 + ) + + +def test_score_records_error_without_changing_exception(fresh_metrics, monkeypatch): + monkeypatch.setattr("rl_engine.executors.stateless_executor.obs_metrics", fresh_metrics) + executor = StatelessForwardExecutor( + _AlwaysFailingModel(), + StatelessForwardConfig(mode="reference", attention_backend="model_default"), + ) + + with pytest.raises(RuntimeError, match="score exploded"): + executor.score(_make_score_inputs()) + + assert ( + _sample_value( + fresh_metrics, + "rl_kernel_requests_total", + {"stage": "score", "status": "error"}, + ) + == 1 + ) + + +def test_score_records_error_on_input_validation_failure(fresh_metrics, monkeypatch): + # Regression: a shape/device mismatch is rejected before the forward pass, + # so it must still be recorded as an errored stage rather than escaping + # uninstrumented. + monkeypatch.setattr("rl_engine.executors.stateless_executor.obs_metrics", fresh_metrics) + executor = StatelessForwardExecutor( + _TinyLogitsModel(), + StatelessForwardConfig(mode="reference", attention_backend="model_default"), + ) + good = _make_score_inputs(batch_size=2, seq_len=6) + mismatched = StatelessForwardInputs( + input_ids=good.input_ids, + attention_mask=torch.ones(2, 7, dtype=torch.long), # wrong shape + completion_mask=good.completion_mask, + ) + + with pytest.raises(ValueError, match="attention_mask shape"): + executor.score(mismatched) + + assert ( + _sample_value( + fresh_metrics, + "rl_kernel_requests_total", + {"stage": "score", "status": "error"}, + ) + == 1 + ) + + +def test_paged_kv_gauges(fresh_metrics, monkeypatch): + monkeypatch.setattr("rl_engine.executors.paged_kv_baseline.obs_metrics", fresh_metrics) + from rl_engine.executors.paged_kv_baseline import collect_paged_kv_metrics + from rl_engine.executors.stateless_executor import summarize_tensor_tree + + inputs = _make_score_inputs(batch_size=2, seq_len=6) + config = PagedKVScoringConfig(block_size=4, kv_cache_blocks=8) + reservation = reserve_paged_kv_cache(inputs, config) + collect_paged_kv_metrics( + inputs, + reservation, + config=config, + elapsed_seconds=0.1, + use_cache_passed=True, + cuda_tracking=False, + model_kv_cache_summary=summarize_tensor_tree(None), + ) + + # 2 sequences x 6 tokens with block_size=4 -> 2 blocks each. + assert _sample_value(fresh_metrics, "rl_kernel_kv_cache_blocks", {"kind": "required"}) == 4 + assert _sample_value(fresh_metrics, "rl_kernel_kv_cache_blocks", {"kind": "reserved"}) == 8 + assert _sample_value(fresh_metrics, "rl_kernel_kv_cache_fragmentation_ratio") == 0.5 + + +def test_training_observability_helper(fresh_metrics, monkeypatch): + monkeypatch.setattr("rl_engine.executors.deepspeed_trainer.obs_metrics", fresh_metrics) + from rl_engine.executors.deepspeed_trainer import _record_training_observability + + _record_training_observability( + {"loss": 0.75, "active_tokens": 12}, + device=torch.device("cpu"), + ) + + assert _sample_value(fresh_metrics, "rl_kernel_training_loss") == 0.75 + assert ( + _sample_value(fresh_metrics, "rl_kernel_output_tokens_total", {"stage": "train_step"}) == 12 + ) + + +def test_nvtx_range_noops_without_cuda(): + with nvtx_range("rlk::test"): + value = 41 + 1 + assert value == 42 + + +def test_nvtx_requires_explicit_runtime_opt_in(monkeypatch): + nvtx_module = importlib.import_module("rl_engine.observability.nvtx") + monkeypatch.delenv(nvtx_module.RL_KERNEL_NVTX, raising=False) + nvtx_module._nvtx_enabled.cache_clear() + nvtx_module._nvtx_module.cache_clear() + assert nvtx_module._nvtx_module() is None + + monkeypatch.setenv(nvtx_module.RL_KERNEL_NVTX, "1") + nvtx_module._nvtx_enabled.cache_clear() + assert nvtx_module._nvtx_enabled() is True + nvtx_module._nvtx_enabled.cache_clear() + nvtx_module._nvtx_module.cache_clear() + + +# --------------------------------------------------------------------------- +# 6. Dashboard contract +# --------------------------------------------------------------------------- + + +def test_dashboard_references_only_schema_metrics(): + dashboard = json.loads(DASHBOARD_PATH.read_text()) + + exprs = [ + target["expr"] + for panel in dashboard["panels"] + for target in panel.get("targets", []) + if "expr" in target + ] + assert exprs, "dashboard defines no queries" + + referenced = set() + for expr in exprs: + referenced.update(re.findall(r"rl_kernel_[a-z0-9_]+", expr)) + assert referenced, "dashboard queries reference no rl_kernel_ metrics" + + for name in referenced: + base = re.sub(r"_(bucket|sum|count)$", "", name) + assert base in SCHEMA_METRIC_NAMES, f"dashboard references unknown metric {name}" + + +def test_dashboard_has_no_hardcoded_datasource_uid(): + dashboard = json.loads(DASHBOARD_PATH.read_text()) + text = DASHBOARD_PATH.read_text() + assert "${DS_PROMETHEUS}" in text + for panel in dashboard["panels"]: + uid = panel.get("datasource", {}).get("uid") + assert uid == "${DS_PROMETHEUS}" + + +# --------------------------------------------------------------------------- +# 7. Isolation between registries +# --------------------------------------------------------------------------- + + +def test_metrics_registries_are_isolated(): + pytest.importorskip("prometheus_client") + first = MetricsRegistry() + second = MetricsRegistry() + + first.set_training_loss(1.0) + second.set_training_loss(2.0) + + assert first.collector_registry is not second.collector_registry + assert _sample_value(first, "rl_kernel_training_loss") == 1.0 + assert _sample_value(second, "rl_kernel_training_loss") == 2.0