Sfoglia il codice sorgente

feat: cluster selection with occupancy awareness

- Add occupied_nodes() — queries all nodes via SSH for running sm121
  containers, returns set of occupied node names
- select_cluster() now accepts occupied set, filters clusters by
  free node count (>= required_nodes), picks smallest total cluster
- resolve_cluster() propagates occupied set through
- serve() queries occupancy before cluster resolution, filters
  cluster nodes to available ones before starting containers
- All other call sites pass empty HashSet for backward compat
- Fix select_cluster node count filter: >= not ==
Alexander Korolev 4 mesi fa
parent
commit
3774e765e5
7 ha cambiato i file con 89 aggiunte e 28 eliminazioni
  1. 24 17
      src/config.rs
  2. 43 1
      src/docker/mod.rs
  3. 2 1
      src/doctor.rs
  4. 4 2
      src/exec.rs
  5. 4 2
      src/main.rs
  6. 9 4
      src/tmux/ops.rs
  7. 3 1
      src/tmux/serve.rs

+ 24 - 17
src/config.rs

@@ -4,6 +4,7 @@ pub use crate::pre_serve::{
 use anyhow::{Context, Result};
 use serde::Deserialize;
 use std::collections::HashMap;
+use std::collections::HashSet;
 use tracing::info;
 #[derive(Clone, Deserialize)]
 pub struct Node {
@@ -160,11 +161,12 @@ fn default_tensor_parallel_size() -> u32 {
     1
 }
 /// Selects a cluster + NCCL config by name or auto-picks the smallest
-/// cluster with NCCL config satisfying `required_nodes`.
+/// cluster with NCCL config and enough free nodes.
 pub fn select_cluster<'c>(
     config: &'c Config,
     cluster_name: &str,
     required_nodes: Option<u32>,
+    occupied: &HashSet<String>,
 ) -> Result<(Cluster, &'c NcclConfig)> {
     if cluster_name != "auto" {
         let cluster = resolve_cluster_name(&config.clusters, cluster_name)?;
@@ -174,19 +176,26 @@ pub fn select_cluster<'c>(
             .context(fluent_i18n::t!("nccl-config-missing"))?;
         return Ok((cluster.clone(), nccl));
     }
+    let free = |c: &Cluster| -> u32 {
+        c.nodes
+            .iter()
+            .filter(|n| !occupied.contains(&n.name))
+            .count() as u32
+    };
     let mut candidates: Vec<(&Cluster, &NcclConfig)> = config
         .clusters
         .iter()
         .filter_map(|c| Some((c, config.nccl.get(&c.name)?)))
-        .filter(|(c, _)| required_nodes.is_none_or(|r| (c.nodes.len() as u32) >= r))
+        .filter(|(c, _)| required_nodes.is_none_or(|r| free(c) >= r))
         .collect();
     candidates.sort_by_key(|(c, _)| c.nodes.len());
     if let Some((cluster, nccl)) = candidates.into_iter().next() {
+        let avail = free(cluster);
         info!(
             "{}",
             fluent_i18n::t!("cluster-selected", {
                 "name" => cluster.name.to_string(),
-                "nodes" => cluster.nodes.len(),
+                "nodes" => avail,
             })
         );
         return Ok((cluster.clone(), nccl));
@@ -202,6 +211,7 @@ pub fn resolve_cluster<'c>(
     config: &'c Config,
     recipe: &Recipe,
     cluster_name: &str,
+    occupied: &HashSet<String>,
 ) -> Result<(Cluster, &'c NcclConfig)> {
     let choice = if cluster_name != "auto" {
         cluster_name.to_string()
@@ -210,7 +220,7 @@ pub fn resolve_cluster<'c>(
     } else {
         "auto".to_string()
     };
-    select_cluster(config, &choice, recipe.required_nodes)
+    select_cluster(config, &choice, recipe.required_nodes, occupied)
 }
 /// Resolve the path to a recipe's `mods/` directory, if it exists.
 /// Checks cwd first, then the default config dir.
@@ -299,16 +309,14 @@ mod tests {
     }
     #[test]
     fn test_recipe_deserialize_minimal() {
-        let toml_str = "\nname = \"Minimal\"\nmodel = \"test/model\"\n";
-        let recipe: Recipe = toml::from_str(toml_str).unwrap();
+        let recipe: Recipe =
+            toml::from_str("\nname = \"Minimal\"\nmodel = \"test/model\"\n").unwrap();
         assert_eq!(recipe.name, "Minimal");
         assert_eq!(recipe.model, "test/model");
         assert_eq!(recipe.engine, "vllm");
         assert_eq!(recipe.port, 30000);
-        assert!(recipe.args.to_vec().is_empty());
-        assert!(recipe.env.is_empty());
+        assert!(recipe.args.to_vec().is_empty() && recipe.env.is_empty());
         assert_eq!(recipe.distributed_backend, None);
-        assert_eq!(recipe.tensor_parallel_size, 1);
     }
     #[test]
     fn test_recipe_deserialize_full() {
@@ -400,7 +408,8 @@ ib_disable = 0
             global_mods_dir: None,
         };
         config.nccl.insert("test-cluster".to_string(), test_nccl());
-        let (cluster, nccl) = select_cluster(&config, "test-cluster", None).unwrap();
+        let (cluster, nccl) =
+            select_cluster(&config, "test-cluster", None, &HashSet::new()).unwrap();
         assert_eq!(cluster.name, "test-cluster");
         assert_eq!(nccl.ib_hca, "rocep1s0f0");
     }
@@ -419,21 +428,21 @@ ib_disable = 0
             global_mods_dir: None,
         };
         config.nccl.insert("test-cluster".to_string(), test_nccl());
