launcher/src/docker/tests.rs

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),
        tensor_parallel_size: 1,
        args: vec!["--x".into(), "4096".into()],
        env: HashMap::from([("K".into(), "V".into())]),
        cluster: None,
        entrypoint: vec![],
        distributed_backend: "mp".into(),
        before_serve: vec![],
    }
}

fn run(recipe: &Recipe, node: &Node, head: bool) -> String {
    build_docker_command(
        &tcls(),
        &tnccl(),
        recipe,
        "r",
        node,
        head,
        2,
        2,
        "/home/user/.cache/huggingface",
        false,
        25000,
    )
    .join(" ")
}

fn run_keep(recipe: &Recipe, node: &Node, head: bool) -> String {
    build_docker_command(
        &tcls(),
        &tnccl(),
        recipe,
        "r",
        node,
        head,
        2,
        2,
        "/home/user/.cache/huggingface",
        true,
        25000,
    )
    .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("--master-port 25000"));
    assert!(cmd.contains("--distributed-executor-backend mp"));
    assert!(!cmd.contains("RAY_ADDRESS"));
    assert!(cmd.contains("--rm"), "default --rm should be present");
    // New: defaults & cluster env
    assert!(cmd.contains("--shm-size 10g"), "default shm-size missing");
    assert!(cmd.contains("--ulimit memlock=-1"));
    assert!(cmd.contains("--ulimit stack=67108864"));
    assert!(
        cmd.contains("-e NCCL_CUMEM_ENABLE=0"),
        "default NCCL_CUMEM_ENABLE missing"
    );
    assert!(
        cmd.contains("-e OMP_NUM_THREADS=4"),
        "default OMP_NUM_THREADS missing"
    );
    // Head node: no --headless
    assert!(!cmd.contains("--headless"));
}

#[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"));
    // New: vLLM workers run --headless
    assert!(cmd.contains("--headless"), "vLLM worker should be headless");
}

#[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"
    );
}

// ── SGLang engine tests ──────────────────────────────────────

fn sglang_recipe() -> Recipe {
    Recipe {
        engine: "sglang".into(),
        name: "S".into(),
        docker_image: "sglang-img".into(),
        model: "org/sgl-model".into(),
        port: 8000,
        required_nodes: Some(4),
        pipeline_parallel_size: None,
        tensor_parallel_size: 4,
        args: vec!["--mem-fraction-static".into(), "0.8".into()],
        env: HashMap::from([("TORCH_NCCL_ASYNC_ERROR_HANDLING".into(), "1".into())]),
        cluster: None,
        entrypoint: vec!["sglang".into(), "serve".into()],
        distributed_backend: "mp".into(),
        before_serve: vec![],
    }
}

fn srun(recipe: &Recipe, node: &Node, head: bool) -> String {
    build_docker_command(
        &tcls(),
        &tnccl(),
        recipe,
        "sg",
        node,
        head,
        4,
        4,
        "/home/user/.cache/huggingface",
        false,
        25000,
    )
    .join(" ")
}

#[test]
fn sglang_host_ip() {
    let cmd = srun(&sglang_recipe(), &tcls().nodes[0], true);
    assert!(cmd.contains("-e SGLANG_WORKER_IP=1"));
    assert!(!cmd.contains("VLLM_HOST_IP"));
}

#[test]
fn sglang_head() {
    let cmd = srun(&sglang_recipe(), &tcls().nodes[0], true);
    assert!(cmd.contains("--model-path"));
    assert!(cmd.contains("org/sgl-model"));
    assert!(cmd.contains("--tp 4"));
    assert!(cmd.contains("--nnodes 4"));
    assert!(cmd.contains("--node-rank 0"));
    assert!(cmd.contains("--dist-init-addr 1:25000"));
    assert!(!cmd.contains("--master-addr"));
    assert!(!cmd.contains("--tensor-parallel-size"));
    assert!(!cmd.contains("--pipeline-parallel-size"));
    assert!(!cmd.contains("--distributed-executor-backend"));
    assert!(cmd.contains("-p 8000:8000"));
    assert!(cmd.contains("-e TORCH_NCCL_ASYNC_ERROR_HANDLING=1"));
    assert!(cmd.contains("--port 8000"));
    assert!(cmd.contains("-e HF_HOME=/cache/huggingface"));
    assert!(cmd.contains("-v /home/user/.cache/huggingface:/cache/huggingface"));
    assert!(
        cmd.contains("--rm"),
        "SGLang default --rm should be present"
    );
}

#[test]
fn sglang_worker() {
    let cmd = srun(&sglang_recipe(), &tcls().nodes[1], false);
    assert!(!cmd.contains("-p 8000:8000"));
    assert!(cmd.contains("--node-rank 1"));
    assert!(cmd.contains("--port 8000"));
}

#[test]
fn sglang_entrypoint_model_path_order() {
    let cmd = srun(&sglang_recipe(), &tcls().nodes[0], true);
    let ep = cmd.find("sglang serve").unwrap();
    let mp = cmd.find("--model-path").unwrap();
    let mdl = cmd.find("org/sgl-model").unwrap();
    let tp = cmd.find("--tp 4").unwrap();
    assert!(
        ep < mp && mp < mdl,
        "entrypoint before model-path before model"
    );
    assert!(mdl < tp, "model before --tp");
    let port = cmd.find("--port 8000").unwrap();
    assert!(tp < port, "--tp before --port");
}

