|
|
@@ -0,0 +1,351 @@
|
|
|
+//! Assembly of the `docker run` command per node — flags, env vars, mounts,
|
|
|
+//! and the vLLM/SGLang tail (engine flags + distributed rendezvous).
|
|
|
+//!
|
|
|
+//! [`build_docker_command`] is the single entry point; everything else here is
|
|
|
+//! private to this module.
|
|
|
+
|
|
|
+use crate::config::{Cluster, Config, NcclConfig, Node, Recipe};
|
|
|
+use anyhow::Result;
|
|
|
+use std::borrow::Cow;
|
|
|
+
|
|
|
+use super::{mods, tail};
|
|
|
+
|
|
|
+/// Shared context for [`build_docker_command`] — constructed once, cloned per node.
|
|
|
+#[derive(Clone)]
|
|
|
+pub(crate) struct DockerCommandCtx {
|
|
|
+ pub cluster: Cluster,
|
|
|
+ pub nccl: NcclConfig,
|
|
|
+ pub recipe: Recipe,
|
|
|
+ pub recipe_name: String,
|
|
|
+ pub pp_size: u32,
|
|
|
+ pub tp_size: u32,
|
|
|
+ pub nnodes: u32,
|
|
|
+ pub cache_dir: String,
|
|
|
+ pub keep: bool,
|
|
|
+ pub init_port: u16,
|
|
|
+ pub mods_remote_path: String,
|
|
|
+ /// Recipe `extra_mounts` with host paths already resolved (alias lookup +
|
|
|
+ /// `~` expansion) against `Config.mount_aliases` / `ssh_user`.
|
|
|
+ pub extra_mounts: Vec<ResolvedMount>,
|
|
|
+}
|
|
|
+
|
|
|
+/// A recipe `ExtraMount` after host-path resolution (alias + tilde), ready to
|
|
|
+/// turn into a `-v` flag / `-e` env pair.
|
|
|
+#[derive(Clone)]
|
|
|
+pub(crate) struct ResolvedMount {
|
|
|
+ pub host: String,
|
|
|
+ pub container: String,
|
|
|
+ pub read_only: bool,
|
|
|
+ pub env: Option<String>,
|
|
|
+}
|
|
|
+
|
|
|
+/// Resolve every `extra_mount` in a recipe to a concrete host path using the
|
|
|
+/// config alias registry and the SSH user (for `~` expansion). Done once per
|
|
|
+/// serve; the resulting `Vec<ResolvedMount>` is cloned into each node's
|
|
|
+/// [`DockerCommandCtx`].
|
|
|
+pub(crate) fn resolve_extra_mounts(config: &Config, recipe: &Recipe) -> Result<Vec<ResolvedMount>> {
|
|
|
+ recipe
|
|
|
+ .extra_mounts
|
|
|
+ .iter()
|
|
|
+ .map(|m| {
|
|
|
+ let host = crate::config::resolve_mount_host(
|
|
|
+ m.alias.as_deref(),
|
|
|
+ m.host.as_deref(),
|
|
|
+ &config.mount_aliases,
|
|
|
+ &config.ssh_user,
|
|
|
+ &m.container,
|
|
|
+ )?;
|
|
|
+ Ok(ResolvedMount {
|
|
|
+ host,
|
|
|
+ container: m.container.clone(),
|
|
|
+ read_only: m.read_only,
|
|
|
+ env: m.env.clone(),
|
|
|
+ })
|
|
|
+ })
|
|
|
+ .collect()
|
|
|
+}
|
|
|
+
|
|
|
+/// Collect NCCL cluster env vars for multi-node deployments.
|
|
|
+fn collect_nccl_env<'a>(ctx: &'a DockerCommandCtx, node: &'a Node) -> Vec<(&'a str, Cow<'a, str>)> {
|
|
|
+ let is_sglang = ctx.recipe.engine == "sglang";
|
|
|
+ let sock = ctx.nccl.socket_ifname.as_str();
|
|
|
+ let gloo_if = ctx.nccl.gloo_socket_ifname.as_deref().unwrap_or(sock);
|
|
|
+ let ucx_dev = ctx.nccl.ucx_net_devices.as_deref().unwrap_or(sock);
|
|
|
+ let mut env: Vec<(&str, Cow<str>)> = vec![
|
|
|
+ ("NCCL_IB_DISABLE", format!("{}", ctx.nccl.ib_disable).into()),
|
|
|
+ ("NCCL_IB_HCA", ctx.nccl.ib_hca.as_str().into()),
|
|
|
+ ("NCCL_SOCKET_IFNAME", sock.into()),
|
|
|
+ ("GLOO_SOCKET_IFNAME", gloo_if.into()),
|
|
|
+ ("TP_SOCKET_IFNAME", gloo_if.into()),
|
|
|
+ (
|
|
|
+ "NCCL_IB_GID_INDEX",
|
|
|
+ ctx.nccl.ib_gid_index.to_string().into(),
|
|
|
+ ),
|
|
|
+ ("MN_IF_NAME", gloo_if.into()),
|
|
|
+ ("UCX_NET_DEVICES", ucx_dev.into()),
|
|
|
+ ("OMPI_MCA_btl_tcp_if_include", gloo_if.into()),
|
|
|
+ ("OMP_NUM_THREADS", "4".into()),
|
|
|
+ ("NCCL_IGNORE_CPU_AFFINITY", "1".into()),
|
|
|
+ ("NCCL_NET", "IB".into()),
|
|
|
+ ("NCCL_DEBUG", "WARN".into()),
|
|
|
+ ];
|
|
|
+ if let Some(v) = ctx.nccl.nccl_cross_nic {
|
|
|
+ env.push(("NCCL_CROSS_NIC", v.to_string().into()));
|
|
|
+ }
|
|
|
+ if let Some(ref v) = ctx.nccl.nccl_proto {
|
|
|
+ env.push(("NCCL_PROTO", v.as_str().into()));
|
|
|
+ }
|
|
|
+ if let Some(ref v) = ctx.nccl.nccl_p2p_level {
|
|
|
+ env.push(("NCCL_P2P_LEVEL", v.as_str().into()));
|
|
|
+ }
|
|
|
+ let worker_ip = Cow::Borrowed(node.ib_ip.as_str());
|
|
|
+ if is_sglang {
|
|
|
+ env.push(("SGLANG_WORKER_IP", worker_ip));
|
|
|
+ } else {
|
|
|
+ env.push(("VLLM_HOST_IP", worker_ip));
|
|
|
+ }
|
|
|
+ env
|
|
|
+}
|
|
|
+
|
|
|
+/// Build the list of `-e KEY=VALUE` Docker env var pairs for a node.
|
|
|
+fn collect_docker_env(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> Vec<String> {
|
|
|
+ let mut pairs: Vec<(&str, Cow<str>)> =
|
|
|
+ vec![("HF_HOME", Cow::Borrowed("/root/.cache/huggingface"))];
|
|
|
+ if ctx.nnodes > 1 {
|
|
|
+ pairs.extend(collect_nccl_env(ctx, node));
|
|
|
+ }
|
|
|
+ if is_head {
|
|
|
+ pairs.push(("VLLM_LOGGING_COLOR", "1".into()));
|
|
|
+ pairs.push(("TERM", "xterm-256color".into()));
|
|
|
+ }
|
|
|
+ if ctx.recipe.distributed_backend == Some(crate::config::DistributedBackend::Ray) {
|
|
|
+ pairs.push((
|
|
|
+ "RAY_ADDRESS",
|
|
|
+ format!("{}:6379", ctx.cluster.master_ib_ip).into(),
|
|
|
+ ));
|
|
|
+ pairs.push(("RAY_PORT", "6379".into()));
|
|
|
+ pairs.push(("RAY_NODE_IP_ADDRESS", node.ib_ip.as_str().into()));
|
|
|
+ pairs.push(("RAY_OVERRIDE_NODE_IP_ADDRESS", node.ib_ip.as_str().into()));
|
|
|
+ }
|
|
|
+ // extra_mount `env` defaults (point at the container path) are injected
|
|
|
+ // before recipe [env] so the recipe can still override them (last wins).
|
|
|
+ for m in &ctx.extra_mounts {
|
|
|
+ if let Some(name) = &m.env {
|
|
|
+ pairs.push((name.as_str(), m.container.as_str().into()));
|
|
|
+ }
|
|
|
+ }
|
|
|
+ for (k, v) in &ctx.recipe.env {
|
|
|
+ pairs.push((k.as_str(), v.as_str().into()));
|
|
|
+ }
|
|
|
+ pairs
|
|
|
+ .into_iter()
|
|
|
+ .map(|(k, v)| format!("-e {k}={v}"))
|
|
|
+ .collect()
|
|
|
+}
|
|
|
+
|
|
|
+/// Build the list of Docker run flags (--gpus, --ipc, --ulimit, etc.).
|
|
|
+fn build_docker_flags(ctx: &DockerCommandCtx) -> Vec<String> {
|
|
|
+ let mut flags: Vec<String> = vec![
|
|
|
+ "--gpus all".to_string(),
|
|
|
+ "--ipc=host".to_string(),
|
|
|
+ "--privileged".to_string(),
|
|
|
+ "--network=host".to_string(),
|
|
|
+ "--init".to_string(),
|
|
|
+ "--ulimit memlock=-1".to_string(),
|
|
|
+ "--ulimit stack=67108864".to_string(),
|
|
|
+ "--ulimit nofile=1048576".to_string(),
|
|
|
+ ];
|
|
|
+ if !ctx.recipe.shm_size.is_empty() {
|
|
|
+ flags.push(format!("--shm-size {}", ctx.recipe.shm_size.as_str()));
|
|
|
+ }
|
|
|
+ if !ctx.keep {
|
|
|
+ flags.push("--rm".to_string());
|
|
|
+ }
|
|
|
+ flags
|
|
|
+}
|
|
|
+
|
|
|
+/// Build the full Docker command line for a node.
|
|
|
+pub(crate) fn build_docker_command(
|
|
|
+ ctx: &DockerCommandCtx,
|
|
|
+ node: &Node,
|
|
|
+ is_head: bool,
|
|
|
+) -> Vec<String> {
|
|
|
+ let flags = build_docker_flags(ctx);
|
|
|
+ let mut cmd = vec![
|
|
|
+ format!("docker run -d {}", flags.join(" ")),
|
|
|
+ format!("-v {}:/root/.cache/huggingface", ctx.cache_dir),
|
|
|
+ ];
|
|
|
+ if !ctx.mods_remote_path.is_empty() && !ctx.recipe.mods.commands.is_empty() {
|
|
|
+ cmd.push(format!("-v {}:/mods:ro", ctx.mods_remote_path));
|
|
|
+ }
|
|
|
+ for m in &ctx.extra_mounts {
|
|
|
+ let mut spec = format!("-v {}:{}", m.host, m.container);
|
|
|
+ if m.read_only {
|
|
|
+ spec.push_str(":ro");
|
|
|
+ }
|
|
|
+ cmd.push(spec);
|
|
|
+ }
|
|
|
+ cmd.extend(collect_docker_env(ctx, node, is_head));
|
|
|
+ cmd.push(format!("--name sm121_{}_{}", ctx.recipe_name, node.rank));
|
|
|
+ cmd.push(format!("--label sm121.recipe={}", ctx.recipe_name));
|
|
|
+ cmd.push(format!("--label sm121.rank={}", node.rank));
|
|
|
+ cmd.push(format!("--label sm121.cluster={}", ctx.cluster.name));
|
|
|
+ // Add solo-node label for single-node deployments
|
|
|
+ if ctx.cluster.nodes.len() == 1 && ctx.cluster.extra_entities.is_empty() {
|
|
|
+ cmd.push(format!("--label sm121.solo-node={}", node.name));
|
|
|
+ }
|
|
|
+ let tctx = tail::TailCtx {
|
|
|
+ recipe: &ctx.recipe,
|
|
|
+ tp_size: ctx.tp_size,
|
|
|
+ pp_size: ctx.pp_size,
|
|
|
+ nnodes: ctx.nnodes,
|
|
|
+ init_port: ctx.init_port,
|
|
|
+ master_ib_ip: &ctx.cluster.master_ib_ip,
|
|
|
+ };
|
|
|
+ let tail = tail::build_tail(&tctx, node);
|
|
|
+ if ctx.recipe.mods.commands.is_empty() {
|
|
|
+ cmd.push(ctx.recipe.docker_image.clone());
|
|
|
+ cmd.extend(tail);
|
|
|
+ } else {
|
|
|
+ let script = mods::build_mod_script(&ctx.recipe.mods.commands, &tail);
|
|
|
+ cmd.extend([
|
|
|
+ "--entrypoint sh".into(),
|
|
|
+ ctx.recipe.docker_image.clone(),
|
|
|
+ "-c".into(),
|
|
|
+ script,
|
|
|
+ ]);
|
|
|
+ }
|
|
|
+ cmd
|
|
|
+}
|
|
|
+
|
|
|
+#[cfg(test)]
|
|
|
+mod tests {
|
|
|
+ use super::*;
|
|
|
+ use crate::config::{
|
|
|
+ Cluster, DistributedBackend, ExtraMount, NcclConfig, Node, Recipe, RecipeArgs,
|
|
|
+ };
|
|
|
+ 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,
|
|
|
+ entity: None,
|
|
|
+ }],
|
|
|
+ extra_entities: vec![],
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ fn tnccl() -> NcclConfig {
|
|
|
+ NcclConfig {
|
|
|
+ ib_hca: "r".into(),
|
|
|
+ socket_ifname: "e".into(),
|
|
|
+ ib_gid_index: 3,
|
|
|
+ ib_disable: 0,
|
|
|
+ gloo_socket_ifname: None,
|
|
|
+ ucx_net_devices: None,
|
|
|
+ nccl_cross_nic: None,
|
|
|
+ nccl_proto: None,
|
|
|
+ nccl_p2p_level: None,
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ fn minimal_config() -> Config {
|
|
|
+ Config {
|
|
|
+ ssh_user: "u".into(),
|
|
|
+ ssh_key_path: "/tmp/key".into(),
|
|
|
+ clusters: vec![],
|
|
|
+ nccl: HashMap::new(),
|
|
|
+ hf_cache_dir: None,
|
|
|
+ init_port: 25000,
|
|
|
+ before_serve: HashMap::new(),
|
|
|
+ exec_aliases: HashMap::new(),
|
|
|
+ home_assistant: None,
|
|
|
+ global_mods_dir: None,
|
|
|
+ drop_caches: true,
|
|
|
+ mount_aliases: HashMap::from([("models".to_string(), "/mnt/nas/models".to_string())]),
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ fn run_with(recipe: &Recipe) -> String {
|
|
|
+ let ctx = DockerCommandCtx {
|
|
|
+ cluster: tcls(),
|
|
|
+ nccl: tnccl(),
|
|
|
+ recipe: recipe.clone(),
|
|
|
+ recipe_name: "r".to_string(),
|
|
|
+ pp_size: 2,
|
|
|
+ tp_size: 1,
|
|
|
+ nnodes: 2,
|
|
|
+ cache_dir: "/home/user/.cache/huggingface".to_string(),
|
|
|
+ keep: false,
|
|
|
+ init_port: 25000,
|
|
|
+ mods_remote_path: String::new(),
|
|
|
+ extra_mounts: resolve_extra_mounts(&minimal_config(), recipe).unwrap_or_default(),
|
|
|
+ };
|
|
|
+ build_docker_command(&ctx, &tcls().nodes[0], true).join(" ")
|
|
|
+ }
|
|
|
+
|
|
|
+ fn base_recipe() -> 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: RecipeArgs::List(vec!["--x".into(), "4096".into()]),
|
|
|
+ env: HashMap::new(),
|
|
|
+ cluster: None,
|
|
|
+ entrypoint: vec![],
|
|
|
+ distributed_backend: Some(DistributedBackend::Mp),
|
|
|
+ before_serve: vec![],
|
|
|
+ mods: Default::default(),
|
|
|
+ shm_size: String::new(),
|
|
|
+ sync_image: true,
|
|
|
+ docker_image_dir: None,
|
|
|
+ extra_mounts: vec![],
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn extra_mounts_are_mounted() {
|
|
|
+ let mut r = base_recipe();
|
|
|
+ r.extra_mounts.push(ExtraMount {
|
|
|
+ alias: Some("models".into()),
|
|
|
+ host: Some("~/.cache/models".into()),
|
|
|
+ container: "/models:ro".into(),
|
|
|
+ read_only: true,
|
|
|
+ env: Some("MODEL_DIR".into()),
|
|
|
+ });
|
|
|
+ let cmd = run_with(&r);
|
|
|
+ // alias present in minimal_config() → alias host wins
|
|
|
+ assert!(cmd.contains("-v /mnt/nas/models:/models:ro:ro"));
|
|
|
+ assert!(cmd.contains("-v /home/user/.cache/huggingface:/root/.cache/huggingface"));
|
|
|
+ // env injected pointing at the container path, before recipe env
|
|
|
+ assert!(cmd.contains("-e MODEL_DIR=/models:ro"));
|
|
|
+ }
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn extra_mounts_fallback_to_host_when_alias_absent() {
|
|
|
+ let mut r = base_recipe();
|
|
|
+ r.extra_mounts.push(ExtraMount {
|
|
|
+ alias: Some("no_such_alias".into()),
|
|
|
+ host: Some("~/.cache/jit".into()),
|
|
|
+ container: "/root/.cache/jit".into(),
|
|
|
+ read_only: false,
|
|
|
+ env: None,
|
|
|
+ });
|
|
|
+ let cmd = run_with(&r);
|
|
|
+ // alias not in config → host default used with ~ expansion (ssh_user=u)
|
|
|
+ assert!(cmd.contains("-v /home/u/.cache/jit:/root/.cache/jit"));
|
|
|
+ assert!(!cmd.contains(":ro"));
|
|
|
+ }
|
|
|
+}
|