Ver código fonte

feat(serve): add solo mode and fix single-node launch bug

Two coupled changes to the serve path: a bug fix for single-node recipes
on multi-node clusters, and a new solo mode for running one container on
the first free node across the fleet.

Bug fix (docker/mod.rs start_cluster): a recipe with required_nodes=1
(or any N < cluster size) launched a container on EVERY node in the
cluster, each told `--nnodes <cluster_size> --tp <cluster_size>` — a
broken rendezvous that never coordinated. start_cluster now trims the
launch set to required_nodes when --node is not given, so required_nodes
finally means "launch on exactly N nodes" (it already meant "cluster
must have N free"). prepare_cluster guarantees >= required_nodes are
free, so taking the first N is safe and deterministic.

Solo mode: a single container on the first free node, with no
distributed flags. Triggered by CLI `--solo` (overrides recipe) or a
recipe `solo = true` field. When solo and no explicit --node, the
launcher selects the smallest cluster with a free node and keeps only
that node (required_nodes is forced to 1 for selection). An explicit
--node still wins, but the run stays single-node.

Single-node flags (docker/tail.rs build_tail): when nnodes == 1 (solo,
required_nodes=1, or a single --node), emit NO --nnodes/--tp/--pp/
--node-rank/--master-addr/--master-port/--headless/
--distributed-executor-backend — the engine defaults to single-node.
Previously these were always emitted, which is redundant at best and
rejected by some engines at worst (--nnodes 1 with --master-addr).
Multi-node runs are unchanged (regression test guards this).

Tests: +4 (116 -> 120) for single-node flag omission, multi-node
regression, worker --headless behavior, and solo field deserialization.
Split into src/docker/solo_tests.rs to keep docker/tests.rs under the
500-line limit. AGENTS.md documents --solo, the solo field, and the
required_nodes cap.
Alexander Korolev 2 meses atrás
pai
commit
23ae7a6484

+ 8 - 1
AGENTS.md

@@ -22,6 +22,12 @@ serve/stop/logs/health/doctor/exec/power/tmux/init-host-keys
 
 Cluster: exact name, prefix (`4`→`4nodes-switch`), or `auto` (smallest with NCCL + matching `required_nodes`).
 
+`serve --solo` overrides cluster selection: finds the first free node across
+all clusters and runs one single-node container there (no `--nnodes`/`--tp`/
+`--master-addr` flags — the engine defaults to single-node). `--no-solo`
+forces multi-node even if the recipe sets `solo = true`. A recipe can also
+set `solo = true` as the default.
+
 > ⚠️ **Cluster may be serving live traffic.** A model (often the one serving this
 > very agent) may be running on the cluster at any time. Before running `stop`,
 > `power off`, or `power switch`, ALWAYS first run `health` and confirm with the
@@ -36,7 +42,8 @@ engine = "vllm"                     # or "sglang"
 docker_image = "img:tag"
 model = "org/model"
 port = 30000                        # default
-required_nodes = 2                  # optional
+required_nodes = 2                  # optional; also caps the launch set (N nodes, not all)
+solo = false                       # optional; single-node on first free node, no dist flags
 tensor_parallel_size = 2
 pipeline_parallel_size = 1          # vLLM only, 0 = auto
 distributed_backend = "mp"          # or "ray"

+ 1 - 0
locales/en/main.ftl

@@ -95,6 +95,7 @@ power-off-container-stop-failed = ⚠️ Failed to stop containers on node { $na
 after-stop-failed = after_stop: node { $name } failed: { $error }
 after-stop-run = { $node } ({ $tag }) :: { $cmd }
 after-stop-ignored = ⚠ { $cmd } ignored: { $error }
+solo-node-selected = 🎯 Solo: selected node { $name } in cluster { $cluster }
 
 # --- Strings that previously bypassed the i18n system (hardcoded) ---
 boot-wait = ⏳ Waiting { $secs }s for nodes to boot...

+ 1 - 0
src/after_stop.rs

@@ -197,6 +197,7 @@ mod tests {
             distributed_backend: Some(crate::config::DistributedBackend::Mp),
             before_serve: vec![],
             after_stop,
+            solo: None,
             mods: Default::default(),
             sync_image: true,
             docker_image_dir: None,

+ 2 - 0
src/before_serve.rs

@@ -284,6 +284,7 @@ mod tests {
         let recipe = Recipe {
             before_serve: vec![],
             after_stop: vec![],
+            solo: None,
             ..test_recipe()
         };
         assert!(combined_commands(&cfg, &test_cluster(), &recipe).is_empty());
@@ -306,6 +307,7 @@ mod tests {
             distributed_backend: Some(crate::config::DistributedBackend::Mp),
             before_serve: vec![],
             after_stop: vec![],
+            solo: None,
             mods: Default::default(),
             sync_image: true,
             docker_image_dir: None,

+ 6 - 0
src/cli.rs

@@ -58,6 +58,12 @@ pub enum Commands {
         /// Force rebuild even if checksum matches.
         #[arg(long)]
         force_build: bool,
+
+        /// Serve single-node on the first free node found (overrides recipe
+        /// `solo`). Use `--no-solo` to force multi-node even if the recipe
+        /// sets `solo = true`.
+        #[arg(long)]
+        solo: Option<bool>,
     },
 
     /// Show container logs

+ 7 - 0
src/config/types.rs

@@ -146,6 +146,13 @@ pub struct Recipe {
     /// services that `before_serve` stopped.
     #[serde(default)]
     pub after_stop: Vec<PreServeCommand>,
+    /// Whether this recipe should always be served single-node on the first
+    /// free node found (solo mode). When set, `sm121 serve` ignores
+    /// `required_nodes` / multi-node cluster selection and launches exactly
+    /// one container with no distributed flags. The CLI `--solo`/`--no-solo`
+    /// flag overrides this per-invocation.
+    #[serde(default)]
+    pub solo: Option<bool>,
     #[serde(default)]
     pub mods: ModsTable,
     #[serde(default = "default_sync_image")]

+ 1 - 0
src/docker/command.rs

@@ -384,6 +384,7 @@ mod tests {
             distributed_backend: Some(DistributedBackend::Mp),
             before_serve: vec![],
             after_stop: vec![],
+            solo: None,
             mods: Default::default(),
             shm_size: String::new(),
             sync_image: true,

+ 51 - 4
src/docker/mod.rs

@@ -26,6 +26,9 @@ pub struct ServeOpts {
     pub sync_image_source: Option<String>,
     pub build_node: Option<String>,
     pub force_build: bool,
+    /// CLI override for recipe `solo`: `Some(true)` (--solo), `Some(false)`
+    /// (--no-solo), `None` (unset → use recipe `solo` field).
+    pub solo: Option<bool>,
 }
 /// Shell-quote a string for safe interpolation into a remote `sh -c`
 /// command (single-quote wrapping, embedded quotes escaped). Re-exported
@@ -55,14 +58,35 @@ pub async fn serve(
     } else {
         None
     };
-    let (cluster, nccl) = if opts.dry_run {
+    // Effective solo: CLI flag overrides recipe field. When solo and no
+    // explicit --node, launch on exactly one node (the first free node in
+    // the smallest cluster that has one) with no distributed flags. When
+    // solo and --node is given, the explicit node wins but we still run
+    // single-node (start_cluster sees one node → no distributed flags).
+    let effective_solo = opts.solo.unwrap_or(recipe.solo.unwrap_or(false));
+    // For solo, ignore the recipe's required_nodes — we always want one node.
+    // Build a shallow clone with required_nodes overridden so cluster
+    // selection picks the smallest cluster with >= 1 free node.
+    let recipe_for_select = if effective_solo {
+        let mut r = recipe.clone();
+        r.required_nodes = Some(1);
+        r
+    } else {
+        recipe.clone()
+    };
+    let (mut cluster, nccl) = if opts.dry_run {
         // Dry-run: skip all SSH probing/power/pinging so it works even when the
         // cluster is off or unreachable (treat every configured node as free).
-        crate::orchestrate::prepare_cluster_dry_run(config, &recipe, cluster_name, node_name)?
+        crate::orchestrate::prepare_cluster_dry_run(
+            config,
+            &recipe_for_select,
+            cluster_name,
+            node_name,
+        )?
     } else {
         prepare_cluster(
             config,
-            &recipe,
+            &recipe_for_select,
             cluster_name,
             node_name,
             &PrepareRequest {
@@ -75,6 +99,18 @@ pub async fn serve(
         )
         .await?
     };
+    if effective_solo && node_name.is_none() && !cluster.nodes.is_empty() {
+        // Solo: keep only the first node. `prepare_cluster` already selected
+        // the smallest cluster with >= required_nodes free (required_nodes is
+        // ignored for solo — we always want one node), so the first free node
+        // is the solo target.
+        let solo_node = cluster.nodes[0].clone();
+        info!(
+            "{}",
+            fluent_i18n::t!("solo-node-selected", { "name" => &solo_node.name, "cluster" => &cluster.name })
+        );
+        cluster.nodes = vec![solo_node];
+    }
     crate::before_serve::run(
         config,
         &cluster,
@@ -259,7 +295,18 @@ async fn start_cluster(
         }
         v
     } else {
-        cluster.nodes.iter().collect()
+        // Trim to `required_nodes` when set: a recipe that needs N nodes
+        // must launch on exactly N nodes, not on every free node in the
+        // cluster. Previously a `required_nodes = 1` recipe on a 4-node
+        // cluster launched 4 containers, each told `--nnodes 4 --tp 4` —
+        // a broken rendezvous that never coordinated. `prepare_cluster`
+        // guarantees >= required_nodes are free, so taking the first N is
+        // safe and deterministic.
+        let all = cluster.nodes.iter().collect::<Vec<&Node>>();
+        match recipe.required_nodes {
+            Some(n) if (n as usize) < all.len() => all.into_iter().take(n as usize).collect(),
+            _ => all,
+        }
     };
     let nnodes = nodes.len() as u32;
     let pp_size = if recipe.pipeline_parallel_size.unwrap_or(0) == 0

+ 104 - 0
src/docker/solo_tests.rs

@@ -0,0 +1,104 @@
+//! Solo / single-node flag-emission tests — split from `tests.rs` to keep
+//! that file under the 500-line limit. Uses the parent module's `trec`,
+//! `tcls`, `make_dm_ctx`, and `build_docker_command` helpers.
+
+use super::*;
+use crate::config::Recipe;
+
+#[test]
+fn single_node_run_omits_distributed_flags() {
+    // A solo or required_nodes=1 run must NOT emit --nnodes/--tp/--master-addr
+    // etc. — the engine defaults to single-node. (Regression for the bug
+    // where a 1-node recipe on a multi-node cluster launched N containers each
+    // told `--nnodes N --tp N`.)
+    let mut r = trec();
+    r.required_nodes = Some(1);
+    r.tensor_parallel_size = 1;
+    r.pipeline_parallel_size = Some(1);
+    // nnodes=1 is what start_cluster computes after trimming to required_nodes.
+    let ctx = make_dm_ctx(&r, 1, 1, 1, false);
+    let cmd = build_docker_command(&ctx, &tcls().nodes[0], true).join(" ");
+    assert!(
+        !cmd.contains("--nnodes"),
+        "single-node must not pass --nnodes: {cmd}"
+    );
+    assert!(
+        !cmd.contains("--tensor-parallel-size"),
+        "single-node must not pass --tp: {cmd}"
+    );
+    assert!(
+        !cmd.contains("--pipeline-parallel-size"),
+        "single-node must not pass --pp: {cmd}"
+    );
+    assert!(
+        !cmd.contains("--master-addr"),
+        "single-node must not pass --master-addr: {cmd}"
+    );
+    assert!(
+        !cmd.contains("--master-port"),
+        "single-node must not pass --master-port: {cmd}"
+    );
+    assert!(
+        !cmd.contains("--node-rank"),
+        "single-node must not pass --node-rank: {cmd}"
+    );
+    assert!(
+        !cmd.contains("--headless"),
+        "single-node must not pass --headless: {cmd}"
+    );
+    // But the engine still gets --host/--port.
+    assert!(cmd.contains("--host"));
+    assert!(cmd.contains("--port"));
+}
+
+#[test]
+fn multi_node_run_still_emits_distributed_flags() {
+    // Regression guard: the single-node fix must not break the multi-node path.
+    let r = trec(); // required_nodes=2, tp collapses to nnodes=2
+    let ctx = make_dm_ctx(&r, 2, 2, 2, false);
+    let cmd = build_docker_command(&ctx, &tcls().nodes[0], true).join(" ");
+    assert!(cmd.contains("--nnodes 2"));
+    assert!(cmd.contains("--tensor-parallel-size 2"));
+    assert!(cmd.contains("--pipeline-parallel-size 2"));
+    assert!(cmd.contains("--master-addr"));
+    assert!(cmd.contains("--master-port"));
+    assert!(cmd.contains("--node-rank 0"));
+}
+
+#[test]
+fn solo_worker_node_emits_headless_no_rendezvous() {
+    // Multi-node: a worker (rank>0) gets --headless; single-node never does.
+    let r = trec();
+    let ctx = make_dm_ctx(&r, 2, 2, 2, false);
+    let worker_cmd = build_docker_command(&ctx, &tcls().nodes[1], false).join(" ");
+    assert!(
+        worker_cmd.contains("--headless"),
+        "multi-node worker needs --headless: {worker_cmd}"
+    );
+    // Single-node worker (n/a really, but guard): headless never appears.
+    let mut r1 = r.clone();
+    r1.required_nodes = Some(1);
+    let ctx1 = make_dm_ctx(&r1, 1, 1, 1, false);
+    let solo_cmd = build_docker_command(&ctx1, &tcls().nodes[1], false).join(" ");
+    assert!(!solo_cmd.contains("--headless"));
+}
+
+#[test]
+fn recipe_solo_field_deserializes() {
+    let toml_str = r#"
+name = "Solo"
+model = "org/m"
+solo = true
+required_nodes = 4
+"#;
+    let r: Recipe = toml::from_str(toml_str).unwrap();
+    assert_eq!(r.solo, Some(true));
+    // unset → None (default multi-node)
+    let toml2 = r#"
+name = "Multi"
+model = "org/m"
+required_nodes = 4
+"#;
+    let r2: Recipe = toml::from_str(toml2).unwrap();
+    assert_eq!(r2.solo, None);
+}

+ 36 - 27
src/docker/tail.rs

@@ -40,41 +40,50 @@ pub fn build_tail(ctx: &TailCtx<'_>, node: &Node) -> Vec<String> {
     tail.push(ctx.recipe.model.clone());
     tail.extend(ctx.recipe.args.to_vec());
 
-    let mut flags: Vec<String> = match ctx.recipe.distributed_backend {
-        Some(DistributedBackend::Ray | DistributedBackend::Mp) if is_sglang => vec![
-            "--tp".into(),
-            ctx.tp_size.to_string(),
-            "--nnodes".into(),
-            ctx.nnodes.to_string(),
-            "--node-rank".into(),
-            node.rank.to_string(),
-            "--dist-init-addr".into(),
-            format!("{}:{}", ctx.master_ib_ip, ctx.init_port),
-        ],
-        Some(DistributedBackend::Ray | DistributedBackend::Mp) => {
-            let mut fp = vec![
-                "--tensor-parallel-size".into(),
+    // Single-node run (solo, required_nodes=1, or a single --node): emit
+    // NO distributed flags — the engine defaults to single-node. Launching
+    // one container with `--nnodes 1 --master-addr ...` is redundant at best
+    // and rejected by some engines at worst. Multi-node runs get the full
+    // rendezvous flag set below.
+    let mut flags: Vec<String> = if ctx.nnodes > 1 {
+        match ctx.recipe.distributed_backend {
+            Some(DistributedBackend::Ray | DistributedBackend::Mp) if is_sglang => vec![
+                "--tp".into(),
                 ctx.tp_size.to_string(),
-                "--pipeline-parallel-size".into(),
-                ctx.pp_size.to_string(),
                 "--nnodes".into(),
                 ctx.nnodes.to_string(),
                 "--node-rank".into(),
                 node.rank.to_string(),
-                "--master-addr".into(),
-                ctx.master_ib_ip.to_string(),
-                "--master-port".into(),
-                ctx.init_port.to_string(),
-            ];
-            if ctx.recipe.distributed_backend == Some(DistributedBackend::Ray) {
-                fp.push("--distributed-executor-backend".into());
-                fp.push("ray".into());
+                "--dist-init-addr".into(),
+                format!("{}:{}", ctx.master_ib_ip, ctx.init_port),
+            ],
+            Some(DistributedBackend::Ray | DistributedBackend::Mp) => {
+                let mut fp = vec![
+                    "--tensor-parallel-size".into(),
+                    ctx.tp_size.to_string(),
+                    "--pipeline-parallel-size".into(),
+                    ctx.pp_size.to_string(),
+                    "--nnodes".into(),
+                    ctx.nnodes.to_string(),
+                    "--node-rank".into(),
+                    node.rank.to_string(),
+                    "--master-addr".into(),
+                    ctx.master_ib_ip.to_string(),
+                    "--master-port".into(),
+                    ctx.init_port.to_string(),
+                ];
+                if ctx.recipe.distributed_backend == Some(DistributedBackend::Ray) {
+                    fp.push("--distributed-executor-backend".into());
+                    fp.push("ray".into());
+                }
+                fp
             }
-            fp
+            _ => vec![],
         }
-        _ => vec![],
+    } else {
+        vec![]
     };
-    if !is_sglang && node.rank > 0 && is_distributed {
+    if ctx.nnodes > 1 && !is_sglang && node.rank > 0 && is_distributed {
         flags.push("--headless".into());
     }
     tail.extend(flags);

+ 5 - 0
src/docker/tests.rs

@@ -60,6 +60,7 @@ fn trec() -> Recipe {
         distributed_backend: Some(DistributedBackend::Mp),
         before_serve: vec![],
         after_stop: vec![],
+        solo: None,
         mods: Default::default(),
         shm_size: String::new(),
         sync_image: true,
@@ -246,6 +247,7 @@ fn sglang_recipe() -> Recipe {
         distributed_backend: Some(DistributedBackend::Mp),
         before_serve: vec![],
         after_stop: vec![],
+        solo: None,
         mods: Default::default(),
         shm_size: String::new(),
         sync_image: true,
@@ -475,3 +477,6 @@ fn entrypoint_script_non_sh_passes_through() {
     let cmd = run(&r, &tcls().nodes[0], true);
     assert!(cmd.contains("docker run"));
 }
+
+#[path = "solo_tests.rs"]
+mod solo_tests;

+ 1 - 0
src/doctor.rs

@@ -370,6 +370,7 @@ mod tests {
             distributed_backend: Some(crate::config::DistributedBackend::Mp),
             before_serve: vec![],
             after_stop: vec![],
+            solo: None,
             mods: Default::default(),
             sync_image: true,
             docker_image_dir: None,

+ 3 - 0
src/main.rs

@@ -95,6 +95,7 @@ async fn main() -> Result<()> {
             sync_image_source,
             build_node,
             force_build,
+            solo,
         } => {
             handle_serve(
                 &config,
@@ -110,6 +111,7 @@ async fn main() -> Result<()> {
                     sync_image_source,
                     build_node,
                     force_build,
+                    solo,
                 },
                 follow,
                 dry_run,
@@ -231,6 +233,7 @@ async fn handle_build_cmd(
         sync_image_source: None,
         build_node: Some(build_node.to_string()),
         force_build,
+        solo: None,
     };
     docker::build::handle_build(config, &cluster, recipe, &recipe_obj, &opts, build_node).await?;
     Ok(())