From 4346cbab04090e340b17cc5df9a7b6fddb44a68c Mon Sep 17 00:00:00 2001 From: Billy1900 Date: Sun, 26 Jul 2026 23:10:18 -0400 Subject: [PATCH 1/2] feat(observability): NVTX kernel tracing + Prometheus metrics endpoint Implements #72: NVTX ranges around every csrc/ops.cpp operator for per-op nsys visibility, plus an opt-in Prometheus /metrics endpoint tracking kernel throughput, backend-fallback rate, and paged-KV-cache fragmentation. Includes a sample Grafana dashboard and docs for both the nsys profiling workflow and the metrics/dashboards setup. --- csrc/ops.cpp | 112 +++-- csrc/utils/nvtx_utils.h | 62 +++ docs/.nav.yml | 2 + .../getting_started/metrics-and-dashboards.md | 97 +++++ docs/getting_started/nsys-profiling.md | 92 +++++ examples/grafana/rl_kernel_dashboard.json | 129 ++++++ pyproject.toml | 1 + rl_engine/executors/paged_kv_baseline.py | 6 + rl_engine/executors/rollout.py | 4 + rl_engine/kernels/registry.py | 20 +- rl_engine/observability/__init__.py | 2 + rl_engine/observability/metrics.py | 198 +++++++++ setup.py | 7 +- tests/test_observability_metrics.py | 386 ++++++++++++++++++ 14 files changed, 1089 insertions(+), 29 deletions(-) create mode 100644 csrc/utils/nvtx_utils.h create mode 100644 docs/getting_started/metrics-and-dashboards.md create mode 100644 docs/getting_started/nsys-profiling.md create mode 100644 examples/grafana/rl_kernel_dashboard.json create mode 100644 rl_engine/observability/__init__.py create mode 100644 rl_engine/observability/metrics.py create mode 100644 tests/test_observability_metrics.py diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 61ba4a3b..723e92f5 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); @@ -93,6 +95,7 @@ at::Tensor prefix_shared_attention( const at::Tensor& K, const at::Tensor& V) { + RL_KERNEL_NVTX_RANGE("rl_kernel::prefix_shared_attention"); TORCH_CHECK(Q.dim() == 4, "Q must be [bs, G, len_q, DIM]"); TORCH_CHECK(K.dim() == 3, "K must be [bs, len_kv, DIM]"); TORCH_CHECK(V.dim() == 3, "V must be [bs, len_kv, DIM]"); @@ -125,49 +128,106 @@ 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", ::rl_kernel::traced("rl_kernel::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", + ::rl_kernel::traced("rl_kernel::fused_logp_sm90", &fused_logp_sm90_forward), + "TMA-accelerated Online Softmax Fused LogP"); + m.def("fused_linear_logp_sm90", + ::rl_kernel::traced("rl_kernel::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", + ::rl_kernel::traced("rl_kernel::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", + ::rl_kernel::traced("rl_kernel::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", + ::rl_kernel::traced("rl_kernel::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", + ::rl_kernel::traced("rl_kernel::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", + ::rl_kernel::traced("rl_kernel::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", + ::rl_kernel::traced("rl_kernel::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_", + ::rl_kernel::traced("rl_kernel::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_, + ::rl_kernel::traced("rl_kernel::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", + ::rl_kernel::traced("rl_kernel::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", + ::rl_kernel::traced("rl_kernel::fused_logp_forward_out", &fused_logp_forward_out), + "Fused logp out"); + m.def("fused_logp_forward_fp32", + ::rl_kernel::traced("rl_kernel::fused_logp_forward_fp32", &fused_logp_forward_fp32), + "Fused logp fp32"); + m.def("fused_logp_forward_indexed_out", + ::rl_kernel::traced("rl_kernel::fused_logp_forward_indexed_out", + &fused_logp_forward_indexed_out), + "Fused logp indexed out"); + m.def("fused_logp_forward_indexed_fp32", + ::rl_kernel::traced("rl_kernel::fused_logp_forward_indexed_fp32", + &fused_logp_forward_indexed_fp32), + "Fused logp indexed fp32"); + m.def("fused_logp_forward_online_out", + ::rl_kernel::traced("rl_kernel::fused_logp_forward_online_out", + &fused_logp_forward_online_out), + "Fused logp online out"); + m.def("fused_logp_forward_online_fp32", + ::rl_kernel::traced("rl_kernel::fused_logp_forward_online_fp32", + &fused_logp_forward_online_fp32), + "Fused logp online fp32"); + m.def("fused_logp_forward_online_indexed_out", + ::rl_kernel::traced("rl_kernel::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", + ::rl_kernel::traced("rl_kernel::fused_logp_forward_online_indexed_fp32", + &fused_logp_forward_online_indexed_fp32), + "Fused logp online indexed fp32"); + m.def("deterministic_logp", + ::rl_kernel::traced("rl_kernel::deterministic_logp", &deterministic_logp_forward), + "Batch-invariant deterministic logp"); + m.def("deterministic_logp_forward_out", + ::rl_kernel::traced("rl_kernel::deterministic_logp_forward_out", + &deterministic_logp_forward_out), + "Batch-invariant deterministic logp out"); + m.def("deterministic_logp_forward_fp32", + ::rl_kernel::traced("rl_kernel::deterministic_logp_forward_fp32", + &deterministic_logp_forward_fp32), + "Batch-invariant deterministic logp fp32"); + m.def("deterministic_logp_forward_indexed_out", + ::rl_kernel::traced("rl_kernel::deterministic_logp_forward_indexed_out", + &deterministic_logp_forward_indexed_out), + "Batch-invariant deterministic logp indexed out"); + m.def("deterministic_logp_forward_indexed_fp32", + ::rl_kernel::traced("rl_kernel::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", &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..dd314a04 --- /dev/null +++ b/csrc/utils/nvtx_utils.h @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors + +#pragma once + +// NVTX ranges are only meaningful -- and only safe to include -- on CUDA +// builds. csrc/ops.cpp also compiles under the ROCm/HIP build (its +// unconditional `fused_logp` binding has no #if guard), so this header must +// degrade to a true no-op there rather than failing to find . +// ROCm/roctx tracing is explicit future work, not in scope here. +#if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) || defined(KERNEL_ALIGN_WITH_SM90) + +#include + +#include + +namespace rl_kernel { + +// RAII scoped NVTX range. Uses the classic API, which links +// against libnvToolsExt (see the `-lnvToolsExt` link flag added in setup.py) +// rather than nvtx3's dlopen-based injection layer. Calls are a cheap no-op +// when no profiler (nsys/ncu) is attached to the process. +class NvtxRange { + public: + explicit NvtxRange(const char* name) { nvtxRangePushA(name); } + ~NvtxRange() { nvtxRangePop(); } + NvtxRange(const NvtxRange&) = delete; + NvtxRange& operator=(const NvtxRange&) = delete; +}; + +// Wraps a free-function pointer so pybind11 can bind the wrapper in place of +// the raw pointer; each call is bracketed by an NVTX range named `name`, so +// nsys shows one labeled block per RL-Kernel op regardless of how many CUDA +// kernels the op launches internally. +template +auto traced(const char* name, Ret (*fn)(Args...)) { + return [name, fn](Args... args) -> Ret { + NvtxRange range(name); + return fn(std::forward(args)...); + }; +} + +} // namespace rl_kernel + +#define RL_KERNEL_NVTX_RANGE(name) ::rl_kernel::NvtxRange _rl_kernel_nvtx_range(name) + +#else // Not a CUDA build (e.g. ROCm-only): compile out entirely. + +namespace rl_kernel { + +template +auto traced(const char* /*name*/, Ret (*fn)(Args...)) { + return fn; +} + +} // namespace rl_kernel + +#define RL_KERNEL_NVTX_RANGE(name) \ + do { \ + } while (0) + +#endif diff --git a/docs/.nav.yml b/docs/.nav.yml index cda8f872..0a367729 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -12,6 +12,8 @@ nav: - getting_started/installation.md - getting_started/faq.md - Hardware Profiling Guide: getting_started/hardware-profiling.md + - NVTX & Nsight Profiling Guide: getting_started/nsys-profiling.md + - Metrics & Dashboards Guide: getting_started/metrics-and-dashboards.md - Operators: - operators/README.md - operators/activation.md diff --git a/docs/getting_started/metrics-and-dashboards.md b/docs/getting_started/metrics-and-dashboards.md new file mode 100644 index 00000000..e3b13ac0 --- /dev/null +++ b/docs/getting_started/metrics-and-dashboards.md @@ -0,0 +1,97 @@ +# Metrics & Dashboards Guide + +This guide explains how to expose RL-Kernel's live Prometheus `/metrics` endpoint and load the +sample Grafana dashboard, for cluster-level monitoring of kernel throughput, backend-fallback +rate, and KV-cache fragmentation across a training/rollout deployment. + +For per-op kernel-launch tracing inside a single process (an `nsys` timeline), see the +[NVTX & Nsight Profiling Guide](nsys-profiling.md) instead — that is a micro-level, offline +trace; this page covers the macro-level, always-on metrics surface. + +## 1. Install + +Prometheus support is an optional dependency: + +```bash +pip install -e .[observability] +``` + +Without it, every metrics function in `rl_engine.observability.metrics` degrades to a no-op and +logs a one-time warning — no other RL-Kernel functionality is affected. + +## 2. Environment Variables + +| Variable | Default | Purpose | +| --- | --- | --- | +| `RL_KERNEL_ENABLE_OP_METRICS` | off | Opt-in: wrap `KernelRegistry.get_op(...)` results to record per-op call count and latency. Off by default because several tests assert on the concrete op class returned by `get_op(...)`. | +| `RL_KERNEL_ENABLE_METRICS_SERVER` | off | Opt-in: auto-start the `/metrics` HTTP endpoint from `RolloutExecutor` on kernel init. | +| `RL_KERNEL_METRICS_PORT` | `9400` | Base port for the `/metrics` endpoint. The actual bind port is `RL_KERNEL_METRICS_PORT + RANK` (falls back to `LOCAL_RANK`, then `0`), so multiple ranks on one node do not collide. | + +Backend-fallback and KV-cache-fragmentation recording require no opt-in beyond having +`prometheus_client` installed — they never change any function's return type, so they are +always active once the dependency is present. + +## 3. Start a Worker and Scrape It + +```bash +RL_KERNEL_ENABLE_METRICS_SERVER=1 RL_KERNEL_ENABLE_OP_METRICS=1 \ + python examples/grpo_single_gpu.py --device cuda --steps 2 \ + --num-prompts 1 --samples-per-prompt 2 --prompt-len 2 --completion-len 3 \ + --vocab-size 16 --hidden-dim 8 +``` + +In another shell: + +```bash +curl http://localhost:9400/metrics +``` + +Confirm the response contains: + +- `rlkernel_op_calls_total` +- `rlkernel_op_latency_seconds_bucket` +- `rlkernel_op_fallbacks_total` +- `rlkernel_kv_cache_fragmentation_ratio` + +You can also start the server directly from Python without any environment variable, for +notebooks or ad hoc scripts: + +```python +from rl_engine.observability.metrics import start_metrics_server + +start_metrics_server(port=9400) +``` + +## 4. Point Prometheus at It + +```yaml +scrape_configs: + - job_name: rl-kernel + static_configs: + - targets: ["localhost:9400"] +``` + +For a multi-rank node, add one target per rank's resolved port +(`RL_KERNEL_METRICS_PORT + rank`). + +## 5. Load the Sample Dashboard + +Import `examples/grafana/rl_kernel_dashboard.json` into Grafana (**Dashboards → New → Import**), +and select your Prometheus datasource when prompted. It ships five panels: + +- Scrape Target Up +- KV-Cache Fragmentation +- Op Throughput (calls/sec) +- Op Fallback Rate +- Op Latency p50 / p95 / p99 + +## Reporting Guidance + +When sharing a dashboard screenshot or a metrics snapshot, include: + +- The RL-Kernel commit and the exact command used to start the worker. +- Whether `RL_KERNEL_ENABLE_OP_METRICS` was set (call-count/latency panels are empty otherwise). +- The number of ranks/workers scraped and their resolved ports. + +Keep committed docs focused on process and configuration. Point-in-time metrics snapshots and +dashboard screenshots should stay outside the repository. diff --git a/docs/getting_started/nsys-profiling.md b/docs/getting_started/nsys-profiling.md new file mode 100644 index 00000000..68d1cdf9 --- /dev/null +++ b/docs/getting_started/nsys-profiling.md @@ -0,0 +1,92 @@ +# NVTX & Nsight Profiling Guide + +This guide explains how to see RL-Kernel's compiled C++/CUDA operators as distinct, named +blocks in an NVIDIA Nsight Systems (`nsys`) timeline. It is a profiling checklist, not a static +report — regenerate the trace on your target machine when you need fresh data, and avoid +committing generated `.nsys-rep` files to the repository. + +## Scope + +Every operator bound in `csrc/ops.cpp` (`fused_logp`, `deterministic_logp`, +`prefix_shared_attention`, the SM90 TMA/WGMMA linear-logp family, and their `*_out`/`*_fp32`/ +`*_indexed`/`*_online` variants) is wrapped in an NVTX range named `rl_kernel::` via +`csrc/utils/nvtx_utils.h`. This lets `nsys` draw one labeled block per RL-Kernel op call, +grouped above the raw CUDA kernel launches that op triggers internally. + +This is a micro-level, kernel-boundary trace. For macro-level throughput, fallback-rate, and +cache-fragmentation metrics across a training/rollout cluster, see the +[Metrics & Dashboards Guide](metrics-and-dashboards.md) instead. + +## 1. Build the Extension + +NVTX ranges link against `libnvToolsExt` (via `-lnvToolsExt`, already wired into `setup.py`'s +CUDA build), which ships with every CUDA toolkit -- no extra install step is required beyond a +normal build: + +```bash +MAX_JOBS=2 python setup.py build_ext --inplace +``` + +## 2. Record a Trace + +Wrap any RL-Kernel entry point with `nsys profile`. The GRPO single-GPU example is a convenient +smoke workload: + +```bash +CUDA_VISIBLE_DEVICES= nsys profile -o rlkernel_report \ + python examples/grpo_single_gpu.py \ + --device cuda \ + --require-fused-logp \ + --steps 2 \ + --num-prompts 1 \ + --samples-per-prompt 2 \ + --prompt-len 2 \ + --completion-len 3 \ + --vocab-size 16 \ + --hidden-dim 8 +``` + +This produces `rlkernel_report.nsys-rep` in the current directory. + +## 3. Inspect the Timeline + +Open the report in the Nsight Systems UI (`nsys-ui rlkernel_report.nsys-rep`), or summarize it +on the command line: + +```bash +nsys stats --report nvtx_sum rlkernel_report.nsys-rep +``` + +Confirm that named ranges appear as distinct, non-overlapping blocks on the **NVTX** row, +directly above the correlated CUDA HW kernel rows. Expect names matching the op(s) actually +exercised by the workload, for example: + +- `rl_kernel::fused_logp` +- `rl_kernel::deterministic_logp` +- `rl_kernel::prefix_shared_attention` +- `rl_kernel::fused_logp_sm90` / `rl_kernel::fused_linear_logp_sm90` (Hopper-only, requires a + build with `KERNEL_ALIGN_FORCE_SM90=1`) + +Each block's duration should track the wall-clock time of that op's C++ entry point, including +every CUDA kernel it launches internally — this is what distinguishes an "op-level" NVTX block +from the finer-grained individual kernel-launch rows `nsys` already draws on its own. + +## 4. Manual Verification Only + +Confirming that labeled blocks render correctly in Nsight Systems is a **manual verification +step**. There is no GPU + `nsys` available in standard CI, and no meaningful unit test can +assert "nsys drew a labeled block" — so this guide, not an automated test, is the source of +truth for validating NVTX coverage. Treat a successful walkthrough of this checklist as the +acceptance bar, not a green CI run. + +## Reporting Guidance + +When sharing a trace or a screenshot from Nsight Systems, include: + +- GPU model and compute capability. +- Driver, CUDA runtime, and `nsys` version (`nsys --version`). +- The exact `CUDA_VISIBLE_DEVICES` mapping and command used to record the trace. +- Which NVTX-named ranges were visible and whether their nesting/order matched expectations. + +Keep committed docs focused on the process and commands. Generated `.nsys-rep` files and +point-in-time screenshots should stay outside the repository. diff --git a/examples/grafana/rl_kernel_dashboard.json b/examples/grafana/rl_kernel_dashboard.json new file mode 100644 index 00000000..08c3d2e0 --- /dev/null +++ b/examples/grafana/rl_kernel_dashboard.json @@ -0,0 +1,129 @@ +{ + "title": "RL-Kernel Observability", + "uid": "rl-kernel-observability", + "schemaVersion": 39, + "version": 1, + "editable": true, + "timezone": "browser", + "refresh": "30s", + "time": { "from": "now-1h", "to": "now" }, + "tags": ["rl-kernel", "observability"], + "panels": [ + { + "id": 1, + "title": "Scrape Target Up", + "type": "stat", + "gridPos": { "h": 4, "w": 6, "x": 0, "y": 0 }, + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "targets": [ + { + "expr": "up{job=\"rl-kernel\"}", + "legendFormat": "{{instance}}", + "refId": "A" + } + ], + "fieldConfig": { + "defaults": { + "mappings": [ + { "type": "value", "options": { "0": { "text": "DOWN", "color": "red" } } }, + { "type": "value", "options": { "1": { "text": "UP", "color": "green" } } } + ] + } + } + }, + { + "id": 2, + "title": "KV-Cache Fragmentation", + "type": "stat", + "gridPos": { "h": 4, "w": 6, "x": 6, "y": 0 }, + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "targets": [ + { + "expr": "rlkernel_kv_cache_fragmentation_ratio", + "legendFormat": "{{baseline_kind}}", + "refId": "A" + } + ], + "fieldConfig": { + "defaults": { + "unit": "percentunit", + "min": 0, + "max": 1, + "thresholds": { + "mode": "absolute", + "steps": [ + { "color": "green", "value": null }, + { "color": "yellow", "value": 0.3 }, + { "color": "red", "value": 0.6 } + ] + } + } + } + }, + { + "id": 3, + "title": "Op Throughput (calls/sec)", + "type": "timeseries", + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 4 }, + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "targets": [ + { + "expr": "sum(rate(rlkernel_op_calls_total[5m])) by (op_type, backend)", + "legendFormat": "{{op_type}} ({{backend}})", + "refId": "A" + } + ], + "fieldConfig": { "defaults": { "unit": "cps" } } + }, + { + "id": 4, + "title": "Op Fallback Rate", + "type": "timeseries", + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 4 }, + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "targets": [ + { + "expr": "sum(rate(rlkernel_op_fallbacks_total[5m])) by (op_type, failed_backend)", + "legendFormat": "{{op_type}} <- {{failed_backend}}", + "refId": "A" + } + ], + "fieldConfig": { "defaults": { "unit": "cps" } } + }, + { + "id": 5, + "title": "Op Latency p50 / p95 / p99", + "type": "timeseries", + "gridPos": { "h": 8, "w": 24, "x": 0, "y": 12 }, + "datasource": { "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, + "targets": [ + { + "expr": "histogram_quantile(0.50, sum(rate(rlkernel_op_latency_seconds_bucket[5m])) by (le, op_type, backend))", + "legendFormat": "p50 {{op_type}} ({{backend}})", + "refId": "A" + }, + { + "expr": "histogram_quantile(0.95, sum(rate(rlkernel_op_latency_seconds_bucket[5m])) by (le, op_type, backend))", + "legendFormat": "p95 {{op_type}} ({{backend}})", + "refId": "B" + }, + { + "expr": "histogram_quantile(0.99, sum(rate(rlkernel_op_latency_seconds_bucket[5m])) by (le, op_type, backend))", + "legendFormat": "p99 {{op_type}} ({{backend}})", + "refId": "C" + } + ], + "fieldConfig": { "defaults": { "unit": "s" } } + } + ], + "templating": { + "list": [ + { + "name": "DS_PROMETHEUS", + "type": "datasource", + "query": "prometheus", + "label": "Prometheus datasource" + } + ] + } +} diff --git a/pyproject.toml b/pyproject.toml index b0b80a1a..6bfc8852 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ dependencies = [ cuda = ["flashinfer-python>=0.1.6", "nvidia-ml-py"] rocm = ["aiter"] vllm = ["vllm>=0.6.0"] +observability = ["prometheus-client>=0.19"] dev = ["pytest", "black", "isort", "ruff", "mypy", "pre-commit"] [tool.setuptools.packages.find] diff --git a/rl_engine/executors/paged_kv_baseline.py b/rl_engine/executors/paged_kv_baseline.py index 0c465720..70cc68e6 100644 --- a/rl_engine/executors/paged_kv_baseline.py +++ b/rl_engine/executors/paged_kv_baseline.py @@ -23,6 +23,7 @@ score_rewards, summarize_tensor_tree, ) +from rl_engine.observability.metrics import record_kv_cache_fragmentation @dataclass(frozen=True) @@ -306,6 +307,11 @@ def collect_paged_kv_metrics( device = input_ids.device metrics["peak_allocated_mb"] = torch.cuda.max_memory_allocated(device) / 1_048_576.0 metrics["peak_reserved_mb"] = torch.cuda.max_memory_reserved(device) / 1_048_576.0 + record_kv_cache_fragmentation( + reservation.required_blocks, + reservation.reserved_blocks, + baseline_kind=str(metrics["baseline_kind"]), + ) return metrics diff --git a/rl_engine/executors/rollout.py b/rl_engine/executors/rollout.py index 1832ae14..63e88ae0 100644 --- a/rl_engine/executors/rollout.py +++ b/rl_engine/executors/rollout.py @@ -16,6 +16,7 @@ ) 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_server_enabled, start_metrics_server from rl_engine.utils.logger import logger @@ -137,6 +138,9 @@ def _prepare_kernels(self): f" Attn: {type(self.attn_op).__name__}" ) + if metrics_server_enabled(): + start_metrics_server() + def _prepare_sampler(self) -> VLLMSharedPrefixSampler: """ Lazily construct the vLLM-backed sampler. diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 2b9b0c30..e78b63db 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 InstrumentedOp, op_metrics_enabled, record_fallback from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -311,7 +312,7 @@ def get_op(self, op_type: str) -> Any: for backend in candidates: if backend.name in self._instance_cache: - return self._instance_cache[backend.name] + return self._instrument(self._instance_cache[backend.name], op_type, backend.name) if backend.name in self._failed_backends: continue @@ -321,15 +322,30 @@ def get_op(self, op_type: str) -> Any: try: op_instance = op_class() self._instance_cache[backend.name] = op_instance - return op_instance + return self._instrument(op_instance, op_type, backend.name) except Exception as e: logger.error(f"Failed to instantiate {backend.name}: {e}") self._failed_backends.add(backend.name) + record_fallback(op_type, backend.name) else: self._failed_backends.add(backend.name) + record_fallback(op_type, backend.name) raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + @staticmethod + def _instrument(op_instance: Any, op_type: str, backend_name: str) -> Any: + """Wrap a resolved op instance for latency/throughput metrics, opt-in only. + + `_instance_cache` keeps storing the raw, unwrapped instance; only the value returned + from `get_op(...)` is wrapped, freshly, per call -- so the same cached instance is + always labeled with the `op_type` of the *current* call, not whichever op_type first + populated the cache. + """ + if not op_metrics_enabled(): + return op_instance + return InstrumentedOp(op_instance, op_type=op_type, backend=backend_name) + 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..86cf4c9d --- /dev/null +++ b/rl_engine/observability/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/rl_engine/observability/metrics.py b/rl_engine/observability/metrics.py new file mode 100644 index 00000000..89fb35fd --- /dev/null +++ b/rl_engine/observability/metrics.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Prometheus metrics for RL-Kernel: op throughput, backend-fallback rate, and +KV-cache fragmentation, exposed over an opt-in `/metrics` HTTP endpoint. + +`prometheus_client` is an optional dependency (install via `pip install -e .[observability]`). +Every public function in this module degrades to a safe no-op when it is unavailable, mirroring +the `_C`/`_EXT_AVAILABLE` optional-extension pattern in `rl_engine/kernels/ops/base.py`. +""" + +from __future__ import annotations + +import os +import time +from typing import Any, Optional + +from rl_engine.utils.logger import logger + +try: + from prometheus_client import Counter, Gauge, Histogram, start_http_server + + _PROMETHEUS_AVAILABLE = True +except ImportError as e: + logger.warning( + f"prometheus_client unavailable: {e}. Observability metrics will be no-ops. " + "Install with `pip install -e .[observability]` to enable them." + ) + _PROMETHEUS_AVAILABLE = False + Counter = Gauge = Histogram = start_http_server = None # type: ignore[assignment] + +_OP_METRICS_ENV = "RL_KERNEL_ENABLE_OP_METRICS" +_METRICS_SERVER_ENV = "RL_KERNEL_ENABLE_METRICS_SERVER" +_METRICS_PORT_ENV = "RL_KERNEL_METRICS_PORT" +_DEFAULT_METRICS_PORT = 9400 + +_TRUE_VALUES = {"1", "true", "yes", "on"} + + +_FALSE_VALUES = {"0", "false", "no", "off"} + + +def _env_flag(name: str, default: bool) -> bool: + value = os.environ.get(name) + if value is None or value.strip() == "": + return default + normalized = value.strip().lower() + if normalized in _TRUE_VALUES: + return True + if normalized in _FALSE_VALUES: + return False + logger.warn_once( + f"{name}={value!r} is not a recognized boolean flag " + f"(expected one of {sorted(_TRUE_VALUES | _FALSE_VALUES)}); using default={default}." + ) + return default + + +def metrics_enabled() -> bool: + """Whether the always-on recorders (fallback, KV-cache fragmentation) are active. + + These recorders never change any caller's return type or value, so they are gated only on + dependency availability -- free once `prometheus_client` is installed. + """ + return _PROMETHEUS_AVAILABLE + + +def op_metrics_enabled() -> bool: + """Whether `KernelRegistry.get_op(...)` should wrap its return value in `InstrumentedOp`. + + Defaults OFF and requires explicit opt-in via `RL_KERNEL_ENABLE_OP_METRICS=1`, separately + from `metrics_enabled()`, because wrapping changes `get_op(...)`'s return type: several + existing tests assert `isinstance(kernel_registry.get_op(op_type), ConcreteOpClass)` (e.g. + tests/test_embedding.py, tests/test_swiglu.py, tests/test_lm_head.py, tests/test_attention.py, + tests/test_kv_cache_attention.py), and wrapping by default would break all of them. + """ + return _PROMETHEUS_AVAILABLE and _env_flag(_OP_METRICS_ENV, default=False) + + +def metrics_server_enabled() -> bool: + """Whether a worker should auto-start the `/metrics` HTTP endpoint on init.""" + return _env_flag(_METRICS_SERVER_ENV, default=False) + + +if _PROMETHEUS_AVAILABLE: + OP_CALLS_TOTAL = Counter( + "rlkernel_op_calls_total", + "Total kernel operator invocations.", + ["op_type", "backend", "method"], + ) + OP_LATENCY_SECONDS = Histogram( + "rlkernel_op_latency_seconds", + "Kernel operator invocation latency, seconds.", + ["op_type", "backend", "method"], + ) + OP_FALLBACKS_TOTAL = Counter( + "rlkernel_op_fallbacks_total", + "Times a preferred backend failed to load/instantiate and dispatch fell back.", + ["op_type", "failed_backend"], + ) + KV_CACHE_FRAGMENTATION_RATIO = Gauge( + "rlkernel_kv_cache_fragmentation_ratio", + "1 - (required_blocks / reserved_blocks) for the most recent paged-KV reservation.", + ["baseline_kind"], + ) +else: + OP_CALLS_TOTAL = None + OP_LATENCY_SECONDS = None + OP_FALLBACKS_TOTAL = None + KV_CACHE_FRAGMENTATION_RATIO = None + + +def record_fallback(op_type: str, failed_backend: str) -> None: + """Record that `failed_backend` could not serve `op_type` and dispatch moved on.""" + if not metrics_enabled(): + return + OP_FALLBACKS_TOTAL.labels(op_type=op_type, failed_backend=failed_backend).inc() + + +def record_kv_cache_fragmentation( + required_blocks: int, reserved_blocks: int, *, baseline_kind: str +) -> None: + """Record the reserved-but-unused fraction of a paged-KV-cache block reservation.""" + if not metrics_enabled() or reserved_blocks <= 0: + return + fragmentation = 1.0 - (required_blocks / reserved_blocks) + KV_CACHE_FRAGMENTATION_RATIO.labels(baseline_kind=baseline_kind).set(fragmentation) + + +class InstrumentedOp: + """Transparent latency/throughput proxy around a resolved kernel op instance. + + Wraps whatever `KernelRegistry.get_op(...)` resolves. `__call__` and every callable + attribute access (`.forward`, `.forward_fp32`, ...) are timed under the same + `(op_type, backend)` label pair plus a `method` label, so the underlying op classes need no + changes. Non-callable attributes (e.g. `.op_class`) pass through untouched. + + Does not fake `isinstance` against the wrapped type -- callers needing type identity should + read `.op_class` or avoid opting into `RL_KERNEL_ENABLE_OP_METRICS`. + """ + + def __init__(self, wrapped: Any, *, op_type: str, backend: str): + object.__setattr__(self, "_wrapped", wrapped) + object.__setattr__(self, "_op_type", op_type) + object.__setattr__(self, "_backend", backend) + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + return self._timed(self._wrapped, "__call__", args, kwargs) + + def __getattr__(self, name: str) -> Any: + attr = getattr(self._wrapped, name) + if not callable(attr): + return attr + + def _call(*args: Any, __attr: Any = attr, __name: str = name, **kwargs: Any) -> Any: + return self._timed(__attr, __name, args, kwargs) + + return _call + + def _timed(self, fn: Any, method: str, args: tuple, kwargs: dict) -> Any: + start = time.perf_counter() + try: + return fn(*args, **kwargs) + finally: + elapsed = time.perf_counter() - start + OP_CALLS_TOTAL.labels(op_type=self._op_type, backend=self._backend, method=method).inc() + OP_LATENCY_SECONDS.labels( + op_type=self._op_type, backend=self._backend, method=method + ).observe(elapsed) + + +_server_started = False + + +def start_metrics_server(port: Optional[int] = None) -> Optional[int]: + """Idempotent per-process `/metrics` HTTP server start. + + Wraps `prometheus_client.start_http_server` (stdlib `http.server` based, so this adds no + Flask/FastAPI dependency). Returns the bound port, or None if metrics are unavailable or the + server is already running in this process. + """ + global _server_started + if not metrics_enabled() or _server_started: + return None + if port is None: + base_port = int(os.environ.get(_METRICS_PORT_ENV, str(_DEFAULT_METRICS_PORT))) + rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0"))) + port = base_port + rank + try: + start_http_server(port) + except OSError as e: + # Observability must never take down real training/rollout work -- a bound port, + # a sandboxed network namespace, etc. should degrade to "no /metrics", not a crash. + logger.warning(f"Failed to start Prometheus /metrics server on :{port}: {e}") + return None + _server_started = True + logger.info(f"Prometheus /metrics server listening on :{port}") + return port diff --git a/setup.py b/setup.py index 6f94e040..d83a652e 100644 --- a/setup.py +++ b/setup.py @@ -134,7 +134,11 @@ def get_extensions(): nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] - extra_link_args = list(torch_rpath) + # csrc/utils/nvtx_utils.h uses the classic API (nvtxRangePushA/Pop), + # which resolves its symbols at link time rather than via nvtx3's dlopen-based + # injection layer -- link libnvToolsExt explicitly so `rl_kernel::traced(...)` in + # csrc/ops.cpp loads cleanly. + extra_link_args = list(torch_rpath) + ["-lnvToolsExt"] sm90_srcs = [ "csrc/cuda/fused_logp_sm90.cu", @@ -187,6 +191,7 @@ def get_cmdclass(): "cuda": ["flashinfer"], "rocm": ["aiter"], "vllm": ["vllm>=0.6.0"], + "observability": ["prometheus-client>=0.19"], }, python_requires=">=3.10", include_package_data=True, diff --git a/tests/test_observability_metrics.py b/tests/test_observability_metrics.py new file mode 100644 index 00000000..3508d9f8 --- /dev/null +++ b/tests/test_observability_metrics.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU-only tests for rl_engine.observability.metrics. + +NVTX correctness (whether csrc/ops.cpp's ranges render as labeled, distinct blocks in an +`nsys` timeline) is NOT covered here and is not unit-testable in CI without a GPU + nsys. That +is a manual verification step documented in docs/getting_started/nsys-profiling.md, not a +Python-testable concern. It has been manually verified on real H100 hardware, including that +the RAII NVTX range still closes correctly when the wrapped op raises (`nsys stats +--report nvtx_sum` shows matched, non-orphaned ranges for both successful and raising calls). +""" + +from __future__ import annotations + +import socket +import urllib.request + +import pytest +import torch + +import rl_engine.observability.metrics as metrics_module +from rl_engine.executors.paged_kv_baseline import ( + PagedKVScoringConfig, + collect_paged_kv_metrics, + reserve_paged_kv_cache, +) +from rl_engine.executors.stateless_executor import StatelessForwardInputs, TensorTreeSummary +from rl_engine.kernels.registry import KernelRegistry, OpBackend +from rl_engine.observability.metrics import ( + InstrumentedOp, + metrics_enabled, + op_metrics_enabled, + record_fallback, + record_kv_cache_fragmentation, + start_metrics_server, +) + + +def _sample_value(metric, name: str, **labels) -> float | None: + """Read a Prometheus sample by exact name + labels via the public collect() API.""" + for family in metric.collect(): + for sample in family.samples: + if sample.name == name and sample.labels == labels: + return sample.value + return None + + +def test_metrics_noop_without_prometheus(monkeypatch): + monkeypatch.setattr(metrics_module, "_PROMETHEUS_AVAILABLE", False) + + assert metrics_enabled() is False + assert op_metrics_enabled() is False + + # Must not raise even though the underlying metric objects are unusable stand-ins. + record_fallback("some_op", "SOME_BACKEND") + record_kv_cache_fragmentation(1, 2, baseline_kind="unit_test") + assert start_metrics_server() is None + + +def test_op_metrics_enabled_requires_explicit_opt_in(monkeypatch): + pytest.importorskip("prometheus_client") + monkeypatch.delenv("RL_KERNEL_ENABLE_OP_METRICS", raising=False) + assert op_metrics_enabled() is False + + monkeypatch.setenv("RL_KERNEL_ENABLE_OP_METRICS", "1") + assert op_metrics_enabled() is True + + monkeypatch.setenv("RL_KERNEL_ENABLE_OP_METRICS", "0") + assert op_metrics_enabled() is False + + +class _DummyOp: + op_class = "dummy" + + def __call__(self, x: int) -> int: + return x + 1 + + def forward(self, x: int) -> int: + return x * 2 + + +def test_instrumented_op_records_calls_and_latency(): + pytest.importorskip("prometheus_client") + op_type, backend = "unit_test_dummy_op", "UNIT_TEST_DUMMY_BACKEND" + wrapped = InstrumentedOp(_DummyOp(), op_type=op_type, backend=backend) + + # Non-callable attributes pass through untouched. + assert wrapped.op_class == "dummy" + + before_call = ( + _sample_value( + metrics_module.OP_CALLS_TOTAL, + "rlkernel_op_calls_total", + op_type=op_type, + backend=backend, + method="__call__", + ) + or 0.0 + ) + before_forward = ( + _sample_value( + metrics_module.OP_CALLS_TOTAL, + "rlkernel_op_calls_total", + op_type=op_type, + backend=backend, + method="forward", + ) + or 0.0 + ) + + assert wrapped(3) == 4 + assert wrapped.forward(3) == 6 + + after_call = _sample_value( + metrics_module.OP_CALLS_TOTAL, + "rlkernel_op_calls_total", + op_type=op_type, + backend=backend, + method="__call__", + ) + after_forward = _sample_value( + metrics_module.OP_CALLS_TOTAL, + "rlkernel_op_calls_total", + op_type=op_type, + backend=backend, + method="forward", + ) + latency_count = _sample_value( + metrics_module.OP_LATENCY_SECONDS, + "rlkernel_op_latency_seconds_count", + op_type=op_type, + backend=backend, + method="__call__", + ) + + assert after_call == before_call + 1 + assert after_forward == before_forward + 1 + assert latency_count is not None and latency_count >= 1 + + +def test_registry_fallback_increments_counter(monkeypatch): + pytest.importorskip("prometheus_client") + + registry = KernelRegistry() + op_type = "unit_test_fallback_op_type" + failing_backend = OpBackend.PYTORCH_NATIVE + working_backend = OpBackend.PYTORCH_NATIVE_SILU + for platform in ("cpu", "cuda", "rocm"): + registry._priority_map[platform][op_type] = [failing_backend, working_backend] + + original_load_backend = KernelRegistry._load_backend + + def fake_load_backend(self, backend): + if backend is failing_backend: + return None + return original_load_backend(self, backend) + + monkeypatch.setattr(KernelRegistry, "_load_backend", fake_load_backend) + + before = ( + _sample_value( + metrics_module.OP_FALLBACKS_TOTAL, + "rlkernel_op_fallbacks_total", + op_type=op_type, + failed_backend=failing_backend.name, + ) + or 0.0 + ) + + op = registry.get_op(op_type) + + after = _sample_value( + metrics_module.OP_FALLBACKS_TOTAL, + "rlkernel_op_fallbacks_total", + op_type=op_type, + failed_backend=failing_backend.name, + ) + + assert op is not None + assert after == before + 1 + + +def test_registry_get_op_returns_raw_instance_by_default(monkeypatch): + monkeypatch.delenv("RL_KERNEL_ENABLE_OP_METRICS", raising=False) + + registry = KernelRegistry() + registry._priority_map["cpu"]["silu"] = [OpBackend.PYTORCH_NATIVE_SILU] + registry._priority_map["cuda"]["silu"] = [OpBackend.PYTORCH_NATIVE_SILU] + registry._priority_map["rocm"]["silu"] = [OpBackend.PYTORCH_NATIVE_SILU] + + op = registry.get_op("silu") + + assert not isinstance(op, InstrumentedOp) + + +def test_registry_get_op_wraps_when_op_metrics_enabled(monkeypatch): + pytest.importorskip("prometheus_client") + monkeypatch.setenv("RL_KERNEL_ENABLE_OP_METRICS", "1") + + registry = KernelRegistry() + registry._priority_map["cpu"]["silu"] = [OpBackend.PYTORCH_NATIVE_SILU] + registry._priority_map["cuda"]["silu"] = [OpBackend.PYTORCH_NATIVE_SILU] + registry._priority_map["rocm"]["silu"] = [OpBackend.PYTORCH_NATIVE_SILU] + + op = registry.get_op("silu") + + assert isinstance(op, InstrumentedOp) + # Still delegates correctly to the underlying implementation. + x = torch.randn(2, 4) + assert torch.allclose(op(x), torch.nn.functional.silu(x)) + + +def test_kv_cache_fragmentation_recorded(): + pytest.importorskip("prometheus_client") + + inputs = StatelessForwardInputs( + input_ids=torch.tensor([[0, 1, 2, 3, 0], [0, 2, 1, 4, 5]], dtype=torch.long), + attention_mask=torch.tensor( + [[True, True, True, True, True], [True, True, True, False, False]] + ), + completion_mask=torch.tensor( + [[False, False, True, True, False], [False, True, True, False, False]] + ), + ) + config = PagedKVScoringConfig( + num_layers=2, + num_kv_heads=2, + head_dim=4, + block_size=2, + kv_cache_dtype=torch.float32, + kv_cache_blocks=8, + ) + reservation = reserve_paged_kv_cache(inputs, config) + assert reservation.required_blocks == 5 + assert reservation.reserved_blocks == 8 + + collect_paged_kv_metrics( + inputs, + reservation, + config=config, + elapsed_seconds=0.0, + use_cache_passed=True, + cuda_tracking=False, + model_kv_cache_summary=TensorTreeSummary(tensor_count=0, total_bytes=0), + ) + + value = _sample_value( + metrics_module.KV_CACHE_FRAGMENTATION_RATIO, + "rlkernel_kv_cache_fragmentation_ratio", + baseline_kind="generation_engine_paged_kv_reservation", + ) + assert value == pytest.approx(1.0 - 5 / 8) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def test_metrics_endpoint_serves_expected_names(monkeypatch): + pytest.importorskip("prometheus_client") + monkeypatch.setattr(metrics_module, "_server_started", False) + + port = _free_port() + bound_port = start_metrics_server(port=port) + assert bound_port == port + + with urllib.request.urlopen(f"http://127.0.0.1:{port}/metrics", timeout=5) as response: + assert response.status == 200 + body = response.read().decode("utf-8") + + for name in ( + "rlkernel_op_calls_total", + "rlkernel_op_latency_seconds", + "rlkernel_op_fallbacks_total", + "rlkernel_kv_cache_fragmentation_ratio", + ): + assert name in body + + +def test_start_metrics_server_returns_none_on_port_conflict(monkeypatch): + pytest.importorskip("prometheus_client") + monkeypatch.setattr(metrics_module, "_server_started", False) + + occupied_port = _free_port() + blocker = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + blocker.bind(("127.0.0.1", occupied_port)) + blocker.listen(1) + try: + # Must not raise OSError -- observability failures must never crash the caller. + result = start_metrics_server(port=occupied_port) + assert result is None + # A failed bind must not mark the server as started, so a later call (e.g. on a + # different port) can still succeed. + assert metrics_module._server_started is False + finally: + blocker.close() + + +@pytest.mark.parametrize("raw_value", ["1yes", "enabled", "TRUE_ISH", "2"]) +def test_env_flag_warns_and_falls_back_to_default_on_unrecognized_value(monkeypatch, raw_value): + monkeypatch.setenv("RL_KERNEL_UNIT_TEST_BOOL_FLAG", raw_value) + + warnings = [] + monkeypatch.setattr( + metrics_module.logger, "warn_once", lambda msg, *a: warnings.append(msg % a if a else msg) + ) + + assert metrics_module._env_flag("RL_KERNEL_UNIT_TEST_BOOL_FLAG", default=False) is False + assert metrics_module._env_flag("RL_KERNEL_UNIT_TEST_BOOL_FLAG", default=True) is True + assert len(warnings) == 2 + assert all(raw_value in w for w in warnings) + + +def test_registry_all_backends_fail_raises_and_records_every_fallback(monkeypatch): + pytest.importorskip("prometheus_client") + + registry = KernelRegistry() + op_type = "unit_test_all_fail_op_type" + candidates = [OpBackend.PYTORCH_NATIVE, OpBackend.PYTORCH_NATIVE_SILU] + for platform in ("cpu", "cuda", "rocm"): + registry._priority_map[platform][op_type] = list(candidates) + + monkeypatch.setattr(KernelRegistry, "_load_backend", lambda self, backend: None) + + before = { + backend.name: _sample_value( + metrics_module.OP_FALLBACKS_TOTAL, + "rlkernel_op_fallbacks_total", + op_type=op_type, + failed_backend=backend.name, + ) + or 0.0 + for backend in candidates + } + + with pytest.raises(RuntimeError, match=op_type): + registry.get_op(op_type) + + for backend in candidates: + after = _sample_value( + metrics_module.OP_FALLBACKS_TOTAL, + "rlkernel_op_fallbacks_total", + op_type=op_type, + failed_backend=backend.name, + ) + assert after == before[backend.name] + 1 + + +def test_instrumented_op_still_raises_and_records_metrics_on_failure(): + pytest.importorskip("prometheus_client") + + class _FailingOp: + def __call__(self, *args, **kwargs): + raise ValueError("boom") + + op_type, backend = "unit_test_failing_op", "UNIT_TEST_FAILING_BACKEND" + wrapped = InstrumentedOp(_FailingOp(), op_type=op_type, backend=backend) + + before = ( + _sample_value( + metrics_module.OP_CALLS_TOTAL, + "rlkernel_op_calls_total", + op_type=op_type, + backend=backend, + method="__call__", + ) + or 0.0 + ) + + with pytest.raises(ValueError, match="boom"): + wrapped() + + after = _sample_value( + metrics_module.OP_CALLS_TOTAL, + "rlkernel_op_calls_total", + op_type=op_type, + backend=backend, + method="__call__", + ) + # The exception must propagate untouched (no swallowing), but the `finally` block should + # still have recorded the call -- a failed invocation still consumed real time and is + # still useful signal for throughput/error-rate dashboards. + assert after == before + 1 From 1ccbf30e678449eecdec55b142732d2bc76e6c40 Mon Sep 17 00:00:00 2001 From: Henry Date: Mon, 27 Jul 2026 11:14:53 +0800 Subject: [PATCH 2/2] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- rl_engine/observability/metrics.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/rl_engine/observability/metrics.py b/rl_engine/observability/metrics.py index 89fb35fd..b043c13a 100644 --- a/rl_engine/observability/metrics.py +++ b/rl_engine/observability/metrics.py @@ -183,8 +183,23 @@ def start_metrics_server(port: Optional[int] = None) -> Optional[int]: if not metrics_enabled() or _server_started: return None if port is None: - base_port = int(os.environ.get(_METRICS_PORT_ENV, str(_DEFAULT_METRICS_PORT))) - rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0"))) + try: + base_port = int(os.environ.get(_METRICS_PORT_ENV, str(_DEFAULT_METRICS_PORT))) + except (TypeError, ValueError): + logger.warn_once( + "%s=%r is not a valid integer port; using default=%d.", + _METRICS_PORT_ENV, + os.environ.get(_METRICS_PORT_ENV), + _DEFAULT_METRICS_PORT, + ) + base_port = _DEFAULT_METRICS_PORT + try: + rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0"))) + except (TypeError, ValueError): + logger.warn_once( + "RANK/LOCAL_RANK env var is not a valid integer; using rank=0.", + ) + rank = 0 port = base_port + rank try: start_http_server(port)