fix: callback registration, fix metric name access, ensure ordered vec, etc... (#3541)

Signed-off-by: Keiven Chang <[email protected]>
Co-authored-by: Keiven Chang <[email protected]>

Author: keivenchang

lib/bindings/python/examples/metrics/README.md

@@ -216,7 +216,7 @@ Method 2 supports all standard Prometheus metric types:
 - **CounterVec**: `CounterVec` (float with labels), `IntCounterVec` (integer with labels)
 - **Histograms**: `Histogram`
 
-All metrics are imported from `dynamo._prometheus_metrics`.
+All metrics are imported from `dynamo.prometheus_metrics`.
 
 #### Adding/Changing Metrics in Method 2
 

lib/bindings/python/examples/metrics/server_with_callback.py

@@ -10,6 +10,12 @@
 2. Register them with an endpoint
 3. Update their values using type-safe methods (set for gauges, inc for counters)
 4. The metrics are automatically served via the /metrics endpoint
+
+Usage:
+    DYN_SYSTEM_ENABLED=true DYN_SYSTEM_PORT=8081 ./server_with_callback.py
+
+    # In another terminal, query the metrics:
+    curl http://localhost:8081/metrics
 """
 
 import asyncio
@@ -18,8 +24,8 @@
 
 # Note that these imports are for type hints only. They cannot be instantiated directly.
 # You can instantiate them using the endpoint.metrics.create_*() methods.
-from dynamo._prometheus_metrics import Gauge, IntCounter, IntGauge, IntGaugeVec
-from dynamo.runtime import DistributedRuntime, dynamo_worker
+from dynamo.prometheus_metrics import Gauge, IntCounter, IntGauge, IntGaugeVec
+from dynamo.runtime import Component, DistributedRuntime, Endpoint, dynamo_worker
 
 
 @dynamo_worker()
@@ -29,10 +35,10 @@ async def worker(runtime: DistributedRuntime) -> None:
 
 async def init(runtime: DistributedRuntime):
     # Create component and endpoint
-    component = runtime.namespace("ns556").component("cp556")
+    component: Component = runtime.namespace("ns556").component("cp556")
     await component.create_service()
 
-    endpoint = component.endpoint("ep556")
+    endpoint: Endpoint = component.endpoint("ep556")
 
     # Step 1: Create metrics using the endpoint's metrics property
     print("[python] Creating metrics...")
@@ -59,11 +65,10 @@ async def init(runtime: DistributedRuntime):
         [("update_method", "callback")],
     )
 
-    print(f"[python] Created IntGauge: {request_total_slots.name}")
-    print(f"[python] Created Gauge: {gpu_cache_usage_perc.name}")
-    print(f"[python] Created IntGaugeVec: {worker_active_requests.name}")
-    print(f"[python] Created IntCounter with constant labels: {update_count.name}")
-    print(f"[python]   Const labels: {update_count.const_labels}")
+    print(f"[python] Created IntGauge: {request_total_slots.name()}")
+    print(f"[python] Created Gauge: {gpu_cache_usage_perc.name()}")
+    print(f"[python] Created IntGaugeVec: {worker_active_requests.name()}")
+    print(f"[python] Created IntCounter: {update_count.name()}")
     print("[python] Metrics automatically registered with endpoint!")
 
     # Step 2: Register a callback to update metrics on-demand

lib/bindings/python/examples/metrics/server_with_loop.py

@@ -10,6 +10,12 @@
 2. A background thread continuously updates metrics in a loop
 3. No callback is used - metrics are updated directly by the thread
 4. The metrics are automatically served via the /metrics endpoint
+
+Usage:
+    DYN_SYSTEM_ENABLED=true DYN_SYSTEM_PORT=8081 ./server_with_loop.py
+
+    # In another terminal, query the metrics:
+    curl http://localhost:8081/metrics
 """
 
 import asyncio
@@ -18,8 +24,8 @@
 
 import uvloop
 
-from dynamo._prometheus_metrics import Gauge, IntCounter, IntGauge, IntGaugeVec
-from dynamo.runtime import DistributedRuntime, dynamo_worker
+from dynamo.prometheus_metrics import Gauge, IntCounter, IntGauge, IntGaugeVec
+from dynamo.runtime import Component, DistributedRuntime, Endpoint, dynamo_worker
 
 
 def metrics_updater_thread(
@@ -60,10 +66,10 @@ async def worker(runtime: DistributedRuntime) -> None:
 
 async def init(runtime: DistributedRuntime):
     # Create component and endpoint
-    component = runtime.namespace("ns557").component("cp557")
+    component: Component = runtime.namespace("ns557").component("cp557")
     await component.create_service()
 
-    endpoint = component.endpoint("ep557")
+    endpoint: Endpoint = component.endpoint("ep557")
 
     # Create metrics using the endpoint's metrics property
     print("[python] Creating metrics...")
@@ -87,10 +93,10 @@ async def init(runtime: DistributedRuntime):
         [("update_method", "background_thread")],
     )
 
-    print(f"[python] Created IntGauge: {request_total_slots.name}")
-    print(f"[python] Created Gauge: {gpu_cache_usage_perc.name}")
-    print(f"[python] Created IntGaugeVec: {worker_active_requests.name}")
-    print(f"[python] Created IntCounter: {update_count.name}")
+    print(f"[python] Created IntGauge: {request_total_slots.name()}")
+    print(f"[python] Created Gauge: {gpu_cache_usage_perc.name()}")
+    print(f"[python] Created IntGaugeVec: {worker_active_requests.name()}")
+    print(f"[python] Created IntCounter: {update_count.name()}")
     print("[python] Metrics automatically registered with endpoint!")
 
     # Set initial values

lib/bindings/python/rust/lib.rs

@@ -185,7 +185,7 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
     engine::add_to_module(m)?;
     parsers::add_to_module(m)?;
 
-    m.add_class::<prometheus_metrics::PyRuntimeMetrics>()?;
+    m.add_class::<prometheus_metrics::RuntimeMetrics>()?;
     let prometheus_metrics = PyModule::new(m.py(), "prometheus_metrics")?;
     prometheus_metrics::add_to_module(&prometheus_metrics)?;
     m.add_submodule(&prometheus_metrics)?;
@@ -637,8 +637,8 @@ impl Component {
 
     /// Get a RuntimeMetrics helper for creating Prometheus metrics
     #[getter]
-    fn metrics(&self) -> prometheus_metrics::PyRuntimeMetrics {
-        prometheus_metrics::PyRuntimeMetrics::from_component(self.inner.clone())
+    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
+        prometheus_metrics::RuntimeMetrics::from_component(self.inner.clone())
     }
 }
 
@@ -726,8 +726,8 @@ impl Endpoint {
 
     /// Get a RuntimeMetrics helper for creating Prometheus metrics
     #[getter]
-    fn metrics(&self) -> prometheus_metrics::PyRuntimeMetrics {
-        prometheus_metrics::PyRuntimeMetrics::from_endpoint(self.inner.clone())
+    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
+        prometheus_metrics::RuntimeMetrics::from_endpoint(self.inner.clone())
     }
 }
 
@@ -743,8 +743,8 @@ impl Namespace {
 
     /// Get a RuntimeMetrics helper for creating Prometheus metrics
     #[getter]
-    fn metrics(&self) -> prometheus_metrics::PyRuntimeMetrics {
-        prometheus_metrics::PyRuntimeMetrics::from_namespace(self.inner.clone())
+    fn metrics(&self) -> prometheus_metrics::RuntimeMetrics {
+        prometheus_metrics::RuntimeMetrics::from_namespace(self.inner.clone())
     }
 }
 

lib/bindings/python/rust/prometheus_metrics.rs

@@ -18,6 +18,37 @@ use std::sync::Arc;
 
 use crate::rs;
 
+/// Helper function to order label values according to variable_labels declaration.
+/// This ensures labels are passed to with_label_values() in the correct order.
+///
+/// # Arguments
+/// * `variable_labels` - The ordered list of label names as declared in the metric
+/// * `labels` - The HashMap of label name-value pairs from Python
+///
+/// # Returns
+/// * `Ok(Vec<&str>)` - Ordered vector of label values matching variable_labels order
+/// * `Err(PyErr)` - If a required label is missing
+fn collect_ordered_label_values<'a>(
+    variable_labels: &[String],
+    labels: &'a HashMap<String, String>,
+) -> PyResult<Vec<&'a str>> {
+    let mut ordered_values = Vec::with_capacity(variable_labels.len());
+    for label_name in variable_labels {
+        match labels.get(label_name) {
+            Some(value) => ordered_values.push(value.as_str()),
+            None => {
+                return Err(pyo3::exceptions::PyValueError::new_err(format!(
+                    "Missing required label '{}'. Expected labels: {:?}, Got: {:?}",
+                    label_name,
+                    variable_labels,
+                    labels.keys().collect::<Vec<_>>()
+                )));
+            }
+        }
+    }
+    Ok(ordered_values)
+}
+
 // Python wrappers for Prometheus metric types.
 //
 // These wrapper structs are necessary because Prometheus types from the external `prometheus` crate
