Browse Source

refactor: replace #[allow] with structs for start_cluster and build_docker_command

- StartClusterOpts: groups cache_dir, keep, node_name, dry_run
- DockerCommandCtx: clusters all shared context for docker command building
  (cluster, nccl, recipe, sizes, cache_dir, keep, init_port)
- build_docker_command reduced from 12 to 3 params
- start_cluster reduced from 9 to 6 params
- No #[allow(clippy::too_many_arguments)] remaining
Alexander Korolev 4 tháng trước cách đây
mục cha
commit
f55cd70082
2 tập tin đã thay đổi với 111 bổ sung và 142 xóa
  1. 88 82
      src/docker/mod.rs
  2. 23 60
      src/docker/tests.rs

+ 88 - 82
src/docker/mod.rs

@@ -85,10 +85,12 @@ pub async fn serve(
         &nccl,
         recipe_name,
         &recipe,
-        cache_dir,
-        keep,
-        node_name,
-        dry_run,
+        StartClusterOpts {
+            cache_dir,
+            keep,
+            node_name: node_name.map(String::from),
+            dry_run,
+        },
     )
     .await?;
     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);
 }
 
-#[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(
     config: &Config,
     cluster: &Cluster,
     nccl: &NcclConfig,
     recipe_name: &str,
     recipe: &Recipe,
-    cache_dir: std::path::PathBuf,
-    keep: bool,
-    node_name: Option<&str>,
-    dry_run: bool,
+    opts: StartClusterOpts,
 ) -> 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() {
             anyhow::bail!("Node '{}' not found in cluster '{}'", nn, cluster.name);
         }
@@ -258,39 +279,32 @@ async fn start_cluster(
     let mut handles = Vec::new();
     for node in &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 dry = opts.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 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 {
             if !dry {
                 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;
             (node.name, node.rank, engine, result)
         }));
@@ -308,7 +322,7 @@ async fn start_cluster(
             Err(e) => anyhow::bail!("Task join error: {}", e),
         }
     }
