improve node grouping and dissolution - link to task, speed up groups api, add config name to metrics (#445)

Author: JannikSt

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

@@ -264,4 +264,53 @@ mod tests {
         assert_eq!(tasks[1].image, "test2");
         assert_eq!(tasks[2].image, "test1");
     }
+
+    #[actix_web::test]
+    async fn test_create_task_with_metadata() {
+        let app_state = create_test_app_state().await;
+        let app = test::init_service(
+            App::new()
+                .app_data(app_state.clone())
+                .route("/tasks", post().to(create_task)),
+        )
+        .await;
+
+        let mut labels = std::collections::HashMap::new();
+        labels.insert("model".to_string(), "qwen3-4b".to_string());
+        labels.insert("dataset".to_string(), "intellect-2-rl-dataset".to_string());
+        labels.insert("version".to_string(), "v1".to_string());
+
+        let payload = TaskRequest {
+            image: "primeintellect/prime-rl:main".to_string(),
+            name: "Qwen3-4B:INTELLECT-2-RL-Dataset".to_string(),
+            metadata: Some(shared::models::task::TaskMetadata {
+                labels: Some(labels),
+            }),
+            ..Default::default()
+        };
+
+        let req = test::TestRequest::post()
+            .uri("/tasks")
+            .set_json(payload)
+            .to_request();
+        let resp = test::call_service(&app, req).await;
+        assert_eq!(resp.status(), StatusCode::OK);
+
+        let body = test::read_body(resp).await;
+        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
+        assert_eq!(json["success"], serde_json::Value::Bool(true));
+        assert!(json["task"]["id"].is_string());
+        assert_eq!(json["task"]["image"], "primeintellect/prime-rl:main");
+        assert_eq!(json["task"]["name"], "Qwen3-4B:INTELLECT-2-RL-Dataset");
+
+        // Verify metadata is preserved
+        assert!(json["task"]["metadata"].is_object());
+        assert!(json["task"]["metadata"]["labels"].is_object());
+        assert_eq!(json["task"]["metadata"]["labels"]["model"], "qwen3-4b");
+        assert_eq!(
+            json["task"]["metadata"]["labels"]["dataset"],
+            "intellect-2-rl-dataset"
+        );
+        assert_eq!(json["task"]["metadata"]["labels"]["version"], "v1");
+    }
 }

crates/orchestrator/src/metrics/mod.rs

