launcher/src/docker.rs

use anyhow::{Context, Result};
use tokio::time::{Duration, sleep};
use tracing::info;

use crate::config::*;
use crate::doctor::try_repaired_image;
use crate::ssh::*;

/// Launches containers across the cluster for a recipe.
pub async fn serve(config: &Config, recipe_name: &str, cluster_name: &str, dry_run: bool) -> Result<()> {
    let recipe = load_recipe(recipe_name).await?;
    info!("🚀 Запуск рецепта: {} ({}) [backend: {}]", recipe.name, recipe.engine, recipe.distributed_backend);
    let (cluster, nccl) = resolve_cluster(config, &recipe, cluster_name)?;

    if let Some(required) = recipe.required_nodes {
        let actual = cluster.nodes.len() as u32;
        if actual != required {
            anyhow::bail!("Recipe '{}' requires {} nodes, but cluster '{}' has {} nodes", recipe.name, required, cluster.name, actual);
        }
    }

    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 {
        info!("\n📜 Последние логи контейнеров:");
        show_logs(config, &cluster, recipe_name, false, dry_run).await?;
    }
    Ok(())
}

/// Shows Docker logs for a recipe across all nodes (or follows head node).
pub async fn show_logs(config: &Config, cluster: &Cluster, recipe_name: &str, follow: bool, dry_run: bool) -> Result<()> {
    if follow {
        let head_node = cluster.nodes.first().context("No nodes in cluster")?;
        info!("📜 Following logs for '{}' (head node)...", head_node.name);
        let cmd = format!("docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank=0\" | xargs -r docker logs -f", recipe_name);
        run_ssh_follow(config, head_node, cmd, dry_run).await?;
    } else {
        let flag = "--tail=100";
        info!("📜 Логи кластера '{}' ({} узлов):", cluster.name, cluster.nodes.len());
        let mut handles = Vec::new();
        for node in &cluster.nodes {
            let cmd = format!("docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank={}\" | xargs -r docker logs {} 2>&1", recipe_name, node.rank, flag);
            let config = config.clone(); let node = node.clone(); let dry = dry_run;
            let name = node.name.clone();
            handles.push(tokio::spawn(async move { (name, run_ssh_with_retry(&config, &node, cmd, dry, 1).await) }));
        }
        for result in futures::future::join_all(handles).await {
            if let Ok((name, Ok(output))) = result && !output.trim().is_empty() {
                println!("\n--- {} ---", name); println!("{}", output);
            }
        }
    }
    Ok(())
}

/// Stops and removes containers for a recipe on all nodes.
pub async fn stop_cluster(config: &Config, cluster: &Cluster, recipe_name: &str, dry_run: bool) -> Result<()> {
    info!("🛑 Останавливаем кластер '{}'...", cluster.name);
    let mut failures = Vec::new();

    for node in &cluster.nodes {
        let stop_cmd = format!("docker ps -q --filter \"label=sm121.recipe={}\" | xargs -r docker stop", recipe_name);
        if let Err(e) = run_ssh(config, node, stop_cmd, dry_run).await {
            tracing::warn!("Stop failed on {}: {}", node.name, e);
            failures.push((node.name.clone(), "stop".to_string(), e));
        }
        let rm_cmd = format!("docker ps -qa --filter \"label=sm121.recipe={}\" | xargs -r docker rm", recipe_name);
        if let Err(e) = run_ssh(config, node, rm_cmd, dry_run).await {
            tracing::warn!("Remove failed on {}: {}", node.name, e);
            failures.push((node.name.clone(), "rm".to_string(), e));
        }
    }

    if !failures.is_empty() {
        for (name, action, e) in &failures { tracing::warn!("Failed to {} on {}: {}", action, name, e); }
        anyhow::bail!("Failed to stop containers on {} node(s)", failures.len());
    }
    info!("✅ Все контейнеры остановлены.");
    Ok(())
}

