Pārlūkot izejas kodu

refactor: integrate health.rs, deduplicate select_cluster, fix Ray backend, add 21 tests

- health.rs: removed dead_code silencers, refactored check_cluster_health to
  be a pure data function (no println), used as backend for tmux_status()
- tmux.rs: removed duplicate select_cluster (now uses crate::select_cluster),
  integrated health::check_cluster_health into tmux_status
- main.rs: added #[serde(default)] for nccl in Config; fixed Ray backend flags
  (added --pipeline-parallel-size, --nnodes, --distributed-executor-backend ray)
- Added 21 unit tests: docker command generation (5), select_cluster (4),
  recipe/config deserialization (6), generate_repair_dockerfile (3),
  defaults (2), select_cluster_by_name (1)
- Fixed clippy warnings: needless_borrows, nonminimal_bool, collapsible_if,
  let_and_return, manual_flatten, len_zero, single_match, eq_op
- AGENTS.md: updated project layout, commands, recipe fields, dependencies
Alexander Korolev 4 mēneši atpakaļ
vecāks
revīzija
750a89980c
5 mainītis faili ar 602 papildinājumiem un 184 dzēšanām
  1. 93 9
      AGENTS.md
  2. 19 52
      src/health.rs
  3. 437 13
      src/main.rs
  4. 2 2
      src/ssh.rs
  5. 51 108
      src/tmux.rs

+ 93 - 9
AGENTS.md

@@ -1,18 +1,20 @@
 # AGENTS.md — vllm-spark-launcher
 
-Rust CLI tool for launching vLLM/SGLang containers across a DGX Spark cluster. SSH-based provisioning over `russh`.
+Rust CLI tool for launching vLLM/SGLang containers across a DGX Spark cluster. SSH-based provisioning over `russh`. Supports two runtimes: **Docker** (default) and **tmux** (alternative).
 
 ## Project Layout
 
 - `Cargo.toml` — binary crate `vllm-spark-launcher`. Edition 2024.
-- `src/main.rs` — single-file binary: CLI parsing, SSH, docker orchestration.
+- `src/main.rs` — CLI parsing, Docker orchestration, Doctor mode, health check, log streaming.
+- `src/ssh.rs` — SSH client session management, key loading, host key verification, retry logic.
+- `src/tmux.rs` — tmux session manager: start/stop/status/logs/attach across a cluster.
+- `src/health.rs` — tmux health-check helpers (unused, see below).
 - `config.toml` — runtime config: SSH key, cluster definitions, per-cluster NCCL settings.
 - `recipes/<name>/recipe.toml` — recipe files.
 
-No `build.rs`, no workspace, no additional modules. Everything is in `main.rs`.
-
 ## Commands
 
+### Docker runtime
 ```
 cargo run serve --recipe glm51                             # Auto-cluster selection
 cargo run serve --recipe glm51 --cluster 4nodes-switch     # Explicit cluster
@@ -20,8 +22,21 @@ cargo run -- -d serve --recipe glm51                       # Dry-run (print only
 cargo run stop --recipe glm51                              # Stop containers on all nodes
 cargo run health                                           # Health check all clusters
 cargo run health --cluster 2nodes-dac                      # Single cluster
+cargo run health --fingerprint                             # Show host keys for config.toml
 cargo run logs --recipe glm51                              # Last 100 lines per node
 cargo run logs --recipe glm51 --follow                     # Follow head node logs
+cargo run doctor --recipe glm51                            # Check node prerequisites
+cargo run doctor --recipe glm51 --repair                   # Auto-fix (build repaired image)
+```
+
+### Tmux runtime
+```
+cargo run tmux serve --recipe glm51                        # Launch via tmux (no Docker)
+cargo run tmux stop --recipe glm51                         # Stop tmux sessions
+cargo run tmux status --recipe glm51                       # Detailed health per node
+cargo run tmux logs --recipe glm51                         # Capture pane output
+cargo run tmux logs --recipe glm51 --follow                # Stream pane output
+cargo run tmux attach --recipe glm51                       # Interactive attach (head node)
 ```
 
 All subprocess calls support `-d` / `--dry-run` as a global flag.
@@ -38,6 +53,7 @@ ssh_key_path = "/absolute/path/to/key"
 name = "my-cluster"
 master_ib_ip = "1.2.3.4"
 nodes = [
+    # host_key = "SHA256:..."  # run 'cargo run health --fingerprint' to discover
     { name = "dgx1.lan", ssh_ip = "10.0.0.1", ib_ip = "1.2.3.4", rank = 0 },
     { name = "dgx2.lan", ssh_ip = "10.0.0.2", ib_ip = "1.2.3.5", rank = 1 },
 ]
