Release/0.3.3 (#549)

* refactor helm charts (#527)

* orchestrator: refactor helm chart

* validator: refactor helm chart

* discovery: create new helm chart

* fix helm chart (#537)

* fix orchestrator validation

* discovery: match existing secrets name

* charts: add quotes (#538)

* charts: add quotes

* charts/validator: fix conditions for redis (#539)

* fix validator redis chart (#540)

* bump version (#541)

* remove deprecated e2e folder (#543)

* fix(validator) dockerfile env causing invalidations not to be triggered(#536)

* imp(validator): ability to pick invalidation type for rejection (#542)

* ability to pick invalidation type for rejection

* imp(orchestrator): optimize heartbeat & metrics storage (#545)

* optimize heartbeat & metrics store

* imp(orchestrator): optimize orchestrator nodes / groups performance (#546)

* optimize orchestrator nodes / groups performance

* bump version (#548)

* add missing toploc invalidation type

* Merge pull request #552 from PrimeIntellect-ai/fix/migration-only-on-processor

imp(orchestrator): restrict redis migrations to processor

* improve node metrics migration (#553)

* validator: fix fullname template (#554)

* bump version to 0.3.3

---------

Co-authored-by: Aiman Ismail <m.aimanismail@gmail.com>
Co-authored-by: Aiman Ismail <aiman@primeintellect.ai>

Author: JannikSt

Cargo.lock

@@ -40,7 +40,7 @@ iroh = "0.34.1"
 rand_v8 = { package = "rand", version = "0.8.5", features = ["std"] }
 rand_core_v6 = { package = "rand_core", version = "0.6.4", features = ["std"] }
 [workspace.package]
-version = "0.3.2"
+version = "0.3.3"
 edition = "2021"
 
 [workspace.features]

Cargo.toml

@@ -95,46 +95,54 @@ async fn heartbeat(
         error!("Heartbeat Error: {}", e);
     }
     if let Some(metrics) = heartbeat.metrics.clone() {
-        // Get all previously reported metrics for this node
-        let previous_metrics = match app_state
+        // Get current metric keys for this node efficiently using HKEYS
+        let previous_metric_keys = match app_state
             .store_context
             .metrics_store
-            .get_metrics_for_node(node_address)
+            .get_metric_keys_for_node(node_address)
             .await
         {
-            Ok(metrics) => metrics,
+            Ok(keys) => keys,
             Err(e) => {
-                error!("Error getting metrics for node: {}", e);
-                Default::default()
+                error!("Error getting metric keys for node: {}", e);
+                Vec::new()
             }
         };
 
-        // Create a HashSet of new metrics for efficient lookup
-        let new_metrics_set: HashSet<_> = metrics
+        // Create a HashSet of new metric keys for efficient lookup
+        let new_metrics_set: HashSet<String> = metrics
             .iter()
-            .map(|metric| (&metric.key.task_id, &metric.key.label))
+            .map(|metric| {
+                let task_id = if metric.key.task_id.is_empty() {
+                    "manual".to_string()
+                } else {
+                    metric.key.task_id.clone()
+                };
+                format!("{}:{}", task_id, metric.key.label.replace(':', ""))
+            })
             .collect();
 
-        // Clean up stale metrics from Redis only
-        // The sync service will handle all Prometheus updates
-        for (task_id, task_metrics) in previous_metrics {
-            for (label, _value) in task_metrics {
-                let prev_key = (&task_id, &label);
-                if !new_metrics_set.contains(&prev_key) {
-                    // Remove from Redis metrics store
-                    if let Err(e) = app_state
-                        .store_context
-                        .metrics_store
-                        .delete_metric(&task_id, &label, &node_address.to_string())
-                        .await
-                    {
-                        error!("Error deleting metric: {}", e);
-                    }
+        // Find stale metrics to delete
+        let stale_metrics: Vec<String> = previous_metric_keys
+            .into_iter()
+            .filter(|key| !new_metrics_set.contains(key))
+            .collect();
+
+        // Delete stale metrics efficiently
+        for metric_key in stale_metrics {
+            if let Some((task_id, label)) = metric_key.split_once(':') {
+                if let Err(e) = app_state
+                    .store_context
+                    .metrics_store
+                    .delete_metric(task_id, label, &node_address.to_string())
+                    .await
+                {
+                    error!("Error deleting metric: {}", e);
                 }
             }
         }
 
-        // Store new metrics in Redis only
+        // Store new metrics in Redis
         if let Err(e) = app_state
             .store_context
             .metrics_store

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

@@ -68,22 +68,34 @@ async fn get_nodes(query: Query<NodeQuery>, app_state: Data<AppState>) -> HttpRe
     if let Some(node_groups_plugin) = &app_state.node_groups_plugin {
         let mut nodes_with_groups = Vec::new();
 
-        for node in &nodes {
-            let mut node_json = json!(node);
+        let node_addresses: Vec<String> =
+            nodes.iter().map(|node| node.address.to_string()).collect();
 
-            if let Ok(Some(group)) = node_groups_plugin
-                .get_node_group(&node.address.to_string())
-                .await
-            {
-                node_json["group"] = json!({
-                    "id": group.id,
-                    "size": group.nodes.len(),
-                    "created_at": group.created_at,
-                    "topology_config": group.configuration_name
-                });
-            }
+        match node_groups_plugin
+            .get_node_groups_batch(&node_addresses)
+            .await
+        {
+            Ok(node_groups) => {
+                for node in &nodes {
+                    let mut node_json = json!(node);
 
-            nodes_with_groups.push(node_json);
+                    if let Some(Some(group)) = node_groups.get(&node.address.to_string()) {
+                        node_json["group"] = json!({
+                            "id": group.id,
+                            "size": group.nodes.len(),
+                            "created_at": group.created_at,
+                            "topology_config": group.configuration_name
+                        });
+                    }
+
+                    nodes_with_groups.push(node_json);
+                }
+            }
+            Err(e) => {
+                error!("Error getting node groups batch: {}", e);
+                // Fall back to nodes without group information
+                nodes_with_groups = nodes.iter().map(|node| json!(node)).collect();
+            }
         }
 
         response["nodes"] = json!(nodes_with_groups);

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

@@ -126,6 +126,51 @@ struct Args {
     max_healthy_nodes_with_same_endpoint: u32,
 }
 
+async fn run_inactive_node_metric_migration(store_context: Arc<StoreContext>) -> Result<()> {
+    info!("Starting migration of inactive node metrics to new data model");
+
+    let node_addresses = match store_context.node_store.get_nodes().await {
+        Ok(nodes) => {
+            let addresses: Vec<_> = nodes.into_iter().map(|node| node.address).collect();
+            info!("Found {} nodes to migrate", addresses.len());
+            addresses
+        }
+        Err(e) => {
+            error!("Error getting all nodes for migration: {}", e);
+            return Ok(()); // Don't fail startup if migration can't get nodes
+        }
+    };
+
+    let mut migrated_count = 0;
+    let mut error_count = 0;
+
+    for node_address in node_addresses {
+        match store_context
+            .metrics_store
+            .migrate_node_metrics_if_needed(node_address)
+            .await
+        {
+            Ok(()) => {
+                migrated_count += 1;
+                if migrated_count % 100 == 0 {
+                    info!("Migrated {} nodes so far...", migrated_count);
+                }
+            }
+            Err(e) => {
+                error!("Error migrating metrics for node {}: {}", node_address, e);
+                error_count += 1;
+                // Continue with other nodes even if one fails
+            }
+        }
+    }
+
+    info!(
+        "Migration completed. Successfully migrated {} nodes, {} errors",
+        migrated_count, error_count
+    );
+    Ok(())
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
     let args = Args::parse();
@@ -176,6 +221,17 @@ async fn main() -> Result<()> {
     let store = Arc::new(RedisStore::new(&args.redis_store_url));
     let store_context = Arc::new(StoreContext::new(store.clone()));
 
+    // Run one-time migration for inactive nodes
+    let migration_store_context = store_context.clone();
+    if matches!(server_mode, ServerMode::ProcessorOnly | ServerMode::Full) {
+        let migration_store_context = migration_store_context.clone();
+        tokio::spawn(async move {
+            if let Err(e) = run_inactive_node_metric_migration(migration_store_context).await {
+                error!("Failed to run inactive node metric migration: {}", e);
+            }
+        });
+    }
+
     let p2p_client = Arc::new(P2PClient::new(wallet.clone()).await.unwrap());
 
     let contracts = ContractBuilder::new(wallet.provider())
@@ -256,6 +312,27 @@ async fn main() -> Result<()> {
                     Some(node_groups_heartbeats.clone()),
                     Some(webhook_plugins.clone()),
                 );
+
+                // Run groups index migration on startup
+                if matches!(server_mode, ServerMode::ProcessorOnly | ServerMode::Full) {
+                    match group_plugin.migrate_groups_index().await {
+                        Ok(count) => {
+                            if count > 0 {
+                                info!(
+                                    "Groups index migration completed: {} groups migrated",
+                                    count
+                                );
+                            } else {
+                                info!("Groups index migration: no groups to migrate");
+                            }
+                        }
+                        Err(e) => {
+                            error!("Groups index migration failed: {}", e);
+                            return Err(e);
+                        }
+                    }
+                }
+
                 let status_group_plugin = group_plugin.clone();
                 let group_plugin_for_server = group_plugin.clone();
                 node_groups_plugin = Some(Arc::new(group_plugin_for_server));