|
|
@@ -7,21 +7,37 @@ use crate::doctor::try_repaired_image;
|
|
|
use crate::ssh::*;
|
|
|
|
|
|
/// Launches containers across the cluster for a recipe.
|
|
|
-pub async fn serve(config: &Config, recipe_name: &str, cluster_name: &str, dry_run: bool) -> Result<()> {
|
|
|
+pub async fn serve(
|
|
|
+ config: &Config,
|
|
|
+ recipe_name: &str,
|
|
|
+ cluster_name: &str,
|
|
|
+ dry_run: bool,
|
|
|
+) -> Result<()> {
|
|
|
let recipe = load_recipe(recipe_name).await?;
|
|
|
- info!("🚀 Запуск рецепта: {} ({}) [backend: {}]", recipe.name, recipe.engine, recipe.distributed_backend);
|
|
|
+ info!(
|
|
|
+ "🚀 Запуск рецепта: {} ({}) [backend: {}]",
|
|
|
+ recipe.name, recipe.engine, recipe.distributed_backend
|
|
|
+ );
|
|
|
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);
|
|
|
+ anyhow::bail!(
|
|
|
+ "Recipe '{}' requires {} nodes, but cluster '{}' has {} nodes",
|
|
|
+ recipe.name,
|
|
|
+ required,
|
|
|
+ cluster.name,
|
|
|
+ actual
|
|
|
+ );
|
|
|
}
|
|
|
}
|
|
|
|
|
|
let image = try_repaired_image(config, &cluster, &recipe).await?;
|
|
|
let mut recipe = recipe;
|
|
|
- if image != recipe.docker_image { recipe.docker_image = image.clone(); }
|
|
|
+ if image != recipe.docker_image {
|
|
|
+ recipe.docker_image = image.clone();
|
|
|
+ }
|
|
|
|
|
|
start_cluster(config, &cluster, nccl, recipe_name, &recipe, dry_run).await?;
|
|
|
if !dry_run {
|
|
|
@@ -32,25 +48,48 @@ pub async fn serve(config: &Config, recipe_name: &str, cluster_name: &str, dry_r
|
|
|
}
|
|
|
|
|
|
/// 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<()> {
|
|
|
+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!("📜 Following logs for '{}' (head node)...", head_node.name);
|
|
|
- let cmd = format!("docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank=0\" | xargs -r docker logs -f", recipe_name);
|
|
|
+ let cmd = format!(
|
|
|
+ "docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank=0\" | xargs -r docker logs -f",
|
|
|
+ recipe_name
|
|
|
+ );
|
|
|
run_ssh_follow(config, head_node, cmd, dry_run).await?;
|
|
|
} else {
|
|
|
let flag = "--tail=100";
|
|
|
- info!("📜 Логи кластера '{}' ({} узлов):", cluster.name, cluster.nodes.len());
|
|
|
+ info!(
|
|
|
+ "📜 Логи кластера '{}' ({} узлов):",
|
|
|
+ cluster.name,
|
|
|
+ cluster.nodes.len()
|
|
|
+ );
|
|
|
let mut handles = Vec::new();
|
|
|
for node in &cluster.nodes {
|
|
|
- let cmd = format!("docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank={}\" | xargs -r docker logs {} 2>&1", recipe_name, node.rank, flag);
|
|
|
- let config = config.clone(); let node = node.clone(); let dry = dry_run;
|
|
|
+ let cmd = format!(
|
|
|
+ "docker ps -q --filter \"label=sm121.recipe={}\" --filter \"label=sm121.rank={}\" | xargs -r docker logs {} 2>&1",
|
|
|
+ recipe_name, node.rank, flag
|
|
|
+ );
|
|
|
+ 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) }));
|
|
|
+ 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!("\n--- {} ---", name); println!("{}", output);
|
|
|
+ if let Ok((name, Ok(output))) = result
|
|
|
+ && !output.trim().is_empty()
|
|
|
+ {
|
|
|
+ println!("\n--- {} ---", name);
|
|
|
+ println!("{}", output);
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
@@ -58,17 +97,28 @@ pub async fn show_logs(config: &Config, cluster: &Cluster, recipe_name: &str, fo
|
|
|
}
|
|
|
|
|
|
/// 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<()> {
|
|
|
+pub async fn stop_cluster(
|
|
|
+ config: &Config,
|
|
|
+ cluster: &Cluster,
|
|
|
+ recipe_name: &str,
|
|
|
+ dry_run: bool,
|
|
|
+) -> Result<()> {
|
|
|
info!("🛑 Останавливаем кластер '{}'...", 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);
|
|
|
+ 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!("Stop failed on {}: {}", node.name, e);
|
|
|
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);
|
|
|
+ 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!("Remove failed on {}: {}", node.name, e);
|
|
|
failures.push((node.name.clone(), "rm".to_string(), e));
|
|
|
@@ -76,7 +126,9 @@ pub async fn stop_cluster(config: &Config, cluster: &Cluster, recipe_name: &str,
|
|
|
}
|
|
|
|
|
|
if !failures.is_empty() {
|
|
|
- for (name, action, e) in &failures { tracing::warn!("Failed to {} on {}: {}", action, name, e); }
|
|
|
+ for (name, action, e) in &failures {
|
|
|
+ tracing::warn!("Failed to {} on {}: {}", action, name, e);
|
|
|
+ }
|
|
|
anyhow::bail!("Failed to stop containers on {} node(s)", failures.len());
|
|
|
}
|
|
|
info!("✅ Все контейнеры остановлены.");
|
|
|
@@ -88,9 +140,19 @@ pub async fn health_check(config: &Config, cluster: &Cluster, dry_run: bool) ->
|
|
|
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;
|
|
|
+ let config = config.clone();
|
|
|
+ let node = node.clone();
|
|
|
+ let dry = dry_run;
|
|
|
handles.push(tokio::spawn(async move {
|
|
|
- let result = run_ssh_with_retry(&config, &node, "docker ps --filter \"label=sm121\" --format '{{.Names}} {{.Status}}' 2>&1".to_string(), dry, 1).await;
|
|
|
+ let result = run_ssh_with_retry(
|
|
|
+ &config,
|
|
|
+ &node,
|
|
|
+ "docker ps --filter \"label=sm121\" --format '{{.Names}} {{.Status}}' 2>&1"
|
|
|
+ .to_string(),
|
|
|
+ dry,
|
|
|
+ 1,
|
|
|
+ )
|
|
|
+ .await;
|
|
|
(node.name, node.ib_ip, result)
|
|
|
}));
|
|
|
}
|
|
|
@@ -99,41 +161,86 @@ pub async fn health_check(config: &Config, cluster: &Cluster, dry_run: bool) ->
|
|
|
match result {
|
|
|
Ok((name, ib_ip, Ok(check))) => {
|
|
|
print!("{} ({}): ", name, ib_ip);
|
|
|
- if check.is_empty() && dry_run { continue; }
|
|
|
+ if check.is_empty() && dry_run {
|
|
|
+ continue;
|
|
|
+ }
|
|
|
if check.starts_with("docker:") || check.contains("permission denied") {
|
|
|
println!("✗ Docker недоступен");
|
|
|
- if let Some(line) = check.lines().next() { println!(" └─ {}", line); }
|
|
|
+ if let Some(line) = check.lines().next() {
|
|
|
+ println!(" └─ {}", line);
|
|
|
+ }
|
|
|
all_healthy = false;
|
|
|
} else {
|
|
|
let containers: Vec<&str> = check.lines().filter(|l| !l.is_empty()).collect();
|
|
|
- if containers.is_empty() { println!("✗ Нет контейнеров"); all_healthy = false; }
|
|
|
- else { println!("✓ {} контейнер(ов)", containers.len()); for c in containers { println!(" └─ {}", c); } }
|
|
|
+ if containers.is_empty() {
|
|
|
+ println!("✗ Нет контейнеров");
|
|
|
+ all_healthy = false;
|
|
|
+ } else {
|
|
|
+ println!("✓ {} контейнер(ов)", containers.len());
|
|
|
+ for c in containers {
|
|
|
+ println!(" └─ {}", c);
|
|
|
+ }
|
|
|
+ }
|
|
|
}
|
|
|
}
|
|
|
- Ok((name, ib_ip, Err(e))) => { println!("{} ({}): ✗ SSH ошибка: {}", name, ib_ip, e); all_healthy = false; }
|
|
|
- Err(join_err) => { println!("✗ Task join error: {}", join_err); all_healthy = false; }
|
|
|
+ Ok((name, ib_ip, Err(e))) => {
|
|
|
+ println!("{} ({}): ✗ SSH ошибка: {}", name, ib_ip, e);
|
|
|
+ all_healthy = false;
|
|
|
+ }
|
|
|
+ Err(join_err) => {
|
|
|
+ println!("✗ Task join error: {}", join_err);
|
|
|
+ all_healthy = false;
|
|
|
+ }
|
|
|
}
|
|
|
}
|
|
|
- if !dry_run && all_healthy { info!("✅ Кластер здоров!"); } else if !dry_run { info!("⚠️ Кластер имеет проблемы"); }
|
|
|
+ if !dry_run && all_healthy {
|
|
|
+ info!("✅ Кластер здоров!");
|
|
|
+ } else if !dry_run {
|
|
|
+ info!("⚠️ Кластер имеет проблемы");
|
|
|
+ }
|
|
|
Ok(())
|
|
|
}
|
|
|
|
|
|
/// Starts containers on all nodes in parallel.
|
|
|
-async fn start_cluster(config: &Config, cluster: &Cluster, nccl: &NcclConfig, recipe_name: &str, recipe: &Recipe, dry_run: bool) -> Result<()> {
|
|
|
+async fn start_cluster(
|
|
|
+ config: &Config,
|
|
|
+ cluster: &Cluster,
|
|
|
+ nccl: &NcclConfig,
|
|
|
+ recipe_name: &str,
|
|
|
+ recipe: &Recipe,
|
|
|
+ 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 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 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 cleanup = format!(
|
|
|
+ "docker rm -f sm121_{}_{} 2>/dev/null; true",
|
|
|
+ recipe_name, node.rank
|
|
|
+ );
|
|
|
|
|
|
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).join(" ");
|
|
|
+ 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,
|
|
|
+ )
|
|
|
+ .join(" ");
|
|
|
let result = run_ssh_with_retry(&config, &node, cmd, dry, 2).await;
|
|
|
(node.name, node.rank, engine, result)
|
|
|
}));
|
|
|
@@ -141,26 +248,67 @@ async fn start_cluster(config: &Config, cluster: &Cluster, nccl: &NcclConfig, re
|
|
|
|
|
|
for result in futures::future::join_all(handles).await {
|
|
|
match result {
|
|
|
- Ok((name, rank, engine, Ok(_))) => info!("✅ {} запущен на {} (rank {})", engine.to_uppercase(), name, rank),
|
|
|
- Ok((name, rank, engine, Err(e))) => anyhow::bail!("Failed to start {} on {} (rank {}): {}", engine.to_uppercase(), name, rank, e),
|
|
|
+ Ok((name, rank, engine, Ok(_))) => info!(
|
|
|
+ "✅ {} запущен на {} (rank {})",
|
|
|
+ engine.to_uppercase(),
|
|
|
+ name,
|
|
|
+ rank
|
|
|
+ ),
|
|
|
+ Ok((name, rank, engine, Err(e))) => anyhow::bail!(
|
|
|
+ "Failed to start {} on {} (rank {}): {}",
|
|
|
+ engine.to_uppercase(),
|
|
|
+ name,
|
|
|
+ rank,
|
|
|
+ e
|
|
|
+ ),
|
|
|
Err(e) => anyhow::bail!("Task join error: {}", e),
|
|
|
}
|
|
|
}
|
|
|
|
|
|
- if !dry_run { sleep(Duration::from_secs(10)).await; info!("🎉 Кластер успешно запущен! API → http://{}:{}", cluster.master_ib_ip, recipe.port); }
|
|
|
+ if !dry_run {
|
|
|
+ sleep(Duration::from_secs(10)).await;
|
|
|
+ info!(
|
|
|
+ "🎉 Кластер успешно запущен! API → http://{}:{}",
|
|
|
+ cluster.master_ib_ip, 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) -> Vec<String> {
|
|
|
+fn build_docker_command(
|
|
|
+ cluster: &Cluster,
|
|
|
+ nccl: &NcclConfig,
|
|
|
+ recipe: &Recipe,
|
|
|
+ recipe_name: &str,
|
|
|
+ node: &Node,
|
|
|
+ is_head: bool,
|
|
|
+ pp_size: u32,
|
|
|
+ nnodes: u32,
|
|
|
+) -> Vec<String> {
|
|
|
+ let is_sglang = recipe.engine == "sglang";
|
|
|
+
|
|
|
let mut cmd = vec![
|
|
|
"docker run -d --gpus all --ipc=host --privileged --network=host".to_string(),
|
|
|
- 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 NCCL_IB_GID_INDEX={}", nccl.ib_gid_index),
|
|
|
- format!("-e VLLM_HOST_IP={}", node.ib_ip),
|
|
|
+ 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 NCCL_IB_GID_INDEX={}", nccl.ib_gid_index),
|
|
|
];
|
|
|
- for (k, v) in &recipe.env { cmd.push(format!("-e {}={}", k, v)); }
|
|
|
- if is_head { cmd.push(format!("-p {}:{}", recipe.port, recipe.port)); }
|
|
|
+
|
|
|
+ // 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));
|
|
|
+ }
|
|
|
+
|
|
|
+ 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));
|
|
|
@@ -169,22 +317,56 @@ fn build_docker_command(cluster: &Cluster, nccl: &NcclConfig, recipe: &Recipe, r
|
|
|
cmd.push(format!("-e RAY_PORT={}", recipe.port));
|
|
|
}
|
|
|
cmd.push(recipe.docker_image.clone());
|
|
|
- if !recipe.entrypoint.is_empty() { cmd.extend(recipe.entrypoint.clone()); }
|
|
|
- cmd.push(recipe.model.clone());
|
|
|
+ if !recipe.entrypoint.is_empty() {
|
|
|
+ cmd.extend(recipe.entrypoint.clone());
|
|
|
+ }
|
|
|
+
|
|
|
+ // 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 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(),
|
|
|
+ "--master-addr".into(),
|
|
|
+ cluster.master_ib_ip.clone(),
|
|
|
+ ],
|
|
|
"ray" | "mp" => vec![
|
|
|
- "--tensor-parallel-size".into(), "1".into(),
|
|
|
- "--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(),
|
|
|
- "--distributed-executor-backend".into(), recipe.distributed_backend.clone(),
|
|
|
+ "--tensor-parallel-size".into(),
|
|
|
+ "1".into(),
|
|
|
+ "--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(),
|
|
|
+ "--distributed-executor-backend".into(),
|
|
|
+ recipe.distributed_backend.clone(),
|
|
|
],
|
|
|
_ => vec![],
|
|
|
};
|
|
|
cmd.extend(flags);
|
|
|
+
|
|
|
+ // Port flag — auto-generated for SGLang (all nodes need to bind)
|
|
|
+ if is_sglang {
|
|
|
+ cmd.push("--port".to_string());
|
|
|
+ cmd.push(recipe.port.to_string());
|
|
|
+ }
|
|
|
+
|
|
|
cmd
|
|
|
}
|
|
|
|
|
|
@@ -192,43 +374,181 @@ fn build_docker_command(cluster: &Cluster, nccl: &NcclConfig, recipe: &Recipe, r
|
|
|
mod tests {
|
|
|
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),
|
|
|
- args: vec!["--x".into(), "4096".into()], env: HashMap::from([("K".into(), "V".into())]),
|
|
|
- cluster: None, entrypoint: vec![], distributed_backend: "mp".into(),
|
|
|
- } }
|
|
|
+ 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(),
|
|
|
+ }
|
|
|
+ }
|
|
|
fn run(recipe: &Recipe, node: &Node, head: bool) -> String {
|
|
|
build_docker_command(&tcls(), &tnccl(), recipe, "r", node, head, 2, 2).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("--distributed-executor-backend mp")); assert!(!cmd.contains("RAY_ADDRESS"));
|
|
|
+ #[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("--distributed-executor-backend mp"));
|
|
|
+ assert!(!cmd.contains("RAY_ADDRESS"));
|
|
|
}
|
|
|
- #[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"));
|
|
|
+ #[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"));
|
|
|
}
|
|
|
- #[test] fn ray_backend() { let mut r = trec(); r.distributed_backend = "ray".into();
|
|
|
+ #[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"));
|
|
|
+ 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 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()];
|
|
|
+ #[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");
|
|
|
+ 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(),
|
|
|
+ }
|
|
|
+ }
|
|
|
+ fn srun(recipe: &Recipe, node: &Node, head: bool) -> String {
|
|
|
+ build_docker_command(&tcls(), &tnccl(), recipe, "sg", node, head, 4, 4).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("--master-addr 1"));
|
|
|
+ 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"));
|
|
|
+ }
|
|
|
+
|
|
|
+ #[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");
|
|
|
}
|
|
|
}
|