/// Checks Docker container health on a cluster.
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 {
            let result = run_ssh_with_retry(&config, &node, "docker ps --filter \"label=sm121\" --format '{{.Names}} {{.Status}}' 2>&1".to_string(), dry, 1).await;
            (node.name, node.ib_ip, result)
        }));
    }
    let mut all_healthy = true;
    for result in futures::future::join_all(handles).await {
        match result {
            Ok((name, ib_ip, Ok(check))) => {
                print!("{} ({}): ", name, ib_ip);
                if check.is_empty() && dry_run { continue; }
                if check.starts_with("docker:") || check.contains("permission denied") {
                    println!("✗ Docker недоступен");
                    if let Some(line) = check.lines().next() { println!("  └─ {}", line); }
                    all_healthy = false;
                } else {
                    let containers: Vec<&str> = check.lines().filter(|l| !l.is_empty()).collect();
                    if containers.is_empty() { println!("✗ Нет контейнеров"); all_healthy = false; }
                    else { println!("✓ {} контейнер(ов)", containers.len()); for c in containers { println!("  └─ {}", c); } }
                }
            }
            Ok((name, ib_ip, Err(e))) => { println!("{} ({}): ✗ SSH ошибка: {}", name, ib_ip, e); all_healthy = false; }
            Err(join_err) => { println!("✗ Task join error: {}", join_err); all_healthy = false; }
        }
    }
    if !dry_run && all_healthy { info!("✅ Кластер здоров!"); } else if !dry_run { info!("⚠️ Кластер имеет проблемы"); }
    Ok(())
}

