launcher/src/engine_flags.rs

//! Shared engine-specific argument builder — used by both the Docker
//! (`docker/tail.rs`) and tmux (`tmux/commands.rs`) deployment paths to
//! produce the model, recipe args, and distributed-rendezvous flags.
//!
//! This consolidates the duplicated flag logic that previously lived in
//! `tail.rs` and `commands.rs` separately. Each deployment path still
//! adds its own entrypoint and host/port flags on top.

use crate::config::{DistributedBackend, Node, Recipe};

/// Context for [`build_engine_args`] — the subset of `DockerCommandCtx`
/// / `TmuxManager` state that the shared flag builder needs.
pub struct EngineFlagsCtx<'a> {
    pub recipe: &'a Recipe,
    pub tp_size: u32,
    pub pp_size: u32,
    pub nnodes: u32,
    pub init_port: u16,
    pub master_ib_ip: &'a str,
}

/// Build the engine-specific arguments for a recipe: model path, recipe
/// args, and distributed-rendezvous flags.
///
/// Does **not** include the entrypoint, `--host`, or `--port` — those are
/// deployment-specific and added by each path:
/// - Docker (`tail.rs`): always adds `--host 0.0.0.0 --port <port>`.
/// - tmux (`commands.rs`): adds `--port <port>` on the head node (or all
///   nodes for SGLang).
///
/// Single-node runs emit **no** distributed flags — the engine defaults
/// to single-node. Multi-node runs get the full rendezvous flag set.
pub fn build_engine_args(ctx: &EngineFlagsCtx<'_>, node: &Node) -> Vec<String> {
    let is_sglang = ctx.recipe.engine == "sglang";
    let is_distributed = matches!(
        ctx.recipe.distributed_backend,
        Some(DistributedBackend::Mp | DistributedBackend::Ray)
    );

    let mut args = Vec::new();

    // Model — SGLang uses --model-path, vLLM uses the bare model arg.
    if is_sglang {
        args.push("--model-path".to_string());
    }
    args.push(ctx.recipe.model.clone());

    // Recipe-specific args (e.g. --kv-cache-dtype, --max-model-len).
    args.extend(ctx.recipe.args.to_vec());

    // Distributed flags — only for multi-node runs.
    if ctx.nnodes > 1 {
        match ctx.recipe.distributed_backend {
            Some(DistributedBackend::Ray | DistributedBackend::Mp) if is_sglang => {
                args.extend(vec![
                    "--tp".into(),
                    ctx.tp_size.to_string(),
                    "--nnodes".into(),
                    ctx.nnodes.to_string(),
                    "--node-rank".into(),
                    node.rank.to_string(),
                    "--dist-init-addr".into(),
                    format!("{}:{}", ctx.master_ib_ip, ctx.init_port),
                ]);
            }
            Some(DistributedBackend::Ray | DistributedBackend::Mp) => {
                let mut fp = vec![
                    "--tensor-parallel-size".into(),
                    ctx.tp_size.to_string(),
                    "--pipeline-parallel-size".into(),
                    ctx.pp_size.to_string(),
                    "--nnodes".into(),
                    ctx.nnodes.to_string(),
                    "--node-rank".into(),
                    node.rank.to_string(),
                    "--master-addr".into(),
                    ctx.master_ib_ip.to_string(),
                    "--master-port".into(),
                    ctx.init_port.to_string(),
                ];
                if ctx.recipe.distributed_backend == Some(DistributedBackend::Ray) {
                    fp.push("--distributed-executor-backend".into());
                    fp.push("ray".into());
                }
                args.extend(fp);
            }
            _ => {}
        }
    }

    // vLLM workers in distributed mode run headless (no port binding).
    if ctx.nnodes > 1 && !is_sglang && node.rank > 0 && is_distributed {
        args.push("--headless".into());
    }

    args
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::{DeployTarget, RecipeArgs, VenvSetup};
    use std::collections::HashMap;

    fn base_recipe() -> Recipe {
        Recipe {
            engine: "vllm".into(),
            name: "T".into(),
            docker_image: "img".into(),
            model: "mdl".into(),
            port: 8000,
            required_nodes: None,
            pipeline_parallel_size: Some(1),
            tensor_parallel_size: 1,
            args: RecipeArgs::List(vec!["--x".into(), "4096".into()]),
            env: HashMap::new(),
            cluster: None,
            entrypoint: vec!["vllm".into(), "serve".into()],
            distributed_backend: Some(DistributedBackend::Mp),
            before_serve: vec![],
            after_stop: vec![],
            solo: None,
            deploy: DeployTarget::Docker,
            venv: None,
            venv_setup: VenvSetup::default(),
            venv_setup_script: None,
            mods: Default::default(),
            shm_size: String::new(),
            sync_image: true,
            docker_image_dir: None,
            extra_mounts: vec![],
            build: Default::default(),
            composed_model: None,
        }
    }

    fn tnode(rank: u32) -> Node {
        Node {
            name: format!("n{}", rank),
            ssh_ip: "10.0.0.1".into(),
            ib_ip: Some("10.0.0.1".into()),
            rank,
            host_key: None,
            entity: None,
        }
    }

    #[test]
    fn single_node_no_distributed_flags() {
        let r = base_recipe();
        let ctx = EngineFlagsCtx {
            recipe: &r,
            tp_size: 1,
            pp_size: 1,
            nnodes: 1,
            init_port: 25000,
            master_ib_ip: "10.0.0.1",
        };
        let args = build_engine_args(&ctx, &tnode(0));
        // model + recipe args, no distributed flags
        assert_eq!(args, vec!["mdl", "--x", "4096"]);
    }

    #[test]
    fn multi_node_vllm_mp_flags() {
        let r = base_recipe();
        let ctx = EngineFlagsCtx {
            recipe: &r,
            tp_size: 2,
            pp_size: 1,
            nnodes: 2,
            init_port: 25000,
            master_ib_ip: "10.0.0.1",
        };
        let head_args = build_engine_args(&ctx, &tnode(0));
        assert!(head_args.contains(&"--tensor-parallel-size".into()));
        assert!(head_args.contains(&"2".into()));
        assert!(head_args.contains(&"--nnodes".into()));
        assert!(head_args.contains(&"2".into()));
        assert!(head_args.contains(&"--master-addr".into()));
        assert!(head_args.contains(&"10.0.0.1".into()));
        assert!(head_args.contains(&"--master-port".into()));
        assert!(head_args.contains(&"25000".into()));
        // head node: no --headless
        assert!(!head_args.contains(&"--headless".into()));

        let worker_args = build_engine_args(&ctx, &tnode(1));
        assert!(worker_args.contains(&"--headless".into()));
    }

    #[test]
    fn multi_node_sglang_flags() {
        let mut r = base_recipe();
        r.engine = "sglang".into();
        let ctx = EngineFlagsCtx {
            recipe: &r,
            tp_size: 2,
            pp_size: 1,
            nnodes: 2,
            init_port: 25000,
            master_ib_ip: "10.0.0.1",
        };
        let args = build_engine_args(&ctx, &tnode(0));
        // SGLang uses --model-path, --tp, --dist-init-addr
        assert!(args.contains(&"--model-path".into()));
        assert!(args.contains(&"--tp".into()));
        assert!(args.contains(&"--dist-init-addr".into()));
        assert!(args.contains(&"10.0.0.1:25000".into()));
        // SGLang never gets --headless
        assert!(!args.contains(&"--headless".into()));
    }

    #[test]
    fn ray_backend_adds_executor() {
        let mut r = base_recipe();
        r.distributed_backend = Some(DistributedBackend::Ray);
        let ctx = EngineFlagsCtx {
            recipe: &r,
            tp_size: 1,
            pp_size: 1,
            nnodes: 2,
            init_port: 25000,
            master_ib_ip: "10.0.0.1",
        };
        let args = build_engine_args(&ctx, &tnode(0));
        assert!(args.contains(&"--distributed-executor-backend".into()));
        assert!(args.contains(&"ray".into()));
    }
}