@@ -49,6 +65,14 @@ ib_gid_index = 3
 ib_disable = 0
 ```
 
+### Recipe-level cluster override
+
+A recipe can pin itself to a specific cluster (bypasses `--cluster` / auto-pick):
+
+```toml
+cluster = "4nodes-switch"
+```
+
 ## recipes/
 
 TOML format. Directory name is intentionally misspelled. Each recipe lives in its own subdirectory:
@@ -59,21 +83,36 @@ name = "Human Readable Name"
 docker_image = "my-image:tag"
 model = "org/model-name"
 port = 8000
-required_nodes = 4                     # 0 = any cluster size
+required_nodes = 4                     # optional: 0 = any cluster size, omitted = no check
 pipeline_parallel_size = 0             # 0 = auto (equals nnodes)
+distributed_backend = "mp"             # "mp" (multiprocessing) or "ray"
+cluster = "4nodes-switch"              # optional: pin to a specific cluster
 
 args = [
     "--kv-cache-dtype", "fp8",
     "--max-model-len", "32768",
 ]
 
+# Optional: override docker entrypoint (appended after image name)
+entrypoint = ["python", "-m", "vllm.entrypoints.openai.api_server"]
+
 [env]
 SOME_VAR = "value"
 ```
 
-## How launch works
+### Engine support
+
+Both `vllm` and `sglang` engines are supported. Engine name is used at runtime for host IP variables:
+- `vllm` → `VLLM_HOST_IP`
+- `sglang` → `SGLANG_WORKER_IP`
+
+## How Docker launch works
 
-`build_docker_command` in `main.rs` constructs the docker run command for each node. It always appends these flags (cannot be overridden by recipe):
+`build_docker_command` in `main.rs` constructs the `docker run` command for each node. Flavour depends on `distributed_backend`:
+
+### MP backend (default)
+
+Always appends these flags (cannot be overridden by recipe):
 
 ```
 --tensor-parallel-size 1
@@ -84,15 +123,60 @@ SOME_VAR = "value"
 --distributed-executor-backend mp
 ```
 
-NCCL env vars come from the cluster's `[nccl."<name>"]` section in `config.toml`.
+### Ray backend
+
+Appends:
+
+```
+--tensor-parallel-size <nnodes>
+--node-rank <rank>
+--master-addr <cluster.master_ib_ip>
+```
+
+Plus `-e RAY_ADDRESS=<master_ib_ip>:6379` and `-e RAY_PORT=<port>`.
+
+### Common to all runs
+
+NCCL env vars come from the cluster's `[nccl."<name>"]` section in `config.toml`. Port mapping (`-p`) is applied on **all nodes** (not just head).
+
+Containers are labelled with:
+- `sm121.recipe=<name>`
+- `sm121.rank=<rank>`
 
 `--cluster auto` picks the first cluster in `config.toml` that has an NCCL section.
 
+## How tmux launch works
+
+`tmux serve` builds sessions on each node with a 3-window layout:
+
+| Window | Content |
+|--------|---------|
+| 0: `model` | Main inference process (vLLM/SGLang) |
+| 1: `gpu` | `nvidia-smi -l 5` GPU monitor |
+| 2: `shell` | Debug shell |
+
+For **Ray backend**, nodes start sequentially: head node first (polled for API readiness up to 150s), then workers. For **MP backend**, all nodes start in parallel.
+
+## Doctor mode
+
+`cargo run doctor --recipe <name>` checks each node for:
+- Docker availability
+- NVIDIA GPU (nvidia-smi)
+- NCCL libraries
+- Engine version inside the image (`vllm version` / `sglang version`)
+- Ray availability (if backend is "ray")
+- Container entrypoint
+
+With `--repair`, generates a Dockerfile that installs missing packages and rebuilds the image with a `.repaired` tag. The repaired image is auto-detected on subsequent `serve` runs.
+
 ## Dependencies
 
-- `clap` 4.6 — derive-based CLI
+- `clap` 4.6 + `derive` — derive-based CLI
 - `russh` 0.60 + `russh-keys` 0.49 — SSH client
 - `tokio` full — async runtime
 - `toml` 1 + `serde` 1 — config/recipe deserialization
 - `tracing` + `tracing-subscriber` — logging
 - `anyhow` — error handling
+- `futures` 0.3 — `join_all` for concurrent SSH
+- `sha2` 0.10 — SHA-256 for host key fingerprints
+- `base64` 0.22 — Base64 encoding for fingerprints

+ 19 - 52
src/health.rs

@@ -1,12 +1,10 @@
 use anyhow::Result;
-use tracing::{info, warn};
 
 use crate::tmux::{SessionHealth, TmuxManager};
 use crate::{Cluster, Config, Recipe};
 
 /// Health status of a node's tmux session
 #[derive(Debug, Clone, PartialEq)]