/// Starts containers on all nodes in parallel.
async fn start_cluster(config: &Config, cluster: &Cluster, nccl: &NcclConfig, recipe_name: &str, recipe: &Recipe, dry_run: bool) -> Result<()> {
    let nnodes = cluster.nodes.len() as u32;
    let pp_size = if recipe.pipeline_parallel_size.unwrap_or(0) == 0 { nnodes } else { recipe.pipeline_parallel_size.unwrap() };
    let mut handles = Vec::new();

    for node in &cluster.nodes {
        let node = node.clone(); let cluster = cluster.clone(); let nccl = nccl.clone();
        let recipe = recipe.clone(); let config = config.clone();
        let rn = recipe_name.to_string(); let dry = dry_run; let engine = recipe.engine.clone();
        let is_head = node.rank == 0;
        let cleanup = format!("docker rm -f sm121_{}_{} 2>/dev/null; true", recipe_name, node.rank);

        handles.push(tokio::spawn(async move {
            if !dry { let _ = run_ssh(&config, &node, cleanup, dry).await; }
            let cmd = build_docker_command(&cluster, &nccl, &recipe, &rn, &node, is_head, pp_size, nnodes).join(" ");
            let result = run_ssh_with_retry(&config, &node, cmd, dry, 2).await;
            (node.name, node.rank, engine, result)
        }));
    }

    for result in futures::future::join_all(handles).await {
        match result {
            Ok((name, rank, engine, Ok(_))) => info!("✅ {} запущен на {} (rank {})", engine.to_uppercase(), name, rank),
            Ok((name, rank, engine, Err(e))) => anyhow::bail!("Failed to start {} on {} (rank {}): {}", engine.to_uppercase(), name, rank, e),
            Err(e) => anyhow::bail!("Task join error: {}", e),
        }
    }

    if !dry_run { sleep(Duration::from_secs(10)).await; info!("🎉 Кластер успешно запущен! API → http://{}:{}", cluster.master_ib_ip, recipe.port); }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn build_docker_command(cluster: &Cluster, nccl: &NcclConfig, recipe: &Recipe, recipe_name: &str, node: &Node, is_head: bool, pp_size: u32, nnodes: u32) -> Vec<String> {
    let mut cmd = vec![
        "docker run -d --gpus all --ipc=host --privileged --network=host".to_string(),
        format!("-e NCCL_IB_DISABLE={}", nccl.ib_disable), format!("-e NCCL_IB_HCA={}", nccl.ib_hca),
        format!("-e NCCL_SOCKET_IFNAME={}", nccl.socket_ifname), format!("-e NCCL_IB_GID_INDEX={}", nccl.ib_gid_index),
        format!("-e VLLM_HOST_IP={}", node.ib_ip),
    ];
    for (k, v) in &recipe.env { cmd.push(format!("-e {}={}", k, v)); }
    if is_head { cmd.push(format!("-p {}:{}", recipe.port, recipe.port)); }
    cmd.push(format!("--name sm121_{}_{}", recipe_name, node.rank));
    cmd.push(format!("--label sm121.recipe={}", recipe_name));
    cmd.push(format!("--label sm121.rank={}", node.rank));
    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());
    if !recipe.entrypoint.is_empty() { cmd.extend(recipe.entrypoint.clone()); }
    cmd.push(recipe.model.clone());
    cmd.extend(recipe.args.clone());

    let flags: Vec<String> = match recipe.distributed_backend.as_str() {
        "ray" | "mp" => vec![
            "--tensor-parallel-size".into(), "1".into(),
            "--pipeline-parallel-size".into(), pp_size.to_string(),
            "--nnodes".into(), nnodes.to_string(),
            "--node-rank".into(), node.rank.to_string(),
            "--master-addr".into(), cluster.master_ib_ip.clone(),
            "--distributed-executor-backend".into(), recipe.distributed_backend.clone(),
        ],
        _ => vec![],
    };
    cmd.extend(flags);
    cmd
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;
    fn tcls() -> Cluster { Cluster { name: "t".into(), master_ib_ip: "1".into(), nodes: vec![
        Node { name: "n0".into(), ssh_ip: "10.0.0.1".into(), ib_ip: "1".into(), rank: 0, host_key: None },
        Node { name: "n1".into(), ssh_ip: "10.0.0.2".into(), ib_ip: "2".into(), rank: 1, host_key: None },
    ] } }
    fn tnccl() -> NcclConfig { NcclConfig { ib_hca: "r".into(), socket_ifname: "e".into(), ib_gid_index: 3, ib_disable: 0 } }
    fn trec() -> Recipe { Recipe {
        engine: "vllm".into(), name: "T".into(), docker_image: "img".into(), model: "mdl".into(), port: 8000,
        required_nodes: Some(2), pipeline_parallel_size: Some(2),
        args: vec!["--x".into(), "4096".into()], env: HashMap::from([("K".into(), "V".into())]),
        cluster: None, entrypoint: vec![], distributed_backend: "mp".into(),
    } }
    fn run(recipe: &Recipe, node: &Node, head: bool) -> String {
        build_docker_command(&tcls(), &tnccl(), recipe, "r", node, head, 2, 2).join(" ")
    }

    #[test] fn mp() { let cmd = run(&trec(), &tcls().nodes[0], true);
        assert!(cmd.contains("docker run -d --gpus all")); assert!(cmd.contains("-e NCCL_IB_HCA=r"));
        assert!(cmd.contains("-e K=V")); assert!(cmd.contains("-p 8000:8000")); assert!(cmd.contains("--name sm121_r_0"));
        assert!(cmd.contains("--label sm121.recipe=r")); assert!(cmd.contains("--label sm121.rank=0"));
        assert!(cmd.contains("img")); assert!(cmd.contains("mdl")); assert!(cmd.contains("--tensor-parallel-size 1"));
        assert!(cmd.contains("--pipeline-parallel-size 2")); assert!(cmd.contains("--nnodes 2"));
        assert!(cmd.contains("--node-rank 0")); assert!(cmd.contains("--master-addr 1"));
        assert!(cmd.contains("--distributed-executor-backend mp")); assert!(!cmd.contains("RAY_ADDRESS"));
    }
    #[test] fn worker_no_port() { let cmd = run(&trec(), &tcls().nodes[1], false);
        assert!(!cmd.contains("-p 8000:8000")); assert!(cmd.contains("--node-rank 1"));
        assert!(cmd.contains("--label sm121.rank=1")); assert!(cmd.contains("--name sm121_r_1"));
    }
    #[test] fn ray_backend() { let mut r = trec(); r.distributed_backend = "ray".into();
        let cmd = run(&r, &tcls().nodes[0], true);
        assert!(cmd.contains("-e RAY_ADDRESS=1:6379")); assert!(cmd.contains("-e RAY_PORT=8000"));
        assert!(cmd.contains("--node-rank 0")); assert!(cmd.contains("--distributed-executor-backend ray"));
    }
    #[test] fn auto_pp_size() { let cmd = run(&trec(), &tcls().nodes[0], true); assert!(cmd.contains("--pipeline-parallel-size 2")); }
    #[test] fn entrypoint() { let mut r = trec(); r.entrypoint = vec!["py".into(), "-m".into(), "mod".into()];
        let cmd = run(&r, &tcls().nodes[0], true);
        let ip = cmd.find("img").unwrap(); let mp = cmd.find("mdl").unwrap(); let ep = cmd.find("py -m mod").unwrap();
        assert!(ip < ep && ep < mp, "entrypoint should be between image and model");
    }
}