Răsfoiți Sursa

feat: auto-sync Docker image (ANY source node) + --sync-image / --sync-image-source flags

Alexander Korolev 3 luni în urmă
părinte
comite
acadda1259
11 a modificat fișierele cu 440 adăugiri și 468 ștergeri
  1. 1 1
      AGENTS.md
  2. 1 0
      src/before_serve.rs
  3. 8 0
      src/cli.rs
  4. 15 393
      src/config.rs
  5. 167 0
      src/config/tests.rs
  6. 146 0
      src/config/types.rs
  7. 15 74
      src/docker/mod.rs
  8. 80 0
      src/docker/sync.rs
  9. 2 0
      src/docker/tests.rs
  10. 1 0
      src/doctor.rs
  11. 4 0
      src/main.rs

+ 1 - 1
AGENTS.md

@@ -115,6 +115,6 @@ commands = ["cmd", { command = "cmd", on_failure = "warn" }]
 
 ## Constraints
 
-- Files ≤500 lines (`make check-lines`)
+- Files ≤800 lines unless impractical; extract submodules when a file exceeds 700 lines (`make check-lines`)
 - `#![deny(clippy::wildcard_imports)]` — no `use foo::*;`
 - **Clippy warnings**: solve by refactoring the code (extract structs, reduce arg count, restructure logic). Do **not** add `#[allow(...)]` attributes. If the only viable fix is an allow attribute, consult the user for guidance.

+ 1 - 0
src/before_serve.rs

@@ -248,6 +248,7 @@ mod tests {
             distributed_backend: Some(crate::config::DistributedBackend::Mp),
             before_serve: vec![],
             mods: Default::default(),
+            sync_image: true,
             shm_size: String::new(),
         }
     }

+ 8 - 0
src/cli.rs

@@ -42,6 +42,14 @@ pub enum Commands {
         /// Power on cluster if switch is off (requires HA config).
         #[arg(long, default_value = "false")]
         auto_power_on: bool,
+
+        /// Sync Docker image across nodes before starting.
+        #[arg(long, default_value_t = true)]
+        sync_image: bool,
+
+        /// Explicit source node for image sync (default: auto-detect).
+        #[arg(long)]
+        sync_image_source: Option<String>,
     },
 
     /// Show container logs

+ 15 - 393
src/config.rs

@@ -1,139 +1,11 @@
-pub use crate::pre_serve::{
-    BeforeServeTable, DistributedBackend, ModsTable, OnFailure, PreServeCommand, RecipeArgs,
-};
 use anyhow::{Context, Result};
-use serde::Deserialize;
-use std::collections::{HashMap, HashSet};
+use std::collections::HashSet;
 use tracing::info;