-        let (cluster, nccl) = select_cluster(&config, "auto", None).unwrap();
+        let (cluster, nccl) = select_cluster(&config, "auto", None, &HashSet::new()).unwrap();
         assert_eq!(cluster.name, "test-cluster");
         assert_eq!(nccl.ib_hca, "rocep1s0f0");
     }
     #[test]
     fn test_select_cluster_auto_no_match() {
         let config = test_config();
-        let result = select_cluster(&config, "auto", None);
+        let result = select_cluster(&config, "auto", None, &HashSet::new());
         assert!(result.is_err());
     }
     #[test]
     fn test_select_cluster_not_found() {
         let mut config = test_config();
         config.clusters = vec![];
-        let result = select_cluster(&config, "nonexistent", None);
+        let result = select_cluster(&config, "nonexistent", None, &HashSet::new());
         assert!(result.is_err());
     }
     #[test]
@@ -456,12 +465,10 @@ ib_disable = 0
     fn test_select_cluster_by_prefix() {
         let mut config = test_config();
         config.nccl.insert("test-cluster".to_string(), test_nccl());
-        // Unique prefix works
         let cluster = select_cluster_by_name(&config, "test-c").unwrap();
         assert_eq!(cluster.name, "test-cluster");
-        // No match errors
-        let result = select_cluster_by_name(&config, "no-such");
-        assert!(result.is_err());
+        let result = select_cluster_by_name(&config, "no-such").is_err();
+        assert!(result);
     }
     #[test]
     fn test_russian_pluralization() {

+ 43 - 1
src/docker/mod.rs

@@ -8,6 +8,7 @@ use crate::{
 };
 use anyhow::{Context, Result};
 use std::borrow::Cow;
+use std::collections::HashSet;
 use tokio::time::{Duration, sleep};
 use tracing::info;
 
@@ -55,7 +56,15 @@ pub async fn serve(
             .context(format!("Node '{}' not found in any cluster", nn))?;
         (solo_c, solo_n)
     } else {
-        let (c, n) = resolve_cluster(config, &recipe, cluster_name)?;
+        let occupied = occupied_nodes(config).await?;
+        let (c, n) = resolve_cluster(config, &recipe, cluster_name, &occupied)?;
+        // Use only available (non-occupied) nodes from the cluster.
+        let avail: Vec<Node> = c
+            .nodes
+            .into_iter()
+            .filter(|n| !occupied.contains(&n.name))
+            .collect();
+        let c = Cluster { nodes: avail, ..c };
         (c, n.clone())
     };
     if let Some(required) = recipe.required_nodes {
@@ -236,6 +245,39 @@ pub async fn resolve_cluster_by_container_label<'c>(
     anyhow::bail!("No running containers found for recipe '{}'", recipe_name);
 }
 
+/// Query running containers across all unique nodes and return the set of
+/// occupied node names.  A node is occupied if any sm121 container is
+/// running on it.
+pub async fn occupied_nodes(config: &Config) -> Result<HashSet<String>> {
+    let mut seen = HashSet::new();
+    let mut handles = Vec::new();
+    for cluster in &config.clusters {
+        for node in &cluster.nodes {
+            if seen.contains(&node.ssh_ip) {
+                continue;
+            }
+            seen.insert(node.ssh_ip.clone());
+            let cmd = "docker ps -q --filter label=sm121.cluster --format '{{.Names}}' | head -1"
+                .to_string();
+            let config = config.clone();
+            let node = node.clone();
+            let name = node.name.clone();
+            handles.push(tokio::spawn(async move {
+                (name, run_ssh(&config, &node, cmd, false).await)
+            }));
+        }
+    }
+    let mut occupied = HashSet::new();
+    for h in futures::future::join_all(handles).await {
+        if let Ok((name, Ok(output))) = h
+            && !output.trim().is_empty()
+        {
+            occupied.insert(name);
+        }
+    }
+    Ok(occupied)
+}
+
 /// Runtime options for [`start_cluster`].
 struct StartClusterOpts {
     cache_dir: std::path::PathBuf,

+ 2 - 1
src/doctor.rs

@@ -1,4 +1,5 @@
 use anyhow::{Context, Result};
+use std::collections::HashSet;
 use tracing::info;
 
 use crate::config::{Cluster, Config, Node, Recipe, load_recipe, resolve_cluster};
@@ -56,7 +57,7 @@ pub async fn doctor(
         })
     );
 
