Add py bindings for otel metrics

monarch_hyperactor/src/telemetry.rs

@@ -13,6 +13,7 @@ use std::cell::Cell;
 use hyperactor::clock::ClockKind;
 use hyperactor::clock::RealClock;
 use hyperactor::clock::SimClock;
+use hyperactor_telemetry::opentelemetry;
 use hyperactor_telemetry::swap_telemetry_clock;
 use pyo3::prelude::*;
 use pyo3::types::PyTraceback;
@@ -134,6 +135,73 @@ impl PySpan {
     }
 }
 
+/// Add to a counter with the given name and attributes
+#[pyfunction]
+#[pyo3(signature = (name, value, attributes = None))]
+pub fn add_to_counter(
+    name: String,
+    value: u64,
+    attributes: Option<Vec<(String, String)>>,
+) -> PyResult<()> {
+    let kv_pairs: Vec<opentelemetry::KeyValue> = attributes
+        .unwrap_or_default()
+        .into_iter()
+        .map(|(k, v)| opentelemetry::KeyValue::new(k, v))
+        .collect();
+
+    println!("Added {} to counter {}", value, name);
+    let counter = hyperactor_telemetry::meter("python")
+        .u64_counter(name)
+        .build();
+
+    counter.add(value, &kv_pairs);
+    Ok(())
+}
+
+/// Add to an up/down counter with the given name and attributes
+#[pyfunction]
+#[pyo3(signature = (name, value, attributes = None))]
+pub fn add_to_up_down_counter(
+    name: String,
+    value: i64,
+    attributes: Option<Vec<(String, String)>>,
+) -> PyResult<()> {
+    let kv_pairs: Vec<opentelemetry::KeyValue> = attributes
+        .unwrap_or_default()
+        .into_iter()
+        .map(|(k, v)| opentelemetry::KeyValue::new(k, v))
+        .collect();
+
+    let counter = hyperactor_telemetry::meter("python")
+        .i64_up_down_counter(name)
+        .build();
+
+    counter.add(value, &kv_pairs);
+    Ok(())
+}
+
+/// Record a value to a gauge with the given name and attributes
+#[pyfunction]
+#[pyo3(signature = (name, value, attributes = None))]
+pub fn add_to_gauge(
+    name: String,
+    value: f64,
+    attributes: Option<Vec<(String, String)>>,
+) -> PyResult<()> {
+    let kv_pairs: Vec<opentelemetry::KeyValue> = attributes
+        .unwrap_or_default()
+        .into_iter()
+        .map(|(k, v)| opentelemetry::KeyValue::new(k, v))
+        .collect();
+
+    let gauge = hyperactor_telemetry::meter("python")
+        .f64_gauge(name)
+        .build();
+
+    gauge.record(value, &kv_pairs);
+    Ok(())
+}
+
 use pyo3::Bound;
 use pyo3::types::PyModule;
 
@@ -182,6 +250,28 @@ pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
     )?;
     module.add_function(use_sim_clock_fn)?;
 
+    // Register telemetry functions
+    let add_to_counter_fn = wrap_pyfunction!(add_to_counter, module)?;
+    add_to_counter_fn.setattr(
+        "__module__",
+        "monarch._rust_bindings.hyperactor_extension.telemetry",
+    )?;
+    module.add_function(add_to_counter_fn)?;
+
+    let add_to_up_down_counter_fn = wrap_pyfunction!(add_to_up_down_counter, module)?;
+    add_to_up_down_counter_fn.setattr(
+        "__module__",
+        "monarch._rust_bindings.hyperactor_extension.telemetry",
+    )?;
+    module.add_function(add_to_up_down_counter_fn)?;
+
+    let add_to_gauge_fn = wrap_pyfunction!(add_to_gauge, module)?;
+    add_to_gauge_fn.setattr(
+        "__module__",
+        "monarch._rust_bindings.hyperactor_extension.telemetry",
+    )?;
+    module.add_function(add_to_gauge_fn)?;
+
     module.add_class::<PySpan>()?;
     Ok(())
 }

python/monarch/_rust_bindings/monarch_hyperactor/telemetry.pyi

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import logging
+from typing import Optional
 
 def forward_to_tracing(record: logging.LogRecord) -> None:
     """
@@ -86,6 +87,60 @@ def use_sim_clock() -> None:
     """
     ...
 
+def add_to_counter(
+    name: str, value: int, attributes: Optional[list[tuple[str, str]]] = None
+) -> None:
+    """
+    Add a value to a counter with the given name and attributes.
+
+    This function creates or uses an existing counter with the specified name
+    and increments it by the given value with the provided attributes.
+
+    Args:
+    - name (str): The name of the counter metric.
+    - value (int): The value to add to the counter (must be non-negative).
+    - attributes (Optional[list[tuple[str, str]]]): Optional list of key-value pairs to use as metric attributes.
+      These attributes allow you to create different time series for the same counter.
+      If None, no attributes will be added.
+    """
+    ...
+
+def add_to_up_down_counter(
+    name: str, value: int, attributes: Optional[list[tuple[str, str]]] = None
+) -> None:
+    """
+    Add a value to an up/down counter with the given name and attributes.
+
+    This function creates or uses an existing up/down counter with the specified name
+    and increments or decrements it by the given value with the provided attributes.
+
+    Args:
+    - name (str): The name of the up/down counter metric.
+    - value (int): The value to add to the counter (can be positive or negative).
+    - attributes (Optional[list[tuple[str, str]]]): Optional list of key-value pairs to use as metric attributes.
+      These attributes allow you to create different time series for the same counter.
+      If None, no attributes will be added.
+    """
+    ...
+
+def add_to_gauge(
+    name: str, value: float, attributes: Optional[list[tuple[str, str]]] = None
+) -> None:
+    """
+    Record a value to a gauge with the given name and attributes.
+
+    This function creates or uses an existing gauge with the specified name
+    and sets it to the given value with the provided attributes.
+
+    Args:
+    - name (str): The name of the gauge metric.
+    - value (float): The value to record to the gauge.
+    - attributes (Optional[list[tuple[str, str]]]): Optional list of key-value pairs to use as metric attributes.
+      These attributes allow you to create different time series for the same gauge.
+      If None, no attributes will be added.
+    """
+    ...
+
 class PySpan:
     def __init__(self, name: str) -> None:
         """