Эх сурвалжийг харах

feat: cluster defaults, configurable init_port, --headless for workers

Three improvements for vLLM/SGLang clusters, learned from sparkrun:

1. Default Docker flags now include --shm-size 10g (was unset, causing
   silent /dev/shm shortages for PyTorch DataLoader).
2. Cluster-wide env defaults NCCL_CUMEM_ENABLE=0 and OMP_NUM_THREADS=4
   injected before recipe env so user-supplied values override.
3. vLLM workers (rank > 0, both mp and ray backends) now pass --headless
   so they don't each try to bind the API server.

Other plumbing:
- Config.init_port (default 25000, configurable via config.toml) feeds
  vLLM --master-port and SGLang --dist-init-addr port.
- vLLM cluster now emits --master-port alongside --master-addr, matching
  the vLLM CLI.

Tests:
- 'mp' test now covers shm-size, ulimits, default cluster env, headless
  absence on rank 0.
- 'worker_no_port' asserts --headless for vLLM workers.
- New 'recipe_env_overrides_cluster_defaults' verifies recipe env wins.
- New 'custom_init_port' verifies the port propagates to --master-port.
Alexander Korolev 4 сар өмнө
parent
commit
2b0be29648
3 өөрчлөгдсөн 104 нэмэгдсэн , 3 устгасан
  1. 14 0
      src/config.rs
  2. 26 3
      src/docker/mod.rs
  3. 64 0
      src/docker/tests.rs

+ 14 - 0
src/config.rs

