export node task metrics to prometheus (#367)

* orchestrator now exports a /metrics/prometheus endpoint that includes the decentralized node's reported metrics

Author: JannikSt

Cargo.lock

@@ -17,6 +17,7 @@ google-cloud-auth = "0.18.0"
 google-cloud-storage = "0.24.0"
 hex = { workspace = true }
 log = { workspace = true }
+prometheus = "0.14.0"
 rand = "0.9.0"
 redis = { workspace = true }
 redis-test = { workspace = true }

crates/orchestrator/Cargo.toml

@@ -41,10 +41,21 @@ async fn heartbeat(
     }
 
     app_state.store_context.heartbeat_store.beat(&heartbeat);
-    app_state
-        .store_context
-        .metrics_store
-        .store_metrics(heartbeat.metrics.clone(), node_address);
+    if let Some(metrics) = heartbeat.metrics.clone() {
+        app_state
+            .store_context
+            .metrics_store
+            .store_metrics(Some(metrics.clone()), node_address);
+
+        for metric in metrics {
+            app_state.metrics.record_compute_task_gauge(
+                &node_address.to_string(),
+                &metric.key.task_id,
+                &metric.key.label,
+                metric.value,
+            );
+        }
+    }
 
     let current_task = app_state.scheduler.get_task_for_node(node_address);
     match current_task {
@@ -73,10 +84,13 @@ pub fn heartbeat_routes() -> Scope {
 mod tests {
     use super::*;
     use crate::api::tests::helper::create_test_app_state;
+
     use actix_web::http::StatusCode;
     use actix_web::test;
     use actix_web::App;
     use serde_json::json;
+    use shared::models::metric::MetricEntry;
+    use shared::models::metric::MetricKey;
     use shared::models::task::TaskRequest;
 
     #[actix_web::test]
@@ -90,7 +104,10 @@ mod tests {
         .await;
 
         let address = "0x0000000000000000000000000000000000000000".to_string();
-        let req_payload = json!({"address": address});
+        let req_payload = json!({"address": address, "metrics": [
+            {"key": {"task_id": "long-task-1234", "label": "performance/batch_avg_seq_length"}, "value": 1.0},
+            {"key": {"task_id": "long-task-1234", "label": "performance/batch_min_seq_length"}, "value": 5.0}
+        ]});
 
         let req = test::TestRequest::post()
             .uri("/heartbeat")
@@ -116,12 +133,31 @@ mod tests {
                 address: "0x0000000000000000000000000000000000000000".to_string(),
                 task_id: None,
                 task_state: None,
-                metrics: None,
+                metrics: Some(vec![
+                    MetricEntry {
+                        key: MetricKey {
+                            task_id: "long-task-1234".to_string(),
+                            label: "performance/batch_avg_seq_length".to_string(),
+                        },
+                        value: 1.0,
+                    },
+                    MetricEntry {
+                        key: MetricKey {
+                            task_id: "long-task-1234".to_string(),
+                            label: "performance/batch_min_seq_length".to_string(),
+                        },
+                        value: 5.0,
+                    }
+                ]),
                 version: None,
                 timestamp: None,
                 p2p_id: None,
             })
         );
+
+        let metrics = app_state.metrics.export_metrics().unwrap();
+        assert!(metrics.contains("performance/batch_avg_seq_length"));
+        assert!(metrics.contains("performance/batch_min_seq_length"));
     }
 
     #[actix_web::test]

crates/orchestrator/src/api/routes/heartbeat.rs

@@ -17,6 +17,7 @@ struct DeleteMetricRequest {
     label: String,
     address: String,
 }
+
 async fn get_metrics(app_state: Data<AppState>) -> HttpResponse {
     let metrics = app_state
         .store_context
@@ -30,6 +31,18 @@ async fn get_all_metrics(app_state: Data<AppState>) -> HttpResponse {
     HttpResponse::Ok().json(json!({"success": true, "metrics": metrics}))
 }
 
+async fn get_prometheus_metrics(app_state: Data<AppState>) -> HttpResponse {
+    match app_state.metrics.export_metrics() {
+        Ok(metrics) => HttpResponse::Ok()
+            .content_type("text/plain; version=0.0.4")
+            .body(metrics),
+        Err(e) => HttpResponse::InternalServerError().json(json!({
+            "success": false,
+            "error": format!("Failed to export metrics: {}", e)
+        })),
+    }
+}
+
 // for potential backup restore purposes
 async fn create_metric(
     app_state: Data<AppState>,
@@ -62,6 +75,7 @@ pub fn metrics_routes() -> Scope {
     web::scope("/metrics")
         .route("", get().to(get_metrics))
         .route("/all", get().to(get_all_metrics))
+        .route("/prometheus", get().to(get_prometheus_metrics))
         .route("", post().to(create_metric))
         .route("/{task_id}", delete().to(delete_metric))
 }

crates/orchestrator/src/api/routes/metrics.rs

@@ -2,6 +2,7 @@ use crate::api::routes::nodes::nodes_routes;
 use crate::api::routes::storage::storage_routes;
 use crate::api::routes::task::tasks_routes;
 use crate::api::routes::{heartbeat::heartbeat_routes, metrics::metrics_routes};
+use crate::metrics::MetricsContext;
 use crate::models::node::NodeStatus;
 use crate::plugins::node_groups::NodeGroupsPlugin;
 use crate::scheduler::Scheduler;
@@ -32,6 +33,7 @@ pub struct AppState {
     pub pool_id: u32,
     pub scheduler: Scheduler,
     pub node_groups_plugin: Option<Arc<NodeGroupsPlugin>>,
+    pub metrics: Arc<MetricsContext>,
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -50,6 +52,7 @@ pub async fn start_server(
     server_mode: ServerMode,
     scheduler: Scheduler,
     node_groups_plugin: Option<Arc<NodeGroupsPlugin>>,
+    metrics: Arc<MetricsContext>,
 ) -> Result<(), Error> {
     info!("Starting server at http://{}:{}", host, port);
     let app_state = Data::new(AppState {
@@ -63,6 +66,7 @@ pub async fn start_server(
         pool_id,
         scheduler,
         node_groups_plugin,
+        metrics,
     });
     let node_store = app_state.store_context.node_store.clone();
     let node_store_clone = node_store.clone();

crates/orchestrator/src/api/server.rs

@@ -19,7 +19,10 @@ use url::Url;
 pub async fn create_test_app_state() -> Data<AppState> {
     use shared::utils::MockStorageProvider;
 
-    use crate::{scheduler::Scheduler, utils::loop_heartbeats::LoopHeartbeats, ServerMode};
+    use crate::{
+        metrics::MetricsContext, scheduler::Scheduler, utils::loop_heartbeats::LoopHeartbeats,
+        ServerMode,
+    };
 
     let store = Arc::new(RedisStore::new_test());
     let mut con = store
@@ -40,6 +43,7 @@ pub async fn create_test_app_state() -> Data<AppState> {
 
     let mock_storage = MockStorageProvider::new();
     let storage_provider = Arc::new(mock_storage);
+    let metrics = Arc::new(MetricsContext::new(1.to_string()));
 
     Data::new(AppState {
         store_context: store_context.clone(),
@@ -58,6 +62,7 @@ pub async fn create_test_app_state() -> Data<AppState> {
         redis_store: store.clone(),
         scheduler,
         node_groups_plugin: None,
+        metrics,
     })
 }
 
@@ -66,6 +71,7 @@ pub async fn create_test_app_state_with_nodegroups() -> Data<AppState> {
     use shared::utils::MockStorageProvider;
 
     use crate::{
+        metrics::MetricsContext,
         plugins::node_groups::{NodeGroupConfiguration, NodeGroupsPlugin},
         scheduler::Scheduler,
         utils::loop_heartbeats::LoopHeartbeats,
@@ -105,6 +111,7 @@ pub async fn create_test_app_state_with_nodegroups() -> Data<AppState> {
 
     let mock_storage = MockStorageProvider::new();
     let storage_provider = Arc::new(mock_storage);
+    let metrics = Arc::new(MetricsContext::new(1.to_string()));
 
     Data::new(AppState {
         store_context: store_context.clone(),
@@ -123,6 +130,7 @@ pub async fn create_test_app_state_with_nodegroups() -> Data<AppState> {
         redis_store: store.clone(),
         scheduler,
         node_groups_plugin,
+        metrics,
     })
 }
 

crates/orchestrator/src/api/tests/helper.rs

@@ -1,6 +1,7 @@
 mod api;
 mod discovery;
 mod events;
+mod metrics;
 mod models;
 mod node;
 mod plugins;
@@ -24,6 +25,7 @@ use log::debug;
 use log::error;
 use log::info;
 use log::LevelFilter;
+use metrics::MetricsContext;
 use plugins::node_groups::NodeGroupConfiguration;
 use plugins::node_groups::NodeGroupsPlugin;
 use plugins::webhook::WebhookPlugin;
@@ -145,6 +147,8 @@ async fn main() -> Result<()> {
     debug!("Log level: {}", log_level);
     debug!("Server mode: {:?}", server_mode);
 
+    let metrics_context = Arc::new(MetricsContext::new(args.compute_pool_id.to_string()));
+
     let heartbeats = Arc::new(LoopHeartbeats::new(&server_mode));
 
     let compute_pool_id = args.compute_pool_id;
@@ -327,6 +331,7 @@ async fn main() -> Result<()> {
             server_mode,
             scheduler,
             node_groups_plugin,
+            metrics_context,
         ) => {
             if let Err(e) = res {
                 error!("Server error: {}", e);

crates/orchestrator/src/main.rs

@@ -0,0 +1,47 @@
+use prometheus::{GaugeVec, Opts, Registry, TextEncoder};
+
+pub struct MetricsContext {
+    pub compute_task_gauges: GaugeVec,
+    pub pool_id: String,
+    pub registry: Registry,
+}
+
+impl MetricsContext {
+    pub fn new(pool_id: String) -> Self {
+        // For current state/rate metrics
+        let compute_task_gauges = GaugeVec::new(
+            Opts::new("compute_gauges", "Compute task gauge metrics"),
+            &["node_address", "task_id", "label", "pool_id"],
+        )
+        .unwrap();
+        let registry = Registry::new();
+        let _ = registry.register(Box::new(compute_task_gauges.clone()));
+
+        Self {
+            compute_task_gauges,
+            pool_id,
+            registry,
+        }
+    }
+
+    pub fn record_compute_task_gauge(
+        &self,
+        node_address: &str,
+        task_id: &str,
+        label: &str,
+        value: f64,
+    ) {
+        println!("record_compute_task_gauge: {:?}", node_address);
+        let vals = vec![node_address, task_id, label, &self.pool_id];
+        println!("vals: {:?}", vals);
+        self.compute_task_gauges
+            .with_label_values(&[node_address, task_id, label, &self.pool_id])
+            .set(value);
+    }
+
+    pub fn export_metrics(&self) -> Result<String, prometheus::Error> {
+        let encoder = TextEncoder::new();
+        let metric_families = self.registry.gather();
+        encoder.encode_to_string(&metric_families)
+    }
+}