@@ -205,21 +236,24 @@ impl CounterVec {
 
     /// Increment counter by 1 with labels
     fn inc(&self, labels: HashMap<String, String>) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.counter.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.counter.with_label_values(&label_values).inc();
         Ok(())
     }
 
     /// Increment counter by value with labels
     fn inc_by(&self, labels: HashMap<String, String>, value: f64) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.counter.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.counter.with_label_values(&label_values).inc_by(value);
         Ok(())
     }
 
     /// Get counter value with labels
     fn get(&self, labels: HashMap<String, String>) -> PyResult<f64> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.counter.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         Ok(self.counter.with_label_values(&label_values).get())
     }
 }
@@ -257,21 +291,24 @@ impl IntCounterVec {
 
     /// Increment counter by 1 with labels
     fn inc(&self, labels: HashMap<String, String>) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.counter.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.counter.with_label_values(&label_values).inc();
         Ok(())
     }
 
     /// Increment counter by value with labels
     fn inc_by(&self, labels: HashMap<String, String>, value: u64) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.counter.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.counter.with_label_values(&label_values).inc_by(value);
         Ok(())
     }
 
     /// Get counter value with labels
     fn get(&self, labels: HashMap<String, String>) -> PyResult<u64> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.counter.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         Ok(self.counter.with_label_values(&label_values).get())
     }
 }
