Merge pull request #564 from PrimeIntellect-ai/release/v.0.3.5

Release: v.0.3.5

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.4"
+version = "0.3.5"
 edition = "2021"
 
 [workspace.features]

Cargo.toml

@@ -158,8 +158,14 @@ async fn restart_node_task(node_id: web::Path<String>, app_state: Data<AppState>
         }));
     }
 
-    let p2p_id = node.worker_p2p_id.as_ref().unwrap();
-    let p2p_addresses = node.worker_p2p_addresses.as_ref().unwrap();
+    let p2p_id = node
+        .worker_p2p_id
+        .as_ref()
+        .expect("worker_p2p_id should be present");
+    let p2p_addresses = node
+        .worker_p2p_addresses
+        .as_ref()
+        .expect("worker_p2p_addresses should be present");
 
     match app_state
         .p2p_client
@@ -227,8 +233,24 @@ async fn get_node_logs(node_id: web::Path<String>, app_state: Data<AppState>) ->
         }));
     }
 
-    let p2p_id = node.worker_p2p_id.as_ref().unwrap();
-    let p2p_addresses = node.worker_p2p_addresses.as_ref().unwrap();
+    let p2p_id = match node.worker_p2p_id.as_ref() {
+        Some(id) => id,
+        None => {
+            return HttpResponse::BadRequest().json(json!({
+                "success": false,
+                "error": "Node does not have worker p2p id"
+            }));
+        }
+    };
+    let p2p_addresses = match node.worker_p2p_addresses.as_ref() {
+        Some(addresses) => addresses,
+        None => {
+            return HttpResponse::BadRequest().json(json!({
+                "success": false,
+                "error": "Node does not have worker p2p addresses"
+            }));
+        }
+    };
 
     match app_state
         .p2p_client

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

@@ -335,6 +335,7 @@ fn generate_file_name(template: &str, original_name: &str) -> String {
 mod tests {
 
     use std::collections::HashMap;
+    use std::sync::Arc;
 
     use super::*;
     use crate::plugins::StatusUpdatePlugin;
@@ -580,13 +581,14 @@ mod tests {
             compute_requirements: None,
         };
 
-        let plugin = NodeGroupsPlugin::new(
+        let plugin = Arc::new(NodeGroupsPlugin::new(
             vec![config],
             app_state.redis_store.clone(),
             app_state.store_context.clone(),
             None,
             None,
-        );
+        ));
+        let _ = plugin.clone().register_observer().await;
 
         let _ = plugin
             .handle_status_change(&node, &NodeStatus::Healthy)

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

@@ -41,6 +41,7 @@ use shared::web3::contracts::core::builder::ContractBuilder;
 use shared::web3::contracts::structs::compute_pool::PoolStatus;
 use shared::web3::wallet::Wallet;
 use std::sync::Arc;
+use std::time::Duration;
 use tokio::task::JoinSet;
 use url::Url;
 
@@ -176,6 +177,32 @@ async fn main() -> Result<()> {
     let store = Arc::new(RedisStore::new(&args.redis_store_url));
     let store_context = Arc::new(StoreContext::new(store.clone()));
 
+    if server_mode == ServerMode::ProcessorOnly || server_mode == ServerMode::Full {
+        info!("Migrating node store to hash");
+        if let Err(e) = store_context.node_store.migrate_json_to_hash().await {
+            error!("Error migrating node store to hash: {}", e);
+        }
+        info!("Node store migrated to hash");
+    } else {
+        // Wait for all nodes to be migrated to hash format
+        loop {
+            match store_context.node_store.count_non_hash_format_nodes().await {
+                Ok(0) => {
+                    info!("All nodes are in hash format");
+                    break;
+                }
+                Ok(count) => {
+                    info!("Waiting for {} nodes to be migrated to hash format", count);
+                    tokio::time::sleep(Duration::from_secs(1)).await;
+                }
+                Err(e) => {
+                    error!("Error counting non-hash format nodes: {}", e);
+                    tokio::time::sleep(Duration::from_secs(1)).await;
+                }
+            }
+        }
+    }
+
     let p2p_client = Arc::new(P2PClient::new(wallet.clone()).await.unwrap());
 
     let contracts = ContractBuilder::new(wallet.provider())
@@ -259,7 +286,14 @@ async fn main() -> Result<()> {
 
                 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));
+                let group_plugin_arc = Arc::new(group_plugin_for_server);
+
+                // Register the plugin as a task observer
+                if let Err(e) = group_plugin_arc.clone().register_observer().await {
+                    error!("Failed to register node groups plugin as observer: {}", e);
+                }
+
+                node_groups_plugin = Some(group_plugin_arc);
                 scheduler_plugins.push(Box::new(group_plugin));
                 status_update_plugins.push(Box::new(status_group_plugin));
                 info!("Plugin: Node group plugin initialized");

crates/orchestrator/src/main.rs

@@ -65,22 +65,12 @@ impl From<DiscoveryNode> for OrchestratorNode {
     }
 }
 
-impl OrchestratorNode {
-    pub fn from_string(s: &str) -> Self {
-        let mut node: Self = serde_json::from_str(s).unwrap();
-        if node.status == NodeStatus::Dead || node.status == NodeStatus::Ejected {
-            node.task_id = None;
-            node.task_state = None;
-        }
-        node
-    }
-}
-
 impl fmt::Display for OrchestratorNode {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "{}", serde_json::to_string(self).unwrap())
     }
 }
