Explorar o código

feat: doctor command with repair, distributed_backend ray/mp switch

New features:
- doctor command: checks nodes for docker/nvidia-smi/nccl libs and
  containers for engine (vllm/sglang), ray (if ray backend), and
  entrypoint compatibility. Runs concurrently across all nodes.
- doctor --repair: auto-generates Dockerfile from base image, installs
  missing packages (ray, engine), and builds a repaired image tagged
  as <original>.repaired on the head node.
- serve auto-detects repaired images: checks if <image>.repaired exists
  on head node before starting, uses it if available.
- distributed_backend recipe field (default: 'ray'):
  - ray: sets --tensor-parallel-size=nnodes, --node-rank, --master-addr,
    sets RAY_ADDRESS env var
  - mp: sets --tensor-parallel-size=1, --pipeline-parallel-size,
    --nnodes, --node-rank, --master-addr, --distributed-executor-backend=mp
Alexander Korolev hai 6 meses
pai
achega
abcb899ef7
Modificáronse 1 ficheiros con 369 adicións e 15 borrados
  1. 369 15
      src/main.rs

+ 369 - 15
src/main.rs

@@ -61,6 +61,16 @@ enum Commands {
         #[arg(long, help = "Показать хост-ключи узлов (для добавления в config.toml)")]
         fingerprint: bool,
     },
+
+    /// Проверить и исправить рецепт (doctor mode)
+    Doctor {
+        #[arg(long, required = true, help = "Имя рецепта для проверки")]
+        recipe: String,
+        #[arg(long, default_value = "auto", help = "Кластер (если не auto — использует recipe.cluster)")]
+        cluster: String,
+        #[arg(long, help = "Автоматически исправить проблемы (создать Dockerfile + собрать образ)")]
+        repair: bool,
+    },
 }
 
 #[derive(Clone, Deserialize)]
@@ -119,12 +129,18 @@ struct Recipe {
     cluster: Option<String>,
     #[serde(default)]
     entrypoint: Vec<String>,
+    #[serde(default = "default_distributed_backend")]
+    distributed_backend: String,
 }
 
 fn default_engine() -> String {
     "vllm".to_string()
 }
 
+fn default_distributed_backend() -> String {
+    "ray".to_string()
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
     tracing_subscriber::fmt::init();
@@ -178,6 +194,9 @@ async fn main() -> Result<()> {
                 }
             }
         },
+        Commands::Doctor { recipe, cluster, repair } => {
+            doctor(&config, &recipe, &cluster, repair, dry_run).await?;
+        },
     }
 
     Ok(())
@@ -193,7 +212,7 @@ async fn serve(
     let recipe_path = format!("recipes/{}/recipe.toml", recipe_name);
     let recipe: Recipe = toml::from_str(&fs::read_to_string(&recipe_path).await?)?;
 
-    info!("🚀 Запуск рецепта: {} ({})", recipe.name, recipe.engine);
+    info!("🚀 Запуск рецепта: {} ({}) [backend: {}]", recipe.name, recipe.engine, recipe.distributed_backend);
 
     let cluster_choice = if cluster_name != "auto" {
         cluster_name.to_string()
@@ -214,6 +233,13 @@ async fn serve(
         }
     }
 
+    // Check for repaired image
+    let image = try_repaired_image(config, &cluster, &recipe).await?;
+    let mut recipe = recipe;
+    if image != recipe.docker_image {
+        recipe.docker_image = image.clone();
+    }
+
     start_cluster(config, &cluster, nccl, recipe_name, &recipe, dry_run).await?;
 
     if !dry_run {
@@ -523,6 +549,12 @@ fn build_docker_command(
     cmd.push(format!("--label sm121.recipe={}", recipe_name));
     cmd.push(format!("--label sm121.rank={}", node.rank));
 
+    // Ray backend: need RAY_ADDRESS + custom entrypoint that starts Ray before vLLM
+    if recipe.distributed_backend == "ray" {
+        cmd.push(format!("-e RAY_ADDRESS={}:6379", cluster.master_ib_ip));
+        cmd.push(format!("-e RAY_PORT={}", recipe.port));
+    }
+
     cmd.push(recipe.docker_image.clone());
 
     // Engine entrypoint (optional — empty for vLLM)
@@ -532,20 +564,37 @@ fn build_docker_command(
 
     cmd.push(recipe.model.clone());
     cmd.extend(recipe.args.clone());
-    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(),
-        "mp".to_string(),
-    ]);
+
+    // Backend-specific flags
+    match recipe.distributed_backend.as_str() {
+        "ray" => {
+            cmd.extend(vec![
+                "--tensor-parallel-size".to_string(),
+                nnodes.to_string(),
+                "--node-rank".to_string(),
+                node.rank.to_string(),
+                "--master-addr".to_string(),
+                cluster.master_ib_ip.clone(),
+            ]);
+        }
+        "mp" => {
+            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(),
+                "mp".to_string(),
+            ]);
+        }
+        _ => {}
+    }
 
     cmd
 }
