Procházet zdrojové kódy

feat: auto-select smallest cluster matching required_nodes

Previously, auto-select picked the first cluster in config.toml
that had an NCCL config.  Now it picks the smallest cluster (fewest
nodes) whose node count exactly matches the recipe's required_nodes.

- select_cluster() now takes a required_nodes: Option<u32> parameter
  (exact match, since required_nodes = 2 means exactly 2 nodes).
- When cluster_name is 'auto', builds a sorted list of candidates
  and picks the smallest.
- Explicit cluster names bypass the filter (same as before).
- All callers updated: resolve_cluster, tmux_* ops, exec recipe path.

Example: with a 2-node and 4-node cluster in config, a recipe with
required_nodes=2 auto-selects the 2-node cluster (previously would
always pick whichever came first in the config).
Alexander Korolev před 4 měsíci
rodič
revize
b0c139ea80
5 změnil soubory, kde provedl 46 přidání a 34 odebrání
  1. 15 3
      recipes/deepseek-v4-flash/recipe.toml
  2. 25 25
      src/config.rs
  3. 1 1
      src/exec.rs
  4. 4 4
      src/tmux/ops.rs
  5. 1 1
      src/tmux/serve.rs

+ 15 - 3
recipes/deepseek-v4-flash/recipe.toml

@@ -11,6 +11,8 @@ distributed_backend = "mp"
 # default 10g.  KV cache + activations for 1M tokens fills up fast.
 shm_size = "64gb"
 
