|
@@ -85,10 +85,12 @@ pub async fn serve(
|
|
|
&nccl,
|
|
&nccl,
|
|
|
recipe_name,
|
|
recipe_name,
|
|
|
&recipe,
|
|
&recipe,
|
|
|
- cache_dir,
|
|
|
|
|
- keep,
|
|
|
|
|
- node_name,
|
|
|
|
|
- dry_run,
|
|
|
|
|
|
|
+ StartClusterOpts {
|
|
|
|
|
+ cache_dir,
|
|
|
|
|
+ keep,
|
|
|
|
|
+ node_name: node_name.map(String::from),
|
|
|
|
|
+ dry_run,
|
|
|
|
|
+ },
|
|
|
)
|
|
)
|
|
|
.await?;
|
|
.await?;
|
|
|
Ok(cluster)
|
|
Ok(cluster)
|
|
@@ -221,20 +223,39 @@ pub async fn resolve_cluster_by_container_label<'c>(
|
|
|
anyhow::bail!("No running containers found for recipe '{}'", recipe_name);
|
|
anyhow::bail!("No running containers found for recipe '{}'", recipe_name);
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
-#[allow(clippy::too_many_arguments)]
|
|
|
|
|
|
|
+/// Runtime options for [`start_cluster`].
|
|
|
|
|
+struct StartClusterOpts {
|
|
|
|
|
+ cache_dir: std::path::PathBuf,
|
|
|
|
|
+ keep: bool,
|
|
|
|
|
+ node_name: Option<String>,
|
|
|
|
|
+ dry_run: bool,
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// Shared context for [`build_docker_command`] — constructed once, cloned per node.
|
|
|
|
|
+#[derive(Clone)]
|
|
|
|
|
+struct DockerCommandCtx {
|
|
|
|
|
+ cluster: Cluster,
|
|
|
|
|
+ nccl: NcclConfig,
|
|
|
|
|
+ recipe: Recipe,
|
|
|
|
|
+ recipe_name: String,
|
|
|
|
|
+ pp_size: u32,
|
|
|
|
|
+ tp_size: u32,
|
|
|
|
|
+ nnodes: u32,
|
|
|
|
|
+ cache_dir: String,
|
|
|
|
|
+ keep: bool,
|
|
|
|
|
+ init_port: u16,
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
async fn start_cluster(
|
|
async fn start_cluster(
|
|
|
config: &Config,
|
|
config: &Config,
|
|
|
cluster: &Cluster,
|
|
cluster: &Cluster,
|
|
|
nccl: &NcclConfig,
|
|
nccl: &NcclConfig,
|
|
|
recipe_name: &str,
|
|
recipe_name: &str,
|
|
|
recipe: &Recipe,
|
|
recipe: &Recipe,
|
|
|
- cache_dir: std::path::PathBuf,
|
|
|
|
|
- keep: bool,
|
|
|
|
|
- node_name: Option<&str>,
|
|
|
|
|
- dry_run: bool,
|
|
|
|
|
|
|
+ opts: StartClusterOpts,
|
|
|
) -> Result<()> {
|
|
) -> Result<()> {
|
|
|
- let nodes: Vec<&Node> = if let Some(nn) = node_name {
|
|
|
|
|
- let v: Vec<&Node> = cluster.nodes.iter().filter(|n| n.name == nn).collect();
|
|
|
|
|
|
|
+ let nodes: Vec<&Node> = if let Some(ref nn) = opts.node_name {
|
|
|
|
|
+ let v: Vec<&Node> = cluster.nodes.iter().filter(|n| n.name == *nn).collect();
|
|
|
if v.is_empty() {
|
|
if v.is_empty() {
|
|
|
anyhow::bail!("Node '{}' not found in cluster '{}'", nn, cluster.name);
|
|
anyhow::bail!("Node '{}' not found in cluster '{}'", nn, cluster.name);
|
|
|
}
|
|
}
|
|
@@ -258,39 +279,32 @@ async fn start_cluster(
|
|
|
let mut handles = Vec::new();
|
|
let mut handles = Vec::new();
|
|
|
for node in &nodes {
|
|
for node in &nodes {
|
|
|
let node = (*node).clone();
|
|
let node = (*node).clone();
|
|
|
- let cluster = cluster.clone();
|
|
|
|
|
- let nccl = nccl.clone();
|
|
|
|
|
- let recipe = recipe.clone();
|
|
|
|
|
let config = config.clone();
|
|
let config = config.clone();
|
|
|
let rn = recipe_name.to_string();
|
|
let rn = recipe_name.to_string();
|
|
|
- let dry = dry_run;
|
|
|
|
|
|
|
+ let dry = opts.dry_run;
|
|
|
let engine = recipe.engine.clone();
|
|
let engine = recipe.engine.clone();
|
|
|
let is_head = node.rank == 0;
|
|
let is_head = node.rank == 0;
|
|
|
let cleanup = format!(
|
|
let cleanup = format!(
|
|
|
"docker rm -f sm121_{}_{} 2>/dev/null; true",
|
|
"docker rm -f sm121_{}_{} 2>/dev/null; true",
|
|
|
recipe_name, node.rank
|
|
recipe_name, node.rank
|
|
|
);
|
|
);
|
|
|
- let cd = cache_dir.clone();
|
|
|
|
|
- let cd_s = cd.to_string_lossy().into_owned();
|
|
|
|
|
|
|
+ let dm_ctx = DockerCommandCtx {
|
|
|
|
|
+ cluster: cluster.clone(),
|
|
|
|
|
+ nccl: nccl.clone(),
|
|
|
|
|
+ recipe: recipe.clone(),
|
|
|
|
|
+ recipe_name: rn.clone(),
|
|
|
|
|
+ pp_size,
|
|
|
|
|
+ tp_size,
|
|
|
|
|
+ nnodes,
|
|
|
|
|
+ cache_dir: opts.cache_dir.to_string_lossy().into_owned(),
|
|
|
|
|
+ keep: opts.keep,
|
|
|
|
|
+ init_port: config.init_port,
|
|
|
|
|
+ };
|
|
|
handles.push(tokio::spawn(async move {
|
|
handles.push(tokio::spawn(async move {
|
|
|
if !dry {
|
|
if !dry {
|
|
|
let _ = run_ssh(&config, &node, cleanup, dry).await;
|
|
let _ = run_ssh(&config, &node, cleanup, dry).await;
|
|
|
}
|
|
}
|
|
|
- let cmd = build_docker_command(
|
|
|
|
|
- &cluster,
|
|
|
|
|
- &nccl,
|
|
|
|
|
- &recipe,
|
|
|
|
|
- &rn,
|
|
|
|
|
- &node,
|
|
|
|
|
- is_head,
|
|
|
|
|
- pp_size,
|
|
|
|
|
- tp_size,
|
|
|
|
|
- nnodes,
|
|
|
|
|
- &cd_s,
|
|
|
|
|
- keep,
|
|
|
|
|
- config.init_port,
|
|
|
|
|
- )
|
|
|
|
|
- .join(" ");
|
|
|
|
|
|
|
+ let cmd = build_docker_command(&dm_ctx, &node, is_head).join(" ");
|
|
|
let result = run_ssh_with_retry(&config, &node, cmd, dry, 2).await;
|
|
let result = run_ssh_with_retry(&config, &node, cmd, dry, 2).await;
|
|
|
(node.name, node.rank, engine, result)
|
|
(node.name, node.rank, engine, result)
|
|
|
}));
|
|
}));
|
|
@@ -308,7 +322,7 @@ async fn start_cluster(
|
|
|
Err(e) => anyhow::bail!("Task join error: {}", e),
|
|
Err(e) => anyhow::bail!("Task join error: {}", e),
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
|
- if !dry_run {
|
|
|
|
|
|
|
+ if !opts.dry_run {
|
|
|
sleep(Duration::from_secs(10)).await;
|
|
sleep(Duration::from_secs(10)).await;
|
|
|
info!(
|
|
info!(
|
|
|
"{}",
|
|
"{}",
|
|
@@ -317,22 +331,8 @@ async fn start_cluster(
|
|
|
}
|
|
}
|
|
|
Ok(())
|
|
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,
|
|
|
|
|
- tp_size: u32,
|
|
|
|
|
- nnodes: u32,
|
|
|
|
|
- cache_dir: &str,
|
|
|
|
|
- keep: bool,
|
|
|
|
|
- init_port: u16,
|
|
|
|
|
-) -> Vec<String> {
|
|
|
|
|
- let is_sglang = recipe.engine == "sglang";
|
|
|
|
|
|
|
+fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> Vec<String> {
|
|
|
|
|
+ let is_sglang = ctx.recipe.engine == "sglang";
|
|
|
let mut flags: Vec<String> = vec![
|
|
let mut flags: Vec<String> = vec![
|
|
|
"--gpus all".to_string(),
|
|
"--gpus all".to_string(),
|
|
|
"--ipc=host".to_string(),
|
|
"--ipc=host".to_string(),
|
|
@@ -341,34 +341,37 @@ fn build_docker_command(
|
|
|
"--ulimit memlock=-1".to_string(),
|
|
"--ulimit memlock=-1".to_string(),
|
|
|
"--ulimit stack=67108864".to_string(),
|
|
"--ulimit stack=67108864".to_string(),
|
|
|
];
|
|
];
|
|
|
- if !recipe.shm_size.is_empty() {
|
|
|
|
|
- flags.push(format!("--shm-size {}", recipe.shm_size.as_str()));
|
|
|
|
|
|
|
+ if !ctx.recipe.shm_size.is_empty() {
|
|
|
|
|
+ flags.push(format!("--shm-size {}", ctx.recipe.shm_size.as_str()));
|
|
|
}
|
|
}
|
|
|
- if !keep {
|
|
|
|
|
|
|
+ if !ctx.keep {
|
|
|
flags.push("--rm".to_string());
|
|
flags.push("--rm".to_string());
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
let mut cmd = vec![
|
|
let mut cmd = vec![
|
|
|
format!("docker run -d {}", flags.join(" ")),
|
|
format!("docker run -d {}", flags.join(" ")),
|
|
|
- format!("-v {}:/root/.cache/huggingface", cache_dir),
|
|
|
|
|
|
|
+ format!("-v {}:/root/.cache/huggingface", ctx.cache_dir),
|
|
|
];
|
|
];
|
|
|
|
|
|
|
|
let mut env = vec![("HF_HOME", Cow::Borrowed("/root/.cache/huggingface"))];
|
|
let mut env = vec![("HF_HOME", Cow::Borrowed("/root/.cache/huggingface"))];
|
|
|
- if nnodes > 1 {
|
|
|
|
|
|
|
+ if ctx.nnodes > 1 {
|
|
|
let worker_ip = Cow::Borrowed(node.ib_ip.as_str());
|
|
let worker_ip = Cow::Borrowed(node.ib_ip.as_str());
|
|
|
if is_sglang {
|
|
if is_sglang {
|
|
|
env.push(("SGLANG_WORKER_IP", worker_ip));
|
|
env.push(("SGLANG_WORKER_IP", worker_ip));
|
|
|
} else {
|
|
} else {
|
|
|
env.push(("VLLM_HOST_IP", worker_ip));
|
|
env.push(("VLLM_HOST_IP", worker_ip));
|
|
|
}
|
|
}
|
|
|
- let nccl_socket_ifname = nccl.socket_ifname.as_str();
|
|
|
|
|
|
|
+ let nccl_socket_ifname = ctx.nccl.socket_ifname.as_str();
|
|
|
env.extend([
|
|
env.extend([
|
|
|
- ("NCCL_IB_DISABLE", format!("{}", nccl.ib_disable).into()),
|
|
|
|
|
- ("NCCL_IB_HCA", nccl.ib_hca.as_str().into()),
|
|
|
|
|
|
|
+ ("NCCL_IB_DISABLE", format!("{}", ctx.nccl.ib_disable).into()),
|
|
|
|
|
+ ("NCCL_IB_HCA", ctx.nccl.ib_hca.as_str().into()),
|
|
|
("NCCL_SOCKET_IFNAME", nccl_socket_ifname.into()),
|
|
("NCCL_SOCKET_IFNAME", nccl_socket_ifname.into()),
|
|
|
("GLOO_SOCKET_IFNAME", nccl_socket_ifname.into()),
|
|
("GLOO_SOCKET_IFNAME", nccl_socket_ifname.into()),
|
|
|
("TP_SOCKET_IFNAME", nccl_socket_ifname.into()),
|
|
("TP_SOCKET_IFNAME", nccl_socket_ifname.into()),
|
|
|
- ("NCCL_IB_GID_INDEX", nccl.ib_gid_index.to_string().into()),
|
|
|
|
|
|
|
+ (
|
|
|
|
|
+ "NCCL_IB_GID_INDEX",
|
|
|
|
|
+ ctx.nccl.ib_gid_index.to_string().into(),
|
|
|
|
|
+ ),
|
|
|
("MN_IF_NAME", nccl_socket_ifname.into()),
|
|
("MN_IF_NAME", nccl_socket_ifname.into()),
|
|
|
("UCX_NET_DEVICES", nccl_socket_ifname.into()),
|
|
("UCX_NET_DEVICES", nccl_socket_ifname.into()),
|
|
|
("OMPI_MCA_btl_tcp_if_include", nccl_socket_ifname.into()),
|
|
("OMPI_MCA_btl_tcp_if_include", nccl_socket_ifname.into()),
|
|
@@ -382,17 +385,17 @@ fn build_docker_command(
|
|
|
env.push(("VLLM_LOGGING_COLOR", "1".into()));
|
|
env.push(("VLLM_LOGGING_COLOR", "1".into()));
|
|
|
env.push(("TERM", "xterm-256color".into()));
|
|
env.push(("TERM", "xterm-256color".into()));
|
|
|
}
|
|
}
|
|
|
- for (k, v) in &recipe.env {
|
|
|
|
|
|
|
+ for (k, v) in &ctx.recipe.env {
|
|
|
env.push((k.as_str(), v.as_str().into()));
|
|
env.push((k.as_str(), v.as_str().into()));
|
|
|
}
|
|
}
|
|
|
- cmd.push(format!("--name sm121_{}_{}", recipe_name, node.rank));
|
|
|
|
|
- cmd.push(format!("--label sm121.recipe={}", recipe_name));
|
|
|
|
|
|
|
+ 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.rank={}", node.rank));
|
|
|
- cmd.push(format!("--label sm121.cluster={}", cluster.name));
|
|
|
|
|
- if recipe.distributed_backend == "ray" {
|
|
|
|
|
|
|
+ cmd.push(format!("--label sm121.cluster={}", ctx.cluster.name));
|
|
|
|
|
+ if ctx.recipe.distributed_backend == "ray" {
|
|
|
env.push((
|
|
env.push((
|
|
|
"RAY_ADDRESS",
|
|
"RAY_ADDRESS",
|
|
|
- format!("{}:6379", cluster.master_ib_ip).into(),
|
|
|
|
|
|
|
+ format!("{}:6379", ctx.cluster.master_ib_ip).into(),
|
|
|
));
|
|
));
|
|
|
env.push(("RAY_PORT", "6379".into()));
|
|
env.push(("RAY_PORT", "6379".into()));
|
|
|
env.push(("RAY_NODE_IP_ADDRESS", node.ib_ip.as_str().into()));
|
|
env.push(("RAY_NODE_IP_ADDRESS", node.ib_ip.as_str().into()));
|
|
@@ -402,10 +405,10 @@ fn build_docker_command(
|
|
|
cmd.push(format!("-e {k}={v}"));
|
|
cmd.push(format!("-e {k}={v}"));
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
- cmd.push(recipe.docker_image.clone());
|
|
|
|
|
|
|
+ cmd.push(ctx.recipe.docker_image.clone());
|
|
|
|
|
|
|
|
- if !recipe.entrypoint.is_empty() {
|
|
|
|
|
- let mut ep = recipe.entrypoint.clone();
|
|
|
|
|
|
|
+ if !ctx.recipe.entrypoint.is_empty() {
|
|
|
|
|
+ let mut ep = ctx.recipe.entrypoint.clone();
|
|
|
if ep.len() >= 3 && ep[ep.len() - 2] == "-c" {
|
|
if ep.len() >= 3 && ep[ep.len() - 2] == "-c" {
|
|
|
let last = ep.len() - 1;
|
|
let last = ep.len() - 1;
|
|
|
ep[last] = shq(&ep[last]);
|
|
ep[last] = shq(&ep[last]);
|
|
@@ -416,52 +419,55 @@ fn build_docker_command(
|
|
|
if is_sglang {
|
|
if is_sglang {
|
|
|
cmd.push("--model-path".to_string());
|
|
cmd.push("--model-path".to_string());
|
|
|
}
|
|
}
|
|
|
- cmd.push(recipe.model.clone());
|
|
|
|
|
|
|
+ cmd.push(ctx.recipe.model.clone());
|
|
|
|
|
|
|
|
- cmd.extend(recipe.args.clone());
|
|
|
|
|
|
|
+ cmd.extend(ctx.recipe.args.clone());
|
|
|
|
|
|
|
|
- let mut flags: Vec<String> = match recipe.distributed_backend.as_str() {
|
|
|
|
|
|
|
+ let mut flags: Vec<String> = match ctx.recipe.distributed_backend.as_str() {
|
|
|
"ray" | "mp" if is_sglang => vec![
|
|
"ray" | "mp" if is_sglang => vec![
|
|
|
"--tp".into(),
|
|
"--tp".into(),
|
|
|
- tp_size.to_string(),
|
|
|
|
|
|
|
+ ctx.tp_size.to_string(),
|
|
|
"--nnodes".into(),
|
|
"--nnodes".into(),
|
|
|
- nnodes.to_string(),
|
|
|
|
|
|
|
+ ctx.nnodes.to_string(),
|
|
|
"--node-rank".into(),
|
|
"--node-rank".into(),
|
|
|
node.rank.to_string(),
|
|
node.rank.to_string(),
|
|
|
"--dist-init-addr".into(),
|
|
"--dist-init-addr".into(),
|
|
|
- format!("{}:{}", cluster.master_ib_ip, init_port),
|
|
|
|
|
|
|
+ format!("{}:{}", ctx.cluster.master_ib_ip, ctx.init_port),
|
|
|
],
|
|
],
|
|
|
"ray" | "mp" => {
|
|
"ray" | "mp" => {
|
|
|
let mut fp = vec![
|
|
let mut fp = vec![
|
|
|
"--tensor-parallel-size".into(),
|
|
"--tensor-parallel-size".into(),
|
|
|
- tp_size.to_string(),
|
|
|
|
|
|
|
+ ctx.tp_size.to_string(),
|
|
|
"--pipeline-parallel-size".into(),
|
|
"--pipeline-parallel-size".into(),
|
|
|
- pp_size.to_string(),
|
|
|
|
|
|
|
+ ctx.pp_size.to_string(),
|
|
|
"--nnodes".into(),
|
|
"--nnodes".into(),
|
|
|
- nnodes.to_string(),
|
|
|
|
|
|
|
+ ctx.nnodes.to_string(),
|
|
|
"--node-rank".into(),
|
|
"--node-rank".into(),
|
|
|
node.rank.to_string(),
|
|
node.rank.to_string(),
|
|
|
"--master-addr".into(),
|
|
"--master-addr".into(),
|
|
|
- cluster.master_ib_ip.clone(),
|
|
|
|
|
|
|
+ ctx.cluster.master_ib_ip.clone(),
|
|
|
"--master-port".into(),
|
|
"--master-port".into(),
|
|
|
- init_port.to_string(),
|
|
|
|
|
|
|
+ ctx.init_port.to_string(),
|
|
|
];
|
|
];
|
|
|
- if recipe.distributed_backend == "ray" {
|
|
|
|
|
|
|
+ if ctx.recipe.distributed_backend == "ray" {
|
|
|
fp.push("--distributed-executor-backend".into());
|
|
fp.push("--distributed-executor-backend".into());
|
|
|
- fp.push(recipe.distributed_backend.clone());
|
|
|
|
|
|
|
+ fp.push(ctx.recipe.distributed_backend.clone());
|
|
|
}
|
|
}
|
|
|
fp
|
|
fp
|
|
|
}
|
|
}
|
|
|
_ => vec![],
|
|
_ => vec![],
|
|
|
};
|
|
};
|
|
|
- if !is_sglang && node.rank > 0 && matches!(recipe.distributed_backend.as_str(), "mp" | "ray") {
|
|
|
|
|
|
|
+ if !is_sglang
|
|
|
|
|
+ && node.rank > 0
|
|
|
|
|
+ && matches!(ctx.recipe.distributed_backend.as_str(), "mp" | "ray")
|
|
|
|
|
+ {
|
|
|
flags.push("--headless".into());
|
|
flags.push("--headless".into());
|
|
|
}
|
|
}
|
|
|
cmd.extend(flags);
|
|
cmd.extend(flags);
|
|
|
cmd.push("--host".to_string());
|
|
cmd.push("--host".to_string());
|
|
|
cmd.push("0.0.0.0".to_string());
|
|
cmd.push("0.0.0.0".to_string());
|
|
|
cmd.push("--port".to_string());
|
|
cmd.push("--port".to_string());
|
|
|
- cmd.push(recipe.port.to_string());
|
|
|
|
|
|
|
+ cmd.push(ctx.recipe.port.to_string());
|
|
|
cmd
|
|
cmd
|
|
|
}
|
|
}
|
|
|
|
|
|