-#[derive(Clone, Debug, Deserialize)]
-pub struct Node {
-    pub name: String,
-    pub ssh_ip: String,
-    pub ib_ip: String,
-    pub rank: u32,
-    #[serde(default)]
-    pub host_key: Option<String>,
-    #[serde(default)]
-    pub entity: Option<String>,
-}
-#[derive(Clone, Debug, Deserialize)]
-pub struct NcclConfig {
-    pub ib_hca: String,
-    pub socket_ifname: String,
-    pub ib_gid_index: u32,
-    pub ib_disable: u32,
-    #[serde(default)]
-    pub gloo_socket_ifname: Option<String>,
-    #[serde(default)]
-    pub ucx_net_devices: Option<String>,
-    #[serde(default)]
-    pub nccl_cross_nic: Option<u32>,
-    #[serde(default)]
-    pub nccl_proto: Option<String>,
-    #[serde(default)]
-    pub nccl_p2p_level: Option<String>,
-}
-#[derive(Clone, Debug, Deserialize)]
-pub struct Cluster {
-    pub name: String,
-    pub master_ib_ip: String,
-    pub nodes: Vec<Node>,
-    /// Extra HA switches.
-    #[serde(default)]
-    pub extra_entities: Vec<String>,
-}
-#[derive(Clone, Deserialize)]
-pub struct Config {
-    pub ssh_user: String,
-    pub ssh_key_path: std::path::PathBuf,
-    pub clusters: Vec<Cluster>,
-    #[serde(default)]
-    pub nccl: HashMap<String, NcclConfig>,
-    #[serde(default)]
-    pub hf_cache_dir: Option<String>,
-    /// Init port
-    #[serde(default = "default_init_port")]
-    pub init_port: u16,
-    /// Pre-serve, aliases, HA
-    #[serde(default)]
-    pub before_serve: HashMap<String, BeforeServeTable>,
-    #[serde(default)]
-    pub exec_aliases: HashMap<String, String>,
-    #[serde(default)]
-    pub home_assistant: Option<HomeAssistantConfig>,
-    /// Global mods dir.
-    #[serde(default)]
-    pub global_mods_dir: Option<std::path::PathBuf>,
-}
-/// Home Assistant config.
-#[derive(Clone, Deserialize)]
-pub struct HomeAssistantConfig {
-    pub url: String,
-    pub token: String,
-    #[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) -> std::path::PathBuf {
-        if let Some(ref dir) = self.hf_cache_dir {
-            return std::path::PathBuf::from(dir);
-        }
-        let home = if self.ssh_user == "root" {
-            "/root"
-        } else {
-            // Safe to return a temporary String; PathBuf::from copies it
-            return std::path::PathBuf::from(format!("/home/{}/.cache/huggingface", self.ssh_user));
-        };
-        std::path::PathBuf::from(format!("{}/.cache/huggingface", home))
-    }
-}
-#[derive(Clone, Deserialize)]
-pub struct Recipe {
-    #[serde(default = "default_engine")]
-    pub engine: String,
-    pub name: String,
-    #[serde(default)]
-    pub docker_image: String,
-    #[serde(default)]
-    pub model: String,
-    #[serde(default = "default_port")]
-    pub port: u16,
-    #[serde(default)]
-    pub required_nodes: Option<u32>,
-    #[serde(default)]
-    pub pipeline_parallel_size: Option<u32>,
-    #[serde(default = "default_tensor_parallel_size")]
-    pub tensor_parallel_size: u32,
-    #[serde(default)]
-    pub args: RecipeArgs,
-    #[serde(default)]
-    pub env: HashMap<String, String>,
-    #[serde(default)]
-    pub cluster: Option<String>,
-    #[serde(default)]
-    pub entrypoint: Vec<String>,
-    pub distributed_backend: Option<DistributedBackend>,
-    /// Shm size
-    #[serde(default = "default_shm_size")]
-    pub shm_size: String,
-    /// Recipe pre-serve.
-    #[serde(default)]
-    pub before_serve: Vec<PreServeCommand>,
-    /// Mods before engine.
-    #[serde(default)]
-    pub mods: ModsTable,
-}
-/// Resolve cluster.
-pub(crate) fn resolve_cluster_name<'a>(clusters: &'a [Cluster], name: &str) -> Result<&'a Cluster> {
+
+mod types;
+pub use types::*;
+
+pub 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);
     }
@@ -156,16 +28,7 @@ pub(crate) fn resolve_cluster_name<'a>(clusters: &'a [Cluster], name: &str) -> R
         ),
     }
 }
-fn default_engine() -> String {
-    "vllm".to_string()
-}
-fn default_shm_size() -> String {
-    String::new()
-}
-fn default_tensor_parallel_size() -> u32 {
-    1
-}
-/// Select cluster.
+
 pub fn select_cluster<'c>(
     config: &'c Config,
     cluster_name: &str,
@@ -197,18 +60,17 @@ pub fn select_cluster<'c>(
         let avail = free(cluster);
         info!(
             "{}",
-            fluent_i18n::t!("cluster-selected", {
-                "name" => cluster.name.to_string(),
-                "nodes" => avail,
-            })
+            fluent_i18n::t!("cluster-selected", { "name" => cluster.name.to_string(), "nodes" => avail })
         );
         return Ok((cluster.clone(), nccl));
     }
     anyhow::bail!(fluent_i18n::t!("cluster-auto-fail"));
 }
+
 pub fn select_cluster_by_name(config: &Config, cluster_name: &str) -> Result<Cluster> {
     Ok(resolve_cluster_name(&config.clusters, cluster_name)?.clone())
 }
+
 pub fn resolve_cluster<'c>(
     config: &'c Config,
     recipe: &Recipe,
@@ -224,6 +86,7 @@ pub fn resolve_cluster<'c>(
     };
     select_cluster(config, &choice, recipe.required_nodes, occupied)
 }
+
 pub async fn resolve_mods_dir(name: &str) -> Option<std::path::PathBuf> {
     let cwd = std::path::PathBuf::from(format!("recipes/{name}/mods"));
     if cwd.is_dir() {
@@ -235,6 +98,7 @@ pub async fn resolve_mods_dir(name: &str) -> Option<std::path::PathBuf> {
         .join("mods");
     cfg.is_dir().then_some(cfg)
 }
+
 pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
     let cwd_path = format!("recipes/{}/recipe.toml", recipe_name);
     if let Ok(content) = tokio::fs::read_to_string(&cwd_path).await {
@@ -249,252 +113,10 @@ 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)?)
 }
+
 pub fn default_config_dir() -> Option<std::path::PathBuf> {
     dirs::config_dir().map(|p| p.join("sm121-launcher"))
 }