-#[allow(dead_code)]
 pub enum SessionStatus {
     Running,
     Stopped,
@@ -16,7 +14,6 @@ pub enum SessionStatus {
 
 /// Health report for a single node
 #[derive(Debug, Clone)]
-#[allow(dead_code)]
 pub struct NodeHealth {
     pub name: String,
     pub ib_ip: String,
@@ -29,7 +26,6 @@ pub struct NodeHealth {
 
 /// Overall health report for a cluster
 #[derive(Debug, Clone)]
-#[allow(dead_code)]
 pub struct HealthReport {
     pub nodes: Vec<NodeHealth>,
     pub overall: HealthStatus,
@@ -37,16 +33,15 @@ pub struct HealthReport {
 
 /// Overall cluster health status
 #[derive(Debug, Clone, PartialEq)]
-#[allow(dead_code)]
 pub enum HealthStatus {
     Healthy,
     Degraded,
     Unhealthy,
 }
 
-#[allow(dead_code)]
 /// Checks the health of all tmux sessions for a recipe across a cluster.
 /// Uses deep health checks: session existence, process liveness, GPU status, API reachability.
+/// Returns a [`HealthReport`] with per-node details — no output is printed by this function.
 pub async fn check_cluster_health(
     manager: &TmuxManager,
     config: &Config,
@@ -54,9 +49,6 @@ pub async fn check_cluster_health(
     recipe: &Recipe,
     dry_run: bool,
 ) -> Result<HealthReport> {
-    println!("\n=== Cluster: {} (deep health) ===", cluster.name);
-
-    let mut all_healthy = true;
     let mut nodes: Vec<NodeHealth> = Vec::new();
 
     let mut handles = Vec::new();
@@ -81,17 +73,6 @@ pub async fn check_cluster_health(
         match result {
             Ok((name, ib_ip, rank, Ok(health))) => {
                 let status = classify_status(&health);
-
-                if status == SessionStatus::Running {
-                    info!("  ✓ {} ({}) — healthy", name, ib_ip);
-                } else if status == SessionStatus::Degraded {
-                    warn!("  ⚠ {} ({}) — degraded", name, ib_ip);
-                    all_healthy = false;
-                } else {
-                    warn!("  ✗ {} ({}) — {}", name, ib_ip, status_msg(&status));
-                    all_healthy = false;
-                }
-
                 nodes.push(NodeHealth {
                     name,
                     ib_ip,
@@ -102,9 +83,7 @@ pub async fn check_cluster_health(
                     status,
                 });
             }
-            Ok((name, ib_ip, rank, Err(e))) => {
-                warn!("  ✗ {} ({}) — SSH error: {}", name, ib_ip, e);
-                all_healthy = false;
+            Ok((name, ib_ip, rank, Err(_))) => {
                 nodes.push(NodeHealth {
                     name,
                     ib_ip,
@@ -115,34 +94,17 @@ pub async fn check_cluster_health(
                     status: SessionStatus::Unknown,
                 });
             }
-            Err(join_err) => {
-                warn!("  ✗ Task join error: {}", join_err);
-                all_healthy = false;
-            }
+            Err(_) => {}
         }
     }
 
-    if !dry_run {
-        if all_healthy {
-            info!("✅ Cluster is healthy!");
-        } else {
-            warn!("⚠️ Cluster has issues");
-        }
-    }
+    let overall = compute_overall_status(&nodes);
 
-    Ok(HealthReport {
-        nodes,
-        overall: if all_healthy {
-            HealthStatus::Healthy
-        } else {
-            HealthStatus::Unhealthy
-        },
-    })
+    Ok(HealthReport { nodes, overall })
 }
 
 /// Classifies overall session status from detailed health data.
-#[allow(dead_code)]
-fn classify_status(health: &SessionHealth) -> SessionStatus {
+pub fn classify_status(health: &SessionHealth) -> SessionStatus {
     if health.is_healthy() {
         SessionStatus::Running
     } else if health.is_degraded() || (health.session_exists && !health.gpu_ok) {
@@ -152,13 +114,18 @@ fn classify_status(health: &SessionHealth) -> SessionStatus {
     }
 }
 
-/// Human-readable message for a session status.
-#[allow(dead_code)]
-fn status_msg(status: &SessionStatus) -> &'static str {
-    match status {
-        SessionStatus::Running => "healthy",
-        SessionStatus::Degraded => "degraded",
-        SessionStatus::Stopped => "not found",
-        SessionStatus::Unknown => "unknown",
+/// Computes overall cluster health from per-node statuses.
+fn compute_overall_status(nodes: &[NodeHealth]) -> HealthStatus {
+    let has_unknown = nodes.iter().any(|n| n.status == SessionStatus::Unknown);
+    let has_stopped = nodes.iter().any(|n| n.status == SessionStatus::Stopped);
+    let has_degraded = nodes.iter().any(|n| n.status == SessionStatus::Degraded);
+    let all_running = nodes.iter().all(|n| n.status == SessionStatus::Running);
+
+    if all_running {
+        HealthStatus::Healthy
+    } else if !has_unknown && !has_stopped && has_degraded {
+        HealthStatus::Degraded
+    } else {
+        HealthStatus::Unhealthy
     }
 }

+ 437 - 13
src/main.rs

@@ -148,6 +148,7 @@ struct Config {
     ssh_user: String,
     ssh_key_path: String,
     clusters: Vec<Cluster>,
+    #[serde(default)]
     nccl: HashMap<String, NcclConfig>,
 }
 
@@ -333,11 +334,11 @@ async fn show_logs(config: &Config, cluster: &Cluster, recipe_name: &str, follow
 
         let results = futures::future::join_all(handles).await;
         for result in results {
-            if let Ok((name, Ok(output))) = result {
-                if !output.trim().is_empty() {
-                    println!("\n--- {} ---", name);
-                    println!("{}", output);
-                }
+            if let Ok((name, Ok(output))) = result
+                && !output.trim().is_empty()
+            {
+                println!("\n--- {} ---", name);
+                println!("{}", output);
             }
         }
     }
@@ -562,6 +563,7 @@ async fn start_cluster(
     Ok(())
 }
 
+#[allow(clippy::too_many_arguments)]
 fn build_docker_command(
     cluster: &Cluster,
     nccl: &NcclConfig,
@@ -615,11 +617,17 @@ fn build_docker_command(
         "ray" => {
             cmd.extend(vec![
                 "--tensor-parallel-size".to_string(),
+                "1".to_string(),
+                "--pipeline-parallel-size".to_string(),
+                pp_size.to_string(),
+                "--nnodes".to_string(),
                 nnodes.to_string(),
                 "--node-rank".to_string(),
                 node.rank.to_string(),
                 "--master-addr".to_string(),
                 cluster.master_ib_ip.clone(),
+                "--distributed-executor-backend".to_string(),
+                "ray".to_string(),
             ]);
         }
         "mp" => {
@@ -709,16 +717,13 @@ async fn doctor(
         let backend = recipe.distributed_backend.clone();
         let image = recipe.docker_image.clone();
         handles.push(tokio::spawn(async move {
-            let result = check_node(&config, &node, &engine, &backend, &image).await;
-            result
+            check_node(&config, &node, &engine, &backend, &image).await
         }));
     }
 
     let results = futures::future::join_all(handles).await;
-    for result in results {
-        if let Ok(r) = result {
-            all_results.push(r);
-        }
+    for r in results.into_iter().flatten() {
+        all_results.push(r);
     }
 
     // Print results
@@ -735,7 +740,7 @@ async fn doctor(
 
     if has_failures {
         if repair {
-            println!("\n🔧 Reparing image...");
+            println!("\n🔧 Repairing image...");
             let repaired = run_cluster_repair(config, &cluster, &recipe, dry_run).await?;
             println!("✅ Repaired image: {}", repaired);
         } else {
@@ -783,7 +788,7 @@ async fn check_node(
 /// Run a command on the node via SSH and return a check result
 async fn node_check(config: &Config, node: &Node, name: &str, command: &str) -> DoctorCheck {
     match run_ssh(config, node, command.to_string(), false).await {
-        Ok(output) if output.trim().len() > 0 => DoctorCheck {
+        Ok(output) if !output.trim().is_empty() => DoctorCheck {
             name: name.to_string(),
             status: DoctorStatus::Ok,
             detail: format!("OK: {}", output.trim().lines().next().unwrap_or(output.trim())),
@@ -948,3 +953,422 @@ async fn try_repaired_image(config: &Config, cluster: &Cluster, recipe: &Recipe)
         _ => Ok(recipe.docker_image.clone()),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::collections::HashMap;
+
+    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,
+                },
+                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,
+                },
+            ],
+        }
+    }
+
+    fn test_nccl() -> NcclConfig {
+        NcclConfig {
+            ib_hca: "rocep1s0f0".to_string(),
+            socket_ifname: "enp1s0f0np0".to_string(),
+            ib_gid_index: 3,
+            ib_disable: 0,
+        }
+    }
+
+    fn test_recipe() -> Recipe {
+        Recipe {
+            engine: "vllm".to_string(),
+            name: "Test Model".to_string(),
+            docker_image: "test/image:latest".to_string(),
+            model: "test/model".to_string(),
+            port: 8000,
+            required_nodes: Some(2),
+            pipeline_parallel_size: Some(2),
+            args: vec!["--max-model-len".to_string(), "4096".to_string()],
+            env: HashMap::from([("NCCL_DEBUG".to_string(), "INFO".to_string())]),
+            cluster: None,
+            entrypoint: vec![],
+            distributed_backend: "mp".to_string(),
+        }
+    }
+
+    // ── build_docker_command tests ──────────────────────────────
+
+    #[test]
+    fn test_build_docker_command_mp_backend() {
+        let cluster = test_cluster();
+        let nccl = test_nccl();
+        let recipe = test_recipe();
+        let node = &cluster.nodes[0];
+
+        let cmd = build_docker_command(&cluster, &nccl, &recipe, "test-recipe", node, true, 2, 2);
+
+        let joined = cmd.join(" ");
+
+        // Basic docker flags
+        assert!(joined.contains("docker run -d --gpus all --ipc=host --privileged --network=host"));
+        // NCCL env
+        assert!(joined.contains("-e NCCL_IB_HCA=rocep1s0f0"));
+        assert!(joined.contains("-e NCCL_SOCKET_IFNAME=enp1s0f0np0"));
+        // Recipe env
+        assert!(joined.contains("-e NCCL_DEBUG=INFO"));
+        // Port only on head
+        assert!(joined.contains("-p 8000:8000"));
+        // Container name & labels
+        assert!(joined.contains("--name sm121_test-recipe_0"));
+        assert!(joined.contains("--label sm121.recipe=test-recipe"));
+        assert!(joined.contains("--label sm121.rank=0"));
+        // Image & model
+        assert!(joined.contains("test/image:latest"));
+        assert!(joined.contains("test/model"));
+        // Args
+        assert!(joined.contains("--max-model-len 4096"));
+        // MP-specific flags
+        assert!(joined.contains("--tensor-parallel-size 1"));
+        assert!(joined.contains("--pipeline-parallel-size 2"));
+        assert!(joined.contains("--nnodes 2"));
+        assert!(joined.contains("--node-rank 0"));
+        assert!(joined.contains("--master-addr 192.168.1.1"));
+        assert!(joined.contains("--distributed-executor-backend mp"));
+        // No Ray env
+        assert!(!joined.contains("RAY_ADDRESS"));
+    }
+
+    #[test]
+    fn test_build_docker_command_worker_no_port() {
+        let cluster = test_cluster();
+        let nccl = test_nccl();
+        let recipe = test_recipe();
+        let node = &cluster.nodes[1]; // rank 1
+
+        let cmd = build_docker_command(&cluster, &nccl, &recipe, "test-recipe", node, false, 2, 2);
+        let joined = cmd.join(" ");
+
+        // No port mapping on worker
+        assert!(!joined.contains("-p 8000:8000"));
+        // Correct rank
+        assert!(joined.contains("--node-rank 1"));
+        assert!(joined.contains("--label sm121.rank=1"));
+        assert!(joined.contains("--name sm121_test-recipe_1"));
+    }
+
+    #[test]
+    fn test_build_docker_command_ray_backend() {
+        let cluster = test_cluster();
+        let nccl = test_nccl();
+        let mut recipe = test_recipe();
+        recipe.distributed_backend = "ray".to_string();
+        let node = &cluster.nodes[0];
+
+        let cmd = build_docker_command(&cluster, &nccl, &recipe, "test-recipe", node, true, 2, 2);
+        let joined = cmd.join(" ");
+
+        // Ray-specific env
+        assert!(joined.contains("-e RAY_ADDRESS=192.168.1.1:6379"));
+        assert!(joined.contains("-e RAY_PORT=8000"));
+        // Ray-specific flags
+        assert!(joined.contains("--tensor-parallel-size 1"));
+        assert!(joined.contains("--pipeline-parallel-size 2"));
+        assert!(joined.contains("--nnodes 2"));
+        assert!(joined.contains("--node-rank 0"));
+        assert!(joined.contains("--master-addr 192.168.1.1"));
+        assert!(joined.contains("--distributed-executor-backend ray"));
+    }
+
+    #[test]
+    fn test_build_docker_command_auto_pp_size() {
+        let cluster = test_cluster();
+        let nccl = test_nccl();
+        let mut recipe = test_recipe();
+        recipe.pipeline_parallel_size = Some(0); // auto
+        let node = &cluster.nodes[0];
+
+        // pp_size = 0 should become nnodes (2) in start_cluster, but build_docker_command
+        // takes pp_size directly. test with pp_size = nnodes = 2.
+        let cmd = build_docker_command(&cluster, &nccl, &recipe, "test-recipe", node, true, 2, 2);
+        let joined = cmd.join(" ");
+
+        assert!(joined.contains("--pipeline-parallel-size 2"));
+    }
+
+    #[test]
+    fn test_build_docker_command_entrypoint() {
+        let cluster = test_cluster();
+        let nccl = test_nccl();
+        let mut recipe = test_recipe();
+        recipe.entrypoint = vec!["python".to_string(), "-m".to_string(), "custom.module".to_string()];
+        let node = &cluster.nodes[0];
+
+        let cmd = build_docker_command(&cluster, &nccl, &recipe, "test-recipe", node, true, 2, 2);
+        let joined = cmd.join(" ");
+
+        // Entrypoint should appear after image name, before model
+        let image_pos = joined.find("test/image:latest").unwrap();
+        let model_pos = joined.find("test/model").unwrap();
+        let entrypoint_pos = joined.find("python -m custom.module").unwrap();
+        assert!(image_pos < entrypoint_pos && entrypoint_pos < model_pos,
+            "entrypoint should be between image and model");
+    }
+
+    // ── select_cluster tests ────────────────────────────────────
+
+    #[test]
+    fn test_select_cluster_by_name() {
+        let mut config = Config {
+            ssh_user: "user".to_string(),
+            ssh_key_path: "/path/to/key".to_string(),
+            clusters: vec![test_cluster()],
+            nccl: HashMap::new(),
+        };
+        config.nccl.insert("test-cluster".to_string(), test_nccl());
+
+        let (cluster, nccl) = select_cluster(&config, "test-cluster").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".to_string(),
+            clusters: vec![test_cluster()],
+            nccl: HashMap::new(),
+        };
+        config.nccl.insert("test-cluster".to_string(), test_nccl());
+
+        let (cluster, nccl) = select_cluster(&config, "auto").unwrap();
+        assert_eq!(cluster.name, "test-cluster");
+        assert_eq!(nccl.ib_hca, "rocep1s0f0");
+    }
+
+    #[test]
+    fn test_select_cluster_auto_no_match() {
+        let config = Config {
+            ssh_user: "user".to_string(),
+            ssh_key_path: "/path/to/key".to_string(),
+            clusters: vec![test_cluster()],
+            nccl: HashMap::new(),
+        };
+        // No NCCL config for test-cluster → auto should fail
+        let result = select_cluster(&config, "auto");
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_select_cluster_not_found() {
+        let config = Config {
+            ssh_user: "user".to_string(),
+            ssh_key_path: "/path/to/key".to_string(),
+            clusters: vec![],
+            nccl: HashMap::new(),
+        };
+        let result = select_cluster(&config, "nonexistent");
+        assert!(result.is_err());
+    }
+
+    // ── Recipe deserialization tests ─────────────────────────────
+
+    #[test]
+    fn test_recipe_deserialize_minimal() {
+        let toml_str = r#"
+name = "Minimal"
+model = "test/model"
+"#;
+        let recipe: Recipe = toml::from_str(toml_str).unwrap();
+        assert_eq!(recipe.name, "Minimal");
+        assert_eq!(recipe.model, "test/model");
+        assert_eq!(recipe.engine, "vllm"); // default
+        assert_eq!(recipe.port, 0); // default
+        assert!(recipe.args.is_empty());
+        assert!(recipe.env.is_empty());
+        assert_eq!(recipe.distributed_backend, "mp"); // default
+    }
+
+    #[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, "ray");
+        assert_eq!(recipe.cluster, Some("my-cluster".to_string()));
+        assert_eq!(recipe.args, 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); // None, not Some(0)
+        assert_eq!(recipe.pipeline_parallel_size, None);
+    }
+
+    // ── Config deserialization tests ───────────────────────────────
+
+    #[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"));
+        assert_eq!(config.nccl["prod"].ib_hca, "mlx5_0");
+    }
+
+    #[test]
+    fn test_config_deserialize_with_host_key() {
+        let toml_str = r#"
+ssh_user = "admin"
+ssh_key_path = "/path/to/key"
+
+[[clusters]]
+name = "secure"
+master_ib_ip = "1.1.1.1"
+nodes = [
+    { name = "n1", ssh_ip = "2.2.2.2", ib_ip = "1.1.1.1", rank = 0, host_key = "SHA256:abc123" },
+]
+"#;
+        let config: Config = toml::from_str(toml_str).unwrap();
+        assert_eq!(
+            config.clusters[0].nodes[0].host_key,
+            Some("SHA256:abc123".to_string())
+        );
+    }
+
+    // ── generate_repair_dockerfile tests ───────────────────────────
+
+    #[test]
+    fn test_generate_repair_dockerfile_vllm() {
+        let recipe = test_recipe();
+        let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);
+
+        assert!(dockerfile.contains("FROM base/image:tag"));
+        assert!(dockerfile.contains("pip install vllm"));
+        assert!(!dockerfile.contains("ray"));
+        assert!(dockerfile.contains("ENTRYPOINT"));
+    }
+
+    #[test]
+    fn test_generate_repair_dockerfile_ray() {
+        let mut recipe = test_recipe();
+        recipe.distributed_backend = "ray".to_string();
+        let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);
+
+        assert!(dockerfile.contains("FROM base/image:tag"));
+        assert!(dockerfile.contains("pip install vllm"));
+        assert!(dockerfile.contains("pip install ray"));
+    }
+
+    #[test]
+    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);
+
+        assert!(dockerfile.contains("pip install sglang"));
+        assert!(!dockerfile.contains("vllm"));
+    }
+
+    // ── select_cluster_by_name tests ───────────────────────────────
+
+    #[test]
+    fn test_select_cluster_by_name_found() {
+        let mut nccl_map = HashMap::new();
+        nccl_map.insert("test-cluster".to_string(), test_nccl());
+        let config = Config {
+            ssh_user: "u".to_string(),
+            ssh_key_path: "/k".to_string(),
+            clusters: vec![test_cluster()],
+            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 = Config {
+            ssh_user: "u".to_string(),
+            ssh_key_path: "/k".to_string(),
+            clusters: vec![test_cluster()],
+            nccl: HashMap::new(),
+        };
+        let result = select_cluster_by_name(&config, "no-such-cluster");
+        assert!(result.is_err());
+    }
+
+    // ── Defaults ───────────────────────────────────────────────────
+
+    #[test]
+    fn test_default_engine() {
+        assert_eq!(default_engine(), "vllm");
+    }
+
+    #[test]
+    fn test_default_distributed_backend() {
+        assert_eq!(default_distributed_backend(), "mp");
+    }
+}

+ 2 - 2
src/ssh.rs

@@ -145,7 +145,7 @@ pub async fn fetch_host_key(config: &Config, node: &Node) -> Result<String> {
     let blob = captured.lock().unwrap().take()
         .context("No server key captured")?;
     let hash = Sha256::digest(&blob);
-    let fp = format!("SHA256:{}", STANDARD.encode(&hash));
+    let fp = format!("SHA256:{}", STANDARD.encode(hash));
     Ok(fp)
 }
 
@@ -170,7 +170,7 @@ impl client::Handler for VerifyingHandler {
         if let Some(ref expected_fp) = expected {
             let key_bytes = server_public_key.to_string().into_bytes();
             let hash = Sha256::digest(&key_bytes);
-            let fp = format!("SHA256:{}", STANDARD.encode(&hash));
+            let fp = format!("SHA256:{}", STANDARD.encode(hash));
             if &fp != expected_fp {
                 return Err(russh::Error::Kex);
             }

+ 51 - 108
src/tmux.rs

@@ -4,7 +4,7 @@ use tokio::time::sleep;
 use tracing::{info, warn};
 
 use crate::ssh::*;
-use crate::{Cluster, Config, NcclConfig, Node, Recipe, TmuxCommands};
+use crate::{health, Cluster, Config, NcclConfig, Node, Recipe, TmuxCommands};
 
 /// Manages tmux sessions for model inference across a DGX Spark cluster.
 #[derive(Clone)]
@@ -80,11 +80,17 @@ impl TmuxManager {
                 // Ray mode: workers connect via RAY_ADDRESS
                 parts.extend(vec![
                     "--tensor-parallel-size".to_string(),
+                    "1".to_string(),
+                    "--pipeline-parallel-size".to_string(),
+                    pp_size.to_string(),
+                    "--nnodes".to_string(),
                     nnodes.to_string(),
                     "--node-rank".to_string(),
                     node.rank.to_string(),
                     "--master-addr".to_string(),
                     cluster.master_ib_ip.clone(),
+                    "--distributed-executor-backend".to_string(),
+                    "ray".to_string(),
                 ]);
             }
             _ => {
@@ -315,7 +321,7 @@ impl TmuxManager {
             match run_ssh(config, node, api_check, dry_run).await {
                 Ok(out) => {
                     let code = out.trim();
-                    code == "200" || /* vLLM openai compat */ code == "200" || !code.contains("fail")
+                    code == "200" || !code.contains("fail")
                 }
                 Err(_) => false,
             }
@@ -551,7 +557,7 @@ async fn tmux_serve(
     } else {
         "auto".to_string()
     };
-    let (cluster, nccl) = select_cluster(config, &cluster_choice)?;
+    let (cluster, nccl) = crate::select_cluster(config, &cluster_choice)?;
 
     if let Some(required) = recipe.required_nodes {
         let actual = cluster.nodes.len() as u32;
@@ -696,15 +702,12 @@ async fn wait_for_head_ready(
             "curl -s -o /dev/null -w '%{{http_code}}' http://127.0.0.1:{}/v1/models 2>/dev/null || echo 'not_ready'",
             recipe.port
         );
-        match run_ssh(config, head_node, api_check, false).await {
-            Ok(out) => {
-                let code = out.trim();
-                if code == "200" {
-                    info!("  Head node API ready (attempt {})", i + 1);
-                    return Ok(true);
-                }
+        if let Ok(out) = run_ssh(config, head_node, api_check, false).await {
+            let code = out.trim();
+            if code == "200" {
+                info!("  Head node API ready (attempt {})", i + 1);
+                return Ok(true);
             }
-            Err(_) => {}
         }
 
         // Also check stderr/output for model load progress
@@ -747,7 +750,7 @@ async fn tmux_stop(
     } else {
         "auto".to_string()
     };
-    let (cluster, _) = select_cluster(config, &cluster_choice)?;
+    let (cluster, _) = crate::select_cluster(config, &cluster_choice)?;
 
     let manager = TmuxManager {
         session_prefix: recipe_name.to_string(),
@@ -797,7 +800,7 @@ async fn tmux_status(
     } else {
         "auto".to_string()
     };
-    let (cluster, _) = select_cluster(config, &cluster_choice)?;
+    let (cluster, _) = crate::select_cluster(config, &cluster_choice)?;
 
     let manager = TmuxManager {
         session_prefix: recipe_name.to_string(),
@@ -806,84 +809,49 @@ async fn tmux_status(
     println!("\n📊 TMUX Status: {} on cluster {}", recipe.name, cluster.name);
     println!("{}", "─".repeat(60));
 
-    let mut handles = Vec::new();
-    for node in &cluster.nodes {
-        let config = config.clone();
-        let node = node.clone();
-        let cluster = cluster.clone();
-        let recipe = recipe.clone();
-        let manager = manager.clone();
-        let dry = dry_run;
-
-        handles.push(tokio::spawn(async move {
-            let sess = manager.session_name(&recipe, node.rank);
-            let health = manager
-                .deep_check_session(&config, &node, &cluster, &sess, &recipe, dry)
-                .await;
-            (node.name, node.rank, node.ib_ip, health)
-        }));
-    }
+    // Use health::check_cluster_health for data collection
+    let report = health::check_cluster_health(&manager, config, &cluster, &recipe, dry_run).await?;
 
-    let mut all_healthy = true;
-    let results = futures::future::join_all(handles).await;
-    for result in results {
-        match result {
-            Ok((name, rank, ib_ip, Ok(health))) => {
-                let status_icon = if health.is_healthy() {
-                    "✓".to_string()
-                } else if health.session_exists {
-                    "⚠".to_string()
-                } else {
-                    "✗".to_string()
-                };
+    for node_health in &report.nodes {
+        let status_icon = if node_health.status == health::SessionStatus::Running {
+            "✓".to_string()
+        } else if node_health.session_exists {
+            "⚠".to_string()
+        } else {
+            "✗".to_string()
+        };
 
-                print!("  {} {} (rank {}, {}): ", status_icon, name, rank, ib_ip);
+        print!("  {} {} (rank {}, {}): ", status_icon, node_health.name, node_health.rank, node_health.ib_ip);
 
-                if health.is_healthy() {
-                    println!("✅ ALL OK");
-                } else {
-                    let mut issues = Vec::new();
-                    if !health.session_exists {
-                        issues.push("no session");
-                    }
-                    if !health.process_alive {
-                        issues.push("process dead");
-                    }
-                    if !health.gpu_ok {
-                        issues.push("GPU issue");
-                    }
-                    if !health.api_ok && rank == 0 {
-                        issues.push("API down");
-                    }
-                    println!("⚠️  {}", issues.join(", "));
-                    all_healthy = false;
-                }
-
-                if health.session_exists {
-                    // Show GPU stats
-                    println!("     Session: alive | Process: {} | GPU: {} | API: {}",
-                        if health.process_alive { "✓" } else { "✗" },
-                        if health.gpu_ok { "✓" } else { "✗" },
-                        if health.api_ok { "✓" } else { "✗" },
-                    );
-                }
+        if node_health.status == health::SessionStatus::Running {
+            println!("✅ ALL OK");
+        } else {
+            let mut issues = Vec::new();
+            if !node_health.session_exists {
+                issues.push("no session");
             }
-            Ok((name, _, _, Err(e))) => {
-                println!("  ✗ {} — SSH error: {}", name, e);
-                all_healthy = false;
+            if !node_health.process_running {
+                issues.push("process dead");
             }
-            Err(join_err) => {
-                println!("  ✗ Task join error: {}", join_err);
-                all_healthy = false;
+            if !node_health.port_accessible && node_health.rank == 0 {
+                issues.push("API down");
             }
+            println!("⚠️  {}", issues.join(", "));
+        }
+
+        if node_health.session_exists {
+            println!("     Session: alive | Process: {} | API: {}",
+                if node_health.process_running { "✓" } else { "✗" },
+                if node_health.port_accessible { "✓" } else { "✗" },
+            );
         }
     }
 
     println!("{}", "─".repeat(60));
-    if all_healthy {
-        info!("✅ Cluster is healthy!");
-    } else {
-        warn!("⚠️ Cluster has issues — check details above.");
+    match report.overall {
+        health::HealthStatus::Healthy => info!("✅ Cluster is healthy!"),
+        health::HealthStatus::Degraded => warn!("⚠️ Cluster is degraded"),
+        health::HealthStatus::Unhealthy => warn!("⚠️ Cluster has issues — check details above."),
     }
 
     Ok(())
@@ -909,7 +877,7 @@ async fn tmux_logs(
     } else {
         "auto".to_string()
     };
-    let (cluster, _) = select_cluster(config, &cluster_choice)?;
+    let (cluster, _) = crate::select_cluster(config, &cluster_choice)?;
 
     let manager = TmuxManager {
         session_prefix: recipe_name.to_string(),
@@ -989,7 +957,7 @@ async fn tmux_attach(
     } else {
         "auto".to_string()
     };
-    let (cluster, _) = select_cluster(config, &cluster_choice)?;
+    let (cluster, _) = crate::select_cluster(config, &cluster_choice)?;
 
     let head_node = cluster
         .nodes
@@ -1010,29 +978,4 @@ async fn tmux_attach(
     Ok(())
 }
 
-// ====================== Helpers ======================
 
-/// Selects a cluster and its NCCL config by name or auto-pick.
-fn select_cluster<'c>(config: &'c Config, cluster_name: &str) -> Result<(Cluster, &'c NcclConfig)> {
-    if cluster_name != "auto" {
-        let cluster = config
-            .clusters
-            .iter()
-            .find(|c| c.name == cluster_name)
-            .context(format!("Cluster '{}' not found", cluster_name))?;
-        let nccl = config
-            .nccl
-            .get(&cluster.name)
-            .context("NCCL config not found")?;
-        return Ok((cluster.clone(), nccl));
-    }
-
-    for cluster in &config.clusters {
-        if let Some(nccl) = config.nccl.get(&cluster.name) {
-            info!("✅ Auto-selected cluster: {}", cluster.name);
-            return Ok((cluster.clone(), nccl));
-        }
-    }
-
-    anyhow::bail!("No suitable cluster found (none have NCCL config)");
-}