Ver código fonte

fix: make IB config optional for single-node clusters

Users running single-node tmux recipes don't have InfiniBand addressing.
Previously `Node.ib_ip` and `Cluster.master_ib_ip` were required fields, so
removing the IB config from config.toml broke loading with "missing field
ib_ip". They were also injected into engine env for single-node runs,
pointing at an unreachable IB address.

- Make `Node.ib_ip` and `Cluster.master_ib_ip` optional (`Option<String>`)
  with `#[serde(default)]`
- Add `Node::effective_ib_ip()` (falls back to `ssh_ip`) and
  `Cluster::effective_master_ip()` (falls back to rank-0 node's effective
  IP)
- Update all call sites (docker/health/command, tmux, orchestrate) to use
  the effective getters
- Single-node renders no `NCCL_*` env (build_nccl_setup already no-ops when
  no `[nccl]` section) and `VLLM_HOST_IP` now uses the reachable ssh_ip
  fallback
Alexander Korolev 2 meses atrás
pai
commit
0181645b9e

+ 2 - 2
src/after_stop.rs

@@ -150,11 +150,11 @@ mod tests {
     fn test_cluster() -> Cluster {
         Cluster {
             name: "c1".to_string(),
-            master_ib_ip: "1.2.3.4".to_string(),
+            master_ib_ip: Some("1.2.3.4".to_string()),
             nodes: vec![Node {
                 name: "n0".to_string(),
                 ssh_ip: "10.0.0.1".to_string(),
-                ib_ip: "1.2.3.4".to_string(),
+                ib_ip: Some("1.2.3.4".to_string()),
                 rank: 0,
                 host_key: None,
                 entity: None,

+ 3 - 3
src/before_serve.rs

@@ -183,12 +183,12 @@ mod tests {
     fn test_cluster() -> Cluster {
         Cluster {
             name: "c1".to_string(),
-            master_ib_ip: "1.2.3.4".to_string(),
+            master_ib_ip: Some("1.2.3.4".to_string()),
             nodes: vec![
                 Node {
                     name: "n0".to_string(),
                     ssh_ip: "10.0.0.1".to_string(),
-                    ib_ip: "1.2.3.4".to_string(),
+                    ib_ip: Some("1.2.3.4".to_string()),
                     rank: 0,
                     host_key: None,
                     entity: None,
@@ -196,7 +196,7 @@ mod tests {
                 Node {
                     name: "n1".to_string(),
                     ssh_ip: "10.0.0.2".to_string(),
-                    ib_ip: "1.2.3.5".to_string(),
+                    ib_ip: Some("1.2.3.5".to_string()),
                     rank: 1,
                     host_key: None,
                     entity: None,

+ 2 - 2
src/cluster_state.rs

@@ -166,7 +166,7 @@ mod tests {
         Node {
             name: name.to_string(),
             ssh_ip: format!("10.0.0.{}", name.chars().last().unwrap_or('0')),
-            ib_ip: "10.0.1.1".to_string(),
+            ib_ip: Some("10.0.1.1".to_string()),
             rank: 0,
             host_key: None,
             entity: None,
@@ -176,7 +176,7 @@ mod tests {
     fn test_cluster(name: &str, node_names: &[&str]) -> Cluster {
         Cluster {
             name: name.to_string(),
-            master_ib_ip: "10.0.1.1".to_string(),
+            master_ib_ip: Some("10.0.1.1".to_string()),
             nodes: node_names.iter().map(|n| test_node(n)).collect(),
             extra_entities: vec![],
         }

+ 24 - 0
src/config.rs

@@ -5,6 +5,30 @@ use tracing::info;
 mod types;
 pub use types::*;
 
+impl Node {
+    /// Host IP for engine exports; falls back to `ssh_ip` when `ib_ip` unset
+    /// (single-node runs need no IB addressing).
+    pub fn effective_ib_ip(&self) -> &str {
+        self.ib_ip.as_deref().unwrap_or(&self.ssh_ip)
+    }
+}
+
+impl Cluster {
+    /// Effective master IP: explicit `master_ib_ip`, else the rank-0 node's
+    /// effective (ib→ssh fallback) IP, else the first node's ssh IP, else empty.
+    pub fn effective_master_ip(&self) -> String {
+        if let Some(ip) = &self.master_ib_ip {
+            return ip.clone();
+        }
+        self.nodes
+            .iter()
+            .find(|n| n.rank == 0)
+            .or_else(|| self.nodes.first())
+            .map(|n| n.effective_ib_ip().to_string())
+            .unwrap_or_default()
+    }
+}
+
 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);

+ 3 - 3
src/config/tests.rs

@@ -5,12 +5,12 @@ use std::collections::{HashMap, HashSet};
 fn tcls() -> Cluster {
     Cluster {
         name: "test-cluster".to_string(),
-        master_ib_ip: "192.168.1.1".to_string(),
+        master_ib_ip: Some("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(),
+                ib_ip: Some("192.168.1.1".to_string()),
                 rank: 0,
                 host_key: None,
                 entity: None,
@@ -18,7 +18,7 @@ fn tcls() -> Cluster {
             Node {
                 name: "node1".to_string(),
                 ssh_ip: "10.0.0.2".to_string(),
-                ib_ip: "192.168.1.2".to_string(),
+                ib_ip: Some("192.168.1.2".to_string()),
                 rank: 1,
                 host_key: None,
                 entity: None,

+ 8 - 19
src/config/types.rs

@@ -10,7 +10,9 @@ pub use crate::pre_serve::{
 pub struct Node {
     pub name: String,
     pub ssh_ip: String,
-    pub ib_ip: String,
+    /// InfiniBand IP; optional (falls back to `ssh_ip`).
+    #[serde(default)]
+    pub ib_ip: Option<String>,
     pub rank: u32,
     #[serde(default)]
     pub host_key: Option<String>,
@@ -39,7 +41,9 @@ pub struct NcclConfig {
 #[derive(Clone, Debug, Deserialize)]
 pub struct Cluster {
     pub name: String,
-    pub master_ib_ip: String,
+    /// Master IB IP; optional (falls back to rank-0 node's effective IP).
+    #[serde(default)]
+    pub master_ib_ip: Option<String>,
     pub nodes: Vec<Node>,
     #[serde(default)]
     pub extra_entities: Vec<String>,
@@ -178,23 +182,8 @@ pub struct Recipe {
     /// before serve; if absent, it runs these commands in order, then
     /// re-checks. A failure aborts the serve.
     ///
-    /// Example (inline):
-    /// ```toml
-    /// venv = "~/venvs/vllm025"
-    /// venv_setup = """
-    ///     uv venv ~/venvs/vllm025 -p 3.12
-    ///     uv pip install vllm==0.25.1 --torch-backend=cu130
-    /// """
-    /// ```
-    ///
-    /// Example (list):
-    /// ```toml
-    /// venv = "~/venvs/vllm025"
-    /// venv_setup = [
-    ///     "uv venv ~/venvs/vllm025 -p 3.12",
-    ///     "uv pip install vllm==0.25.1 --torch-backend=cu130",
-    /// ]
-    /// ```
+    /// Example (inline): `uv venv ~/venvs/vllm025 -p 3.12 && uv pip install ...`
+    /// Example (list): `["uv venv ~/venvs/vllm025 -p 3.12", "..."]`
     #[serde(default)]
     pub venv_setup: VenvSetup,
     /// Path to a shell script (relative to the recipe directory) that creates

+ 9 - 5
src/docker/command.rs

@@ -182,7 +182,7 @@ fn collect_nccl_env<'a>(ctx: &'a DockerCommandCtx, node: &'a Node) -> Vec<(&'a s
     }
     // Worker IP is always set for multi-node (collect_nccl_env is only
     // called when nnodes > 1).
-    let worker_ip = Cow::Borrowed(node.ib_ip.as_str());
+    let worker_ip = Cow::Borrowed(node.effective_ib_ip());
     if is_sglang {
         env.push(("SGLANG_WORKER_IP", worker_ip));
     } else {
@@ -205,11 +205,14 @@ fn collect_docker_env(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> Vec
     if ctx.recipe.distributed_backend == Some(crate::config::DistributedBackend::Ray) {
         pairs.push((
             "RAY_ADDRESS",
-            format!("{}:6379", ctx.cluster.master_ib_ip).into(),
+            format!("{}:6379", ctx.cluster.effective_master_ip()).into(),
         ));
         pairs.push(("RAY_PORT", "6379".into()));
-        pairs.push(("RAY_NODE_IP_ADDRESS", node.ib_ip.as_str().into()));
-        pairs.push(("RAY_OVERRIDE_NODE_IP_ADDRESS", node.ib_ip.as_str().into()));
+        pairs.push(("RAY_NODE_IP_ADDRESS", node.effective_ib_ip().into()));
+        pairs.push((
+            "RAY_OVERRIDE_NODE_IP_ADDRESS",
+            node.effective_ib_ip().into(),
+        ));
     }
     // extra_mount `env` defaults (point at the container path) are injected
     // before recipe [env] so the recipe can still override them (last wins).
@@ -278,13 +281,14 @@ pub(crate) fn build_docker_command(
     if ctx.cluster.nodes.len() == 1 && ctx.cluster.extra_entities.is_empty() {
         cmd.push(format!("--label sm121.solo-node={}", node.name));
     }
+    let master_ip = ctx.cluster.effective_master_ip();
     let tctx = tail::TailCtx {
         recipe: &ctx.recipe,
         tp_size: ctx.tp_size,
         pp_size: ctx.pp_size,
         nnodes: ctx.nnodes,
         init_port: ctx.init_port,
-        master_ib_ip: &ctx.cluster.master_ib_ip,
+        master_ib_ip: &master_ip,
     };
     let tail = tail::build_tail(&tctx, node);
     if ctx.recipe.mods.commands.is_empty() {

+ 2 - 2
src/docker/command_tests.rs

@@ -8,11 +8,11 @@ use std::collections::HashMap;
 fn tcls() -> Cluster {
     Cluster {
         name: "t".into(),
-        master_ib_ip: "1".into(),
+        master_ib_ip: Some("1".into()),
         nodes: vec![Node {
             name: "n0".into(),
             ssh_ip: "10.0.0.1".into(),
-            ib_ip: "1".into(),
+            ib_ip: Some("1".into()),
             rank: 0,
             host_key: None,
             entity: None,

+ 1 - 1
src/docker/health.rs

@@ -62,7 +62,7 @@ pub async fn health_check(config: &Config, cluster: &Cluster, dry_run: bool) ->
         let config = config.clone();
         let node = node.clone();
         let name = node.name.clone();
-        let ib_ip = node.ib_ip.clone();
+        let ib_ip = node.effective_ib_ip().to_string();
         let cl = cluster_label.clone();
         let ns = net_script.clone();
         let docker_cmd = format!(

+ 3 - 3
src/docker/tests.rs

@@ -6,12 +6,12 @@ use std::collections::HashMap;
 fn tcls() -> Cluster {
     Cluster {
         name: "t".into(),
-        master_ib_ip: "1".into(),
+        master_ib_ip: Some("1".into()),
         nodes: vec![
             Node {
                 name: "n0".into(),
                 ssh_ip: "10.0.0.1".into(),
-                ib_ip: "1".into(),
+                ib_ip: Some("1".into()),
                 rank: 0,
                 host_key: None,
                 entity: None,
@@ -19,7 +19,7 @@ fn tcls() -> Cluster {
             Node {
                 name: "n1".into(),
                 ssh_ip: "10.0.0.2".into(),
-                ib_ip: "2".into(),
+                ib_ip: Some("2".into()),
                 rank: 1,
                 host_key: None,
                 entity: None,

+ 1 - 1
src/engine_flags.rs

@@ -138,7 +138,7 @@ mod tests {
         Node {
             name: format!("n{}", rank),
             ssh_ip: "10.0.0.1".into(),
-            ib_ip: "10.0.0.1".into(),
+            ib_ip: Some("10.0.0.1".into()),
             rank,
             host_key: None,
             entity: None,

+ 6 - 1
src/health.rs

@@ -65,7 +65,12 @@ pub async fn check_cluster_health(
             let health = manager
                 .deep_check_session(&config, &node, &cluster, &sess, &recipe, dry)
                 .await;
-            (node.name.clone(), node.ib_ip.clone(), node.rank, health)
+            (
+                node.name.clone(),
+                node.effective_ib_ip().to_string(),
+                node.rank,
+                health,
+            )
         }));
     }
 

+ 2 - 2
src/orchestrate/tests.rs

@@ -42,7 +42,7 @@ fn test_node(name: &str) -> Node {
     Node {
         name: name.to_string(),
         ssh_ip: format!("10.0.0.{}", name.chars().last().unwrap_or('0')),
-        ib_ip: "10.0.1.1".to_string(),
+        ib_ip: Some("10.0.1.1".to_string()),
         rank: 0,
         host_key: None,
         entity: Some(format!("switch.{}", name)),
@@ -52,7 +52,7 @@ fn test_node(name: &str) -> Node {
 fn test_cluster(name: &str, node_names: &[&str]) -> Cluster {
     Cluster {
         name: name.to_string(),
-        master_ib_ip: "10.0.1.1".to_string(),
+        master_ib_ip: Some("10.0.1.1".to_string()),
         nodes: node_names.iter().map(|n| test_node(n)).collect(),
         extra_entities: vec![],
     }

+ 6 - 4
src/tmux/commands.rs

@@ -48,13 +48,14 @@ impl TmuxManager {
         }
 
         // Shared engine args: model, recipe args, distributed flags, headless.
+        let master_ip = cluster.effective_master_ip();
         let ectx = EngineFlagsCtx {
             recipe,
             tp_size: recipe.tensor_parallel_size,
             pp_size,
             nnodes,
             init_port: 25000,
-            master_ib_ip: &cluster.master_ib_ip,
+            master_ib_ip: &master_ip,
         };
         parts.extend(build_engine_args(&ectx, node));
 
@@ -95,7 +96,8 @@ impl TmuxManager {
         {
             cmds.push(format!(
                 "tmux send-keys -t {}:0 'export RAY_ADDRESS={}:6379' Enter",
-                sess, cluster.master_ib_ip
+                sess,
+                cluster.effective_master_ip()
             ));
         }
         for (k, v) in &recipe.env {
@@ -172,11 +174,11 @@ mod tests {
     fn tcluster() -> Cluster {
         Cluster {
             name: "c".into(),
-            master_ib_ip: "10.0.0.1".into(),
+            master_ib_ip: Some("10.0.0.1".into()),
             nodes: vec![Node {
                 name: "n0".into(),
                 ssh_ip: "10.0.0.1".into(),
-                ib_ip: "10.0.0.1".into(),
+                ib_ip: Some("10.0.0.1".into()),
                 rank: 0,
                 host_key: None,
                 entity: None,

+ 5 - 4
src/tmux/mod.rs

@@ -49,12 +49,13 @@ impl TmuxManager {
 
     /// Builds engine-specific host IP environment variable.
     pub fn build_host_ip_export(&self, engine: &str, node: &Node) -> Vec<String> {
+        let ip = node.effective_ib_ip();
         match engine {
-            "vllm" => vec![format!("export VLLM_HOST_IP={}", node.ib_ip)],
-            "sglang" => vec![format!("export SGLANG_WORKER_IP={}", node.ib_ip)],
+            "vllm" => vec![format!("export VLLM_HOST_IP={}", ip)],
+            "sglang" => vec![format!("export SGLANG_WORKER_IP={}", ip)],
             _ => vec![
-                format!("export VLLM_HOST_IP={}", node.ib_ip),
-                format!("export SGLANG_WORKER_IP={}", node.ib_ip),
+                format!("export VLLM_HOST_IP={}", ip),
+                format!("export SGLANG_WORKER_IP={}", ip),
             ],
         }
     }

+ 2 - 1
src/tmux/serve.rs

@@ -186,7 +186,8 @@ pub async fn tmux_serve(
         sleep(Duration::from_secs(5)).await;
         info!(
             "🎉 Cluster launched via tmux! API → http://{}:{}",
-            cluster.master_ib_ip, recipe.port
+            cluster.effective_master_ip(),
+            recipe.port
         );
         info!(
             "📌 To view sessions: `cargo run tmux status --recipe {}`",