-    let (cluster, _nccl) = resolve_cluster(config, &recipe, cluster_name)?;
+    let (cluster, _nccl) = resolve_cluster(config, &recipe, cluster_name, &HashSet::new())?;
 
     // Check each node
     let mut all_results: Vec<DoctorResult> = Vec::new();

+ 4 - 2
src/exec.rs

@@ -2,6 +2,7 @@ use anyhow::{Context, Result, anyhow};
 use tracing::{info, warn};
 
 use crate::config::{Cluster, Config, Recipe, select_cluster};
+use std::collections::HashSet;
 
 /// Resolve the target cluster for `cargo run exec` from `--cluster` or
 /// `--recipe` (exactly one of which must be given).
@@ -24,8 +25,9 @@ pub async fn resolve_target(
         (None, Some(rname)) => {
             let recipe = crate::config::load_recipe(rname).await?;
             let cname = recipe.cluster.as_deref().unwrap_or("auto");
-            let (cluster, _) = select_cluster(config, cname, recipe.required_nodes)
-                .with_context(|| format!("resolving cluster for recipe '{}'", rname))?;
+            let (cluster, _) =
+                select_cluster(config, cname, recipe.required_nodes, &HashSet::new())
+                    .with_context(|| format!("resolving cluster for recipe '{}'", rname))?;
             Ok((cluster, Some(recipe)))
         }
     }

+ 4 - 2
src/main.rs

