launcher/src/docker/mod.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,
    keep: bool,
    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();
    }

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

/// 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 = logs_follow_cmd(recipe_name);
        run_ssh_follow(config, head_node, cmd, dry_run).await?;
    } else {
        info!(
            "📜 Логи кластера '{}' ({} узлов):",
            cluster.name,
            cluster.nodes.len()
        );
        let mut handles = Vec::new();
        for node in &cluster.nodes {
            let cmd = logs_tail_cmd(recipe_name, node.rank, 100);
            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 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,
    cluster: &Cluster,
    nccl: &NcclConfig,
    recipe_name: &str,
    recipe: &Recipe,
    cache_dir: String,
    keep: bool,
    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
        );

        let cd = cache_dir.clone();
        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,
                &cd,
                keep,
                config.init_port,
            )
            .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,
    cache_dir: &str,
    keep: bool,
    init_port: u16,
) -> Vec<String> {
    let is_sglang = recipe.engine == "sglang";

    let mut flags = vec![
        "--gpus all",
        "--ipc=host",
        "--privileged",
        "--network=host",
        "--shm-size 10g",
        "--ulimit memlock=-1",
        "--ulimit stack=67108864",
    ];
    if !keep {
        flags.push("--rm");
    }
    let mut cmd = vec![
        format!("docker run -d {}", flags.join(" ")),
        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 HF_HOME=/cache/huggingface"),
        format!("-v {}:/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));
    }

    // Default cluster env (recipe env below overrides these if duplicated)
    cmd.push("-e NCCL_CUMEM_ENABLE=0".to_string());
    cmd.push("-e OMP_NUM_THREADS=4".to_string());

    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());
    }

    // Model argument - engine-specific
    if is_sglang {
        cmd.push("--model-path".to_string());
        cmd.push(recipe.model.clone());
    } else {
        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(),
            recipe.tensor_parallel_size.to_string(),
            "--nnodes".into(),
            nnodes.to_string(),
            "--node-rank".into(),
            node.rank.to_string(),
            "--dist-init-addr".into(),
            format!("{}:{}", cluster.master_ib_ip, init_port),
        ],
        "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(),
            "--master-port".into(),
            init_port.to_string(),
            "--distributed-executor-backend".into(),
            recipe.distributed_backend.clone(),
        ],
        _ => 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);

    // Port flag — auto-generated for SGLang (all nodes need to bind)
    if is_sglang {
        cmd.push("--port".to_string());
        cmd.push(recipe.port.to_string());
    }

    cmd
}

#[cfg(test)]
mod tests;