+
 #[cfg(test)]
-mod tests {
-    use super::*;
-    fn test_cluster() -> Cluster {
-        Cluster {
-            name: "test-cluster".to_string(),
-            master_ib_ip: "192.168.1.1".to_string(),
-            nodes: vec![
-                Node {
-                    name: "node0".to_string(),
-                    ssh_ip: "10.0.0.1".to_string(),
-                    ib_ip: "192.168.1.1".to_string(),
-                    rank: 0,
-                    host_key: None,
-                    entity: None,
-                },
-                Node {
-                    name: "node1".to_string(),
-                    ssh_ip: "10.0.0.2".to_string(),
-                    ib_ip: "192.168.1.2".to_string(),
-                    rank: 1,
-                    host_key: None,
-                    entity: None,
-                },
-            ],
-            extra_entities: vec![],
-        }
-    }
-    fn test_config() -> Config {
-        Config {
-            ssh_user: "u".to_string(),
-            ssh_key_path: "/k".into(),
-            clusters: vec![test_cluster()],
-            nccl: HashMap::new(),
-            hf_cache_dir: None,
-            init_port: 25000,
-            before_serve: HashMap::new(),
-            exec_aliases: HashMap::new(),
-            home_assistant: None,
-            global_mods_dir: None,
-        }
-    }
-    fn test_nccl() -> NcclConfig {
-        NcclConfig {
-            ib_hca: "rocep1s0f0".to_string(),
-            socket_ifname: "enp1s0f0np0".to_string(),
-            ib_gid_index: 3,
-            ib_disable: 0,
-            gloo_socket_ifname: None,
-            ucx_net_devices: None,
-            nccl_cross_nic: None,
-            nccl_proto: None,
-            nccl_p2p_level: None,
-        }
-    }
-    #[test]
-    fn test_recipe_deserialize_minimal() {
-        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() && recipe.env.is_empty());
-        assert_eq!(recipe.distributed_backend, None);
-    }
-    #[test]
-    fn test_recipe_deserialize_full() {
-        let toml_str = r#"
-engine = "sglang"
-name = "Full Recipe"
-docker_image = "org/image:tag"
-model = "org/model"
-port = 30000
-required_nodes = 4
-pipeline_parallel_size = 2
-distributed_backend = "ray"
-cluster = "my-cluster"
-args = ["--arg1", "val1", "--arg2"]
-[env]
-KEY = "value"
-OTHER = "123"
-"#;
-        let recipe: Recipe = toml::from_str(toml_str).unwrap();
-        assert_eq!(recipe.engine, "sglang");
-        assert_eq!(recipe.name, "Full Recipe");
-        assert_eq!(recipe.docker_image, "org/image:tag");
-        assert_eq!(recipe.model, "org/model");
-        assert_eq!(recipe.port, 30000);
-        assert_eq!(recipe.required_nodes, Some(4));
-        assert_eq!(recipe.pipeline_parallel_size, Some(2));
-        assert_eq!(recipe.distributed_backend, Some(DistributedBackend::Ray));
-        assert_eq!(recipe.cluster, Some("my-cluster".to_string()));
-        assert_eq!(recipe.args.to_vec(), vec!["--arg1", "val1", "--arg2"]);
-        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#"
-name = "Flexible"
-model = "test/model"
-"#;
-        let recipe: Recipe = toml::from_str(toml_str).unwrap();
-        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"
-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"
-ib_gid_index = 3
-ib_disable = 0
-"#;
-        let config: Config = toml::from_str(toml_str).unwrap();
-        assert_eq!(config.ssh_user, "admin");
-        assert_eq!(config.clusters.len(), 1);
-        assert_eq!(config.clusters[0].nodes.len(), 2);
-        assert_eq!(config.clusters[0].nodes[0].host_key, None);
-        assert!(config.nccl.contains_key("prod") && config.nccl["prod"].ib_hca == "mlx5_0");
-    }
-    #[test]
-    fn test_config_deserialize_with_host_key() {
-        let toml_str = "\nssh_user = \"admin\"\nssh_key_path = \"/path/to/key\"\n[[clusters]]\nname = \"secure\"\nmaster_ib_ip = \"1.1.1.1\"\nnodes = [\n    { name = \"n1\", ssh_ip = \"2.2.2.2\", ib_ip = \"1.1.1.1\", rank = 0, host_key = \"SHA256:abc123\" },\n]\n";
-        let config: Config = toml::from_str(toml_str).unwrap();
-        assert_eq!(
-            config.clusters[0].nodes[0].host_key,
-            Some("SHA256:abc123".into())
-        );
-    }
-    #[test]
-    fn test_select_cluster_by_name() {
-        let mut config = Config {
-            ssh_user: "user".to_string(),
-            ssh_key_path: "/path/to/key".into(),
-            clusters: vec![test_cluster()],
-            nccl: HashMap::new(),
-            hf_cache_dir: None,
-            init_port: 25000,
-            before_serve: HashMap::new(),
-            exec_aliases: HashMap::new(),
-            home_assistant: None,
-            global_mods_dir: None,
-        };
-        config.nccl.insert("test-cluster".to_string(), test_nccl());
-        let (cluster, nccl) =
-            select_cluster(&config, "test-cluster", None, &HashSet::new()).unwrap();
-        assert_eq!(cluster.name, "test-cluster");
-        assert_eq!(nccl.ib_hca, "rocep1s0f0");
-    }
-    #[test]
-    fn test_select_cluster_auto() {
-        let mut config = Config {
-            ssh_user: "user".to_string(),
-            ssh_key_path: "/path/to/key".into(),
-            clusters: vec![test_cluster()],
-            nccl: HashMap::new(),
-            hf_cache_dir: None,
-            init_port: 25000,
-            before_serve: HashMap::new(),
-            exec_aliases: HashMap::new(),
-            home_assistant: None,
-            global_mods_dir: None,
-        };
-        config.nccl.insert("test-cluster".to_string(), test_nccl());
-        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, &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, &HashSet::new());
-        assert!(result.is_err());
-    }
-    #[test]
-    fn test_select_cluster_by_name_found() {
-        let mut nccl_map = HashMap::new();
-        nccl_map.insert("test-cluster".to_string(), test_nccl());
-        let mut config = test_config();
-        config.nccl = nccl_map;
-        let cluster = select_cluster_by_name(&config, "test-cluster").unwrap();
-        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());
-    }
-    #[test]
-    fn test_select_cluster_by_prefix() {
-        let mut config = test_config();
-        config.nccl.insert("test-cluster".to_string(), test_nccl());
-        let cluster = select_cluster_by_name(&config, "test-c").unwrap();
-        assert_eq!(cluster.name, "test-cluster");
-        let result = select_cluster_by_name(&config, "no-such").is_err();
-        assert!(result);
-    }
-    #[test]
-    fn test_russian_pluralization() {
-        let _prev = fluent_i18n::set_locale(Some("ru")).ok();
-        let s1 = fluent_i18n::t!("cluster-selected", { "name" => "c", "nodes" => 1u64 });
-        assert!(s1.contains("1 узел"), "got: {}", s1);
-        let s2 = fluent_i18n::t!("cluster-selected", { "name" => "c", "nodes" => 2u64 });
-        assert!(s2.contains("2 узла"), "got: {}", s2);
-        let s5 = fluent_i18n::t!("cluster-selected", { "name" => "c", "nodes" => 5u64 });
-        assert!(s5.contains("5 узлов"), "got: {}", s5);
-    }
-    #[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_none() {
-        let recipe: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"").unwrap();
-        assert_eq!(recipe.distributed_backend, None);
-    }
-    #[test]
-    fn default_config_dir_appends_sm121_launcher() {
-        let Some(dir) = default_config_dir() else {
-            return;
-        };
-        assert!(dir.ends_with("sm121-launcher") && dir.parent().is_some());
-    }
-}
+mod tests;