@@ -37,6 +37,14 @@ pub struct Config {
     pub nccl: HashMap<String, NcclConfig>,
     #[serde(default)]
     pub hf_cache_dir: Option<String>,
+    /// TCP port used for cluster init (vLLM `--master-port`,
+    /// SGLang `--dist-init-addr` port). Default: 25000.
+    #[serde(default = "default_init_port")]
+    pub init_port: u16,
+}
+
+fn default_init_port() -> u16 {
+    25000
 }
 
 impl Config {
@@ -331,6 +339,7 @@ nodes = [
             clusters: vec![test_cluster()],
             nccl: HashMap::new(),
             hf_cache_dir: None,
+            init_port: 25000,
         };
         config.nccl.insert("test-cluster".to_string(), test_nccl());
 
@@ -347,6 +356,7 @@ nodes = [
             clusters: vec![test_cluster()],
             nccl: HashMap::new(),
             hf_cache_dir: None,
+            init_port: 25000,
         };
         config.nccl.insert("test-cluster".to_string(), test_nccl());
 
@@ -363,6 +373,7 @@ nodes = [
             clusters: vec![test_cluster()],
             nccl: HashMap::new(),
             hf_cache_dir: None,
+            init_port: 25000,
         };
         let result = select_cluster(&config, "auto");
         assert!(result.is_err());
@@ -376,6 +387,7 @@ nodes = [
             clusters: vec![],
             nccl: HashMap::new(),
             hf_cache_dir: None,
+            init_port: 25000,
         };
         let result = select_cluster(&config, "nonexistent");
         assert!(result.is_err());
@@ -391,6 +403,7 @@ nodes = [
             clusters: vec![test_cluster()],
             nccl: nccl_map,
             hf_cache_dir: None,
+            init_port: 25000,
         };
         let cluster = select_cluster_by_name(&config, "test-cluster").unwrap();
         assert_eq!(cluster.name, "test-cluster");
@@ -405,6 +418,7 @@ nodes = [
             clusters: vec![test_cluster()],
             nccl: HashMap::new(),
             hf_cache_dir: None,
+            init_port: 25000,
         };
         let result = select_cluster_by_name(&config, "no-such-cluster");
         assert!(result.is_err());

+ 26 - 3
src/docker/mod.rs

@@ -262,7 +262,17 @@ async fn start_cluster(
                 let _ = run_ssh(&config, &node, cleanup, dry).await;
             }
             let cmd = build_docker_command(
-                &cluster, &nccl, &recipe, &rn, &node, is_head, pp_size, nnodes, &cd, keep,
+                &cluster,
+                &nccl,
+                &recipe,
+                &rn,
+                &node,
+                is_head,
+                pp_size,
+                nnodes,
+                &cd,
+                keep,
+                config.init_port,
             )
             .join(" ");
             let result = run_ssh_with_retry(&config, &node, cmd, dry, 2).await;
@@ -311,6 +321,7 @@ fn build_docker_command(
     nnodes: u32,
     cache_dir: &str,
     keep: bool,
+    init_port: u16,
 ) -> Vec<String> {
     let is_sglang = recipe.engine == "sglang";
 
@@ -319,6 +330,7 @@ fn build_docker_command(
         "--ipc=host",
         "--privileged",
         "--network=host",
+        "--shm-size 10g",
         "--ulimit memlock=-1",
         "--ulimit stack=67108864",
     ];
@@ -342,6 +354,10 @@ fn build_docker_command(
         cmd.push(format!("-e VLLM_HOST_IP={}", node.ib_ip));
     }
 
+    // Default cluster env (recipe env below overrides these if duplicated)
+    cmd.push("-e NCCL_CUMEM_ENABLE=0".to_string());
+    cmd.push("-e OMP_NUM_THREADS=4".to_string());
+
     for (k, v) in &recipe.env {
         cmd.push(format!("-e {}={}", k, v));
     }
@@ -371,7 +387,7 @@ fn build_docker_command(
     cmd.extend(recipe.args.clone());
 
     // Distributed flags - engine-specific
-    let flags: Vec<String> = match recipe.distributed_backend.as_str() {
+    let mut flags: Vec<String> = match recipe.distributed_backend.as_str() {
         "ray" | "mp" if is_sglang => vec![
             "--tp".into(),
             recipe.tensor_parallel_size.to_string(),
@@ -380,7 +396,7 @@ fn build_docker_command(
             "--node-rank".into(),
             node.rank.to_string(),
             "--dist-init-addr".into(),
-            format!("{}:{}", cluster.master_ib_ip, 25000),
+            format!("{}:{}", cluster.master_ib_ip, init_port),
         ],
         "ray" | "mp" => vec![
             "--tensor-parallel-size".into(),
@@ -393,11 +409,18 @@ fn build_docker_command(
             node.rank.to_string(),
             "--master-addr".into(),
             cluster.master_ib_ip.clone(),
+            "--master-port".into(),
+            init_port.to_string(),
             "--distributed-executor-backend".into(),
             recipe.distributed_backend.clone(),
         ],
         _ => vec![],
     };
+
+    // vLLM workers run headless (skip API server) when not rank 0
+    if !is_sglang && node.rank > 0 && matches!(recipe.distributed_backend.as_str(), "mp" | "ray") {
+        flags.push("--headless".into());
+    }
     cmd.extend(flags);
 
     // Port flag — auto-generated for SGLang (all nodes need to bind)

+ 64 - 0
src/docker/tests.rs

@@ -63,6 +63,7 @@ fn run(recipe: &Recipe, node: &Node, head: bool) -> String {
         2,
         "/home/user/.cache/huggingface",
         false,
+        25000,
     )
     .join(" ")
 }
@@ -79,6 +80,7 @@ fn run_keep(recipe: &Recipe, node: &Node, head: bool) -> String {
         2,
         "/home/user/.cache/huggingface",
         true,
+        25000,
     )
     .join(" ")
 }
@@ -100,9 +102,24 @@ fn mp() {
     assert!(cmd.contains("--nnodes 2"));
     assert!(cmd.contains("--node-rank 0"));
     assert!(cmd.contains("--master-addr 1"));
+    assert!(cmd.contains("--master-port 25000"));
     assert!(cmd.contains("--distributed-executor-backend mp"));
     assert!(!cmd.contains("RAY_ADDRESS"));
     assert!(cmd.contains("--rm"), "default --rm should be present");
+    // New: defaults & cluster env
+    assert!(cmd.contains("--shm-size 10g"), "default shm-size missing");
+    assert!(cmd.contains("--ulimit memlock=-1"));
+    assert!(cmd.contains("--ulimit stack=67108864"));
+    assert!(
+        cmd.contains("-e NCCL_CUMEM_ENABLE=0"),
+        "default NCCL_CUMEM_ENABLE missing"
+    );
+    assert!(
+        cmd.contains("-e OMP_NUM_THREADS=4"),
+        "default OMP_NUM_THREADS missing"
+    );
+    // Head node: no --headless
+    assert!(!cmd.contains("--headless"));
 }
 
 #[test]
@@ -112,6 +129,8 @@ fn worker_no_port() {
     assert!(cmd.contains("--node-rank 1"));
     assert!(cmd.contains("--label sm121.rank=1"));
     assert!(cmd.contains("--name sm121_r_1"));
+    // New: vLLM workers run --headless
+    assert!(cmd.contains("--headless"), "vLLM worker should be headless");
 }
 
 #[test]
@@ -177,6 +196,7 @@ fn srun(recipe: &Recipe, node: &Node, head: bool) -> String {
         4,
         "/home/user/.cache/huggingface",
         false,
+        25000,
     )
     .join(" ")
 }
@@ -245,3 +265,47 @@ fn keep_disables_rm() {
     assert!(cmd.contains("-e HF_HOME=/cache/huggingface"));
     assert!(cmd.contains("-v /home/user/.cache/huggingface:/cache/huggingface"));
 }
+
+#[test]
+fn recipe_env_overrides_cluster_defaults() {
+    // Recipe-level env should win over the defaults we inject.
+    let mut r = trec();
+    r.env.insert("NCCL_CUMEM_ENABLE".into(), "1".into());
+    r.env.insert("OMP_NUM_THREADS".into(), "16".into());
+    let cmd = run(&r, &tcls().nodes[0], true);
+    // Docker uses the last -e for duplicate keys, so recipe values come after.
+    // rfind returns the last occurrence, which should be the recipe override.
+    let cumem_pos = cmd.rfind("NCCL_CUMEM_ENABLE=").unwrap();
+    assert!(
+        cmd[cumem_pos..].starts_with("NCCL_CUMEM_ENABLE=1"),
+        "recipe override should be last: slice starts at {}",
+        cumem_pos
+    );
+    let omp_pos = cmd.rfind("OMP_NUM_THREADS=").unwrap();
+    assert!(
+        cmd[omp_pos..].starts_with("OMP_NUM_THREADS=16"),
+        "recipe override should be last"
+    );
+}
+
+#[test]
+fn custom_init_port() {
+    // Non-default init_port should propagate to vLLM master-port and
+    // SGLang dist-init-addr.
+    let cmd = build_docker_command(
+        &tcls(),
+        &tnccl(),
+        &trec(),
+        "r",
+        &tcls().nodes[0],
+        true,
+        2,
+        2,
+        "/home/user/.cache/huggingface",
+        false,
+        29500,
+    )
+    .join(" ");
+    assert!(cmd.contains("--master-port 29500"));
+    assert!(cmd.contains("--master-addr 1"));
+}