|
@@ -148,6 +148,7 @@ struct Config {
|
|
|
ssh_user: String,
|
|
ssh_user: String,
|
|
|
ssh_key_path: String,
|
|
ssh_key_path: String,
|
|
|
clusters: Vec<Cluster>,
|
|
clusters: Vec<Cluster>,
|
|
|
|
|
+ #[serde(default)]
|
|
|
nccl: HashMap<String, NcclConfig>,
|
|
nccl: HashMap<String, NcclConfig>,
|
|
|
}
|
|
}
|
|
|
|
|
|
|
@@ -333,11 +334,11 @@ async fn show_logs(config: &Config, cluster: &Cluster, recipe_name: &str, follow
|
|
|
|
|
|
|
|
let results = futures::future::join_all(handles).await;
|
|
let results = futures::future::join_all(handles).await;
|
|
|
for result in results {
|
|
for result in results {
|
|
|
- if let Ok((name, Ok(output))) = result {
|
|
|
|
|
- if !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);
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
@@ -562,6 +563,7 @@ async fn start_cluster(
|
|
|
Ok(())
|
|
Ok(())
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
|
|
+#[allow(clippy::too_many_arguments)]
|
|
|
fn build_docker_command(
|
|
fn build_docker_command(
|
|
|
cluster: &Cluster,
|
|
cluster: &Cluster,
|
|
|
nccl: &NcclConfig,
|
|
nccl: &NcclConfig,
|
|
@@ -615,11 +617,17 @@ fn build_docker_command(
|
|
|
"ray" => {
|
|
"ray" => {
|
|
|
cmd.extend(vec![
|
|
cmd.extend(vec![
|
|
|
"--tensor-parallel-size".to_string(),
|
|
"--tensor-parallel-size".to_string(),
|
|
|
|
|
+ "1".to_string(),
|
|
|
|
|
+ "--pipeline-parallel-size".to_string(),
|
|
|
|
|
+ pp_size.to_string(),
|
|
|
|
|
+ "--nnodes".to_string(),
|
|
|
nnodes.to_string(),
|
|
nnodes.to_string(),
|
|
|
"--node-rank".to_string(),
|
|
"--node-rank".to_string(),
|
|
|
node.rank.to_string(),
|
|
node.rank.to_string(),
|
|
|
"--master-addr".to_string(),
|
|
"--master-addr".to_string(),
|
|
|
cluster.master_ib_ip.clone(),
|
|
cluster.master_ib_ip.clone(),
|
|
|
|
|
+ "--distributed-executor-backend".to_string(),
|
|
|
|
|
+ "ray".to_string(),
|
|
|
]);
|
|
]);
|
|
|
}
|
|
}
|
|
|
"mp" => {
|
|
"mp" => {
|
|
@@ -709,16 +717,13 @@ async fn doctor(
|
|
|
let backend = recipe.distributed_backend.clone();
|
|
let backend = recipe.distributed_backend.clone();
|
|
|
let image = recipe.docker_image.clone();
|
|
let image = recipe.docker_image.clone();
|
|
|
handles.push(tokio::spawn(async move {
|
|
handles.push(tokio::spawn(async move {
|
|
|
- let result = check_node(&config, &node, &engine, &backend, &image).await;
|
|
|
|
|
- result
|
|
|
|
|
|
|
+ check_node(&config, &node, &engine, &backend, &image).await
|
|
|
}));
|
|
}));
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
let results = futures::future::join_all(handles).await;
|
|
let results = futures::future::join_all(handles).await;
|
|
|
- for result in results {
|
|
|
|
|
- if let Ok(r) = result {
|
|
|
|
|
- all_results.push(r);
|
|
|
|
|
- }
|
|
|
|
|
|
|
+ for r in results.into_iter().flatten() {
|
|
|
|
|
+ all_results.push(r);
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
// Print results
|
|
// Print results
|
|
@@ -735,7 +740,7 @@ async fn doctor(
|
|
|
|
|
|
|
|
if has_failures {
|
|
if has_failures {
|
|
|
if repair {
|
|
if repair {
|
|
|
- println!("\n🔧 Reparing image...");
|
|
|
|
|
|
|
+ println!("\n🔧 Repairing image...");
|
|
|
let repaired = run_cluster_repair(config, &cluster, &recipe, dry_run).await?;
|
|
let repaired = run_cluster_repair(config, &cluster, &recipe, dry_run).await?;
|
|
|
println!("✅ Repaired image: {}", repaired);
|
|
println!("✅ Repaired image: {}", repaired);
|
|
|
} else {
|
|
} else {
|
|
@@ -783,7 +788,7 @@ async fn check_node(
|
|
|
/// Run a command on the node via SSH and return a check result
|
|
/// Run a command on the node via SSH and return a check result
|
|
|
async fn node_check(config: &Config, node: &Node, name: &str, command: &str) -> DoctorCheck {
|
|
async fn node_check(config: &Config, node: &Node, name: &str, command: &str) -> DoctorCheck {
|
|
|
match run_ssh(config, node, command.to_string(), false).await {
|
|
match run_ssh(config, node, command.to_string(), false).await {
|
|
|
- Ok(output) if output.trim().len() > 0 => DoctorCheck {
|
|
|
|
|
|
|
+ Ok(output) if !output.trim().is_empty() => DoctorCheck {
|
|
|
name: name.to_string(),
|
|
name: name.to_string(),
|
|
|
status: DoctorStatus::Ok,
|
|
status: DoctorStatus::Ok,
|
|
|
detail: format!("OK: {}", output.trim().lines().next().unwrap_or(output.trim())),
|
|
detail: format!("OK: {}", output.trim().lines().next().unwrap_or(output.trim())),
|
|
@@ -948,3 +953,422 @@ async fn try_repaired_image(config: &Config, cluster: &Cluster, recipe: &Recipe)
|
|
|
_ => Ok(recipe.docker_image.clone()),
|
|
_ => Ok(recipe.docker_image.clone()),
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
|
|
|
+
|
|
|
|
|
+#[cfg(test)]
|
|
|
|
|
+mod tests {
|
|
|
|
|
+ use super::*;
|
|
|
|
|
+ use std::collections::HashMap;
|
|
|
|
|
+
|
|
|
|
|
+ fn test_cluster() -> Cluster {
|
|
|
|
|
+ Cluster {
|
|
|
|
|
+ name: "test-cluster".to_string(),
|
|
|
|
|
+ master_ib_ip: "192.168.1.1".to_string(),
|
|
|
|
|
+ nodes: vec![
|
|
|
|
|
+ Node {
|
|
|
|
|
+ name: "node0".to_string(),
|
|
|
|
|
+ ssh_ip: "10.0.0.1".to_string(),
|
|
|
|
|
+ ib_ip: "192.168.1.1".to_string(),
|
|
|
|
|
+ rank: 0,
|
|
|
|
|
+ host_key: None,
|
|
|
|
|
+ },
|
|
|
|
|
+ Node {
|
|
|
|
|
+ name: "node1".to_string(),
|
|
|
|
|
+ ssh_ip: "10.0.0.2".to_string(),
|
|
|
|
|
+ ib_ip: "192.168.1.2".to_string(),
|
|
|
|
|
+ rank: 1,
|
|
|
|
|
+ host_key: None,
|
|
|
|
|
+ },
|
|
|
|
|
+ ],
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ fn test_nccl() -> NcclConfig {
|
|
|
|
|
+ NcclConfig {
|
|
|
|
|
+ ib_hca: "rocep1s0f0".to_string(),
|
|
|
|
|
+ socket_ifname: "enp1s0f0np0".to_string(),
|
|
|
|
|
+ ib_gid_index: 3,
|
|
|
|
|
+ ib_disable: 0,
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ fn test_recipe() -> Recipe {
|
|
|
|
|
+ Recipe {
|
|
|
|
|
+ engine: "vllm".to_string(),
|
|
|
|
|
+ name: "Test Model".to_string(),
|
|
|
|
|
+ docker_image: "test/image:latest".to_string(),
|
|
|
|
|
+ model: "test/model".to_string(),
|
|
|
|
|
+ port: 8000,
|
|
|
|
|
+ required_nodes: Some(2),
|
|
|
|
|
+ pipeline_parallel_size: Some(2),
|
|
|
|
|
+ args: vec!["--max-model-len".to_string(), "4096".to_string()],
|
|
|
|
|
+ env: HashMap::from([("NCCL_DEBUG".to_string(), "INFO".to_string())]),
|
|
|
|
|
+ cluster: None,
|
|
|
|
|
+ entrypoint: vec![],
|
|
|
|
|
+ distributed_backend: "mp".to_string(),
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // ── build_docker_command tests ──────────────────────────────
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_build_docker_command_mp_backend() {
|
|
|
|
|
+ let cluster = test_cluster();
|
|
|
|
|
+ let nccl = test_nccl();
|
|
|
|
|
+ let recipe = test_recipe();
|
|
|
|
|
+ let node = &cluster.nodes[0];
|
|
|
|
|
+
|
|
|
|
|
+ let cmd = build_docker_command(&cluster, &nccl, &recipe, "test-recipe", node, true, 2, 2);
|
|
|
|
|
+
|
|
|
|
|
+ let joined = cmd.join(" ");
|
|
|
|
|
+
|
|
|
|
|
+ // Basic docker flags
|
|
|
|
|
+ assert!(joined.contains("docker run -d --gpus all --ipc=host --privileged --network=host"));
|
|
|
|
|
+ // NCCL env
|
|
|
|
|
+ assert!(joined.contains("-e NCCL_IB_HCA=rocep1s0f0"));
|
|
|
|
|
+ assert!(joined.contains("-e NCCL_SOCKET_IFNAME=enp1s0f0np0"));
|
|
|
|
|
+ // Recipe env
|
|
|
|
|
+ assert!(joined.contains("-e NCCL_DEBUG=INFO"));
|
|
|
|
|
+ // Port only on head
|
|
|
|
|
+ assert!(joined.contains("-p 8000:8000"));
|
|
|
|
|
+ // Container name & labels
|
|
|
|
|
+ assert!(joined.contains("--name sm121_test-recipe_0"));
|
|
|
|
|
+ assert!(joined.contains("--label sm121.recipe=test-recipe"));
|
|
|
|
|
+ assert!(joined.contains("--label sm121.rank=0"));
|
|
|
|
|
+ // Image & model
|
|
|
|
|
+ assert!(joined.contains("test/image:latest"));
|
|
|
|
|
+ assert!(joined.contains("test/model"));
|
|
|
|
|
+ // Args
|
|
|
|
|
+ assert!(joined.contains("--max-model-len 4096"));
|
|
|
|
|
+ // MP-specific flags
|
|
|
|
|
+ assert!(joined.contains("--tensor-parallel-size 1"));
|
|
|
|
|
+ assert!(joined.contains("--pipeline-parallel-size 2"));
|
|
|
|
|
+ assert!(joined.contains("--nnodes 2"));
|
|
|
|
|
+ assert!(joined.contains("--node-rank 0"));
|
|
|
|
|
+ assert!(joined.contains("--master-addr 192.168.1.1"));
|
|
|
|
|
+ assert!(joined.contains("--distributed-executor-backend mp"));
|
|
|
|
|
+ // No Ray env
|
|
|
|
|
+ assert!(!joined.contains("RAY_ADDRESS"));
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_build_docker_command_worker_no_port() {
|
|
|
|
|
+ let cluster = test_cluster();
|
|
|
|
|
+ let nccl = test_nccl();
|
|
|
|
|
+ let recipe = test_recipe();
|
|
|
|
|
+ let node = &cluster.nodes[1]; // rank 1
|
|
|
|
|
+
|
|
|
|
|
+ let cmd = build_docker_command(&cluster, &nccl, &recipe, "test-recipe", node, false, 2, 2);
|
|
|
|
|
+ let joined = cmd.join(" ");
|
|
|
|
|
+
|
|
|
|
|
+ // No port mapping on worker
|
|
|
|
|
+ assert!(!joined.contains("-p 8000:8000"));
|
|
|
|
|
+ // Correct rank
|
|
|
|
|
+ assert!(joined.contains("--node-rank 1"));
|
|
|
|
|
+ assert!(joined.contains("--label sm121.rank=1"));
|
|
|
|
|
+ assert!(joined.contains("--name sm121_test-recipe_1"));
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_build_docker_command_ray_backend() {
|
|
|
|
|
+ let cluster = test_cluster();
|
|
|
|
|
+ let nccl = test_nccl();
|
|
|
|
|
+ let mut recipe = test_recipe();
|
|
|
|
|
+ recipe.distributed_backend = "ray".to_string();
|
|
|
|
|
+ let node = &cluster.nodes[0];
|
|
|
|
|
+
|
|
|
|
|
+ let cmd = build_docker_command(&cluster, &nccl, &recipe, "test-recipe", node, true, 2, 2);
|
|
|
|
|
+ let joined = cmd.join(" ");
|
|
|
|
|
+
|
|
|
|
|
+ // Ray-specific env
|
|
|
|
|
+ assert!(joined.contains("-e RAY_ADDRESS=192.168.1.1:6379"));
|
|
|
|
|
+ assert!(joined.contains("-e RAY_PORT=8000"));
|
|
|
|
|
+ // Ray-specific flags
|
|
|
|
|
+ assert!(joined.contains("--tensor-parallel-size 1"));
|
|
|
|
|
+ assert!(joined.contains("--pipeline-parallel-size 2"));
|
|
|
|
|
+ assert!(joined.contains("--nnodes 2"));
|
|
|
|
|
+ assert!(joined.contains("--node-rank 0"));
|
|
|
|
|
+ assert!(joined.contains("--master-addr 192.168.1.1"));
|
|
|
|
|
+ assert!(joined.contains("--distributed-executor-backend ray"));
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_build_docker_command_auto_pp_size() {
|
|
|
|
|
+ let cluster = test_cluster();
|
|
|
|
|
+ let nccl = test_nccl();
|
|
|
|
|
+ let mut recipe = test_recipe();
|
|
|
|
|
+ recipe.pipeline_parallel_size = Some(0); // auto
|
|
|
|
|
+ let node = &cluster.nodes[0];
|
|
|
|
|
+
|
|
|
|
|
+ // pp_size = 0 should become nnodes (2) in start_cluster, but build_docker_command
|
|
|
|
|
+ // takes pp_size directly. test with pp_size = nnodes = 2.
|
|
|
|
|
+ let cmd = build_docker_command(&cluster, &nccl, &recipe, "test-recipe", node, true, 2, 2);
|
|
|
|
|
+ let joined = cmd.join(" ");
|
|
|
|
|
+
|
|
|
|
|
+ assert!(joined.contains("--pipeline-parallel-size 2"));
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_build_docker_command_entrypoint() {
|
|
|
|
|
+ let cluster = test_cluster();
|
|
|
|
|
+ let nccl = test_nccl();
|
|
|
|
|
+ let mut recipe = test_recipe();
|
|
|
|
|
+ recipe.entrypoint = vec!["python".to_string(), "-m".to_string(), "custom.module".to_string()];
|
|
|
|
|
+ let node = &cluster.nodes[0];
|
|
|
|
|
+
|
|
|
|
|
+ let cmd = build_docker_command(&cluster, &nccl, &recipe, "test-recipe", node, true, 2, 2);
|
|
|
|
|
+ let joined = cmd.join(" ");
|
|
|
|
|
+
|
|
|
|
|
+ // Entrypoint should appear after image name, before model
|
|
|
|
|
+ let image_pos = joined.find("test/image:latest").unwrap();
|
|
|
|
|
+ let model_pos = joined.find("test/model").unwrap();
|
|
|
|
|
+ let entrypoint_pos = joined.find("python -m custom.module").unwrap();
|
|
|
|
|
+ assert!(image_pos < entrypoint_pos && entrypoint_pos < model_pos,
|
|
|
|
|
+ "entrypoint should be between image and model");
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // ── select_cluster tests ────────────────────────────────────
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_select_cluster_by_name() {
|
|
|
|
|
+ let mut config = Config {
|
|
|
|
|
+ ssh_user: "user".to_string(),
|
|
|
|
|
+ ssh_key_path: "/path/to/key".to_string(),
|
|
|
|
|
+ clusters: vec![test_cluster()],
|
|
|
|
|
+ nccl: HashMap::new(),
|
|
|
|
|
+ };
|
|
|
|
|
+ config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
|
|
|
+
|
|
|
|
|
+ let (cluster, nccl) = select_cluster(&config, "test-cluster").unwrap();
|
|
|
|
|
+ assert_eq!(cluster.name, "test-cluster");
|
|
|
|
|
+ assert_eq!(nccl.ib_hca, "rocep1s0f0");
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_select_cluster_auto() {
|
|
|
|
|
+ let mut config = Config {
|
|
|
|
|
+ ssh_user: "user".to_string(),
|
|
|
|
|
+ ssh_key_path: "/path/to/key".to_string(),
|
|
|
|
|
+ clusters: vec![test_cluster()],
|
|
|
|
|
+ nccl: HashMap::new(),
|
|
|
|
|
+ };
|
|
|
|
|
+ config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
|
|
|
+
|
|
|
|
|
+ let (cluster, nccl) = select_cluster(&config, "auto").unwrap();
|
|
|
|
|
+ assert_eq!(cluster.name, "test-cluster");
|
|
|
|
|
+ assert_eq!(nccl.ib_hca, "rocep1s0f0");
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_select_cluster_auto_no_match() {
|
|
|
|
|
+ let config = Config {
|
|
|
|
|
+ ssh_user: "user".to_string(),
|
|
|
|
|
+ ssh_key_path: "/path/to/key".to_string(),
|
|
|
|
|
+ clusters: vec![test_cluster()],
|
|
|
|
|
+ nccl: HashMap::new(),
|
|
|
|
|
+ };
|
|
|
|
|
+ // No NCCL config for test-cluster → auto should fail
|
|
|
|
|
+ let result = select_cluster(&config, "auto");
|
|
|
|
|
+ assert!(result.is_err());
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_select_cluster_not_found() {
|
|
|
|
|
+ let config = Config {
|
|
|
|
|
+ ssh_user: "user".to_string(),
|
|
|
|
|
+ ssh_key_path: "/path/to/key".to_string(),
|
|
|
|
|
+ clusters: vec![],
|
|
|
|
|
+ nccl: HashMap::new(),
|
|
|
|
|
+ };
|
|
|
|
|
+ let result = select_cluster(&config, "nonexistent");
|
|
|
|
|
+ assert!(result.is_err());
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // ── Recipe deserialization tests ─────────────────────────────
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_recipe_deserialize_minimal() {
|
|
|
|
|
+ let toml_str = r#"
|
|
|
|
|
+name = "Minimal"
|
|
|
|
|
+model = "test/model"
|
|
|
|
|
+"#;
|
|
|
|
|
+ let recipe: Recipe = toml::from_str(toml_str).unwrap();
|
|
|
|
|
+ assert_eq!(recipe.name, "Minimal");
|
|
|
|
|
+ assert_eq!(recipe.model, "test/model");
|
|
|
|
|
+ assert_eq!(recipe.engine, "vllm"); // default
|
|
|
|
|
+ assert_eq!(recipe.port, 0); // default
|
|
|
|
|
+ assert!(recipe.args.is_empty());
|
|
|
|
|
+ assert!(recipe.env.is_empty());
|
|
|
|
|
+ assert_eq!(recipe.distributed_backend, "mp"); // default
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_recipe_deserialize_full() {
|
|
|
|
|
+ let toml_str = r#"
|
|
|
|
|
+engine = "sglang"
|
|
|
|
|
+name = "Full Recipe"
|
|
|
|
|
+docker_image = "org/image:tag"
|
|
|
|
|
+model = "org/model"
|
|
|
|
|
+port = 30000
|
|
|
|
|
+required_nodes = 4
|
|
|
|
|
+pipeline_parallel_size = 2
|
|
|
|
|
+distributed_backend = "ray"
|
|
|
|
|
+cluster = "my-cluster"
|
|
|
|
|
+
|
|
|
|
|
+args = ["--arg1", "val1", "--arg2"]
|
|
|
|
|
+
|
|
|
|
|
+[env]
|
|
|
|
|
+KEY = "value"
|
|
|
|
|
+OTHER = "123"
|
|
|
|
|
+"#;
|
|
|
|
|
+ let recipe: Recipe = toml::from_str(toml_str).unwrap();
|
|
|
|
|
+ assert_eq!(recipe.engine, "sglang");
|
|
|
|
|
+ assert_eq!(recipe.name, "Full Recipe");
|
|
|
|
|
+ assert_eq!(recipe.docker_image, "org/image:tag");
|
|
|
|
|
+ assert_eq!(recipe.model, "org/model");
|
|
|
|
|
+ assert_eq!(recipe.port, 30000);
|
|
|
|
|
+ assert_eq!(recipe.required_nodes, Some(4));
|
|
|
|
|
+ assert_eq!(recipe.pipeline_parallel_size, Some(2));
|
|
|
|
|
+ assert_eq!(recipe.distributed_backend, "ray");
|
|
|
|
|
+ assert_eq!(recipe.cluster, Some("my-cluster".to_string()));
|
|
|
|
|
+ assert_eq!(recipe.args, vec!["--arg1", "val1", "--arg2"]);
|
|
|
|
|
+ assert_eq!(recipe.env.get("KEY").unwrap(), "value");
|
|
|
|
|
+ assert_eq!(recipe.env.get("OTHER").unwrap(), "123");
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_recipe_deserialize_no_required_nodes() {
|
|
|
|
|
+ let toml_str = r#"
|
|
|
|
|
+name = "Flexible"
|
|
|
|
|
+model = "test/model"
|
|
|
|
|
+"#;
|
|
|
|
|
+ let recipe: Recipe = toml::from_str(toml_str).unwrap();
|
|
|
|
|
+ assert_eq!(recipe.required_nodes, None); // None, not Some(0)
|
|
|
|
|
+ assert_eq!(recipe.pipeline_parallel_size, None);
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // ── Config deserialization tests ───────────────────────────────
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_config_deserialize() {
|
|
|
|
|
+ let toml_str = r#"
|
|
|
|
|
+ssh_user = "admin"
|
|
|
|
|
+ssh_key_path = "/home/admin/.ssh/id_rsa"
|
|
|
|
|
+
|
|
|
|
|
+[[clusters]]
|
|
|
|
|
+name = "prod"
|
|
|
|
|
+master_ib_ip = "10.0.0.1"
|
|
|
|
|
+nodes = [
|
|
|
|
|
+ { name = "node1", ssh_ip = "10.0.0.10", ib_ip = "10.0.0.1", rank = 0 },
|
|
|
|
|
+ { name = "node2", ssh_ip = "10.0.0.11", ib_ip = "10.0.0.2", rank = 1 },
|
|
|
|
|
+]
|
|
|
|
|
+
|
|
|
|
|
+[nccl."prod"]
|
|
|
|
|
+ib_hca = "mlx5_0"
|
|
|
|
|
+socket_ifname = "eth0"
|
|
|
|
|
+ib_gid_index = 3
|
|
|
|
|
+ib_disable = 0
|
|
|
|
|
+"#;
|
|
|
|
|
+ let config: Config = toml::from_str(toml_str).unwrap();
|
|
|
|
|
+ assert_eq!(config.ssh_user, "admin");
|
|
|
|
|
+ assert_eq!(config.clusters.len(), 1);
|
|
|
|
|
+ assert_eq!(config.clusters[0].nodes.len(), 2);
|
|
|
|
|
+ assert_eq!(config.clusters[0].nodes[0].host_key, None);
|
|
|
|
|
+ assert!(config.nccl.contains_key("prod"));
|
|
|
|
|
+ assert_eq!(config.nccl["prod"].ib_hca, "mlx5_0");
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_config_deserialize_with_host_key() {
|
|
|
|
|
+ let toml_str = r#"
|
|
|
|
|
+ssh_user = "admin"
|
|
|
|
|
+ssh_key_path = "/path/to/key"
|
|
|
|
|
+
|
|
|
|
|
+[[clusters]]
|
|
|
|
|
+name = "secure"
|
|
|
|
|
+master_ib_ip = "1.1.1.1"
|
|
|
|
|
+nodes = [
|
|
|
|
|
+ { name = "n1", ssh_ip = "2.2.2.2", ib_ip = "1.1.1.1", rank = 0, host_key = "SHA256:abc123" },
|
|
|
|
|
+]
|
|
|
|
|
+"#;
|
|
|
|
|
+ let config: Config = toml::from_str(toml_str).unwrap();
|
|
|
|
|
+ assert_eq!(
|
|
|
|
|
+ config.clusters[0].nodes[0].host_key,
|
|
|
|
|
+ Some("SHA256:abc123".to_string())
|
|
|
|
|
+ );
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // ── generate_repair_dockerfile tests ───────────────────────────
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_generate_repair_dockerfile_vllm() {
|
|
|
|
|
+ let recipe = test_recipe();
|
|
|
|
|
+ let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);
|
|
|
|
|
+
|
|
|
|
|
+ assert!(dockerfile.contains("FROM base/image:tag"));
|
|
|
|
|
+ assert!(dockerfile.contains("pip install vllm"));
|
|
|
|
|
+ assert!(!dockerfile.contains("ray"));
|
|
|
|
|
+ assert!(dockerfile.contains("ENTRYPOINT"));
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_generate_repair_dockerfile_ray() {
|
|
|
|
|
+ let mut recipe = test_recipe();
|
|
|
|
|
+ recipe.distributed_backend = "ray".to_string();
|
|
|
|
|
+ let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);
|
|
|
|
|
+
|
|
|
|
|
+ assert!(dockerfile.contains("FROM base/image:tag"));
|
|
|
|
|
+ assert!(dockerfile.contains("pip install vllm"));
|
|
|
|
|
+ assert!(dockerfile.contains("pip install ray"));
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_generate_repair_dockerfile_sglang() {
|
|
|
|
|
+ let mut recipe = test_recipe();
|
|
|
|
|
+ recipe.engine = "sglang".to_string();
|
|
|
|
|
+ let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);
|
|
|
|
|
+
|
|
|
|
|
+ assert!(dockerfile.contains("pip install sglang"));
|
|
|
|
|
+ assert!(!dockerfile.contains("vllm"));
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // ── select_cluster_by_name tests ───────────────────────────────
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_select_cluster_by_name_found() {
|
|
|
|
|
+ let mut nccl_map = HashMap::new();
|
|
|
|
|
+ nccl_map.insert("test-cluster".to_string(), test_nccl());
|
|
|
|
|
+ let config = Config {
|
|
|
|
|
+ ssh_user: "u".to_string(),
|
|
|
|
|
+ ssh_key_path: "/k".to_string(),
|
|
|
|
|
+ clusters: vec![test_cluster()],
|
|
|
|
|
+ nccl: nccl_map,
|
|
|
|
|
+ };
|
|
|
|
|
+ let cluster = select_cluster_by_name(&config, "test-cluster").unwrap();
|
|
|
|
|
+ assert_eq!(cluster.name, "test-cluster");
|
|
|
|
|
+ assert_eq!(cluster.nodes.len(), 2);
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_select_cluster_by_name_not_found() {
|
|
|
|
|
+ let config = Config {
|
|
|
|
|
+ ssh_user: "u".to_string(),
|
|
|
|
|
+ ssh_key_path: "/k".to_string(),
|
|
|
|
|
+ clusters: vec![test_cluster()],
|
|
|
|
|
+ nccl: HashMap::new(),
|
|
|
|
|
+ };
|
|
|
|
|
+ let result = select_cluster_by_name(&config, "no-such-cluster");
|
|
|
|
|
+ assert!(result.is_err());
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // ── Defaults ───────────────────────────────────────────────────
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_default_engine() {
|
|
|
|
|
+ assert_eq!(default_engine(), "vllm");
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ #[test]
|
|
|
|
|
+ fn test_default_distributed_backend() {
|
|
|
|
|
+ assert_eq!(default_distributed_backend(), "mp");
|
|
|
|
|
+ }
|
|
|
|
|
+}
|