diff --git a/monarch_hyperactor/src/telemetry.rs b/monarch_hyperactor/src/telemetry.rs index 90871232..802491d5 100644 --- a/monarch_hyperactor/src/telemetry.rs +++ b/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>, +) -> PyResult<()> { + let kv_pairs: Vec = 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>, +) -> PyResult<()> { + let kv_pairs: Vec = 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>, +) -> PyResult<()> { + let kv_pairs: Vec = 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::()?; Ok(()) } diff --git a/python/monarch/_rust_bindings/monarch_hyperactor/telemetry.pyi b/python/monarch/_rust_bindings/monarch_hyperactor/telemetry.pyi index a9c6fc87..b14a54cf 100644 --- a/python/monarch/_rust_bindings/monarch_hyperactor/telemetry.pyi +++ b/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: """