Feature: Node group plugin with configurable topologies (#335)

* ability to have multiple node group configurations

* ability to set compute requirements per node group

Author: JannikSt

Makefile

@@ -91,7 +91,7 @@ watch-validator:
 
 watch-orchestrator:
 	set -a; source ${ENV_FILE}; set +a; \
-	cargo watch -w crates/orchestrator/src -x "run --bin orchestrator -- -r $$RPC_URL -k $$POOL_OWNER_PRIVATE_KEY -d 0  -p 8090 -i 10 -u http://localhost:8090 --s3-credentials $$S3_CREDENTIALS --compute-pool-id $$WORKER_COMPUTE_POOL_ID --bucket-name $$BUCKET_NAME -l $${LOG_LEVEL:-info} --hourly-s3-upload-limit $${HOURLY_S3_LIMIT:-3} --with-basic-group-plugin --group-size 1"
+	cargo watch -w crates/orchestrator/src -x "run --bin orchestrator -- -r $$RPC_URL -k $$POOL_OWNER_PRIVATE_KEY -d 0  -p 8090 -i 10 -u http://localhost:8090 --s3-credentials $$S3_CREDENTIALS --compute-pool-id $$WORKER_COMPUTE_POOL_ID --bucket-name $$BUCKET_NAME -l $${LOG_LEVEL:-info} --hourly-s3-upload-limit $${HOURLY_S3_LIMIT:-3} --node-group-configs '[{\"name\": \"test-config\", \"min_group_size\": 1, \"max_group_size\": 1, \"compute_requirements\": null}]'"
 
 build-worker:
 	cargo build --release --bin worker

crates/orchestrator/Dockerfile

@@ -23,8 +23,7 @@ ENV LOG_LEVEL=""
 ENV HOURLY_S3_UPLOAD_LIMIT="2"
 ENV WEBHOOK_URLS=""
 ENV MODE="full"
-ENV WITH_BASIC_GROUP_PLUGIN="false"
-ENV GROUP_SIZE="4"
+ENV NODE_GROUP_CONFIGS=""
 
 RUN echo '#!/bin/sh\n\
 exec /usr/local/bin/orchestrator \
@@ -46,8 +45,7 @@ $([ ! -z "$BUCKET_NAME" ] && echo "--bucket-name $BUCKET_NAME") \
 $([ ! -z "$LOG_LEVEL" ] && echo "--log-level $LOG_LEVEL") \
 $([ ! -z "$HOURLY_S3_UPLOAD_LIMIT" ] && echo "--hourly-s3-upload-limit $HOURLY_S3_UPLOAD_LIMIT") \
 $([ ! -z "$WEBHOOK_URLS" ] && echo "--webhook-urls $WEBHOOK_URLS") \
-$([ "$WITH_BASIC_GROUP_PLUGIN" = "true" ] && echo "--with-basic-group-plugin") \
-$([ ! -z "$GROUP_SIZE" ] && echo "--group-size $GROUP_SIZE") \
+$([ ! -z "$NODE_GROUP_CONFIGS" ] && echo "--node-group-configs $NODE_GROUP_CONFIGS") \
 "$@"' > /entrypoint.sh && \
 chmod +x /entrypoint.sh
 

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

@@ -45,7 +45,8 @@ async fn get_nodes(app_state: Data<AppState>) -> HttpResponse {
                 node_json["group"] = json!({
                     "id": group.id,
                     "size": group.nodes.len(),
-                    "created_at": group.created_at
+                    "created_at": group.created_at,
+                    "topology_config": group.configuration_name
                 });
             }
 

crates/orchestrator/src/main.rs

@@ -28,6 +28,7 @@ use scheduler::plugins::SchedulerPlugin;
 use shared::web3::contracts::core::builder::ContractBuilder;
 use shared::web3::contracts::structs::compute_pool::PoolStatus;
 use shared::web3::wallet::Wallet;
+use status_update::plugins::node_groups::NodeGroupConfiguration;
 use status_update::plugins::node_groups::NodeGroupsPlugin;
 use status_update::plugins::StatusUpdatePlugin;
 use std::sync::Arc;
@@ -115,14 +116,9 @@ struct Args {
     #[arg(long, default_value = "")]
     webhook_urls: Option<String>,
 
-    /// With basic group plugin
-    /// Only temporary setting - will be moved to proper plugin config
+    /// Node group configurations in JSON format
     #[arg(long)]
-    with_basic_group_plugin: bool,
-
-    /// Group size
-    #[arg(long, default_value = "4")]
-    group_size: u32,
+    node_group_configs: Option<String>,
 }
 
 #[tokio::main]
@@ -203,16 +199,26 @@ async fn main() -> Result<()> {
     let mut status_update_plugins: Vec<Box<dyn StatusUpdatePlugin>> = vec![];
     let mut node_groups_plugin: Option<Arc<NodeGroupsPlugin>> = None;
 
-    // Add group plugin if enabled
-    if args.with_basic_group_plugin {
-        let group_size: usize = args.group_size as usize;
-        let group_plugin =
-            NodeGroupsPlugin::new(group_size, group_size, store.clone(), group_store_context);
-        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));
-        scheduler_plugins.push(Box::new(group_plugin));
-        status_update_plugins.push(Box::new(status_group_plugin));
+    // This config loading is pretty ugly atm and should be optimized
+    // Issue: https://github.com/PrimeIntellect-ai/protocol/issues/336
+    if let Some(configs_json) = args.node_group_configs {
+        match serde_json::from_str::<Vec<NodeGroupConfiguration>>(&configs_json) {
+            Ok(configs) if !configs.is_empty() => {
+                let group_plugin =
+                    NodeGroupsPlugin::new(configs, store.clone(), group_store_context);
+                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));
+                scheduler_plugins.push(Box::new(group_plugin));
+                status_update_plugins.push(Box::new(status_group_plugin));
+            }
+            Ok(_) => {
+                info!("No node group configurations provided, skipping plugin setup");
+            }
+            Err(e) => {
+                panic!("Failed to parse node group configurations: {}", e);
+            }
+        }
     }
 
     let scheduler = Scheduler::new(store_context.clone(), scheduler_plugins);