@@ -722,3 +771,308 @@ impl client::Handler for VerifyingHandler {
 fn load_secret_key(path: &str) -> Result<russh::keys::PrivateKey> {
     russh_load_secret_key(path, None).context("Failed to load SSH key")
 }
+
+// ====================== Doctor ======================
+
+/// Doctor check result for a single node/container
+struct DoctorResult {
+    node_name: String,
+    checks: Vec<DoctorCheck>,
+}
+
+struct DoctorCheck {
+    name: String,
+    status: DoctorStatus,
+    detail: String,
+}
+
+#[derive(PartialEq)]
+enum DoctorStatus {
+    Ok,
+    Warn,
+    Fail,
+}
+
+impl std::fmt::Display for DoctorStatus {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            DoctorStatus::Ok => write!(f, "✓"),
+            DoctorStatus::Warn => write!(f, "⚠"),
+            DoctorStatus::Fail => write!(f, "✗"),
+        }
+    }
+}
+
+/// Run doctor checks on the cluster for a recipe
+async fn doctor(
+    config: &Config,
+    recipe_name: &str,
+    cluster_name: &str,
+    repair: bool,
+    dry_run: bool,
+) -> Result<()> {
+    let recipe_path = format!("recipes/{}/recipe.toml", recipe_name);
+    let recipe: Recipe = toml::from_str(&fs::read_to_string(&recipe_path).await?)?;
+
+    println!("\n🩺 Doctor: {} ({}) [backend: {}]", recipe.name, recipe.engine, recipe.distributed_backend);
+
+    // Resolve cluster
+    let cluster_choice = if cluster_name != "auto" {
+        cluster_name.to_string()
+    } else if let Some(ref rc) = recipe.cluster {
+        rc.clone()
+    } else {
+        "auto".to_string()
+    };
+    let (cluster, _nccl) = select_cluster(config, &cluster_choice)?;
+
+    // Check each node
+    let mut all_results: Vec<DoctorResult> = Vec::new();
+    let mut handles = Vec::new();
+
+    for node in &cluster.nodes {
+        let config = config.clone();
+        let node = node.clone();
+        let engine = recipe.engine.clone();
+        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
+        }));
+    }
+
+    let results = futures::future::join_all(handles).await;
+    for result in results {
+        if let Ok(r) = result {
+            all_results.push(r);
+        }
+    }
+
+    // Print results
+    let mut has_failures = false;
+    for result in &all_results {
+        println!("\n--- {} ---", result.node_name);
+        for check in &result.checks {
+            println!("  {} {} {}", check.status, check.name, check.detail);
+            if check.status == DoctorStatus::Fail {
+                has_failures = true;
+            }
+        }
+    }
+
+    if has_failures {
+        if repair {
+            println!("\n🔧 Reparing image...");
+            let repaired = run_cluster_repair(config, &cluster, &recipe, dry_run).await?;
+            println!("✅ Repaired image: {}", repaired);
+        } else {
+            println!("\n⚠ Some checks failed. Run with --repair to auto-fix.");
+            std::process::exit(1);
+        }
+    } else {
+        println!("\n✅ All checks passed!");
+    }
+
+    Ok(())
+}
+
+/// Check a single node for required software
+async fn check_node(
+    config: &Config,
+    node: &Node,
+    engine: &str,
+    backend: &str,
+    image: &str,
+) -> DoctorResult {
+    let mut checks = Vec::new();
+
+    // Node-level checks
+    checks.push(node_check(config, node, "Docker", "docker --version").await);
+    checks.push(node_check(config, node, "NVIDIA GPU", "nvidia-smi").await);
+    checks.push(node_check(config, node, "NCCL libs", "ldconfig -p | grep -q nccl && echo OK").await);
+
+    // Container-level checks
+    checks.push(container_check(config, node, image, "Engine", format!("{} version", engine)).await);
+
+    if backend == "ray" {
+        checks.push(container_check(config, node, image, "Ray", "ray --version".to_string()).await);
+    }
+
+    // Check container entrypoint
+    checks.push(container_entrypoint_check(config, node, image).await);
+
+    DoctorResult {
+        node_name: node.name.clone(),
+        checks,
+    }
+}
+
+/// 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, &Cluster { name: String::new(), master_ib_ip: String::new(), nodes: vec![] }, node, command.to_string(), false).await {
+        Ok(output) if output.trim().len() > 0 => DoctorCheck {
+            name: name.to_string(),
+            status: DoctorStatus::Ok,
+            detail: format!("OK: {}", output.trim().lines().next().unwrap_or(output.trim())),
+        },
+        Ok(_) => DoctorCheck {
+            name: name.to_string(),
+            status: DoctorStatus::Fail,
+            detail: "not found".to_string(),
+        },
+        Err(e) => DoctorCheck {
+            name: name.to_string(),
+            status: DoctorStatus::Fail,
+            detail: format!("error: {}", e),
+        },
+    }
+}
+
+/// Run a command inside a container on the node and return a check result
+async fn container_check(config: &Config, node: &Node, image: &str, name: &str, command: String) -> DoctorCheck {
+    let docker_cmd = format!(
+        "docker run --rm {} {} 2>&1 | head -1",
+        image, command
+    );
+    match run_ssh(config, &Cluster { name: String::new(), master_ib_ip: String::new(), nodes: vec![] }, node, docker_cmd, false).await {
+        Ok(output) if !output.contains("not found") && !output.contains("No such file") => DoctorCheck {
+            name: name.to_string(),
+            status: DoctorStatus::Ok,
+            detail: format!("OK: {}", output.trim()),
+        },
+        Ok(output) => DoctorCheck {
+            name: name.to_string(),
+            status: DoctorStatus::Fail,
+            detail: format!("not found: {}", output.trim()),
+        },
+        Err(e) => DoctorCheck {
+            name: name.to_string(),
+            status: DoctorStatus::Fail,
+            detail: format!("error: {}", e),
+        },
+    }
+}
+
+/// Check the container entrypoint
+async fn container_entrypoint_check(config: &Config, node: &Node, image: &str) -> DoctorCheck {
+    let docker_cmd = format!(
+        "docker inspect --format '{{{{json .Config.Entrypoint}}}}' {} 2>/dev/null || echo 'null'",
+        image
+    );
+    match run_ssh(config, &Cluster { name: String::new(), master_ib_ip: String::new(), nodes: vec![] }, node, docker_cmd, false).await {
+        Ok(output) => {
+            let entrypoint = output.trim();
+            if entrypoint == "null" || entrypoint == "[]" {
+                DoctorCheck {
+                    name: "Entrypoint".to_string(),
+                    status: DoctorStatus::Ok,
+                    detail: "default (shell) entrypoint".to_string(),
+                }
+            } else {
+                DoctorCheck {
+                    name: "Entrypoint".to_string(),
+                    status: DoctorStatus::Warn,
+                    detail: format!("custom: {}", entrypoint),
+                }
+            }
+        }
+        Err(e) => DoctorCheck {
+            name: "Entrypoint".to_string(),
+            status: DoctorStatus::Fail,
+            detail: format!("error: {}", e),
+        },
+    }
+}
+
+/// Generate Dockerfile and build repaired image on head node
+async fn run_cluster_repair(
+    config: &Config,
+    cluster: &Cluster,
+    recipe: &Recipe,
+    dry_run: bool,
+) -> Result<String> {
+    let head_node = cluster.nodes.first().context("No nodes in cluster")?;
+
+    // Generate repaired image name
+    let repaired_tag = format!("{}.repaired", recipe.docker_image);
+
+    // Generate Dockerfile
+    let dockerfile = generate_repair_dockerfile(&repaired_tag, recipe);
+
+    // Write Dockerfile locally
+    let dockerfile_path = format!("recipes/{}/Dockerfile.repaired", "temp");
+    fs::write(&dockerfile_path, &dockerfile).await
+        .context("Failed to write Dockerfile")?;
+
+    // SSH to head node, create Dockerfile, build image
+    let create_dockerfile = format!(
+        "mkdir -p /tmp/sm121_docker && cat > /tmp/sm121_docker/Dockerfile << 'DOCKERFILE_END'\n{}DOCKERFILE_END",
+        dockerfile
+    );
+
+    if !dry_run {
+        run_ssh(config, cluster, head_node, create_dockerfile, false).await?;
+
+        let build_cmd = format!(
+            "cd /tmp/sm121_docker && docker build -t {} .",
+            repaired_tag
+        );
+        let output = run_ssh_with_retry(config, cluster, head_node, build_cmd, false, 1).await
+            .context("Failed to build repaired image")?;
+        println!("Build output: {}", output.lines().last().unwrap_or(&output));
+    } else {
+        println!("  (dry-run) Would build {} on {}", repaired_tag, head_node.name);
+    }
+
+    // Cleanup
+    fs::remove_file(&dockerfile_path).await.ok();
+
+    Ok(repaired_tag)
+}
+
+/// Generate a Dockerfile to repair a container
+fn generate_repair_dockerfile(base_image: &str, recipe: &Recipe) -> String {
+    let mut lines = vec![format!("FROM {}", base_image)];
+
+    // Install ray if needed
+    if recipe.distributed_backend == "ray" {
+        lines.push("RUN pip install ray[default] || pip install ray".to_string());
+    }
+
+    // Install engine if needed
+    let engine_install = match recipe.engine.as_str() {
+        "vllm" => "pip install vllm",
+        "sglang" => "pip install sglang",
+        _ => "",
+    };
+    if !engine_install.is_empty() {
+        lines.push(format!("RUN {} || true", engine_install));
+    }
+
+    // Ensure entrypoint is shell-compatible
+    lines.push("ENTRYPOINT [\"/bin/bash\", \"-c\"]".to_string());
+
+    lines.join("\n")
+}
+
+// ====================== serve() repair image check ======================
+
+async fn try_repaired_image(config: &Config, cluster: &Cluster, recipe: &Recipe) -> Result<String> {
+    let head_node = cluster.nodes.first().context("No nodes in cluster")?;
+    let repaired_tag = format!("{}.repaired", recipe.docker_image);
+
+    // Check if repaired image exists on head node
+    let check_cmd = format!(
+        "docker images -q {} | head -1",
+        repaired_tag
+    );
+
+    match run_ssh(config, cluster, head_node, check_cmd, false).await {
+        Ok(output) if !output.trim().is_empty() => {
+            info!("🩺 Found repaired image: {}", repaired_tag);
+            Ok(repaired_tag)
+        }
+        _ => Ok(recipe.docker_image.clone()),
+    }
+}