@@ -4,6 +4,7 @@ pub mod webhook_sender;
 
 pub struct MetricsContext {
     pub compute_task_gauges: GaugeVec,
+    pub task_info: GaugeVec,
     pub pool_id: String,
     pub registry: Registry,
     pub file_upload_requests_total: CounterVec,
@@ -19,7 +20,21 @@ impl MetricsContext {
         // For current state/rate metrics
         let compute_task_gauges = GaugeVec::new(
             Opts::new("compute_gauges", "Compute task gauge metrics"),
-            &["node_address", "task_id", "task_name", "label", "pool_id"],
+            &[
+                "node_address",
+                "task_id",
+                "task_name",
+                "label",
+                "pool_id",
+                "group_id",
+                "group_config_name",
+            ],
+        )
+        .unwrap();
+
+        let task_info = GaugeVec::new(
+            Opts::new("task_info", "Task information with metadata"),
+            &["task_id", "task_name", "pool_id", "metadata"],
         )
         .unwrap();
 
@@ -77,6 +92,7 @@ impl MetricsContext {
 
         let registry = Registry::new();
         let _ = registry.register(Box::new(compute_task_gauges.clone()));
+        let _ = registry.register(Box::new(task_info.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()));
@@ -86,6 +102,7 @@ impl MetricsContext {
 
         Self {
             compute_task_gauges,
+            task_info,
             pool_id,
             registry,
             file_upload_requests_total,
@@ -97,16 +114,29 @@ impl MetricsContext {
         }
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn record_compute_task_gauge(
         &self,
         node_address: &str,
         task_id: &str,
         task_name: &str,
         label: &str,
         value: f64,
+        group_id: Option<&str>,
+        group_config_name: Option<&str>,
     ) {
+        let group_id_str = group_id.unwrap_or("none");
+        let group_config_name_str = group_config_name.unwrap_or("none");
         self.compute_task_gauges
-            .with_label_values(&[node_address, task_id, task_name, label, &self.pool_id])
+            .with_label_values(&[
+                node_address,
+                task_id,
+                task_name,
+                label,
+                &self.pool_id,
+                group_id_str,
+                group_config_name_str,
+            ])
             .set(value);
     }
 
@@ -151,6 +181,12 @@ impl MetricsContext {
             .set(count);
     }
 
+    pub fn set_task_info(&self, task_id: &str, task_name: &str, metadata: &str) {
+        self.task_info
+            .with_label_values(&[task_id, task_name, &self.pool_id, metadata])
+            .set(1.0);
+    }
+
     pub fn export_metrics(&self) -> Result<String, prometheus::Error> {
         let encoder = TextEncoder::new();
         let metric_families = self.registry.gather();
@@ -170,5 +206,6 @@ impl MetricsContext {
         self.tasks_total.reset();
         self.groups_total.reset();
         self.nodes_per_task.reset();
+        self.task_info.reset();
     }
 }

crates/orchestrator/src/metrics/sync_service.rs

@@ -3,6 +3,7 @@ use crate::plugins::node_groups::NodeGroupsPlugin;
 use crate::store::core::StoreContext;
 use crate::ServerMode;
 use log::{debug, error, info};
+use shared::models::task::Task;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
@@ -33,6 +34,70 @@ impl MetricsSyncService {
         }
     }
 
+    /// Format task metadata into a structured string for Prometheus labels
+    /// Example: "model:qwen3-4b|dataset:intellect-2-rl|version:v1"
+    fn format_task_metadata(task: &Task) -> String {
+        if let Some(metadata) = &task.metadata {
+            if let Some(labels) = &metadata.labels {
+                if !labels.is_empty() {
+                    return labels
+                        .iter()
+                        .map(|(k, v)| format!("{}:{}", k, v))
+                        .collect::<Vec<_>>()
+                        .join("|");
+                }
+            }
+        }
+        "".to_string()
+    }
+
+    /// Efficiently get all node-to-group mappings including both group_id and group_config_name
+    /// Returns a HashMap where key is node_address and value is (group_id, group_config_name)
+    async fn get_all_node_group_info(&self) -> anyhow::Result<HashMap<String, (String, String)>> {
+        if let Some(node_groups_plugin) = &self.node_groups_plugin {
+            // First get all node to group_id mappings
+            let node_to_group_mappings =
+                match node_groups_plugin.get_all_node_group_mappings().await {
+                    Ok(mappings) => mappings,
+                    Err(e) => {
+                        error!("Failed to get node group mappings: {}", e);
+                        return Ok(HashMap::new());
+                    }
+                };
+
+            // Then get all groups to get their configuration names
+            let groups = match node_groups_plugin.get_all_groups().await {
+                Ok(groups) => groups,
+                Err(e) => {
+                    error!("Failed to get all groups: {}", e);
+                    return Ok(HashMap::new());
+                }
+            };
+
+            // Create a mapping from group_id to configuration_name
+            let group_id_to_config: HashMap<String, String> = groups
+                .into_iter()
+                .map(|group| (group.id, group.configuration_name))
+                .collect();
+
+            // Combine the mappings to create node_address -> (group_id, group_config_name)
+            let mut result = HashMap::new();
+            for (node_address, group_id) in node_to_group_mappings {
+                if let Some(config_name) = group_id_to_config.get(&group_id) {
+                    result.insert(node_address, (group_id, config_name.clone()));
+                } else {
+                    // If we can't find the config name, still include the group_id
+                    debug!("No configuration name found for group_id: {}", group_id);
+                    result.insert(node_address, (group_id, "unknown".to_string()));
+                }
+            }
+
+            Ok(result)
+        } else {
+            Ok(HashMap::new())
+        }
+    }
+
     pub async fn run(&self) -> anyhow::Result<()> {
         // Only run the sync service on ProcessorOnly or Full mode instances
         if !matches!(
@@ -86,6 +151,18 @@ impl MetricsSyncService {
             .map(|task| (task.id.to_string(), task.name.clone()))
             .collect();
 
+        let node_to_group_info = if self.node_groups_plugin.is_some() {
+            match self.get_all_node_group_info().await {
+                Ok(info) => info,
+                Err(e) => {
+                    error!("Failed to get node group info: {}", e);
+                    HashMap::new()
+                }
+            }
+        } else {
+            HashMap::new()
+        };
+
         // Clear existing Prometheus metrics
         self.metrics_context.clear_compute_task_metrics();
 
@@ -99,12 +176,19 @@ impl MetricsSyncService {
 
             for (label, node_metrics) in task_metrics {
                 for (node_address, value) in node_metrics {
+                    let (group_id, group_config_name) = node_to_group_info
+                        .get(&node_address)
+                        .map(|(id, config)| (Some(id.as_str()), Some(config.as_str())))
+                        .unwrap_or((None, None));
+
                     self.metrics_context.record_compute_task_gauge(
                         &node_address,
                         &task_id,
                         &task_name,
                         &label,
                         value,
+                        group_id,
+                        group_config_name,
                     );
                     total_metrics += 1;
                 }
@@ -150,6 +234,15 @@ impl MetricsSyncService {
             self.metrics_context.set_tasks_count(total_tasks);
             debug!("Synced task statistics: {} total tasks", total_tasks);
 
+            // Sync task info metrics with metadata
+            for task in &tasks {
+                let task_id = task.id.to_string();
+                let metadata = Self::format_task_metadata(task);
+                self.metrics_context
+                    .set_task_info(&task_id, &task.name, &metadata);
+            }
+            debug!("Synced task info metrics with metadata");
+
             // Sync nodes per task based on node assignments
             // Create task name mapping
             let task_name_map: HashMap<String, String> = tasks
@@ -202,3 +295,86 @@ impl MetricsSyncService {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use shared::models::task::{Task, TaskMetadata, TaskState};
+    use std::collections::HashMap;
+    use uuid::Uuid;
+
+    #[test]
+    fn test_format_task_metadata_with_labels() {
+        let mut labels = HashMap::new();
+        labels.insert("model".to_string(), "qwen3-4b".to_string());
+        labels.insert("dataset".to_string(), "intellect-2-rl-dataset".to_string());
+        labels.insert("version".to_string(), "v1".to_string());
+
+        let task = Task {
+            id: Uuid::new_v4(),
+            image: "test".to_string(),
+            name: "test".to_string(),
+            state: TaskState::PENDING,
+            metadata: Some(TaskMetadata {
+                labels: Some(labels),
+            }),
+            ..Default::default()
+        };
+
+        let formatted = MetricsSyncService::format_task_metadata(&task);
+
+        // The format should be key:value pairs separated by |
+        // Order might vary due to HashMap iteration
+        assert!(formatted.contains("model:qwen3-4b"));
+        assert!(formatted.contains("dataset:intellect-2-rl-dataset"));
+        assert!(formatted.contains("version:v1"));
+        assert_eq!(formatted.matches('|').count(), 2); // Should have 2 separators for 3 labels
+    }
+
+    #[test]
+    fn test_format_task_metadata_empty() {
+        let task = Task {
+            id: Uuid::new_v4(),
+            image: "test".to_string(),
+            name: "test".to_string(),
+            state: TaskState::PENDING,
+            metadata: None,
+            ..Default::default()
+        };
+
+        let formatted = MetricsSyncService::format_task_metadata(&task);
+        assert_eq!(formatted, "");
+    }
+
+    #[test]
+    fn test_format_task_metadata_empty_labels() {
+        let task = Task {
+            id: Uuid::new_v4(),
+            image: "test".to_string(),
+            name: "test".to_string(),
+            state: TaskState::PENDING,
+            metadata: Some(TaskMetadata {
+                labels: Some(HashMap::new()),
+            }),
+            ..Default::default()
+        };
+
+        let formatted = MetricsSyncService::format_task_metadata(&task);
+        assert_eq!(formatted, "");
+    }
+
+    #[test]
+    fn test_format_task_metadata_no_labels() {
+        let task = Task {
+            id: Uuid::new_v4(),
+            image: "test".to_string(),
+            name: "test".to_string(),
+            state: TaskState::PENDING,
+            metadata: Some(TaskMetadata { labels: None }),
+            ..Default::default()
+        };
+
+        let formatted = MetricsSyncService::format_task_metadata(&task);
+        assert_eq!(formatted, "");
+    }
+}