prometheus_metrics.rs

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Python bindings for Prometheus metrics callbacks.
//!
//! This module provides minimal bindings for registering callbacks to integrate
//! external Prometheus metrics (e.g., from vLLM, SGLang, TensorRT-LLM) into
//! Dynamo's metrics endpoint.

use pyo3::prelude::*;
use std::sync::Arc;

use crate::rs;

/// RuntimeMetrics provides utilities for registering metrics callbacks.
/// Exposed as endpoint.metrics in Python.
///
/// Note: Metric creation methods have been removed from the public API.
/// This class only provides callback registration for integrating external metrics.
#[pyclass]
#[derive(Clone)]
pub struct RuntimeMetrics {
    hierarchy: Arc<dyn rs::metrics::MetricsHierarchy>,
}

impl RuntimeMetrics {
    /// Create from Endpoint
    pub fn from_endpoint(endpoint: dynamo_runtime::component::Endpoint) -> Self {
        Self {
            hierarchy: Arc::new(endpoint),
        }
    }
}

#[pymethods]
impl RuntimeMetrics {
    /// Register a Python callback that returns Prometheus exposition text
    /// The returned text will be appended to the /metrics endpoint output
    /// The callback should return a string in Prometheus text exposition format
    fn register_prometheus_expfmt_callback(&self, callback: PyObject, _py: Python) -> PyResult<()> {
        // Create the callback once (Arc allows sharing across registries)
        let callback_arc = Arc::new(move || {
            // Execute the Python callback in the Python event loop
            Python::with_gil(|py| {
                match callback.call0(py) {
                    Ok(result) => {
                        // Try to extract a string from the result
                        match result.extract::<String>(py) {
                            Ok(text) => Ok(text),
                            Err(e) => {
                                tracing::error!(
                                    "Metrics exposition text callback must return a string: {}",
                                    e
                                );
                                Ok(String::new())
                            }
                        }
                    }
                    Err(e) => {
                        tracing::error!("Metrics exposition text callback failed: {}", e);
                        Ok(String::new())
                    }
                }
            })
        });

        // Register the callback at this hierarchy level only.
        // Do NOT register on parent hierarchies - combined scrapes automatically
        // traverse child registries and include their callbacks.
        self.hierarchy
            .get_metrics_registry()
            .add_expfmt_callback(callback_arc.clone());

        Ok(())
    }
}

pub fn add_to_module(_m: &Bound<'_, PyModule>) -> PyResult<()> {
    // No metric type classes to add - only RuntimeMetrics is exposed
    Ok(())
}