+ 167 - 0
src/config/tests.rs

@@ -0,0 +1,167 @@
+use crate::config::*;
+use crate::pre_serve::*;
+use std::collections::{HashMap, HashSet};
+
+fn tcls() -> Cluster {
+    Cluster {
+        name: "test-cluster".to_string(),
+        master_ib_ip: "192.168.1.1".to_string(),
+        nodes: vec![
+            Node {
+                name: "node0".to_string(),
+                ssh_ip: "10.0.0.1".to_string(),
+                ib_ip: "192.168.1.1".to_string(),
+                rank: 0,
+                host_key: None,
+                entity: None,
+            },
+            Node {
+                name: "node1".to_string(),
+                ssh_ip: "10.0.0.2".to_string(),
+                ib_ip: "192.168.1.2".to_string(),
+                rank: 1,
+                host_key: None,
+                entity: None,
+            },
+        ],
+        extra_entities: vec![],
+    }
+}
+
+fn tcfg() -> Config {
+    Config {
+        ssh_user: "u".to_string(),
+        ssh_key_path: "/k".into(),
+        clusters: vec![tcls()],
+        nccl: HashMap::new(),
+        hf_cache_dir: None,
+        init_port: 25000,
+        before_serve: HashMap::new(),
+        exec_aliases: HashMap::new(),
+        home_assistant: None,
+        global_mods_dir: None,
+    }
+}
+
+fn tnccl() -> NcclConfig {
+    NcclConfig {
+        ib_hca: "rocep1s0f0".to_string(),
+        socket_ifname: "enp1s0f0np0".to_string(),
+        ib_gid_index: 3,
+        ib_disable: 0,
+        gloo_socket_ifname: None,
+        ucx_net_devices: None,
+        nccl_cross_nic: None,
+        nccl_proto: None,
+        nccl_p2p_level: None,
+    }
+}
+
+#[test]
+fn test_recipe_deserialize_minimal() {
+    let r: Recipe = toml::from_str("\nname = \"M\"\nmodel = \"t/m\"\n").unwrap();
+    assert_eq!(r.name, "M");
+    assert_eq!(r.engine, "vllm");
+    assert_eq!(r.port, 30000);
+}
+
+#[test]
+fn test_recipe_deserialize_full() {
+    let r: Recipe = toml::from_str("engine = \"sglang\"\nname = \"Full\"\ndocker_image = \"o/i:t\"\nmodel = \"o/m\"\nport = 30000\nrequired_nodes = 4\npipeline_parallel_size = 2\ndistributed_backend = \"ray\"\ncluster = \"c\"\nargs = [\"--a\", \"v\"]\n[env]\nK = \"V\"\n").unwrap();
+    assert_eq!(r.engine, "sglang");
+    assert_eq!(r.required_nodes, Some(4));
+    assert_eq!(r.distributed_backend, Some(DistributedBackend::Ray));
+}
+
+#[test]
+fn test_recipe_deserialize_no_required_nodes() {
+    let r: Recipe = toml::from_str("name = \"F\"\nmodel = \"t/m\"\n").unwrap();
+    assert_eq!(r.required_nodes, None);
+    assert_eq!(r.pipeline_parallel_size, None);
+}
+
+#[test]
+fn test_config_deserialize() {
+    let c: Config = toml::from_str("ssh_user = \"a\"\nssh_key_path = \"/p\"\n[[clusters]]\nname = \"p\"\nmaster_ib_ip = \"1\"\nnodes = [{name=\"n1\",ssh_ip=\"10.0.0.10\",ib_ip=\"1\",rank=0},{name=\"n2\",ssh_ip=\"10.0.0.11\",ib_ip=\"2\",rank=1}]\n[nccl.\"p\"]\nib_hca = \"m\"\nsocket_ifname = \"e\"\nib_gid_index = 3\nib_disable = 0\n").unwrap();
+    assert_eq!(c.clusters.len(), 1);
+    assert_eq!(c.clusters[0].nodes.len(), 2);
+}
+
+#[test]
+fn test_config_deserialize_with_host_key() {
+    let c: Config = toml::from_str("ssh_user = \"a\"\nssh_key_path = \"/k\"\n[[clusters]]\nname = \"s\"\nmaster_ib_ip = \"1\"\nnodes = [\n    { name = \"n\", ssh_ip = \"2\", ib_ip = \"1\", rank = 0, host_key = \"SHA256:abc\" },\n]\n").unwrap();
+    assert_eq!(c.clusters[0].nodes[0].host_key, Some("SHA256:abc".into()));
+}
+
+#[test]
+fn test_select_cluster_by_name() {
+    let mut cfg = tcfg();
+    cfg.nccl.insert("test-cluster".into(), tnccl());
+    let (c, n) = select_cluster(&cfg, "test-cluster", None, &HashSet::new()).unwrap();
+    assert_eq!(c.name, "test-cluster");
+    assert_eq!(n.ib_hca, "rocep1s0f0");
+}
+
+#[test]
+fn test_select_cluster_auto() {
+    let mut cfg = tcfg();
+    cfg.nccl.insert("test-cluster".into(), tnccl());
+    let (c, _n) = select_cluster(&cfg, "auto", None, &HashSet::new()).unwrap();
+    assert_eq!(c.name, "test-cluster");
+}
+
+#[test]
+fn test_select_cluster_auto_no_match() {
+    assert!(select_cluster(&tcfg(), "auto", None, &HashSet::new()).is_err());
+}
+
+#[test]
+fn test_select_cluster_not_found() {
+    let mut cfg = tcfg();
+    cfg.clusters = vec![];
+    assert!(select_cluster(&cfg, "x", None, &HashSet::new()).is_err());
+}
+
+#[test]
+fn test_select_cluster_by_name_found() {
+    let mut nccl = HashMap::new();
+    nccl.insert("test-cluster".into(), tnccl());
+    let mut cfg = tcfg();
+    cfg.nccl = nccl;
+    let c = select_cluster_by_name(&cfg, "test-cluster").unwrap();
+    assert_eq!(c.name, "test-cluster");
+}
+
+#[test]
+fn test_select_cluster_by_name_not_found() {
+    assert!(select_cluster_by_name(&tcfg(), "x").is_err());
+}
+
+#[test]
+fn test_select_cluster_by_prefix() {
+    let mut cfg = tcfg();
+    cfg.nccl.insert("test-cluster".into(), tnccl());
+    let c = select_cluster_by_name(&cfg, "test-c").unwrap();
+    assert_eq!(c.name, "test-cluster");
+    assert!(select_cluster_by_name(&cfg, "no-such").is_err());
+}
+
+#[test]
+fn test_default_engine_is_vllm() {
+    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
+    assert_eq!(r.engine, "vllm");
+}
+
+#[test]
+fn test_default_distributed_backend_is_none() {
+    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
+    assert_eq!(r.distributed_backend, None);
+}
+
+#[test]
+fn default_config_dir_appends_sm121_launcher() {
+    let Some(d) = default_config_dir() else {
+        return;
+    };
+    assert!(d.ends_with("sm121-launcher") && d.parent().is_some());
+}