+entrypoint = ["vllm", "serve"]
+
 args = [
     "--kv-cache-dtype", "fp8",
     "--block-size", "256",
@@ -19,12 +21,14 @@ args = [
     "--max-num-batched-tokens", "8192",
     "--gpu-memory-utilization", "0.82",
     "--enable-prefix-caching",
-    "--speculative-config", '{"method":"mtp","num_speculative_tokens":2}',
+    "--speculative-config", """
+'{"method":"mtp","num_speculative_tokens":2}'""",
     "--tokenizer-mode", "deepseek_v4",
     "--tool-call-parser", "deepseek_v4",
     "--enable-auto-tool-choice",
     "--reasoning-parser", "deepseek_v4",
-    "--reasoning-config", '{"reasoning_parser":"deepseek_v4","reasoning_start_str":"","reasoning_end_str":""}',
+    "--reasoning-config", """
+'{"reasoning_parser":"deepseek_v4","reasoning_start_str":"","reasoning_end_str":""}'""",
     "--default-chat-template-kwargs.thinking", "true",
     "--default-chat-template-kwargs.reasoning_effort", "high",
     "--enable-flashinfer-autotune",
@@ -34,7 +38,6 @@ args = [
 
 [env]
 # HuggingFace / cache config
-HF_HUB_OFFLINE = "0"
 VLLM_CACHE_ROOT = "/cache/huggingface/vllm-cache"
 
 # vLLM features for DeepSeek V4 on Blackwell (GB10 / sm_120)
@@ -45,3 +48,12 @@ VLLM_SPARSE_INDEXER_MAX_LOGITS_MB = "256"
 # Compile-time arch hints for PyTorch and FlashInfer on GB10
 TORCH_CUDA_ARCH_LIST = "12.1a"
 FLASHINFER_CUDA_ARCH_LIST = "12.1a"
+
+CUDA_HOME = "/opt/env/targets/sbsa-linux"
+FLASHINFER_DISABLE_VERSION_CHECK = "1"
+TILELANG_CLEANUP_TEMP_FILES = "1"
+DG_JIT_USE_NVRTC = "0"
+DG_JIT_PRINT_COMPILER_COMMAND = "1"
+VLLM_NCCL_SO_PATH = "/opt/env/lib/python3.12/site-packages/nvidia/nccl/lib/libnccl.so.2"
+HF_HUB_OFFLINE = "1"
+TRANSFORMERS_OFFLINE = "1"

+ 25 - 25
src/config.rs

@@ -16,7 +16,6 @@ pub struct Node {
     #[serde(default)]
     pub entity: Option<String>,
 }
-
 #[derive(Clone, Deserialize)]
 pub struct NcclConfig {
     pub ib_hca: String,
@@ -24,19 +23,15 @@ pub struct NcclConfig {
     pub ib_gid_index: u32,
     pub ib_disable: u32,
 }
-
 #[derive(Clone, Deserialize)]
 pub struct Cluster {
     pub name: String,
     pub master_ib_ip: String,
     pub nodes: Vec<Node>,
-    /// Extra Home Assistant switches to toggle alongside the nodes
-    /// during `power on`/`power off`.  E.g. a network switch that powers
-    /// the cluster.  Toggled together with per-node entities.
+    /// Extra Home Assistant switches to toggle alongside the nodes during power on/off.
     #[serde(default)]
     pub extra_entities: Vec<String>,
 }
-
 /// Top-level config (clusters + per-cluster pre-serve commands).
 #[derive(Clone, Deserialize)]
 pub struct Config {
@@ -144,10 +139,12 @@ fn default_tensor_parallel_size() -> u32 {
     1
 }
 
-/// Selects a cluster + NCCL config by name or auto-picks the first one with NCCL.
+/// Selects a cluster + NCCL config by name or auto-picks the smallest
+/// cluster with NCCL config satisfying `required_nodes`.
 pub fn select_cluster<'c>(
     config: &'c Config,
     cluster_name: &str,
+    required_nodes: Option<u32>,
 ) -> Result<(Cluster, &'c NcclConfig)> {
     if cluster_name != "auto" {
         let cluster = config
@@ -162,11 +159,22 @@ pub fn select_cluster<'c>(
         return Ok((cluster.clone(), nccl));
     }
 
-    for cluster in &config.clusters {
-        if let Some(nccl) = config.nccl.get(&cluster.name) {
-            info!("✅ Автоподбор кластера: {}", cluster.name);
-            return Ok((cluster.clone(), nccl));
-        }
+    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))
+        .collect();
+
+    candidates.sort_by_key(|(c, _)| c.nodes.len());
+
+    if let Some((cluster, nccl)) = candidates.into_iter().next() {
+        info!(
+            "✅ Автоподбор кластера: {} ({} узлов)",
+            cluster.name,
+            cluster.nodes.len()
+        );
+        return Ok((cluster.clone(), nccl));
     }
 
     anyhow::bail!("Не удалось подобрать подходящий кластер");
@@ -195,7 +203,7 @@ pub fn resolve_cluster<'c>(
     } else {
         "auto".to_string()
     };
-    select_cluster(config, &choice)
+    select_cluster(config, &choice, recipe.required_nodes)
 }
 
 /// Loads a recipe (cwd first, then `<default_config_dir>/recipes/<name>/recipe.toml`,
@@ -286,7 +294,6 @@ mod tests {
     }
 
     // ── Deserialization tests ──────────────────────────────────
-
     #[test]
     fn test_recipe_deserialize_minimal() {
         let toml_str = r#"
@@ -398,7 +405,6 @@ nodes = [
             Some("SHA256:abc123".to_string())
         );
     }
-
     #[test]
     fn test_select_cluster_by_name() {
         let mut config = Config {
@@ -414,11 +420,10 @@ nodes = [
         };
         config.nccl.insert("test-cluster".to_string(), test_nccl());
 
-        let (cluster, nccl) = select_cluster(&config, "test-cluster").unwrap();
+        let (cluster, nccl) = select_cluster(&config, "test-cluster", None).unwrap();
         assert_eq!(cluster.name, "test-cluster");
         assert_eq!(nccl.ib_hca, "rocep1s0f0");
     }
-
     #[test]
     fn test_select_cluster_auto() {
         let mut config = Config {
@@ -434,15 +439,14 @@ nodes = [
         };
         config.nccl.insert("test-cluster".to_string(), test_nccl());
 
-        let (cluster, nccl) = select_cluster(&config, "auto").unwrap();
+        let (cluster, nccl) = select_cluster(&config, "auto", None).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");
+        let result = select_cluster(&config, "auto", None);
         assert!(result.is_err());
     }
 
@@ -450,10 +454,9 @@ nodes = [
     fn test_select_cluster_not_found() {
         let mut config = test_config();
         config.clusters = vec![];
-        let result = select_cluster(&config, "nonexistent");
+        let result = select_cluster(&config, "nonexistent", None);
         assert!(result.is_err());
     }
-
     #[test]
     fn test_select_cluster_by_name_found() {
         let mut nccl_map = HashMap::new();
@@ -464,14 +467,12 @@ nodes = [
         assert_eq!(cluster.name, "test-cluster");
         assert_eq!(cluster.nodes.len(), 2);
     }
-
     #[test]
     fn test_select_cluster_by_name_not_found() {
         let config = test_config();
         let result = select_cluster_by_name(&config, "no-such-cluster");
         assert!(result.is_err());
     }
-
     // ── Defaults ───────────────────────────────────────────────
 
     #[test]
@@ -485,7 +486,6 @@ nodes = [
         let recipe: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"").unwrap();
         assert_eq!(recipe.distributed_backend, "mp");
     }
-
     // ── Default config dir ─────────────────────────────────────
 
     #[test]

+ 1 - 1
src/exec.rs

@@ -24,7 +24,7 @@ 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)
+            let (cluster, _) = select_cluster(config, cname, recipe.required_nodes)
                 .with_context(|| format!("resolving cluster for recipe '{}'", rname))?;
             Ok((cluster, Some(recipe)))
         }

+ 4 - 4
src/tmux/ops.rs

@@ -22,7 +22,7 @@ pub async fn tmux_stop(
     } else {
         "auto".to_string()
     };
-    let (cluster, _) = crate::config::select_cluster(config, &choice)?;
+    let (cluster, _) = crate::config::select_cluster(config, &choice, recipe.required_nodes)?;
 
     let manager = TmuxManager {
         session_prefix: recipe_name.to_string(),
@@ -73,7 +73,7 @@ pub async fn tmux_status(
     } else {
         "auto".to_string()
     };
-    let (cluster, _) = crate::config::select_cluster(config, &choice)?;
+    let (cluster, _) = crate::config::select_cluster(config, &choice, recipe.required_nodes)?;
 
     let manager = TmuxManager {
         session_prefix: recipe_name.to_string(),
@@ -162,7 +162,7 @@ pub async fn tmux_logs(
     } else {
         "auto".to_string()
     };
-    let (cluster, _) = crate::config::select_cluster(config, &choice)?;
+    let (cluster, _) = crate::config::select_cluster(config, &choice, recipe.required_nodes)?;
 
     let manager = TmuxManager {
         session_prefix: recipe_name.to_string(),
@@ -236,7 +236,7 @@ pub async fn tmux_attach(
     } else {
         "auto".to_string()
     };
-    let (cluster, _) = crate::config::select_cluster(config, &choice)?;
+    let (cluster, _) = crate::config::select_cluster(config, &choice, recipe.required_nodes)?;
 
     let head_node = cluster
         .nodes

+ 1 - 1
src/tmux/serve.rs

@@ -30,7 +30,7 @@ pub async fn tmux_serve(
     } else {
         "auto".to_string()
     };
-    let (cluster, nccl) = crate::config::select_cluster(config, &choice)?;
+    let (cluster, nccl) = crate::config::select_cluster(config, &choice, recipe.required_nodes)?;
 
     if let Some(required) = recipe.required_nodes {
         let actual = cluster.nodes.len() as u32;