#[test]
fn keep_disables_rm() {
    let cmd = run_keep(&trec(), &tcls().nodes[0], true);
    assert!(!cmd.contains("--rm"), "--keep should suppress --rm");
    assert!(cmd.contains("docker run -d"));
    assert!(cmd.contains("--gpus all"));
    assert!(cmd.contains("-e HF_HOME=/cache/huggingface"));
    assert!(cmd.contains("-v /home/user/.cache/huggingface:/cache/huggingface"));
}

#[test]
fn recipe_env_overrides_cluster_defaults() {
    // Recipe-level env should win over the defaults we inject.
    let mut r = trec();
    r.env.insert("NCCL_CUMEM_ENABLE".into(), "1".into());
    r.env.insert("OMP_NUM_THREADS".into(), "16".into());
    let cmd = run(&r, &tcls().nodes[0], true);
    // Docker uses the last -e for duplicate keys, so recipe values come after.
    // rfind returns the last occurrence, which should be the recipe override.
    let cumem_pos = cmd.rfind("NCCL_CUMEM_ENABLE=").unwrap();
    assert!(
        cmd[cumem_pos..].starts_with("NCCL_CUMEM_ENABLE=1"),
        "recipe override should be last: slice starts at {}",
        cumem_pos
    );
    let omp_pos = cmd.rfind("OMP_NUM_THREADS=").unwrap();
    assert!(
        cmd[omp_pos..].starts_with("OMP_NUM_THREADS=16"),
        "recipe override should be last"
    );
}

#[test]
fn custom_init_port() {
    // Non-default init_port should propagate to vLLM master-port and
    // SGLang dist-init-addr.
    let cmd = build_docker_command(
        &tcls(),
        &tnccl(),
        &trec(),
        "r",
        &tcls().nodes[0],
        true,
        2,
        2,
        "/home/user/.cache/huggingface",
        false,
        29500,
    )
    .join(" ");
    assert!(cmd.contains("--master-port 29500"));
    assert!(cmd.contains("--master-addr 1"));
}

// ── Logs command format ───────────────────────────────────────

#[test]
fn logs_tail_cmd_format() {
    let cmd = super::logs_tail_cmd("test-sleep", 0, 100);
    assert!(cmd.contains("docker ps -q"));
    assert!(cmd.contains("--filter \"label=sm121.recipe=test-sleep\""));
    assert!(cmd.contains("--filter \"label=sm121.rank=0\""));
    assert!(cmd.contains("xargs -r docker logs"));
    assert!(cmd.contains("--timestamps"));
    assert!(cmd.contains("--tail=100"));
    assert!(cmd.ends_with(" 2>&1"));
}

#[test]
fn logs_tail_cmd_per_rank() {
    let cmd0 = super::logs_tail_cmd("glm51", 0, 100);
    let cmd2 = super::logs_tail_cmd("glm51", 2, 100);
    assert!(cmd0.contains("label=sm121.rank=0"));
    assert!(cmd2.contains("label=sm121.rank=2"));
    assert!(!cmd0.contains("label=sm121.rank=2"));
}

#[test]
fn logs_follow_cmd_format() {
    let cmd = super::logs_follow_cmd("test-sleep");
    assert!(cmd.contains("xargs -r docker logs -f"));
    assert!(cmd.contains("--timestamps"));
    // Follow targets the head node only (rank=0).
    assert!(cmd.contains("--filter \"label=sm121.rank=0\""));
    // Follow should not include --tail (we want a live stream from now).
    assert!(!cmd.contains("--tail"));
    // Follow must redirect stderr so docker's own errors are captured
    // (mirrors the behaviour of logs_tail_cmd).
    assert!(
        cmd.contains("2>&1"),
        "follow cmd must capture stderr, got: {cmd}"
    );
}

// ── Shell quoting ─────────────────────────────────────────────

#[test]
fn shq_wraps_simple_strings() {
    assert_eq!(shq("foo"), "'foo'");
    assert_eq!(shq(""), "''");
    assert_eq!(shq("foo bar"), "'foo bar'");
}

#[test]
fn shq_escapes_inner_single_quotes() {
    // "don't" → 'don'\''t'
    assert_eq!(shq("don't"), "'don'\\''t'");
}

#[test]
fn entrypoint_script_is_shell_quoted() {
    // sh -c SCRIPT should wrap the script in single quotes so it survives
    // bash word-splitting on the SSH server.
    let mut r = trec();
    r.entrypoint = vec!["sh".into(), "-c".into(), "echo $i; sleep 1".into()];
    let cmd = run(&r, &tcls().nodes[0], true);
    // The script element should appear wrapped in single quotes
    assert!(
        cmd.contains("'echo $i; sleep 1'"),
        "script should be shell-quoted, got: {cmd}"
    );
}

#[test]
fn entrypoint_script_non_sh_passes_through() {
    // For non `sh -c` entrypoints (e.g. ["sh", "-c", "...", "extra"]) or
    // direct binaries, we leave the args alone — no false positive shq.
    let mut r = trec();
    r.entrypoint = vec!["sh".into(), "-c".into()];
    let cmd = run(&r, &tcls().nodes[0], true);
    assert!(cmd.contains("docker run"));
}