Sfoglia il codice sorgente

refactor(docker): move health_check to submodule

Move health_check function to src/docker/health.rs to keep
docker/mod.rs under the 500-line limit after adding --node
support and solo mode (resolve_solo_node).
Alexander Korolev 4 mesi fa
parent
commit
03c293c4ca
4 ha cambiato i file con 133 aggiunte e 108 eliminazioni
  1. 4 0
      src/cli.rs
  2. 74 0
      src/docker/health.rs
  3. 51 105
      src/docker/mod.rs
  4. 4 3
      src/main.rs

+ 4 - 0
src/cli.rs

@@ -49,6 +49,10 @@ pub enum Commands {
     Stop {
         #[arg(long, required = true)]
         recipe: String,
+
+        /// Target a specific node within the cluster (default: all).
+        #[arg(long)]
+        node: Option<String>,
     },
 
     /// Проверить здоровье кластера

+ 74 - 0
src/docker/health.rs

@@ -0,0 +1,74 @@
+use crate::config::{Cluster, Config};
+use crate::ssh::run_ssh_with_retry;
+use anyhow::Result;
+
+/// Checks Docker health on a cluster — shows all containers and marks ours.
+pub async fn health_check(config: &Config, cluster: &Cluster, dry_run: bool) -> Result<()> {
+    println!("\n=== Cluster: {} ===", cluster.name);
+    let mut handles = Vec::new();
+    for node in &cluster.nodes {
+        let config = config.clone();
+        let node = node.clone();
+        let dry = dry_run;
+        handles.push(tokio::spawn(async move {
+            let result = run_ssh_with_retry(
+                &config,
+                &node,
+                "docker ps --format '{{.Names}}\t{{.Status}}\t{{.Image}}' 2>&1; echo '---SM121---'; docker ps --filter label=sm121.recipe --format '{{.Names}}' 2>&1"
+                    .to_string(),
+                dry,
+                1,
+            )
+            .await;
+            (node.name, node.ib_ip, result)
+        }));
+    }
+    for result in futures::future::join_all(handles).await {
+        match result {
+            Ok((name, ib_ip, Ok(output))) => {
+                println!("{} ({})", name, ib_ip);
+                let parts: Vec<&str> = output.splitn(2, "---SM121---").collect();
+                let all_lines: Vec<&str> = parts
+                    .first()
+                    .map(|s| s.lines().filter(|l| !l.is_empty()).collect())
+                    .unwrap_or_default();
+                let sm121_names: Vec<&str> = parts
+                    .get(1)
+                    .map(|s| s.lines().filter(|l| !l.is_empty()).collect())
+                    .unwrap_or_default();
+                if dry_run && all_lines.is_empty() {
+                    continue;
+                }
+                if output.starts_with("docker:") || output.contains("permission denied") {
+                    println!("  ✗ Docker unavailable");
+                    if let Some(line) = output.lines().next() {
+                        println!("    └─ {}", line);
+                    }
+                } else if all_lines.is_empty() {
+                    println!("  (no containers)");
+                } else {
+                    println!("  Total containers: {}", all_lines.len());
+                    for line in &all_lines {
+                        let name_col = line.split('\t').next().unwrap_or("");
+                        let tag = if sm121_names.contains(&name_col) {
+                            "sm121"
+                        } else {
+                            "other"
+                        };
+                        println!("  {tag:5} {line}");
+                    }
+                    if sm121_names.is_empty() {
+                        println!("  (no sm121 containers)");
+                    }
+                }
+            }
+            Ok((name, ib_ip, Err(e))) => {
+                println!("{} ({}): ✗ SSH error: {}", name, ib_ip, e);
+            }
+            Err(join_err) => {
+                println!("✗ Task join error: {}", join_err);
+            }
+        }
+    }
+    Ok(())
+}

+ 51 - 105
src/docker/mod.rs

@@ -1,10 +1,11 @@
+pub mod health;
+
 use crate::config::{Cluster, Config, NcclConfig, Node, Recipe, load_recipe, resolve_cluster};
 use crate::doctor::try_repaired_image;
 use crate::ssh::{run_ssh, run_ssh_follow, run_ssh_with_retry};
 use anyhow::{Context, Result};
 use tokio::time::{Duration, sleep};
 use tracing::info;
-/// Launches containers across the cluster for a recipe.
 pub async fn serve(
     config: &Config,
     recipe_name: &str,
@@ -22,7 +23,37 @@ pub async fn serve(
         "{}",
         fluent_i18n::t!("recipe-starting", { "name" => rn, "engine" => re, "backend" => rb })
     );
-    let (cluster, nccl) = resolve_cluster(config, &recipe, cluster_name)?;
+    let (cluster, nccl): (Cluster, NcclConfig) = if let Some(nn) = node_name {
+        let (solo_c, solo_n) = config
+            .clusters
+            .iter()
+            .find_map(|c| {
+                c.nodes.iter().find(|n| n.name == nn).map(|node| {
+                    let solo = Cluster {
+                        name: format!("solo/{}", node.name),
+                        master_ib_ip: node.ib_ip.clone(),
+                        nodes: vec![node.clone()],
+                        extra_entities: vec![],
+                    };
+                    let nc = config
+                        .nccl
+                        .get(&c.name)
+                        .cloned()
+                        .unwrap_or_else(|| NcclConfig {
+                            ib_hca: String::new(),
+                            socket_ifname: "lo".to_string(),
+                            ib_gid_index: 3,
+                            ib_disable: 1,
+                        });
+                    (solo, nc)
+                })
+            })
+            .context(format!("Node '{}' not found in any cluster", nn))?;
+        (solo_c, solo_n)
+    } else {
+        let (c, n) = resolve_cluster(config, &recipe, cluster_name)?;
+        (c, n.clone())
+    };
     if let Some(required) = recipe.required_nodes {
         let actual = cluster.nodes.len() as u32;
         if actual != required {
@@ -45,7 +76,7 @@ pub async fn serve(
     start_cluster(
         config,
         &cluster,
-        nccl,
+        &nccl,
         recipe_name,
         &recipe,
         cache_dir,
@@ -60,25 +91,18 @@ pub async fn serve(
     }
     Ok(())
 }
-/// Builds the `docker logs` shell command used to fetch the last `tail`
-/// lines of the container for `recipe_name` on a specific `rank`.
 pub(crate) fn logs_tail_cmd(recipe_name: &str, rank: u32, tail: usize) -> String {
     format!(
         "docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank={}\" | xargs -r docker logs --timestamps --tail={} 2>&1",
         recipe_name, rank, tail
     )
 }
-/// Builds the `docker logs -f` shell command used to follow the head
-/// container of `recipe_name` in real time.
 pub(crate) fn logs_follow_cmd(recipe_name: &str) -> String {
     format!(
         "docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank=0\" | xargs -r docker logs -f --timestamps 2>&1",
         recipe_name
     )
 }
-/// Shell-quote `s` (single-quote + escape embedded `'` as `'\''`).
-/// Used for entrypoint scripts passed to `sh -c` so bash treats them
-/// as a single argv element in the docker run command.
 fn shq(s: &str) -> String {
     let mut out = String::with_capacity(s.len() + 2);
     out.push('\'');
@@ -92,7 +116,6 @@ fn shq(s: &str) -> String {
     out.push('\'');
     out
 }
-/// Shows Docker logs for a recipe across all nodes (or follows head node).
 pub async fn show_logs(
     config: &Config,
     cluster: &Cluster,
@@ -141,14 +164,24 @@ pub async fn stop_cluster(
     config: &Config,
     cluster: &Cluster,
     recipe_name: &str,
+    node_name: Option<&str>,
     dry_run: bool,
 ) -> Result<()> {
     info!(
         "{}",
         fluent_i18n::t!("cluster-stopping", { "name" => &cluster.name })
     );
+    let nodes: Vec<&Node> = if let Some(nn) = node_name {
+        let v: Vec<&Node> = cluster.nodes.iter().filter(|n| n.name == nn).collect();
+        if v.is_empty() {
+            anyhow::bail!("Node '{}' not found in cluster '{}'", nn, cluster.name);
+        }
+        v
+    } else {
+        cluster.nodes.iter().collect()
+    };
     let mut failures = Vec::new();
-    for node in &cluster.nodes {
+    for node in &nodes {
         let stop_cmd = format!(
             "docker ps -q --filter \"label=sm121.recipe={}\" | xargs -r docker stop",
             recipe_name
@@ -184,78 +217,7 @@ pub async fn stop_cluster(
     info!("{}", fluent_i18n::t!("cluster-all-stopped"));
     Ok(())
 }
-/// Checks Docker health on a cluster — shows all containers and marks ours.
-pub async fn health_check(config: &Config, cluster: &Cluster, dry_run: bool) -> Result<()> {
-    println!("\n=== Кластер: {} ===", cluster.name);
-    let mut handles = Vec::new();
-    for node in &cluster.nodes {
-        let config = config.clone();
-        let node = node.clone();
-        let dry = dry_run;
-        handles.push(tokio::spawn(async move {
-            // Query all containers + mark sm121-managed ones
-            let result = run_ssh_with_retry(
-                &config,
-                &node,
-                "docker ps --format '{{.Names}}\t{{.Status}}\t{{.Image}}' 2>&1; echo '---SM121---'; docker ps --filter label=sm121.recipe --format '{{.Names}}' 2>&1"
-                    .to_string(),
-                dry,
-                1,
-            )
-            .await;
-            (node.name, node.ib_ip, result)
-        }));
-    }
-    for result in futures::future::join_all(handles).await {
-        match result {
-            Ok((name, ib_ip, Ok(output))) => {
-                println!("{} ({})", name, ib_ip);
-                let parts: Vec<&str> = output.splitn(2, "---SM121---").collect();
-                let all_lines: Vec<&str> = parts
-                    .first()
-                    .map(|s| s.lines().filter(|l| !l.is_empty()).collect())
-                    .unwrap_or_default();
-                let sm121_names: Vec<&str> = parts
-                    .get(1)
-                    .map(|s| s.lines().filter(|l| !l.is_empty()).collect())
-                    .unwrap_or_default();
-                if dry_run && all_lines.is_empty() {
-                    continue;
-                }
-                if output.starts_with("docker:") || output.contains("permission denied") {
-                    println!("  ✗ Docker недоступен");
-                    if let Some(line) = output.lines().next() {
-                        println!("    └─ {}", line);
-                    }
-                } else if all_lines.is_empty() {
-                    println!("  (no containers)");
-                } else {
-                    println!("  Всего контейнеров: {}", all_lines.len());
-                    for line in &all_lines {
-                        let name_col = line.split('\t').next().unwrap_or("");
-                        let tag = if sm121_names.contains(&name_col) {
-                            "sm121"
-                        } else {
-                            "other"
-                        };
-                        println!("  {tag:5} {line}");
-                    }
-                    if sm121_names.is_empty() {
-                        println!("  (нет контейнеров sm121)");
-                    }
-                }
-            }
-            Ok((name, ib_ip, Err(e))) => {
-                println!("{} ({}): ✗ SSH ошибка: {}", name, ib_ip, e);
-            }
-            Err(join_err) => {
-                println!("✗ Task join error: {}", join_err);
-            }
-        }
-    }
-    Ok(())
-}
-/// Starts containers on all nodes in parallel.
+
 #[allow(clippy::too_many_arguments)]
 async fn start_cluster(
     config: &Config,
@@ -269,17 +231,14 @@ async fn start_cluster(
     dry_run: bool,
 ) -> Result<()> {
     let nodes: Vec<&Node> = if let Some(nn) = node_name {
-        cluster.nodes.iter().filter(|n| n.name == nn).collect()
+        let v: Vec<&Node> = cluster.nodes.iter().filter(|n| n.name == nn).collect();
+        if v.is_empty() {
+            anyhow::bail!("Node '{}' not found in cluster '{}'", nn, cluster.name);
+        }
+        v
     } else {
         cluster.nodes.iter().collect()
     };
-    if nodes.is_empty() {
-        anyhow::bail!(
-            "Node '{:?}' not found in cluster '{}'",
-            node_name,
-            cluster.name
-        );
-    }
     let nnodes = nodes.len() as u32;
     let pp_size = if recipe.pipeline_parallel_size.unwrap_or(0) == 0 {
         nnodes
@@ -387,25 +346,21 @@ fn build_docker_command(
         format!("-e HF_HOME=/root/.cache/huggingface"),
         format!("-v {}:/root/.cache/huggingface", cache_dir),
     ];
-    // Engine-specific host IP env var
     if is_sglang {
         cmd.push(format!("-e SGLANG_WORKER_IP={}", node.ib_ip));
     } else {
         cmd.push(format!("-e VLLM_HOST_IP={}", node.ib_ip));
     }
-    // Network interface env vars from cluster NCCL config
     cmd.push(format!("-e MN_IF_NAME={}", nccl.socket_ifname));
     cmd.push(format!("-e UCX_NET_DEVICES={}", nccl.socket_ifname));
     cmd.push(format!(
         "-e OMPI_MCA_btl_tcp_if_include={}",
         nccl.socket_ifname
     ));
-    // Default cluster env (recipe env below overrides these if duplicated)
     cmd.push("-e OMP_NUM_THREADS=4".to_string());
     cmd.push("-e NCCL_IGNORE_CPU_AFFINITY=1".to_string());
     cmd.push("-e NCCL_NET=IB".to_string());
     cmd.push("-e NCCL_DEBUG=WARN".to_string());
-    // Force colored vLLM logs through Docker pipeline
     if is_head {
         cmd.push("-e VLLM_LOGGING_COLOR=1".to_string());
         cmd.push("-e TERM=xterm-256color".to_string());
@@ -428,17 +383,12 @@ fn build_docker_command(
     cmd.push(recipe.docker_image.clone());
     if !recipe.entrypoint.is_empty() {
         let mut ep = recipe.entrypoint.clone();
-        // If the entrypoint is `sh -c SCRIPT`, shell-quote the script so
-        // bash on the SSH server treats it as one argv element to docker.
-        // Without this, bash word-splits the script and sh -c only sees
-        // the first token, then the container exits immediately.
         if ep.len() >= 3 && ep[ep.len() - 2] == "-c" {
             let last = ep.len() - 1;
             ep[last] = shq(&ep[last]);
         }
         cmd.extend(ep);
     }
-    // Model argument - engine-specific
     if is_sglang {
         cmd.push("--model-path".to_string());
         cmd.push(recipe.model.clone());
@@ -446,7 +396,6 @@ fn build_docker_command(
         cmd.push(recipe.model.clone());
     }
     cmd.extend(recipe.args.clone());
-    // Distributed flags - engine-specific
     let mut flags: Vec<String> = match recipe.distributed_backend.as_str() {
         "ray" | "mp" if is_sglang => vec![
             "--tp".into(),
@@ -473,7 +422,6 @@ fn build_docker_command(
                 "--master-port".into(),
                 init_port.to_string(),
             ];
-            // vLLM auto-detects mp; pass explicitly for ray only.
             if recipe.distributed_backend == "ray" {
                 fp.push("--distributed-executor-backend".into());
                 fp.push(recipe.distributed_backend.clone());
@@ -482,12 +430,10 @@ fn build_docker_command(
         }
         _ => vec![],
     };
-    // vLLM workers run headless (skip API server) when not rank 0
     if !is_sglang && node.rank > 0 && matches!(recipe.distributed_backend.as_str(), "mp" | "ray") {
         flags.push("--headless".into());
     }
     cmd.extend(flags);
-    // Host + port flags — auto-generated for both engines.
     cmd.push("--host".to_string());
     cmd.push("0.0.0.0".to_string());
     cmd.push("--port".to_string());

+ 4 - 3
src/main.rs

@@ -104,9 +104,10 @@ async fn main() -> Result<()> {
         }
         Commands::Stop {
             recipe: recipe_name,
+            node,
         } => {
             let (cluster, _recipe) = resolve_cluster_for_recipe(&config, &recipe_name).await?;
-            docker::stop_cluster(&config, &cluster, &recipe_name, dry_run).await?;
+            docker::stop_cluster(&config, &cluster, &recipe_name, node.as_deref(), dry_run).await?;
         }
         Commands::Health {
             cluster: cluster_name,
@@ -132,11 +133,11 @@ async fn main() -> Result<()> {
                 match cluster_name {
                     Some(ref name) => {
                         let cluster = select_cluster_by_name(&config, name)?;
-                        docker::health_check(&config, &cluster, dry_run).await?;
+                        docker::health::health_check(&config, &cluster, dry_run).await?;
                     }
                     None => {
                         for cluster in &config.clusters {
-                            docker::health_check(&config, cluster, dry_run).await?;
+                            docker::health::health_check(&config, cluster, dry_run).await?;
                         }
                     }
                 }