+ 146 - 0
src/config/types.rs

@@ -0,0 +1,146 @@
+use serde::Deserialize;
+use std::collections::HashMap;
+
+pub use crate::pre_serve::{
+    BeforeServeTable, DistributedBackend, ModsTable, OnFailure, PreServeCommand, RecipeArgs,
+};
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct Node {
+    pub name: String,
+    pub ssh_ip: String,
+    pub ib_ip: String,
+    pub rank: u32,
+    #[serde(default)]
+    pub host_key: Option<String>,
+    #[serde(default)]
+    pub entity: Option<String>,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct NcclConfig {
+    pub ib_hca: String,
+    pub socket_ifname: String,
+    pub ib_gid_index: u32,
+    pub ib_disable: u32,
+    #[serde(default)]
+    pub gloo_socket_ifname: Option<String>,
+    #[serde(default)]
+    pub ucx_net_devices: Option<String>,
+    #[serde(default)]
+    pub nccl_cross_nic: Option<u32>,
+    #[serde(default)]
+    pub nccl_proto: Option<String>,
+    #[serde(default)]
+    pub nccl_p2p_level: Option<String>,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct Cluster {
+    pub name: String,
+    pub master_ib_ip: String,
+    pub nodes: Vec<Node>,
+    #[serde(default)]
+    pub extra_entities: Vec<String>,
+}
+
+#[derive(Clone, Deserialize)]
+pub struct Config {
+    pub ssh_user: String,
+    pub ssh_key_path: std::path::PathBuf,
+    pub clusters: Vec<Cluster>,
+    #[serde(default)]
+    pub nccl: HashMap<String, NcclConfig>,
+    #[serde(default)]
+    pub hf_cache_dir: Option<String>,
+    #[serde(default = "default_init_port")]
+    pub init_port: u16,
+    #[serde(default)]
+    pub before_serve: HashMap<String, BeforeServeTable>,
+    #[serde(default)]
+    pub exec_aliases: HashMap<String, String>,
+    #[serde(default)]
+    pub home_assistant: Option<HomeAssistantConfig>,
+    #[serde(default)]
+    pub global_mods_dir: Option<std::path::PathBuf>,
+}
+
+#[derive(Clone, Deserialize)]
+pub struct HomeAssistantConfig {
+    pub url: String,
+    pub token: String,
+    #[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
+}
+fn default_engine() -> String {
+    "vllm".to_string()
+}
+fn default_shm_size() -> String {
+    String::new()
+}
+fn default_tensor_parallel_size() -> u32 {
+    1
+}
+fn default_sync_image() -> bool {
+    true
+}
+
+#[derive(Clone, Deserialize)]
+pub struct Recipe {
+    #[serde(default = "default_engine")]
+    pub engine: String,
+    pub name: String,
+    #[serde(default)]
+    pub docker_image: String,
+    #[serde(default)]
+    pub model: String,
+    #[serde(default = "default_port")]
+    pub port: u16,
+    #[serde(default)]
+    pub required_nodes: Option<u32>,
+    #[serde(default)]
+    pub pipeline_parallel_size: Option<u32>,
+    #[serde(default = "default_tensor_parallel_size")]
+    pub tensor_parallel_size: u32,
+    #[serde(default)]
+    pub args: RecipeArgs,
+    #[serde(default)]
+    pub env: HashMap<String, String>,
+    #[serde(default)]
+    pub cluster: Option<String>,
+    #[serde(default)]
+    pub entrypoint: Vec<String>,
+    pub distributed_backend: Option<DistributedBackend>,
+    #[serde(default = "default_shm_size")]
+    pub shm_size: String,
+    #[serde(default)]
+    pub before_serve: Vec<PreServeCommand>,
+    #[serde(default)]
+    pub mods: ModsTable,
+    #[serde(default = "default_sync_image")]
+    pub sync_image: bool,
+}
+
+impl Config {
+    pub fn resolved_hf_cache_dir(&self) -> std::path::PathBuf {
+        if let Some(ref dir) = self.hf_cache_dir {
+            return std::path::PathBuf::from(dir);
+        }
+        let home = if self.ssh_user == "root" {
+            "/root"
+        } else {
+            return std::path::PathBuf::from(format!("/home/{}/.cache/huggingface", self.ssh_user));
+        };
+        std::path::PathBuf::from(format!("{}/.cache/huggingface", home))
+    }
+}

+ 15 - 74
src/docker/mod.rs

@@ -3,6 +3,7 @@ pub mod health;
 pub mod logs;
 pub mod mods;
 pub use logs::show_logs;
+pub mod sync;
 pub mod tail;
 use crate::{
     config::{Cluster, Config, NcclConfig, Node, Recipe, load_recipe},
@@ -20,6 +21,8 @@ pub struct ServeOpts {
     pub skip_before_serve: bool,
     pub dry_run: bool,
     pub auto_power_on: bool,
+    pub sync_image: bool,
+    pub sync_image_source: Option<String>,
 }
 pub(crate) fn shq(s: &str) -> String {
     let mut out = String::with_capacity(s.len() + 2);
@@ -76,6 +79,7 @@ pub async fn serve(
         opts.dry_run,
     )
     .await?;
+    use self::sync::sync_image_to_nodes;
     let mods_remote_path = mods::upload_mods_to_cluster(
         config,
         &cluster,
@@ -90,9 +94,17 @@ pub async fn serve(
     if image != recipe.docker_image {
         recipe.docker_image = image.clone();
     }
-    // Sync Docker image from head node to workers if needed.
-    if !opts.dry_run && cluster.nodes.len() > 1 {
-        sync_image_to_nodes(config, &cluster, &recipe.docker_image).await?;
+    // Sync Docker image across nodes if needed.
+    let should_sync =
+        opts.sync_image && recipe.sync_image && !opts.dry_run && cluster.nodes.len() > 1;
+    if should_sync {
+        sync_image_to_nodes(
+            config,
+            &cluster,
+            &recipe.docker_image,
+            opts.sync_image_source.as_deref(),
+        )
+        .await?;
     }
     let cache_dir = config.resolved_hf_cache_dir();
     start_cluster(
@@ -411,76 +423,5 @@ fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> V
     cmd
 }
 
-/// Get the Docker image digest (short ID) from a node.
-async fn get_image_digest(
-    config: &Config,
-    node: &Node,
-    image: &str,
-) -> Result<Option<String>> {
-    let cmd = format!(
-        "docker image ls --format '{{{{.ID}}}}' {image}"
-    );
-    let out = run_ssh(config, node, cmd, false).await?;
-    let digest = out.trim().to_string();
-    if digest.is_empty() {
-        Ok(None)
-    } else {
-        Ok(Some(digest))
-    }
-}
-
-/// Sync Docker image from the head node (rank 0) to all worker nodes.
-/// Skips workers that already have the same image digest.
-pub async fn sync_image_to_nodes(
-    config: &Config,
-    cluster: &Cluster,
-    image: &str,
-) -> Result<()> {
-    let head = cluster.nodes.iter().find(|n| n.rank == 0);
-    let head = match head {
-        Some(h) => h,
-        None => return Ok(()),  // single-node, nothing to sync
-    };
-    let head_digest = get_image_digest(config, head, image).await?;
-    if head_digest.is_none() {
-        anyhow::bail!(
-            "Image '{}' not found on head node '{}'",
-            image,
-            head.name
-        );
-    }
-    let head_digest = head_digest.unwrap();
-    info!("Syncing image '{}' from '{}' to workers...", image, head.name);
-    for worker in cluster.nodes.iter().filter(|n| n.rank > 0) {
-        let worker_digest = get_image_digest(config, worker, image).await?;
-        if worker_digest.as_deref() == Some(&head_digest) {
-            info!("  '{}' OK (digest matches)", worker.name);
-            continue;
-        }
-        info!("  {} -> {}...", worker.name, worker.ssh_ip);
-        // Run docker save | ssh docker load from the head node to the worker.
-        let pipe_cmd = format!(
-            "docker save '{}' | ssh -o StrictHostKeyChecking=no -i {} {}@{} docker load",
-            image,
-            config.ssh_key_path.display(),
-            config.ssh_user,
-            worker.ssh_ip
-        );
-        run_ssh(config, head, pipe_cmd, false).await?;
-        // Verify the image landed with the right digest.
-        let verify = get_image_digest(config, worker, image).await?;
-        if verify.as_deref() != Some(&head_digest) {
-            anyhow::bail!(
-                "Image sync to '{}' failed: digest mismatch (head={}, worker={:?})",
-                worker.name,
-                head_digest,
-                verify
-            );
-        }
-        info!("  {} OK", worker.name);
-    }
-    Ok(())
-}
-
 #[cfg(test)]
 mod tests;

+ 80 - 0
src/docker/sync.rs

@@ -0,0 +1,80 @@
+use crate::config::{Cluster, Config, Node};
+use crate::ssh::run_ssh;
+use anyhow::Result;
+
+use tracing::info;
+/// Get the Docker image digest (short ID) from a node.
+async fn get_image_digest(config: &Config, node: &Node, image: &str) -> Result<Option<String>> {
+    let cmd = format!("docker image ls --format '{{{{.ID}}}}' {image}");
+    let out = run_ssh(config, node, cmd, false).await?;
+    let digest = out.trim().to_string();
+    if digest.is_empty() {
+        Ok(None)
+    } else {
+        Ok(Some(digest))
+    }
+}
+
+/// Sync Docker image across cluster nodes.
+/// Probes all nodes, finds a source (explicit or first with image), then pipes
+/// `docker save | ssh docker load` to every node missing or with a different digest.
+pub async fn sync_image_to_nodes(
+    config: &Config,
+    cluster: &Cluster,
+    image: &str,
+    explicit_source: Option<&str>,
+) -> Result<()> {
+    let mut digests: Vec<(&Node, Option<String>)> = Vec::new();
+    for node in &cluster.nodes {
+        let d = get_image_digest(config, node, image).await?;
+        digests.push((node, d));
+    }
+    let source_idx = if let Some(src_name) = explicit_source {
+        digests
+            .iter()
+            .position(|(n, _)| n.name == src_name)
+            .ok_or_else(|| {
+                anyhow::anyhow!("sync-image-source '{}' not found in cluster", src_name)
+            })?
+    } else {
+        digests
+            .iter()
+            .position(|(_, d)| d.is_some())
+            .ok_or_else(|| anyhow::anyhow!("Image '{}' not found on any cluster node", image))?
+    };
+    let (source, source_digest_opt) = &digests[source_idx];
+    let source_digest = source_digest_opt.as_ref().unwrap();
+    info!(
+        "Syncing image '{}' from '{}' to peers...",
+        image, source.name
+    );
+    for (target, target_digest) in &digests {
+        if target.name == source.name {
+            continue;
+        }
+        if target_digest.as_deref() == Some(source_digest) {
+            info!("  '{}' OK (digest matches)", target.name);
+            continue;
+        }
+        info!("  {} ({}): syncing...", target.name, target.ssh_ip);
+        let pipe_cmd = format!(
+            "docker save '{}' | ssh -o StrictHostKeyChecking=no -i {} {}@{} docker load",
+            image,
+            config.ssh_key_path.display(),
+            config.ssh_user,
+            target.ssh_ip
+        );
+        run_ssh(config, source, pipe_cmd, false).await?;
+        let verify = get_image_digest(config, target, image).await?;
+        if verify.as_deref() != Some(source_digest) {
+            anyhow::bail!(
+                "Image sync to '{}' failed: digest mismatch (source={}, target={:?})",
+                target.name,
+                source_digest,
+                verify
+            );
+        }
+        info!("  {} OK", target.name);
+    }
+    Ok(())
+}

+ 2 - 0
src/docker/tests.rs

@@ -60,6 +60,7 @@ fn trec() -> Recipe {
         before_serve: vec![],
         mods: Default::default(),
         shm_size: String::new(),
+        sync_image: true,
     }
 }
 
@@ -241,6 +242,7 @@ fn sglang_recipe() -> Recipe {
         before_serve: vec![],
         mods: Default::default(),
         shm_size: String::new(),
+        sync_image: true,
     }
 }
 

+ 1 - 0
src/doctor.rs

@@ -367,6 +367,7 @@ mod tests {
             distributed_backend: Some(crate::config::DistributedBackend::Mp),
             before_serve: vec![],
             mods: Default::default(),
+            sync_image: true,
             shm_size: String::new(),
         }
     }

+ 4 - 0
src/main.rs

@@ -88,6 +88,8 @@ async fn main() -> Result<()> {
             skip_before_serve,
             follow,
             auto_power_on,
+            sync_image,
+            sync_image_source,
         } => {
             let cluster = docker::serve(
                 &config,
@@ -99,6 +101,8 @@ async fn main() -> Result<()> {
                     skip_before_serve,
                     dry_run,
                     auto_power_on,
+                    sync_image,
+                    sync_image_source,
                 },
             )
             .await?;