@@ -205,6 +205,8 @@ async fn main() -> Result<()> {
 /// Resolve a cluster for container operations (stop/logs).
 /// When `auto`, tries live container label lookup first, then falls back
 /// to config-based auto-selection (for dry-run / offline scenarios).
+use std::collections::HashSet;
+
 async fn resolve_cluster_for_op(
     config: &Config,
     recipe_name: &str,
@@ -212,7 +214,7 @@ async fn resolve_cluster_for_op(
 ) -> Result<Cluster> {
     if cluster_name != "auto" {
         let recipe = load_recipe(recipe_name).await?;
-        let (cluster, _) = resolve_cluster(config, &recipe, cluster_name)?;
+        let (cluster, _) = resolve_cluster(config, &recipe, cluster_name, &HashSet::new())?;
         return Ok(cluster);
     }
     // Try live container label lookup first
@@ -221,6 +223,6 @@ async fn resolve_cluster_for_op(
     }
     // Fall back to config-based auto (dry-run / containers not running)
     let recipe = load_recipe(recipe_name).await?;
-    let (cluster, _) = resolve_cluster(config, &recipe, "auto")?;
+    let (cluster, _) = resolve_cluster(config, &recipe, "auto", &HashSet::new())?;
     Ok(cluster)
 }

+ 9 - 4
src/tmux/ops.rs

@@ -1,4 +1,5 @@
 use anyhow::{Context, Result};
+use std::collections::HashSet;
 use tracing::{info, warn};
 
 use crate::config::{Config, Recipe};
@@ -22,7 +23,8 @@ pub async fn tmux_stop(
     } else {
         "auto".to_string()
     };
-    let (cluster, _) = crate::config::select_cluster(config, &choice, recipe.required_nodes)?;
+    let (cluster, _) =
+        crate::config::select_cluster(config, &choice, recipe.required_nodes, &HashSet::new())?;
 
     let manager = TmuxManager {
         session_prefix: recipe_name.to_string(),
@@ -73,7 +75,8 @@ pub async fn tmux_status(
     } else {
         "auto".to_string()
     };
-    let (cluster, _) = crate::config::select_cluster(config, &choice, recipe.required_nodes)?;
+    let (cluster, _) =
+        crate::config::select_cluster(config, &choice, recipe.required_nodes, &HashSet::new())?;
 
     let manager = TmuxManager {
         session_prefix: recipe_name.to_string(),
@@ -162,7 +165,8 @@ pub async fn tmux_logs(
     } else {
         "auto".to_string()
     };
-    let (cluster, _) = crate::config::select_cluster(config, &choice, recipe.required_nodes)?;
+    let (cluster, _) =
+        crate::config::select_cluster(config, &choice, recipe.required_nodes, &HashSet::new())?;
 
     let manager = TmuxManager {
         session_prefix: recipe_name.to_string(),
@@ -236,7 +240,8 @@ pub async fn tmux_attach(
     } else {
         "auto".to_string()
     };
-    let (cluster, _) = crate::config::select_cluster(config, &choice, recipe.required_nodes)?;
+    let (cluster, _) =
+        crate::config::select_cluster(config, &choice, recipe.required_nodes, &HashSet::new())?;
 
     let head_node = cluster
         .nodes

+ 3 - 1
src/tmux/serve.rs

@@ -1,4 +1,5 @@
 use anyhow::{Context, Result};
+use std::collections::HashSet;
 use std::time::Duration;
 use tokio::time::sleep;
 use tracing::{info, warn};
@@ -35,7 +36,8 @@ pub async fn tmux_serve(
     } else {
         "auto".to_string()
     };
-    let (cluster, nccl) = crate::config::select_cluster(config, &choice, recipe.required_nodes)?;
+    let (cluster, nccl) =
+        crate::config::select_cluster(config, &choice, recipe.required_nodes, &HashSet::new())?;
 
     if let Some(required) = recipe.required_nodes {
         let actual = cluster.nodes.len() as u32;