@@ -443,41 +480,47 @@ impl GaugeVec {
 
     /// Set gauge value with labels
     fn set(&self, value: f64, labels: HashMap<String, String>) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.gauge.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.gauge.with_label_values(&label_values).set(value);
         Ok(())
     }
 
     /// Get gauge value with labels
     fn get(&self, labels: HashMap<String, String>) -> PyResult<f64> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.gauge.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         Ok(self.gauge.with_label_values(&label_values).get())
     }
 
     /// Increment gauge by 1 with labels
     fn inc(&self, labels: HashMap<String, String>) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.gauge.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.gauge.with_label_values(&label_values).inc();
         Ok(())
     }
 
     /// Decrement gauge by 1 with labels
     fn dec(&self, labels: HashMap<String, String>) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.gauge.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.gauge.with_label_values(&label_values).dec();
         Ok(())
     }
 
     /// Add value to gauge with labels
     fn add(&self, labels: HashMap<String, String>, value: f64) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.gauge.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.gauge.with_label_values(&label_values).add(value);
         Ok(())
     }
 
     /// Subtract value from gauge with labels
     fn sub(&self, labels: HashMap<String, String>, value: f64) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.gauge.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.gauge.with_label_values(&label_values).sub(value);
         Ok(())
     }
@@ -516,41 +559,47 @@ impl IntGaugeVec {
 
     /// Set gauge value with labels
     fn set(&self, value: i64, labels: HashMap<String, String>) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.gauge.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.gauge.with_label_values(&label_values).set(value);
         Ok(())
     }
 
     /// Get gauge value with labels
     fn get(&self, labels: HashMap<String, String>) -> PyResult<i64> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.gauge.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         Ok(self.gauge.with_label_values(&label_values).get())
     }
 
     /// Increment gauge by 1 with labels
     fn inc(&self, labels: HashMap<String, String>) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.gauge.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.gauge.with_label_values(&label_values).inc();
         Ok(())
     }
 
     /// Decrement gauge by 1 with labels
     fn dec(&self, labels: HashMap<String, String>) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.gauge.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.gauge.with_label_values(&label_values).dec();
         Ok(())
     }
 
     /// Add value to gauge with labels
     fn add(&self, labels: HashMap<String, String>, value: i64) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.gauge.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.gauge.with_label_values(&label_values).add(value);
         Ok(())
     }
 
     /// Subtract value from gauge with labels
     fn sub(&self, labels: HashMap<String, String>, value: i64) -> PyResult<()> {
-        let label_values: Vec<&str> = labels.values().map(|s| s.as_str()).collect();
+        let desc = self.gauge.desc();
+        let label_values = collect_ordered_label_values(&desc[0].variable_labels, &labels)?;
         self.gauge.with_label_values(&label_values).sub(value);
         Ok(())
     }
@@ -598,15 +647,15 @@ impl Histogram {
 /// and utilities for registering metrics callbacks.
 /// Exposed as endpoint.metrics, component.metrics, and namespace.metrics in Python.
 ///
-/// NOTE: The create_* methods in PyRuntimeMetrics must stay in sync with the MetricsRegistry trait
+/// NOTE: The create_* methods in RuntimeMetrics must stay in sync with the MetricsRegistry trait
 /// in lib/runtime/src/metrics.rs. When adding new metric types, update both locations.
 #[pyclass]
 #[derive(Clone)]
-pub struct PyRuntimeMetrics {
+pub struct RuntimeMetrics {
     metricsregistry: Arc<dyn rs::metrics::MetricsRegistry>,
 }
 
-impl PyRuntimeMetrics {
+impl RuntimeMetrics {
     /// Create from Endpoint
     pub fn from_endpoint(endpoint: dynamo_runtime::component::Endpoint) -> Self {
         Self {
@@ -668,7 +717,7 @@ impl PyRuntimeMetrics {
 }
 
 #[pymethods]
-impl PyRuntimeMetrics {
+impl RuntimeMetrics {
     /// Register a Python callback to be invoked before metrics are scraped
     /// This callback will be called for this endpoint's metrics hierarchy
     fn register_update_callback(&self, callback: PyObject, _py: Python) -> PyResult<()> {