launcher/src/docker/mod.rs

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,
    cluster_name: &str,
    keep: bool,
    skip_before_serve: bool,
    dry_run: bool,
) -> Result<()> {
    let recipe = load_recipe(recipe_name).await?;
    let rn = recipe.name.clone();
    let re = recipe.engine.clone();
    let rb = recipe.distributed_backend.clone();
    info!(
        "{}",
        fluent_i18n::t!("recipe-starting", { "name" => rn, "engine" => re, "backend" => rb })
    );
    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
            );
        }
    }
    crate::before_serve::run(config, &cluster, &recipe, skip_before_serve, dry_run).await?;
    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!("{}", fluent_i18n::t!("container-log-header"));
        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 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('\'');
    for c in s.chars() {
        if c == '\'' {
            out.push_str(r"'\''");
        } else {
            out.push(c);
        }
    }
    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,
    recipe_name: &str,
    follow: bool,
    dry_run: bool,
) -> Result<()> {
    if follow {
        let head_node = cluster.nodes.first().context("No nodes in cluster")?;
        info!(
            "{}",
            fluent_i18n::t!("logs-follow-head", { "name" => &head_node.name })
        );
        let cmd = logs_follow_cmd(recipe_name);
        run_ssh_follow(config, head_node, cmd, dry_run).await?;
    } else {
        let cn = cluster.name.clone();
        info!(
            "{}",
            fluent_i18n::t!("logs-header", { "name" => cn, "count" => 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!("{}", fluent_i18n::t!("logs-section", { "name" => &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!(
        "{}",
        fluent_i18n::t!("cluster-stopping", { "name" => &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!(
                "{}",
                fluent_i18n::t!("stop-failed", { "name" => &node.name, "error" => e.to_string() })
            );
            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!(
                "{}",
                fluent_i18n::t!("remove-failed", { "name" => &node.name, "error" => e.to_string() })
            );
            failures.push((node.name.clone(), "rm".to_string(), e));
        }
    }
    if !failures.is_empty() {
        for (name, action, e) in &failures {
            tracing::warn!(
                "{}",
                fluent_i18n::t!("failed-to", { "action" => action, "name" => name, "error" => e.to_string() })
            );
        }
        anyhow::bail!("Failed to stop containers on {} node(s)", failures.len());
    }
    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,
    cluster: &Cluster,
    nccl: &NcclConfig,
    recipe_name: &str,
    recipe: &Recipe,
    cache_dir: std::path::PathBuf,
    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();
        let cd_s = cd.to_string_lossy().into_owned();
        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_s,
                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!(
                "{}",
                fluent_i18n::t!("engine-started", { "engine" => engine.to_uppercase(), "name" => &name, "rank" => rank })
            ),
            Ok((name, rank, engine, Err(e))) => anyhow::bail!(
                "{}",
                fluent_i18n::t!("engine-start-failed", { "engine" => engine.to_uppercase(), "name" => &name, "rank" => rank, "error" => e.to_string() })
            ),
            Err(e) => anyhow::bail!("Task join error: {}", e),
        }
    }
    if !dry_run {
        sleep(Duration::from_secs(10)).await;
        info!(
            "{}",
            fluent_i18n::t!("cluster-started", { "ip" => &cluster.master_ib_ip, "port" => 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<String> = vec![
        "--gpus all".to_string(),
        "--ipc=host".to_string(),
        "--privileged".to_string(),
        "--network=host".to_string(),
        "--ulimit memlock=-1".to_string(),
        "--ulimit stack=67108864".to_string(),
    ];
    if !recipe.shm_size.is_empty() {
        flags.push(format!("--shm-size {}", recipe.shm_size.as_str()));
    }
    if !keep {
        flags.push("--rm".to_string());
    }
    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 GLOO_SOCKET_IFNAME={}", nccl.socket_ifname),
        format!("-e TP_SOCKET_IFNAME={}", nccl.socket_ifname),
        format!("-e NCCL_IB_GID_INDEX={}", nccl.ib_gid_index),
        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());
    }
    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(format!("-e RAY_NODE_IP_ADDRESS={}", node.ib_ip));
        cmd.push(format!("-e RAY_OVERRIDE_NODE_IP_ADDRESS={}", node.ib_ip));
    }
    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());
    } 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" => {
            let mut fp = vec![
                "--tensor-parallel-size".into(),
                recipe.tensor_parallel_size.to_string(),
                "--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(),
            ];
            // 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());
            }
            fp
        }
        _ => 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());
    cmd.push(recipe.port.to_string());
    cmd
}
#[cfg(test)]
mod tests;