Переглянути джерело

feat: prefix matching for cluster names (--cluster 2 → 2nodes-dac)

Instead of typing the full cluster name you can use any unambiguous
prefix.  Resolution order: exact match first, then unique prefix.
If the prefix matches zero or multiple clusters, fail fast with a
clear error message listing the conflicting names.

Also applied to select_cluster_by_name (used by power/exec cmds).
Alexander Korolev 4 місяців тому
батько
коміт
17f619fb30
1 змінених файлів з 39 додано та 55 видалено
  1. 39 55
      src/config.rs

+ 39 - 55
src/config.rs

@@ -1,10 +1,8 @@
+pub use crate::pre_serve::{BeforeServeTable, OnFailure, PreServeCommand};
 use anyhow::{Context, Result};
 use serde::Deserialize;
 use std::collections::HashMap;
 use tracing::info;
-
-pub use crate::pre_serve::{BeforeServeTable, OnFailure, PreServeCommand};
-
 #[derive(Clone, Deserialize)]
 pub struct Node {
     pub name: String,
@@ -53,7 +51,6 @@ pub struct Config {
     #[serde(default)]
     pub home_assistant: Option<HomeAssistantConfig>,
 }
-
 /// Home Assistant config for `power on`/`power off`.
 #[derive(Clone, Deserialize)]
 pub struct HomeAssistantConfig {
@@ -62,18 +59,15 @@ pub struct HomeAssistantConfig {
     #[serde(default = "default_shutdown_delay")]
     pub shutdown_delay_seconds: u64,
 }
-
 fn default_shutdown_delay() -> u64 {
     30
 }
-
 fn default_init_port() -> u16 {
     25000
 }
 fn default_port() -> u16 {
     30000
 }
-
 impl Config {
     pub fn resolved_hf_cache_dir(&self) -> String {
         if let Some(ref dir) = self.hf_cache_dir {
@@ -87,7 +81,6 @@ impl Config {
         format!("{}/.cache/huggingface", home)
     }
 }
-
 #[derive(Clone, Deserialize)]
 pub struct Recipe {
     #[serde(default = "default_engine")]
@@ -122,23 +115,42 @@ pub struct Recipe {
     #[serde(default)]
     pub before_serve: Vec<PreServeCommand>,
 }
-
+/// Resolves cluster by exact name, then unique prefix. Errors if none or ambiguous.
+fn resolve_cluster_name<'a>(clusters: &'a [Cluster], name: &str) -> Result<&'a Cluster> {
+    if let Some(c) = clusters.iter().find(|c| c.name == name) {
+        return Ok(c);
+    }
+    let matches: Vec<&Cluster> = clusters
+        .iter()
+        .filter(|c| c.name.starts_with(name))
+        .collect();
+    match matches.len() {
+        0 => anyhow::bail!("Кластер {} не найден", name),
+        1 => Ok(matches[0]),
+        n => anyhow::bail!(
+            "Префикс '{}' соответствует {} кластерам: {}",
+            name,
+            n,
+            matches
+                .iter()
+                .map(|c| c.name.as_str())
+                .collect::<Vec<_>>()
+                .join(", ")
+        ),
+    }
+}
 fn default_engine() -> String {
     "vllm".to_string()
 }
-
 fn default_distributed_backend() -> String {
     "mp".to_string()
 }
-
 fn default_shm_size() -> String {
     String::new()
 }
-
 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`.
 pub fn select_cluster<'c>(
@@ -147,27 +159,20 @@ pub fn select_cluster<'c>(
     required_nodes: Option<u32>,
 ) -> Result<(Cluster, &'c NcclConfig)> {
     if cluster_name != "auto" {
-        let cluster = config
-            .clusters
-            .iter()
-            .find(|c| c.name == cluster_name)
-            .context(format!("Кластер {} не найден", cluster_name))?;
+        let cluster = resolve_cluster_name(&config.clusters, cluster_name)?;
         let nccl = config
             .nccl
             .get(&cluster.name)
             .context("NCCL конфиг не найден")?;
         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!(
             "✅ Автоподбор кластера: {} ({} узлов)",
@@ -176,20 +181,12 @@ pub fn select_cluster<'c>(
         );
         return Ok((cluster.clone(), nccl));
     }
-
     anyhow::bail!("Не удалось подобрать подходящий кластер");
 }
-
-/// Looks up a cluster by name without requiring NCCL config.
+/// Looks up a cluster by name (or unique prefix) without requiring NCCL config.
 pub fn select_cluster_by_name(config: &Config, cluster_name: &str) -> Result<Cluster> {
-    let cluster = config
-        .clusters
-        .iter()
-        .find(|c| c.name == cluster_name)
-        .context(format!("Кластер {} не найден", cluster_name))?;
-    Ok(cluster.clone())
+    Ok(resolve_cluster_name(&config.clusters, cluster_name)?.clone())
 }
-
 /// Resolves the cluster for a recipe, applying recipe-level cluster pinning.
 pub fn resolve_cluster<'c>(
     config: &'c Config,
@@ -205,7 +202,6 @@ pub fn resolve_cluster<'c>(
     };
     select_cluster(config, &choice, recipe.required_nodes)
 }
-
 /// Loads a recipe (cwd first, then `<default_config_dir>/recipes/<name>/recipe.toml`,
 /// falling back to cwd for the error).
 pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
@@ -223,13 +219,11 @@ pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
     let content = tokio::fs::read_to_string(&cwd_path).await?;
     Ok(toml::from_str(&content)?)
 }
-
 /// OS-appropriate default config dir (e.g. `~/.config/sm121-launcher`),
 /// or `None` for sandboxes with no known config dir.
 pub fn default_config_dir() -> Option<std::path::PathBuf> {
     dirs::config_dir().map(|p| p.join("sm121-launcher"))
 }
-
 /// Resolves the cluster for a recipe (stop/logs flow).
 pub async fn resolve_cluster_for_recipe(
     config: &Config,
@@ -239,11 +233,9 @@ pub async fn resolve_cluster_for_recipe(
     let (cluster, _) = resolve_cluster(config, &recipe, "auto")?;
     Ok((cluster, recipe))
 }
-
 #[cfg(test)]
 mod tests {
     use super::*;
-
     fn test_cluster() -> Cluster {
         Cluster {
             name: "test-cluster".to_string(),
@@ -269,7 +261,6 @@ mod tests {
             extra_entities: vec![],
         }
     }
-
     fn test_config() -> Config {
         Config {
             ssh_user: "u".to_string(),
@@ -283,7 +274,6 @@ mod tests {
             home_assistant: None,
         }
     }
-
     fn test_nccl() -> NcclConfig {
         NcclConfig {
             ib_hca: "rocep1s0f0".to_string(),
@@ -292,7 +282,6 @@ mod tests {
             ib_disable: 0,
         }
     }
-
     // ── Deserialization tests ──────────────────────────────────
     #[test]
     fn test_recipe_deserialize_minimal() {
@@ -310,7 +299,6 @@ model = "test/model"
         assert_eq!(recipe.distributed_backend, "mp");
         assert_eq!(recipe.tensor_parallel_size, 1);
     }
-
     #[test]
     fn test_recipe_deserialize_full() {
         let toml_str = r#"
@@ -323,9 +311,7 @@ required_nodes = 4
 pipeline_parallel_size = 2
 distributed_backend = "ray"
 cluster = "my-cluster"
-
 args = ["--arg1", "val1", "--arg2"]
-
 [env]
 KEY = "value"
 OTHER = "123"
@@ -345,7 +331,6 @@ OTHER = "123"
         assert_eq!(recipe.env.get("KEY").unwrap(), "value");
         assert_eq!(recipe.env.get("OTHER").unwrap(), "123");
     }
-
     #[test]
     fn test_recipe_deserialize_no_required_nodes() {
         let toml_str = r#"
@@ -356,13 +341,11 @@ model = "test/model"
         assert_eq!(recipe.required_nodes, None);
         assert_eq!(recipe.pipeline_parallel_size, None);
     }
-
     #[test]
     fn test_config_deserialize() {
         let toml_str = r#"
 ssh_user = "admin"
 ssh_key_path = "/home/admin/.ssh/id_rsa"
-
 [[clusters]]
 name = "prod"
 master_ib_ip = "10.0.0.1"
@@ -370,7 +353,6 @@ nodes = [
     { name = "node1", ssh_ip = "10.0.0.10", ib_ip = "10.0.0.1", rank = 0 },
     { name = "node2", ssh_ip = "10.0.0.11", ib_ip = "10.0.0.2", rank = 1 },
 ]
-
 [nccl."prod"]
 ib_hca = "mlx5_0"
 socket_ifname = "eth0"
@@ -385,13 +367,11 @@ ib_disable = 0
         assert!(config.nccl.contains_key("prod"));
         assert_eq!(config.nccl["prod"].ib_hca, "mlx5_0");
     }
-
     #[test]
     fn test_config_deserialize_with_host_key() {
         let toml_str = r#"
 ssh_user = "admin"
 ssh_key_path = "/path/to/key"
-
 [[clusters]]
 name = "secure"
 master_ib_ip = "1.1.1.1"
@@ -419,7 +399,6 @@ nodes = [
             home_assistant: None,
         };
         config.nccl.insert("test-cluster".to_string(), test_nccl());
-
         let (cluster, nccl) = select_cluster(&config, "test-cluster", None).unwrap();
         assert_eq!(cluster.name, "test-cluster");
         assert_eq!(nccl.ib_hca, "rocep1s0f0");
@@ -438,7 +417,6 @@ nodes = [
             home_assistant: None,
         };
         config.nccl.insert("test-cluster".to_string(), test_nccl());
-
         let (cluster, nccl) = select_cluster(&config, "auto", None).unwrap();
         assert_eq!(cluster.name, "test-cluster");
         assert_eq!(nccl.ib_hca, "rocep1s0f0");
@@ -449,7 +427,6 @@ nodes = [
         let result = select_cluster(&config, "auto", None);
         assert!(result.is_err());
     }
-
     #[test]
     fn test_select_cluster_not_found() {
         let mut config = test_config();
@@ -473,21 +450,28 @@ nodes = [
         let result = select_cluster_by_name(&config, "no-such-cluster");
         assert!(result.is_err());
     }
-    // ── Defaults ───────────────────────────────────────────────
-
+    #[test]
+    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());
+    }
     #[test]
     fn test_default_engine_is_vllm() {
         let recipe: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"").unwrap();
         assert_eq!(recipe.engine, "vllm");
     }
-
     #[test]
     fn test_default_distributed_backend_is_mp() {
         let recipe: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"").unwrap();
         assert_eq!(recipe.distributed_backend, "mp");
     }
     // ── Default config dir ─────────────────────────────────────
-
     #[test]
     fn default_config_dir_appends_sm121_launcher() {
         let Some(dir) = default_config_dir() else {