add additional prometheus operator metrics, add taskname to node metrics (#428)

* add additional prometheus operator metrics, add taskname to node metrics

Author: JannikSt

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

@@ -28,6 +28,11 @@ async fn heartbeat(
         }
     };
 
+    // Track heartbeat request in metrics
+    app_state
+        .metrics
+        .increment_heartbeat_requests(&heartbeat.address);
+
     let node_opt = app_state
         .store_context
         .node_store
@@ -257,13 +262,15 @@ mod tests {
         // Verify Prometheus registry is initially empty (no sync service has run)
         let prometheus_metrics_before = app_state.metrics.export_metrics().unwrap();
         assert!(!prometheus_metrics_before.contains("performance/batch_avg_seq_length"));
+        assert!(prometheus_metrics_before.contains(&format!("orchestrator_heartbeat_requests_total{{node_address=\"0x0000000000000000000000000000000000000000\",pool_id=\"{}\"}} 1", app_state.metrics.pool_id)));
 
         // Create and run sync service manually to test the sync functionality
         let sync_service = MetricsSyncService::new(
             app_state.store_context.clone(),
             app_state.metrics.clone(),
             ServerMode::Full, // Test app uses Full mode
             10,
+            None, // No node groups plugin in test
         );
 
         // Manually trigger a sync operation

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

@@ -274,6 +274,12 @@ async fn request_upload(
                 file_name
             );
 
+            app_state.metrics.increment_file_upload_requests(
+                &request_upload.task_id,
+                &task.name,
+                address,
+            );
+
             #[cfg(test)]
             return HttpResponse::Ok().json(serde_json::json!({
                 "success": true,
@@ -380,6 +386,10 @@ mod tests {
             json["file_name"],
             serde_json::Value::String("model_123/user_uploads/test.parquet".to_string())
         );
+
+        let metrics = app_state.metrics.export_metrics().unwrap();
+        assert!(metrics.contains("orchestrator_file_upload_requests_total"));
+        assert!(metrics.contains(&format!("orchestrator_file_upload_requests_total{{node_address=\"test_address\",pool_id=\"{}\",task_id=\"{}\",task_name=\"test-task\"}} 1", app_state.metrics.pool_id, task.id)));
     }
 
     #[actix_web::test]

crates/orchestrator/src/main.rs

@@ -272,12 +272,14 @@ async fn main() -> Result<()> {
         // Start metrics sync service to centralize metrics from Redis to Prometheus
         let metrics_sync_store_context = store_context.clone();
         let metrics_sync_context = metrics_context.clone();
+        let metrics_sync_node_groups = node_groups_plugin.clone();
         tasks.spawn(async move {
             let sync_service = MetricsSyncService::new(
                 metrics_sync_store_context,
                 metrics_sync_context,
                 server_mode,
                 10,
+                metrics_sync_node_groups,
             );
             sync_service.run().await
         });
@@ -328,8 +330,6 @@ async fn main() -> Result<()> {
 
         let status_update_store_context = store_context.clone();
         let status_update_heartbeats = heartbeats.clone();
-        let status_update_metrics = metrics_context.clone();
-
         tasks.spawn({
             let contracts = contracts.clone();
             async move {
@@ -342,7 +342,6 @@ async fn main() -> Result<()> {
                     args.disable_ejection,
                     status_update_heartbeats.clone(),
                     status_update_plugins,
-                    status_update_metrics,
                 );
                 status_updater.run().await
             }

crates/orchestrator/src/metrics/mod.rs

@@ -1,52 +1,154 @@
-use prometheus::{GaugeVec, Opts, Registry, TextEncoder};
+use prometheus::{CounterVec, GaugeVec, Opts, Registry, TextEncoder};
 pub mod sync_service;
 pub mod webhook_sender;
 
 pub struct MetricsContext {
     pub compute_task_gauges: GaugeVec,
     pub pool_id: String,
     pub registry: Registry,
+    pub file_upload_requests_total: CounterVec,
+    pub nodes_total: GaugeVec,
+    pub tasks_total: GaugeVec,
+    pub groups_total: GaugeVec,
+    pub heartbeat_requests_total: CounterVec,
+    pub nodes_per_task: GaugeVec,
 }
 
 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"],
+            &["node_address", "task_id", "task_name", "label", "pool_id"],
         )
         .unwrap();
+
+        // New metrics for orchestrator statistics
+        let file_upload_requests_total = CounterVec::new(
+            Opts::new(
+                "orchestrator_file_upload_requests_total",
+                "Total number of file upload requests",
+            ),
+            &["task_id", "task_name", "node_address", "pool_id"],
+        )
+        .unwrap();
+
+        let nodes_total = GaugeVec::new(
+            Opts::new(
+                "orchestrator_nodes_total",
+                "Total number of nodes by status",
+            ),
+            &["status", "pool_id"],
+        )
+        .unwrap();
+
+        let tasks_total = GaugeVec::new(
+            Opts::new("orchestrator_tasks_total", "Total number of tasks"),
+            &["pool_id"],
+        )
+        .unwrap();
+
+        let groups_total = GaugeVec::new(
+            Opts::new(
+                "orchestrator_groups_total",
+                "Total number of node groups by configuration",
+            ),
+            &["configuration_name", "pool_id"],
+        )
+        .unwrap();
+
+        let heartbeat_requests_total = CounterVec::new(
+            Opts::new(
+                "orchestrator_heartbeat_requests_total",
+                "Total number of heartbeat requests per node",
+            ),
+            &["node_address", "pool_id"],
+        )
+        .unwrap();
+
+        let nodes_per_task = GaugeVec::new(
+            Opts::new(
+                "orchestrator_nodes_per_task",
+                "Number of nodes actively working on each task",
+            ),
+            &["task_id", "task_name", "pool_id"],
+        )
+        .unwrap();
+
         let registry = Registry::new();
         let _ = registry.register(Box::new(compute_task_gauges.clone()));
+        let _ = registry.register(Box::new(file_upload_requests_total.clone()));
+        let _ = registry.register(Box::new(nodes_total.clone()));
+        let _ = registry.register(Box::new(tasks_total.clone()));
+        let _ = registry.register(Box::new(groups_total.clone()));
+        let _ = registry.register(Box::new(heartbeat_requests_total.clone()));
+        let _ = registry.register(Box::new(nodes_per_task.clone()));
 
         Self {
             compute_task_gauges,
             pool_id,
             registry,
+            file_upload_requests_total,
+            nodes_total,
+            tasks_total,
+            groups_total,
+            heartbeat_requests_total,
+            nodes_per_task,
         }
     }
 
     pub fn record_compute_task_gauge(
         &self,
         node_address: &str,
         task_id: &str,
+        task_name: &str,
         label: &str,
         value: f64,
     ) {
         self.compute_task_gauges
-            .with_label_values(&[node_address, task_id, label, &self.pool_id])
+            .with_label_values(&[node_address, task_id, task_name, label, &self.pool_id])
             .set(value);
     }
 
-    pub fn remove_compute_task_gauge(&self, node_address: &str, task_id: &str, label: &str) {
-        if let Err(e) = self.compute_task_gauges.remove_label_values(&[
-            node_address,
-            task_id,
-            label,
-            &self.pool_id,
-        ]) {
-            println!("Error removing compute task gauge: {}", e);
-        }
+    pub fn increment_file_upload_requests(
+        &self,
+        task_id: &str,
+        task_name: &str,
+        node_address: &str,
+    ) {
+        self.file_upload_requests_total
+            .with_label_values(&[task_id, task_name, node_address, &self.pool_id])
+            .inc();
+    }
+
+    pub fn increment_heartbeat_requests(&self, node_address: &str) {
+        self.heartbeat_requests_total
+            .with_label_values(&[node_address, &self.pool_id])
+            .inc();
+    }
+
+    pub fn set_nodes_count(&self, status: &str, count: f64) {
+        self.nodes_total
+            .with_label_values(&[status, &self.pool_id])
+            .set(count);
+    }
+
+    pub fn set_tasks_count(&self, count: f64) {
+        self.tasks_total
+            .with_label_values(&[&self.pool_id])
+            .set(count);
+    }
+
+    pub fn set_groups_count(&self, configuration_name: &str, count: f64) {
+        self.groups_total
+            .with_label_values(&[configuration_name, &self.pool_id])
+            .set(count);
+    }
+
+    pub fn set_nodes_per_task(&self, task_id: &str, task_name: &str, count: f64) {
+        self.nodes_per_task
+            .with_label_values(&[task_id, task_name, &self.pool_id])
+            .set(count);
     }
 
     pub fn export_metrics(&self) -> Result<String, prometheus::Error> {
@@ -61,4 +163,12 @@ impl MetricsContext {
         // This removes all existing metrics so we can rebuild from Redis
         self.compute_task_gauges.reset();
     }
+
+    /// Clear all orchestrator statistics metrics
+    pub fn clear_orchestrator_statistics(&self) {
+        self.nodes_total.reset();
+        self.tasks_total.reset();
+        self.groups_total.reset();
+        self.nodes_per_task.reset();
+    }
 }