Improvement: adjust node config loading (#346)

* load node group config for orchestrator from env

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 --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}]'"
+	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 --compute-pool-id $$WORKER_COMPUTE_POOL_ID --bucket-name $$BUCKET_NAME -l $${LOG_LEVEL:-info} --hourly-s3-upload-limit $${HOURLY_S3_LIMIT:-3}"
 
 build-worker:
 	cargo build --release --bin worker

crates/orchestrator/Dockerfile

@@ -44,7 +44,6 @@ $([ ! -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") \
-$([ ! -z "$NODE_GROUP_CONFIGS" ] && echo "--node-group-configs $NODE_GROUP_CONFIGS") \
 "$@"' > /entrypoint.sh && \
 chmod +x /entrypoint.sh
 

crates/orchestrator/src/main.rs

@@ -111,10 +111,6 @@ struct Args {
     /// Webhook urls (comma-separated string)
     #[arg(long, default_value = "")]
     webhook_urls: Option<String>,
-
-    /// Node group configurations in JSON format
-    #[arg(long)]
-    node_group_configs: Option<String>,
 }
 
 #[tokio::main]
@@ -195,11 +191,11 @@ async fn main() -> Result<()> {
     let mut status_update_plugins: Vec<Box<dyn StatusUpdatePlugin>> = vec![];
     let mut node_groups_plugin: Option<Arc<NodeGroupsPlugin>> = None;
 
-    // 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 {
+    // Load node group configurations from environment variable
+    if let Ok(configs_json) = std::env::var("NODE_GROUP_CONFIGS") {
         match serde_json::from_str::<Vec<NodeGroupConfiguration>>(&configs_json) {
             Ok(configs) if !configs.is_empty() => {
+                println!("configs for node groups: {:?}", configs);
                 let group_plugin =
                     NodeGroupsPlugin::new(configs, store.clone(), group_store_context);
                 let status_group_plugin = group_plugin.clone();
@@ -209,10 +205,15 @@ async fn main() -> Result<()> {
                 status_update_plugins.push(Box::new(status_group_plugin));
             }
             Ok(_) => {
-                info!("No node group configurations provided, skipping plugin setup");
+                info!(
+                    "No node group configurations provided in environment, skipping plugin setup"
+                );
             }
             Err(e) => {
-                panic!("Failed to parse node group configurations: {}", e);
+                error!(
+                    "Failed to parse node group configurations from environment: {}",
+                    e
+                );
             }
         }
     }

deployment/k8s/orchestrator-chart/templates/orchestrator-deployment.yaml

@@ -67,8 +67,8 @@
       value: "{{ .root.Values.env.WEBHOOK_URLS }}"
     {{- end }}
     {{- if .root.Values.env.NODE_GROUP_CONFIGS }}
-    - name: NODE_GROUP_CONFIGS
-      value: "{{ .root.Values.env.NODE_GROUP_CONFIGS }}"
+      - name: NODE_GROUP_CONFIGS
+        value: "{{ .root.Values.env.NODE_GROUP_CONFIGS | trim | quote }}"
     {{- end }}
 {{- end }}
 

deployment/k8s/orchestrator-chart/values.example.yaml

@@ -10,6 +10,7 @@ env:
   DISCOVERY_REFRESH_INTERVAL: "60"
   BUCKET_NAME: "your-bucket-name"
   LOG_LEVEL: "info"
+  NODE_GROUP_CONFIGS: '[{\"name\": \"h100-config\", \"min_group_size\": 1, \"max_group_size\": 1, \"compute_requirements\": \"gpu:model=h100;gpu:count=8\"}, {\"name\": \"a100-config\", \"min_group_size\": 2, \"max_group_size\": 2, \"compute_requirements\": \"gpu:model=a100;gpu:count=8\"}]'
 
 secrets:
   coordinatorKey: "your-coordinator-private-key"