-    if !dry_run {
+    if !opts.dry_run {
         sleep(Duration::from_secs(10)).await;
         info!(
             "{}",
@@ -317,22 +331,8 @@ async fn start_cluster(
     }
     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![
         "--gpus all".to_string(),
         "--ipc=host".to_string(),
@@ -341,34 +341,37 @@ fn build_docker_command(
         "--ulimit memlock=-1".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());
     }
 
     let mut cmd = vec![
         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"))];
-    if nnodes > 1 {
+    if ctx.nnodes > 1 {
         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));
         }
-        let nccl_socket_ifname = nccl.socket_ifname.as_str();
+        let nccl_socket_ifname = ctx.nccl.socket_ifname.as_str();
         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()),
             ("GLOO_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()),
             ("UCX_NET_DEVICES", 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(("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()));
     }
-    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.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((
             "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_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(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" {
             let last = ep.len() - 1;
             ep[last] = shq(&ep[last]);
@@ -416,52 +419,55 @@ fn build_docker_command(
     if is_sglang {
         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![
             "--tp".into(),
-            tp_size.to_string(),
+            ctx.tp_size.to_string(),
             "--nnodes".into(),
-            nnodes.to_string(),
+            ctx.nnodes.to_string(),
             "--node-rank".into(),
             node.rank.to_string(),
             "--dist-init-addr".into(),
-            format!("{}:{}", cluster.master_ib_ip, init_port),
+            format!("{}:{}", ctx.cluster.master_ib_ip, ctx.init_port),
         ],
         "ray" | "mp" => {
             let mut fp = vec![
                 "--tensor-parallel-size".into(),
-                tp_size.to_string(),
+                ctx.tp_size.to_string(),
                 "--pipeline-parallel-size".into(),
-                pp_size.to_string(),
+                ctx.pp_size.to_string(),
                 "--nnodes".into(),
-                nnodes.to_string(),
+                ctx.nnodes.to_string(),
                 "--node-rank".into(),
                 node.rank.to_string(),
                 "--master-addr".into(),
-                cluster.master_ib_ip.clone(),
+                ctx.cluster.master_ib_ip.clone(),
                 "--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(recipe.distributed_backend.clone());
+                fp.push(ctx.recipe.distributed_backend.clone());
             }
             fp
         }
         _ => 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());
     }
     cmd.extend(flags);
     cmd.push("--host".to_string());
     cmd.push("0.0.0.0".to_string());
     cmd.push("--port".to_string());
-    cmd.push(recipe.port.to_string());
+    cmd.push(ctx.recipe.port.to_string());
     cmd
 }
 

+ 23 - 60
src/docker/tests.rs

@@ -56,27 +56,28 @@ fn trec() -> Recipe {
     }
 }
 
+fn make_dm_ctx(recipe: &Recipe, nnodes: u32, pp: u32, tp: u32, keep: bool) -> DockerCommandCtx {
+    DockerCommandCtx {
+        cluster: tcls(),
+        nccl: tnccl(),
+        recipe: recipe.clone(),
+        recipe_name: "r".to_string(),
+        pp_size: pp,
+        tp_size: tp,
+        nnodes,
+        cache_dir: "/home/user/.cache/huggingface".to_string(),
+        keep,
+        init_port: 25000,
+    }
+}
+
 fn run(recipe: &Recipe, node: &Node, head: bool) -> String {
     let tp = if recipe.required_nodes == Some(recipe.tensor_parallel_size) {
         2
     } else {
         recipe.tensor_parallel_size
     };
-    build_docker_command(
-        &tcls(),
-        &tnccl(),
-        recipe,
-        "r",
-        node,
-        head,
-        2,
-        tp,
-        2,
-        "/home/user/.cache/huggingface",
-        false,
-        25000,
-    )
-    .join(" ")
+    build_docker_command(&make_dm_ctx(recipe, 2, 2, tp, false), node, head).join(" ")
 }
 
 fn run_keep(recipe: &Recipe, node: &Node, head: bool) -> String {
@@ -85,21 +86,7 @@ fn run_keep(recipe: &Recipe, node: &Node, head: bool) -> String {
     } else {
         recipe.tensor_parallel_size
     };
-    build_docker_command(
-        &tcls(),
-        &tnccl(),
-        recipe,
-        "r",
-        node,
-        head,
-        2,
-        tp,
-        2,
-        "/home/user/.cache/huggingface",
-        true,
-        25000,
-    )
-    .join(" ")
+    build_docker_command(&make_dm_ctx(recipe, 2, 2, tp, true), node, head).join(" ")
 }
 
 #[test]
@@ -254,21 +241,9 @@ fn srun(recipe: &Recipe, node: &Node, head: bool) -> String {
     } else {
         recipe.tensor_parallel_size
     };
-    build_docker_command(
-        &tcls(),
-        &tnccl(),
-        recipe,
-        "sg",
-        node,
-        head,
-        4,
-        tp,
-        4,
-        "/home/user/.cache/huggingface",
-        false,
-        25000,
-    )
-    .join(" ")
+    let mut ctx = make_dm_ctx(recipe, 4, 4, tp, false);
+    ctx.recipe_name = "sg".to_string();
+    build_docker_command(&ctx, node, head).join(" ")
 }
 
 #[test]
@@ -396,21 +371,9 @@ fn custom_init_port() {
     } else {
         r.tensor_parallel_size
     };
-    let cmd = build_docker_command(
-        &tcls(),
-        &tnccl(),
-        &r,
-        "r",
-        &tcls().nodes[0],
-        true,
-        2,
-        tp,
-        2,
-        "/home/user/.cache/huggingface",
-        false,
-        29500,
-    )
-    .join(" ");
+    let mut ctx = make_dm_ctx(&r, 2, 2, tp, false);
+    ctx.init_port = 29500;
+    let cmd = build_docker_command(&ctx, &tcls().nodes[0], true).join(" ");
     assert!(cmd.contains("--master-port 29500"));
     assert!(cmd.contains("--master-addr 1"));
 }