refactor(core): extract Docker command helpers, rewrite host_keys with toml_edit, extract per-command handlers

Extract build_docker_flags(), collect_nccl_env(), and collect_docker_env()
from build_docker_command() (reduced from 116 to 50 lines).

Extract per-command handler functions from main() match dispatch,
reducing the dispatch block from ~230 to 40 lines.

Replace fragile line-by-line TOML heuristic in host_keys.rs with
toml_edit round-trip parsing for safe, format-preserving updates.

Additional robustness:
- generate_repair_dockerfile() returns Result instead of panicking
- home_assistant get_state() logs HTTP/JSON errors before returning None
- ha_controller_from_config() logs HA client construction failures

Document design rationale:
- daisy-chain StrictHostKeyChecking risk in image sync
- fire-and-forget SSH shutdown semantics in power off
- optimistic resolve_cluster_simple assumption in orchestrate
- stderr loss on mods Warn/Continue || true
- no-delay power_switch path when disable_nodes is empty

Cargo.lock

@@ -3452,6 +3452,7 @@ dependencies = [
  "sha2 0.11.0",
  "tokio",
  "toml",
+ "toml_edit",
  "tracing",
  "tracing-subscriber",
  "walkdir",
@@ -3797,12 +3798,18 @@ dependencies = [
  "indexmap",
  "serde_core",
  "serde_spanned",
- "toml_datetime",
+ "toml_datetime 1.1.1+spec-1.1.0",
  "toml_parser",
  "toml_writer",
- "winnow",
+ "winnow 1.0.3",
 ]
 
+[[package]]
+name = "toml_datetime"
+version = "0.6.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "22cddaf88f4fbc13c51aebbf5f8eceb5c7c5a9da2ac40a13519eb5b0a0e8f11c"
+
 [[package]]
 name = "toml_datetime"
 version = "1.1.1+spec-1.1.0"
@@ -3812,15 +3819,33 @@ dependencies = [
  "serde_core",
 ]
 
+[[package]]
+name = "toml_edit"
+version = "0.22.27"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a"
+dependencies = [
+ "indexmap",
+ "toml_datetime 0.6.11",
+ "toml_write",
+ "winnow 0.7.15",
+]
+
 [[package]]
 name = "toml_parser"
 version = "1.1.2+spec-1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526"
 dependencies = [
- "winnow",
+ "winnow 1.0.3",
 ]
 
+[[package]]
+name = "toml_write"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801"
+
 [[package]]
 name = "toml_writer"
 version = "1.1.1+spec-1.1.0"
@@ -4616,6 +4641,15 @@ version = "0.53.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650"
 
+[[package]]
+name = "winnow"
+version = "0.7.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "winnow"
 version = "1.0.3"

Cargo.toml

@@ -12,6 +12,7 @@ anyhow = "1.0"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1"
 toml = "1"
+toml_edit = "0.22"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 sha2 = "0.11"

src/docker/mod.rs

@@ -343,8 +343,79 @@ async fn start_cluster(
     }
     Ok(())
 }
-fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> Vec<String> {
+/// Collect NCCL cluster env vars for multi-node deployments.
+fn collect_nccl_env<'a>(ctx: &'a DockerCommandCtx, node: &'a Node) -> Vec<(&'a str, Cow<'a, str>)> {
     let is_sglang = ctx.recipe.engine == "sglang";
+    let sock = ctx.nccl.socket_ifname.as_str();
+    let gloo_if = ctx.nccl.gloo_socket_ifname.as_deref().unwrap_or(sock);
+    let ucx_dev = ctx.nccl.ucx_net_devices.as_deref().unwrap_or(sock);
+    let mut env: Vec<(&str, Cow<str>)> = vec![
+        ("NCCL_IB_DISABLE", format!("{}", ctx.nccl.ib_disable).into()),
+        ("NCCL_IB_HCA", ctx.nccl.ib_hca.as_str().into()),
+        ("NCCL_SOCKET_IFNAME", sock.into()),
+        ("GLOO_SOCKET_IFNAME", gloo_if.into()),
+        ("TP_SOCKET_IFNAME", gloo_if.into()),
+        (
+            "NCCL_IB_GID_INDEX",
+            ctx.nccl.ib_gid_index.to_string().into(),
+        ),
+        ("MN_IF_NAME", gloo_if.into()),
+        ("UCX_NET_DEVICES", ucx_dev.into()),
+        ("OMPI_MCA_btl_tcp_if_include", gloo_if.into()),
+        ("OMP_NUM_THREADS", "4".into()),
+        ("NCCL_IGNORE_CPU_AFFINITY", "1".into()),
+        ("NCCL_NET", "IB".into()),
+        ("NCCL_DEBUG", "WARN".into()),
+    ];
+    if let Some(v) = ctx.nccl.nccl_cross_nic {
+        env.push(("NCCL_CROSS_NIC", v.to_string().into()));
+    }
+    if let Some(ref v) = ctx.nccl.nccl_proto {
+        env.push(("NCCL_PROTO", v.as_str().into()));
+    }
+    if let Some(ref v) = ctx.nccl.nccl_p2p_level {
+        env.push(("NCCL_P2P_LEVEL", v.as_str().into()));
+    }
+    let worker_ip = Cow::Borrowed(node.ib_ip.as_str());
+    if is_sglang {
+        env.push(("SGLANG_WORKER_IP", worker_ip));
+    } else {
+        env.push(("VLLM_HOST_IP", worker_ip));
+    }
+    env
+}
+
+/// Build the list of `-e KEY=VALUE` Docker env var pairs for a node.
+fn collect_docker_env(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> Vec<String> {
+    let mut pairs: Vec<(&str, Cow<str>)> =
+        vec![("HF_HOME", Cow::Borrowed("/root/.cache/huggingface"))];
+    if ctx.nnodes > 1 {
+        pairs.extend(collect_nccl_env(ctx, node));
+    }
+    if is_head {
+        pairs.push(("VLLM_LOGGING_COLOR", "1".into()));
+        pairs.push(("TERM", "xterm-256color".into()));
+    }
+    if ctx.recipe.distributed_backend == Some(crate::config::DistributedBackend::Ray) {
+        pairs.push((
+            "RAY_ADDRESS",
+            format!("{}:6379", ctx.cluster.master_ib_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()));
+    }
+    for (k, v) in &ctx.recipe.env {
+        pairs.push((k.as_str(), v.as_str().into()));
+    }
+    pairs
+        .into_iter()
+        .map(|(k, v)| format!("-e {k}={v}"))
+        .collect()
+}
+
+/// Build the list of Docker run flags (--gpus, --ipc, --ulimit, etc.).
+fn build_docker_flags(ctx: &DockerCommandCtx) -> Vec<String> {
     let mut flags: Vec<String> = vec![
         "--gpus all".to_string(),
         "--ipc=host".to_string(),
@@ -361,6 +432,12 @@ fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> V
     if !ctx.keep {
         flags.push("--rm".to_string());
     }
+    flags
+}
+
+/// Build the full Docker command line for a node.
+fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> Vec<String> {
+    let flags = build_docker_flags(ctx);
     let mut cmd = vec![
         format!("docker run -d {}", flags.join(" ")),
         format!("-v {}:/root/.cache/huggingface", ctx.cache_dir),
@@ -368,53 +445,7 @@ fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> V
     if !ctx.mods_remote_path.is_empty() && !ctx.recipe.mods.commands.is_empty() {
         cmd.push(format!("-v {}:/mods:ro", ctx.mods_remote_path));
     }
-    let mut env = vec![("HF_HOME", Cow::Borrowed("/root/.cache/huggingface"))];
-    if ctx.nnodes > 1 {
-        let worker_ip = Cow::Borrowed(node.ib_ip.as_str());
-        if is_sglang {
-            env.push(("SGLANG_WORKER_IP", worker_ip));
-        } else {
-            env.push(("VLLM_HOST_IP", worker_ip));
-        }
-        let sock = ctx.nccl.socket_ifname.as_str();
-        let gloo_if = ctx.nccl.gloo_socket_ifname.as_deref().unwrap_or(sock);
-        let ucx_dev = ctx.nccl.ucx_net_devices.as_deref().unwrap_or(sock);
-        let mut ncclext: Vec<(&str, Cow<str>)> = vec![
-            ("NCCL_IB_DISABLE", format!("{}", ctx.nccl.ib_disable).into()),
-            ("NCCL_IB_HCA", ctx.nccl.ib_hca.as_str().into()),
-            ("NCCL_SOCKET_IFNAME", sock.into()),
-            ("GLOO_SOCKET_IFNAME", gloo_if.into()),
-            ("TP_SOCKET_IFNAME", gloo_if.into()),
-            (
-                "NCCL_IB_GID_INDEX",
-                ctx.nccl.ib_gid_index.to_string().into(),
-            ),
-            ("MN_IF_NAME", gloo_if.into()),
-            ("UCX_NET_DEVICES", ucx_dev.into()),
-            ("OMPI_MCA_btl_tcp_if_include", gloo_if.into()),
-            ("OMP_NUM_THREADS", "4".into()),
-            ("NCCL_IGNORE_CPU_AFFINITY", "1".into()),
-            ("NCCL_NET", "IB".into()),
-            ("NCCL_DEBUG", "WARN".into()),
-        ];
-        if let Some(v) = ctx.nccl.nccl_cross_nic {
-            ncclext.push(("NCCL_CROSS_NIC", v.to_string().into()));
-        }
-        if let Some(ref v) = ctx.nccl.nccl_proto {
-            ncclext.push(("NCCL_PROTO", v.as_str().into()));
-        }
-        if let Some(ref v) = ctx.nccl.nccl_p2p_level {
-            ncclext.push(("NCCL_P2P_LEVEL", v.as_str().into()));
-        }
-        env.extend(ncclext);
-    }
-    if is_head {
-        env.push(("VLLM_LOGGING_COLOR", "1".into()));
-        env.push(("TERM", "xterm-256color".into()));
-    }
-    for (k, v) in &ctx.recipe.env {
-        env.push((k.as_str(), v.as_str().into()));
-    }
+    cmd.extend(collect_docker_env(ctx, node, is_head));
     cmd.push(format!("--name sm121_{}_{}", ctx.recipe_name, node.rank));
     cmd.push(format!("--label sm121.recipe={}", ctx.recipe_name));
     cmd.push(format!("--label sm121.rank={}", node.rank));
@@ -423,18 +454,6 @@ fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> V
     if ctx.cluster.nodes.len() == 1 && ctx.cluster.extra_entities.is_empty() {
         cmd.push(format!("--label sm121.solo-node={}", node.name));
     }
-    if ctx.recipe.distributed_backend == Some(crate::config::DistributedBackend::Ray) {
-        env.push((
-            "RAY_ADDRESS",
-            format!("{}:6379", ctx.cluster.master_ib_ip).into(),
-        ));
-        env.push(("RAY_PORT", "6379".into()));
-        env.push(("RAY_NODE_IP_ADDRESS", node.ib_ip.as_str().into()));
-        env.push(("RAY_OVERRIDE_NODE_IP_ADDRESS", node.ib_ip.as_str().into()));
-    }
-    for (k, v) in env {
-        cmd.push(format!("-e {k}={v}"));
-    }
     let tctx = tail::TailCtx {
         recipe: &ctx.recipe,
         tp_size: ctx.tp_size,

src/docker/mods.rs

@@ -34,6 +34,10 @@ pub fn build_mod_script(mods: &[PreServeCommand], tail: &[String]) -> String {
                 script.push_str(&format!("{cmd} || exit 1; "));
             }
             OnFailure::Warn | OnFailure::Continue => {
+                // NOTE: `|| true` swallows the exit code and stderr is
+                // lost to the container log.  If mod debugging is needed,
+                // consider redirecting stderr to a temp file, e.g.:
+                //   cmd 2>/tmp/sm121-mod-err.log; true
                 script.push_str(&format!("{cmd} || true; "));
             }
         }

src/docker/sync.rs

@@ -27,6 +27,15 @@ async fn get_image_created(config: &Config, node: &Node, image: &str) -> Option<
 /// Sync Docker image across cluster nodes.
 /// Finds the node with the *most recently built* image (by creation timestamp),
 /// then pipes `docker save | ssh docker load` to nodes missing or with a different digest.
+///
+/// **Security note:** The daisy-chain SSH pipe (`docker save | ssh -o
+/// StrictHostKeyChecking=no ... docker load`) bypasses host key verification for
+/// the *inner* SSH hop.  The outer SSH session (via russh) is verified, but the
+/// piped `ssh` command trusts the target host unconditionally for the duration
+/// of the pipe.  This is a design trade-off: re-using the verified russh session
+/// for the pipe would require piping the data through the agent, which adds
+/// complexity.  Acceptable for trusted-cluster deployments where the config
+/// file is the trust anchor.
 pub async fn sync_image_to_nodes(
     config: &Config,
     cluster: &Cluster,

src/doctor.rs

@@ -246,7 +246,8 @@ async fn run_cluster_repair(
 ) -> Result<String> {
     let head_node = cluster.nodes.first().context("No nodes in cluster")?;
     let repaired_tag = format!("{}.repaired", recipe.docker_image);
-    let dockerfile = generate_repair_dockerfile(&repaired_tag, recipe);
+    let dockerfile = generate_repair_dockerfile(&repaired_tag, recipe)
+        .context("Failed to generate repair Dockerfile")?;
 
     let dockerfile_path = format!("recipes/{}/Dockerfile.repaired", "temp");
     tokio::fs::write(&dockerfile_path, &dockerfile)
@@ -297,9 +298,11 @@ fn is_valid_docker_image(image: &str) -> bool {
 }
 
 /// Generate a Dockerfile to repair a container.
-fn generate_repair_dockerfile(base_image: &str, recipe: &Recipe) -> String {
+/// Returns an error if `base_image` contains unsafe characters that
+/// could lead to Dockerfile injection.
+fn generate_repair_dockerfile(base_image: &str, recipe: &Recipe) -> Result<String> {
     if !is_valid_docker_image(base_image) {
-        panic!("Invalid docker image name: {}", base_image);
+        anyhow::bail!("Invalid docker image name: {}", base_image);
     }
     let mut lines = vec![format!("FROM {}", base_image)];
 
@@ -317,7 +320,7 @@ fn generate_repair_dockerfile(base_image: &str, recipe: &Recipe) -> String {
     }
 
     lines.push("ENTRYPOINT [\"/bin/bash\", \"-c\"]".to_string());
-    lines.join("\n")
+    Ok(lines.join("\n"))
 }
 
 /// Check if a repaired image exists on the head node.
@@ -376,7 +379,7 @@ mod tests {
     #[test]
     fn test_generate_repair_dockerfile_vllm() {
         let recipe = test_recipe();
-        let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);
+        let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe).unwrap();
 
         assert!(dockerfile.contains("FROM base/image:tag"));
         assert!(dockerfile.contains("pip install vllm"));
@@ -388,7 +391,7 @@ mod tests {
     fn test_generate_repair_dockerfile_ray() {
         let mut recipe = test_recipe();
         recipe.distributed_backend = Some(crate::config::DistributedBackend::Ray);
-        let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);
+        let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe).unwrap();
 
         assert!(dockerfile.contains("FROM base/image:tag"));
         assert!(dockerfile.contains("pip install vllm"));
@@ -399,7 +402,7 @@ mod tests {
     fn test_generate_repair_dockerfile_sglang() {
         let mut recipe = test_recipe();
         recipe.engine = "sglang".to_string();
-        let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);
+        let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe).unwrap();
 
         assert!(dockerfile.contains("pip install sglang"));
         assert!(!dockerfile.contains("vllm"));

src/ends.rs

@@ -40,6 +40,7 @@ pub trait NodePinger {
 use crate::home_assistant::Client as HaClient;
 use crate::home_assistant::warn_if_token_looks_short;
 use crate::ssh::run_ssh;
+use tracing::warn;
 
 pub struct SshProber;
 
@@ -81,7 +82,13 @@ impl NodePinger for SshNodePinger {
 /// not configured.
 pub fn ha_controller_from_config(config: &Config) -> Option<HaPowerController> {
     let ha_cfg = config.home_assistant.as_ref()?;
-    let ha = HaClient::new(ha_cfg).ok()?;
+    let ha = match HaClient::new(ha_cfg) {
+        Ok(c) => c,
+        Err(e) => {
+            warn!("Failed to construct HA client: {}", e);
+            return None;
+        }
+    };
     warn_if_token_looks_short(&ha_cfg.token);
     Some(HaPowerController::new(ha))
 }

src/home_assistant.rs

@@ -102,19 +102,29 @@ impl Client {
 
     /// GET `/api/states/<entity_id>` and return the state string.
     /// Returns `None` if the entity doesn't exist or isn't reachable.
+    /// Logs a warning when the HTTP request or JSON parsing fails so
+    /// callers (e.g. `entity_is_on`) can distinguish "entity is off"
+    /// from "HA is unreachable" via the log.
     pub async fn get_state(&self, entity_id: &str) -> Option<String> {
         let url = format!("{}/api/states/{}", self.base_url, entity_id);
-        let resp = self
-            .http
-            .get(&url)
-            .bearer_auth(&self.token)
-            .send()
-            .await
-            .ok()?;
+        let resp = match self.http.get(&url).bearer_auth(&self.token).send().await {
+            Ok(r) => r,
+            Err(e) => {
+                warn!("HA GET {} failed: {}", url, e);
+                return None;
+            }
+        };
         if !resp.status().is_success() {
+            warn!("HA GET {} returned {}", url, resp.status());
             return None;
         }
-        let body: serde_json::Value = resp.json().await.ok()?;
+        let body: serde_json::Value = match resp.json().await {
+            Ok(v) => v,
+            Err(e) => {
+                warn!("HA GET {} JSON parse failed: {}", url, e);
+                return None;
+            }
+        };
         body.get("state").and_then(|s| s.as_str()).map(String::from)
     }
 }

src/host_keys.rs

@@ -1,5 +1,8 @@
 //! Auto-discover SSH host key fingerprints for every node in the config
 //! and rewrite `config.toml` with `host_key = "SHA256:..."` entries.
+//!
+//! Uses `toml_edit` for round-trip-safe TOML manipulation so formatting
+//! and comments in the config file are preserved.
 
 use anyhow::{Context, Result};
 use std::fs;
@@ -31,7 +34,7 @@ pub async fn init_host_keys(
 
     let content = fs::read_to_string(config_path)
         .with_context(|| format!("Failed to read {}", config_path))?;
-    let new_content = update_host_keys_text(&content, &updates);
+    let new_content = update_host_keys_toml(&content, &updates);
 
     if content == new_content {
         println!("\nNo changes needed in {}", config_path);
@@ -68,75 +71,43 @@ async fn collect_updates(
     Ok(out)
 }
 
-/// Walk `content` line by line and rewrite any node definition whose
-/// `name = "..."` matches an entry in `updates`, setting or updating the
-/// `host_key` field.
-fn update_host_keys_text(content: &str, updates: &[(String, String, String)]) -> String {
-    let mut lines: Vec<String> = content.lines().map(String::from).collect();
-    for (name, _ip, fp) in updates {
-        let marker = format!("name = \"{}\"", name);
-        for line in &mut lines {
-            if is_node_line(line) && line.contains(&marker) {
-                *line = update_node_line(line, fp);
-            }
-        }
-    }
-    if content.ends_with('\n') {
-        lines.join("\n") + "\n"
-    } else {
-        lines.join("\n")
-    }
-}
+/// Parse `content` as TOML, walk the `[[clusters]]` tables, find each
+/// node whose `name` matches an entry in `updates`, and set the
+/// `host_key` field on it.  Preserves comments and formatting.
+fn update_host_keys_toml(content: &str, updates: &[(String, String, String)]) -> String {
+    use toml_edit::{DocumentMut, Value};
 
-/// Heuristic: a node definition starts (after indent) with `{`.
-fn is_node_line(line: &str) -> bool {
-    line.trim_start().starts_with('{')
-}
+    let mut doc = content.parse::<DocumentMut>().expect("valid TOML config");
+    let clusters = doc
+        .get_mut("clusters")
+        .and_then(|v| v.as_array_of_tables_mut());
+
+    let Some(clusters) = clusters else {
+        return content.to_string();
+    };
 
-/// Set or replace the `host_key = "..."` field in a single node line.
-///
-/// If the field is already present, its value is swapped for `fp`.
-/// Otherwise the field is inserted before the trailing `}` with a comma
-/// if needed.
-fn update_node_line(line: &str, fp: &str) -> String {
-    let marker = "host_key = \"";
-    if let Some(start) = line.find(marker) {
-        let value_start = start + marker.len();
-        if let Some(rel_end) = line[value_start..].find('"') {
-            let value_end = value_start + rel_end;
-            let mut result = String::with_capacity(line.len());
-            result.push_str(&line[..value_start]);
-            result.push_str(fp);
-            result.push_str(&line[value_end..]);
-            return result;
+    for cluster in clusters.iter_mut() {
+        let Some(nodes) = cluster.get_mut("nodes") else {
+            continue;
+        };
+        let Some(arr) = nodes.as_array_mut() else {
+            continue;
+        };
+        for item in arr.iter_mut() {
+            let Some(table) = item.as_inline_table_mut() else {
+                continue;
+            };
+            let name_val = table.get("name").and_then(|v| v.as_str());
+            let Some(name) = name_val else {
+                continue;
+            };
+            if let Some((_n, _ip, fp)) = updates.iter().find(|(n, _, _)| n == name) {
+                table.insert("host_key", Value::from(fp.as_str()));
+            }
         }
     }
 
-    // No host_key yet — insert before the closing `}`.
-    let close_idx = line.rfind('}').expect("node line must contain '}'");
-    let before_close = &line[..close_idx];
-    let after_close = &line[close_idx..];
-    let trimmed = before_close.trim_end();
-    let needs_comma = !trimmed.is_empty() && !trimmed.ends_with(',');
-    // Original line either ends with `}` or `},`. Detect which so we can
-    // keep the formatting consistent (space before `}` only when a comma
-    // follows it).
-    let has_trailing_comma = after_close.starts_with("},");
-    let mut result = String::with_capacity(line.len() + fp.len() + 20);
-    result.push_str(trimmed);
-    if needs_comma {
-        result.push(',');
-    }
-    result.push(' ');
-    result.push_str(&format!("host_key = \"{}\"", fp));
-    if has_trailing_comma {
-        result.push(' ');
-    }
-    result.push('}');
-    if has_trailing_comma {
-        result.push(',');
-    }
-    result
+    doc.to_string()
 }
 
 #[cfg(test)]
@@ -145,50 +116,110 @@ mod tests {
 
     #[test]
     fn adds_host_key_when_missing() {
-        let line =
-            r#"    { name = "dgx1.lan", ssh_ip = "10.0.0.1", ib_ip = "1.2.3.4", rank = 0 },"#;
-        let result = update_node_line(line, "SHA256:abc");
-        assert_eq!(
-            result,
-            r#"    { name = "dgx1.lan", ssh_ip = "10.0.0.1", ib_ip = "1.2.3.4", rank = 0, host_key = "SHA256:abc" },"#
-        );
+        let content = r#"
+[[clusters]]
+name = "c"
+nodes = [
+    { name = "dgx1.lan", ssh_ip = "10.0.0.1", ib_ip = "1.2.3.4", rank = 0 },
+]
+"#;
+        let updates = vec![(
+            "dgx1.lan".to_string(),
+            "10.0.0.1".to_string(),
+            "SHA256:abc".to_string(),
+        )];
+        let result = update_host_keys_toml(content, &updates);
+        assert!(result.contains(r#"host_key = "SHA256:abc""#));
+        assert!(result.contains("dgx1.lan"));
     }
 
     #[test]
     fn replaces_host_key_value() {
-        let line = r#"    { name = "x", rank = 0, host_key = "OLD" },"#;
-        let result = update_node_line(line, "NEW");
-        assert_eq!(result, r#"    { name = "x", rank = 0, host_key = "NEW" },"#);
+        let content = r#"
+[[clusters]]
+name = "c"
+nodes = [
+    { name = "x", rank = 0, host_key = "OLD" },
+]
+"#;
+        let updates = vec![("x".to_string(), "10.0.0.1".to_string(), "NEW".to_string())];
+        let result = update_host_keys_toml(content, &updates);
+        assert!(result.contains(r#"host_key = "NEW""#));
+        assert!(!result.contains(r#"host_key = "OLD""#));
     }
 
     #[test]
-    fn update_host_keys_text_targets_only_named_lines() {
-        let content = "\
+    fn only_updates_targeted_nodes() {
+        let content = r#"
 [[clusters]]
-name = \"c\"
+name = "c"
 nodes = [
-    { name = \"a\", rank = 0 },
-    { name = \"b\", rank = 1 },
+    { name = "a", rank = 0 },
+    { name = "b", rank = 1 },
 ]
-";
+"#;
         let updates = vec![(
             "a".to_string(),
             "10.0.0.1".to_string(),
             "SHA256:newA".to_string(),
         )];
-        let result = update_host_keys_text(content, &updates);
-        // a was updated to include host_key
-        assert!(result.contains("name = \"a\", rank = 0, host_key = \"SHA256:newA\""));
-        // b was NOT updated — original line preserved, no host_key added
-        assert!(result.contains("name = \"b\", rank = 1 }"));
-        assert!(!result.contains("name = \"b\", rank = 1, host_key"));
+        let result = update_host_keys_toml(content, &updates);
+        assert!(
+            result.contains(r#"host_key = "SHA256:newA""#),
+            "result: {result}"
+        );
+        // b should NOT have host_key added
+        assert!(result.contains(r#"name = "b""#), "result has b: {result}");
+        assert!(
+            !result.contains("host_key") || result.matches("host_key").count() == 1,
+            "only 'a' should have host_key: {result}"
+        );
+    }
+
+    #[test]
+    fn updates_multiple_clusters() {
+        let content = r#"
+[[clusters]]
+name = "c1"
+nodes = [
+    { name = "n1", rank = 0 },
+]
+
+[[clusters]]
+name = "c2"
+nodes = [
+    { name = "n2", rank = 0 },
+]
+"#;
+        let updates = vec![
+            ("n1".to_string(), "10.0.0.1".to_string(), "FP1".to_string()),
+            ("n2".to_string(), "10.0.0.2".to_string(), "FP2".to_string()),
+        ];
+        let result = update_host_keys_toml(content, &updates);
+        assert!(
+            result.contains(r#"host_key = "FP1""#),
+            "result has FP1: {result}"
+        );
+        assert!(
+            result.contains(r#"host_key = "FP2""#),
+            "result has FP2: {result}"
+        );
+    }
+
+    #[test]
+    fn preserves_trailing_newline() {
+        let content = "[[clusters]]\nname = \"c\"\nnodes = []\n";
+        let result = update_host_keys_toml(content, &[]);
+        assert_eq!(result, content);
     }
 
     #[test]
-    fn update_node_line_with_no_trailing_comma() {
-        // Last field already ends with comma — no extra comma needed.
-        let line = r#"    { name = "x", rank = 0 }"#;
-        let result = update_node_line(line, "FP");
-        assert_eq!(result, r#"    { name = "x", rank = 0, host_key = "FP"}"#);
+    fn no_host_key_for_absent_update() {
+        let content = r#"[[clusters]]
+name = "c"
+nodes = [{ name = "n1", rank = 0 }]
+"#;
+        let result = update_host_keys_toml(content, &[]);
+        assert_eq!(result, content);
     }
 }

src/main.rs

@@ -22,7 +22,7 @@ mod pre_serve;
 mod ssh;
 mod tmux;
 
-use crate::cli::{Cli, Commands};
+use crate::cli::{Cli, Commands, PowerAction};
 use crate::config::{
     Cluster, Config, default_config_dir, load_recipe, resolve_cluster, select_cluster_by_name,
 };
@@ -93,11 +93,11 @@ async fn main() -> Result<()> {
             build_node,
             force_build,
         } => {
-            let cluster = docker::serve(
+            handle_serve(
                 &config,
                 &recipe,
                 &cluster,
-                node.as_deref(),
+                node,
                 docker::ServeOpts {
                     keep,
                     skip_before_serve,
@@ -108,13 +108,10 @@ async fn main() -> Result<()> {
                     build_node,
                     force_build,
                 },
+                follow,
+                dry_run,
             )
             .await?;
-            if follow && !dry_run {
-                tracing::info!("{}", fluent_i18n::t!("container-log-header"));
-                docker::show_logs(&config, &cluster, &recipe, true, node.as_deref(), dry_run)
-                    .await?;
-            }
         }
         Commands::Build {
             recipe,
@@ -122,25 +119,13 @@ async fn main() -> Result<()> {
             cluster,
             force_build,
         } => {
-            let recipe_obj = crate::config::load_recipe(&recipe).await?;
-            let cluster = resolve_cluster_for_op(&config, &recipe, &cluster).await?;
-            let opts = docker::ServeOpts {
-                keep: false,
-                skip_before_serve: true,
-                dry_run,
-                auto_power_on: false,
-                sync_image: false,
-                sync_image_source: None,
-                build_node: Some(build_node.clone()),
-                force_build,
-            };
-            docker::build::handle_build(
+            handle_build_cmd(
                 &config,
-                &cluster,
                 &recipe,
-                &recipe_obj,
-                &opts,
                 &build_node,
+                &cluster,
+                force_build,
+                dry_run,
             )
             .await?;
         }
@@ -150,16 +135,7 @@ async fn main() -> Result<()> {
             follow,
             node,
         } => {
-            let cluster = resolve_cluster_for_op(&config, &recipe_name, &cluster).await?;
-            docker::show_logs(
-                &config,
-                &cluster,
-                &recipe_name,
-                follow,
-                node.as_deref(),
-                dry_run,
-            )
-            .await?;
+            handle_logs(&config, &recipe_name, &cluster, follow, node, dry_run).await?;
         }
         Commands::Stop {
             recipe: recipe_name,
@@ -167,89 +143,13 @@ async fn main() -> Result<()> {
             node,
             power_off,
         } => {
-            if !node.is_empty() {
-                // User specified specific nodes — stop only those.
-                let cluster = resolve_cluster_for_op(&config, &recipe_name, &cluster).await?;
-                docker::stop_cluster(&config, &cluster, &recipe_name, &node, dry_run).await?;
-            } else {
-                // Discover running deployments and resolve ambiguity.
-                let variants =
-                    docker::deploy::discover_recipe_deployments(&config, &recipe_name).await?;
-                if variants.is_empty() {
-                    anyhow::bail!("No running containers found for recipe '{}'", recipe_name);
-                }
-                if variants.len() == 1 {
-                    let v = &variants[0];
-                    let cluster = select_cluster_by_name(&config, &v.cluster)?;
-                    docker::stop_cluster(&config, &cluster, &recipe_name, &v.nodes, dry_run)
-                        .await?;
-                } else {
-                    let mut msg = format!(
-                        "Multiple deployments of '{}' found — use --node to select:\n",
-                        recipe_name
-                    );
-                    for (i, v) in variants.iter().enumerate() {
-                        let nodes = v
-                            .nodes
-                            .iter()
-                            .map(|n| format!("--node {}", n))
-                            .collect::<Vec<_>>()
-                            .join(" ");
-                        msg.push_str(&format!(
-                            "  {}: cluster '{}' with {}\n",
-                            i + 1,
-                            v.cluster,
-                            nodes
-                        ));
-                    }
-                    anyhow::bail!(msg);
-                }
-            }
-            // Handle power_off flag
-            if power_off {
-                let cluster = resolve_cluster_for_op(&config, &recipe_name, &cluster).await?;
-                let node_str = if node.is_empty() {
-                    None
-                } else {
-                    Some(node[0].as_str())
-                };
-                crate::power::power_off(&config, &cluster, node_str, dry_run).await?;
-            }
+            handle_stop(&config, &recipe_name, &cluster, &node, power_off, dry_run).await?;
         }
         Commands::Health {
             cluster: cluster_name,
             fingerprint,
         } => {
-            if fingerprint {
-                let clusters_to_check: Vec<Cluster> = match &cluster_name {
-                    Some(name) => vec![select_cluster_by_name(&config, name)?],
-                    None => config.clusters.clone(),
-                };
-                for cluster in &clusters_to_check {
-                    println!("\n=== Кластер: {} ===", cluster.name);
-                    for node in &cluster.nodes {
-                        match fetch_host_key(&config, node).await {
-                            Ok(fp) => {
-                                println!("{} ({}) -> host_key = \"{}\"", node.name, node.ssh_ip, fp)
-                            }
-                            Err(e) => println!("{} ({}) -> error: {}", node.name, node.ssh_ip, e),
-                        }
-                    }
-                }
-            } else {
-                match cluster_name {
-                    Some(ref name) => {
-                        let cluster = select_cluster_by_name(&config, name)?;
-                        docker::health::health_check(&config, &cluster, dry_run).await?;
-                    }
-                    None => {
-                        for cluster in &config.clusters {
-                            docker::health::health_check(&config, cluster, dry_run).await?;
-                        }
-                        docker::health::health_summary(&config).await?;
-                    }
-                }
-            }
+            handle_health(&config, cluster_name, fingerprint).await?;
         }
         Commands::Doctor {
             recipe,
@@ -276,30 +176,213 @@ async fn main() -> Result<()> {
             alias,
             command,
         } => {
-            let (cluster, _recipe) =
-                exec::resolve_target(&config, cluster.as_deref(), recipe.as_deref()).await?;
-            // Resolve command: alias takes precedence, otherwise join
-            // the trailing positional args into a single command line.
-            let cmd_str = exec::resolve_command(&config, alias.as_deref(), &command)?;
-            exec::run(&config, &cluster, node.as_deref(), &cmd_str, dry_run).await?;
+            handle_exec(&config, cluster, recipe, node, alias, command, dry_run).await?;
         }
-        Commands::Power { action } => match action {
-            cli::PowerAction::Off { cluster, node } => {
-                let cluster = select_cluster_by_name(&config, &cluster)?;
-                power::power_off(&config, &cluster, node.as_deref(), dry_run).await?;
+        Commands::Power { action } => {
+            handle_power(&config, action, dry_run).await?;
+        }
+    }
+
+    Ok(())
+}
+
+// ---------------------------------------------------------------------------
+// Per-command handler functions
+// ---------------------------------------------------------------------------
+
+/// Handle `serve` command: launch inference containers on a cluster.
+async fn handle_serve(
+    config: &Config,
+    recipe: &str,
+    cluster: &str,
+    node: Option<String>,
+    opts: docker::ServeOpts,
+    follow: bool,
+    dry_run: bool,
+) -> Result<()> {
+    let cluster = docker::serve(config, recipe, cluster, node.as_deref(), opts).await?;
+    if follow && !dry_run {
+        tracing::info!("{}", fluent_i18n::t!("container-log-header"));
+        docker::show_logs(config, &cluster, recipe, true, node.as_deref(), dry_run).await?;
+    }
+    Ok(())
+}
+
+/// Handle `build` command: build a Docker image on a remote node.
+async fn handle_build_cmd(
+    config: &Config,
+    recipe: &str,
+    build_node: &str,
+    cluster: &str,
+    force_build: bool,
+    dry_run: bool,
+) -> Result<()> {
+    let recipe_obj = crate::config::load_recipe(recipe).await?;
+    let cluster = resolve_cluster_for_op(config, recipe, cluster).await?;
+    let opts = docker::ServeOpts {
+        keep: false,
+        skip_before_serve: true,
+        dry_run,
+        auto_power_on: false,
+        sync_image: false,
+        sync_image_source: None,
+        build_node: Some(build_node.to_string()),
+        force_build,
+    };
+    docker::build::handle_build(config, &cluster, recipe, &recipe_obj, &opts, build_node).await?;
+    Ok(())
+}
+
+/// Handle `logs` command: show container logs.
+async fn handle_logs(
+    config: &Config,
+    recipe_name: &str,
+    cluster: &str,
+    follow: bool,
+    node: Option<String>,
+    dry_run: bool,
+) -> Result<()> {
+    let cluster = resolve_cluster_for_op(config, recipe_name, cluster).await?;
+    docker::show_logs(
+        config,
+        &cluster,
+        recipe_name,
+        follow,
+        node.as_deref(),
+        dry_run,
+    )
+    .await?;
+    Ok(())
+}
+
+/// Handle `stop` command: stop containers and optionally power off.
+async fn handle_stop(
+    config: &Config,
+    recipe_name: &str,
+    cluster: &str,
+    node: &[String],
+    power_off: bool,
+    dry_run: bool,
+) -> Result<()> {
+    if !node.is_empty() {
+        let cluster = resolve_cluster_for_op(config, recipe_name, cluster).await?;
+        docker::stop_cluster(config, &cluster, recipe_name, node, dry_run).await?;
+    } else {
+        let variants = docker::deploy::discover_recipe_deployments(config, recipe_name).await?;
+        if variants.is_empty() {
+            anyhow::bail!("No running containers found for recipe '{}'", recipe_name);
+        }
+        if variants.len() == 1 {
+            let v = &variants[0];
+            let cluster = select_cluster_by_name(config, &v.cluster)?;
+            docker::stop_cluster(config, &cluster, recipe_name, &v.nodes, dry_run).await?;
+        } else {
+            let mut msg = format!(
+                "Multiple deployments of '{}' found — use --node to select:\n",
+                recipe_name
+            );
+            for (i, v) in variants.iter().enumerate() {
+                let nodes = v
+                    .nodes
+                    .iter()
+                    .map(|n| format!("--node {}", n))
+                    .collect::<Vec<_>>()
+                    .join(" ");
+                msg.push_str(&format!(
+                    "  {}: cluster '{}' with {}\n",
+                    i + 1,
+                    v.cluster,
+                    nodes
+                ));
+            }
+            anyhow::bail!(msg);
+        }
+    }
+    if power_off {
+        let cluster = resolve_cluster_for_op(config, recipe_name, cluster).await?;
+        let node_str = if node.is_empty() {
+            None
+        } else {
+            Some(node[0].as_str())
+        };
+        crate::power::power_off(config, &cluster, node_str, dry_run).await?;
+    }
+    Ok(())
+}
+
+/// Handle `health` command: check cluster health or show host key fingerprints.
+async fn handle_health(
+    config: &Config,
+    cluster_name: Option<String>,
+    fingerprint: bool,
+) -> Result<()> {
+    if fingerprint {
+        let clusters_to_check: Vec<Cluster> = match &cluster_name {
+            Some(name) => vec![select_cluster_by_name(config, name)?],
+            None => config.clusters.clone(),
+        };
+        for cluster in &clusters_to_check {
+            println!("\n=== Кластер: {} ===", cluster.name);
+            for node in &cluster.nodes {
+                match fetch_host_key(config, node).await {
+                    Ok(fp) => {
+                        println!("{} ({}) -> host_key = \"{}\"", node.name, node.ssh_ip, fp)
+                    }
+                    Err(e) => println!("{} ({}) -> error: {}", node.name, node.ssh_ip, e),
+                }
             }
-            cli::PowerAction::On { cluster, node } => {
-                let cluster = select_cluster_by_name(&config, &cluster)?;
-                power::power_on(&config, &cluster, node.as_deref(), dry_run).await?;
+        }
+    } else {
+        match cluster_name {
+            Some(ref name) => {
+                let cluster = select_cluster_by_name(config, name)?;
+                docker::health::health_check(config, &cluster, false).await?;
             }
-            cli::PowerAction::Switch { from, to } => {
-                let from = select_cluster_by_name(&config, &from)?;
-                let to = select_cluster_by_name(&config, &to)?;
-                power::power_switch(&config, &from, &to, dry_run).await?;
+            None => {
+                for cluster in &config.clusters {
+                    docker::health::health_check(config, cluster, false).await?;
+                }
+                docker::health::health_summary(config).await?;
             }
-        },
+        }
     }
+    Ok(())
+}
+
+/// Handle `exec` command: run a command on cluster nodes via SSH.
+async fn handle_exec(
+    config: &Config,
+    cluster: Option<String>,
+    recipe: Option<String>,
+    node: Option<String>,
+    alias: Option<String>,
+    command: Vec<String>,
+    dry_run: bool,
+) -> Result<()> {
+    let (cluster, _recipe) =
+        exec::resolve_target(config, cluster.as_deref(), recipe.as_deref()).await?;
+    let cmd_str = exec::resolve_command(config, alias.as_deref(), &command)?;
+    exec::run(config, &cluster, node.as_deref(), &cmd_str, dry_run).await?;
+    Ok(())
+}
 
+/// Handle `power` command: power on/off/switch nodes.
+async fn handle_power(config: &Config, action: PowerAction, dry_run: bool) -> Result<()> {
+    match action {
+        PowerAction::Off { cluster, node } => {
+            let cluster = select_cluster_by_name(config, &cluster)?;
+            power::power_off(config, &cluster, node.as_deref(), dry_run).await?;
+        }
+        PowerAction::On { cluster, node } => {
+            let cluster = select_cluster_by_name(config, &cluster)?;
+            power::power_on(config, &cluster, node.as_deref(), dry_run).await?;
+        }
+        PowerAction::Switch { from, to } => {
+            let from = select_cluster_by_name(config, &from)?;
+            let to = select_cluster_by_name(config, &to)?;
+            power::power_switch(config, &from, &to, dry_run).await?;
+        }
+    }
     Ok(())
 }
 

src/orchestrate/mod.rs

@@ -123,7 +123,8 @@ fn resolve_solo_node(config: &Config, node_name: &str) -> Result<(Cluster, NcclC
 
 /// Resolve cluster by name or auto treating every config node as free
 /// (reachable, not occupied). Only used for the power-on pre-check
-/// before any probe has been made.
+/// before any probe has been made, so being optimistic is safe:
+/// a full probe follows immediately after power-on completes.
 fn resolve_cluster_simple<'c>(
     config: &'c Config,
     recipe: &Recipe,

src/power.rs

@@ -51,6 +51,11 @@ pub async fn power_off(
     }
 
     // 1) SSH shutdown every target in parallel.
+    // NOTE: This is fire-and-forget — if a node is unreachable, its
+    // SSH failure is logged but the flow continues.  The HA `turn_off`
+    // still fires for all targets regardless of SSH success.  This is
+    // by design: already-offline nodes get their HA entity toggled off,
+    // which is idempotent.
     info!(
         "🛑 power off: {} node(s) — sudo shutdown now",
         targets.len()
@@ -311,6 +316,9 @@ pub async fn power_switch(
             fluent_i18n::t!("power-switch-waiting", { "secs" => delay })
         );
         tokio::time::sleep(Duration::from_secs(delay)).await;
+    } else {
+        // No nodes to shut down — only extra_entities differ.  HA
+        // turn_off / turn_on runs immediately without delay.
     }
 
     if let Some(ha_cfg) = &config.home_assistant {