+
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default, ToSchema)]
 pub enum NodeStatus {
     #[default]

crates/orchestrator/src/models/node.rs

@@ -166,22 +166,24 @@ impl NodeGroupsPlugin {
             }
         });
 
-        let plugin = Self {
+        Self {
             configuration_templates: sorted_configs,
             store,
             store_context,
             node_groups_heartbeats,
             webhook_plugins,
             task_switching_policy,
             proximity_optimization_policy,
-        };
+        }
+    }
 
-        plugin
-            .store_context
+    /// Register this plugin as a task observer (async)
+    pub async fn register_observer(self: Arc<Self>) -> Result<()> {
+        self.store_context
             .task_store
-            .add_observer(Arc::new(plugin.clone()));
-
-        plugin
+            .add_observer(self.clone())
+            .await;
+        Ok(())
     }
 
     /// Check if a node is compatible with a configuration's compute requirements
@@ -776,19 +778,93 @@ impl NodeGroupsPlugin {
         let mut total_nodes = BTreeSet::new();
         let mut groups_to_dissolve = Vec::new();
 
-        // Select groups for merging
-        for group in compatible_groups {
-            if total_nodes.len() + group.nodes.len() <= config.max_group_size {
-                merge_batch.push(group.clone());
-                total_nodes.extend(group.nodes.iter().cloned());
-                groups_to_dissolve.push(group.id.clone());
+        // If proximity optimization is enabled, try to use location-based selection
+        if self.proximity_optimization_policy.enabled {
+            // Get node information for location data
+            let nodes = self.store_context.node_store.get_nodes().await?;
+            let node_map: HashMap<String, &OrchestratorNode> = nodes
+                .iter()
+                .map(|node| (node.address.to_string(), node))
+                .collect();
+
+            // Try to find a seed group with location data
+            let seed_group = compatible_groups.iter().find(|group| {
+                group
+                    .nodes
+                    .iter()
+                    .next()
+                    .and_then(|addr| node_map.get(addr))
+                    .and_then(|node| node.location.as_ref())
+                    .is_some()
+            });
+
+            if let Some(seed) = seed_group {
+                // Found a seed with location, use proximity-based selection
+                let seed_node = node_map.get(seed.nodes.iter().next().unwrap()).unwrap();
 
-                if total_nodes.len() >= config.max_group_size {
-                    break;
+                merge_batch.push(seed.clone());
+                total_nodes.extend(seed.nodes.iter().cloned());
+                groups_to_dissolve.push(seed.id.clone());
+
+                // Create a sorted list of remaining groups by proximity
+                let mut remaining_with_distance: Vec<(f64, &NodeGroup)> = compatible_groups
+                    .iter()
+                    .filter(|g| g.id != seed.id)
+                    .filter_map(|group| {
+                        let node_addr = group.nodes.iter().next()?;
+                        let node = node_map.get(node_addr)?;
+                        let node_loc = node.location.as_ref()?;
+                        let seed_loc = seed_node.location.as_ref()?;
+                        let distance = Self::calculate_distance(seed_loc, node_loc);
+                        Some((distance, group))
+                    })
+                    .collect();
+
+                remaining_with_distance
+                    .sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
+
+                // Add closest groups first
+                for (_, group) in remaining_with_distance {
+                    if total_nodes.len() + group.nodes.len() <= config.max_group_size {
+                        merge_batch.push(group.clone());
+                        total_nodes.extend(group.nodes.iter().cloned());
+                        groups_to_dissolve.push(group.id.clone());
+
+                        if total_nodes.len() >= config.max_group_size {
+                            break;
+                        }
+                    }
                 }
             }
         }
 
+        // If no proximity-based selection happened or we still need more groups, use original logic
+        if merge_batch.is_empty()
+            || (total_nodes.len() < config.max_group_size
+                && total_nodes.len() < config.min_group_size)
+        {
+            // Reset if we didn't get enough nodes
+            if total_nodes.len() < config.min_group_size {
+                merge_batch.clear();
+                total_nodes.clear();
+                groups_to_dissolve.clear();
+            }
+
+            // Original selection logic
+            for group in compatible_groups {
+                if !groups_to_dissolve.contains(&group.id)
+                    && total_nodes.len() + group.nodes.len() <= config.max_group_size
+                {
+                    merge_batch.push(group.clone());
+                    total_nodes.extend(group.nodes.iter().cloned());
+                    groups_to_dissolve.push(group.id.clone());
+
+                    if total_nodes.len() >= config.max_group_size {
+                        break;
+                    }
+                }
+            }
+        }
         // Validate merge conditions
         if !self
             .is_merge_beneficial(&merge_batch, total_nodes.len())