Selaa lähdekoodia

refactor: replace String distributed_backend with DistributedBackend enum

- Add DistributedBackend enum (Mp/Ray/None_) with serde + Display
- Make it optional in Recipe — defaults to Mp when absent
- Replace all string comparisons with enum matching
- Remove default_distributed_backend() helper
Alexander Korolev 4 kuukautta sitten
vanhempi
sitoutus
dbdbc7ed62
9 muutettua tiedostoa jossa 58 lisäystä ja 47 poistoa
  1. 1 1
      src/before_serve.rs
  2. 8 9
      src/config.rs
  3. 2 3
      src/docker/mod.rs
  4. 10 11
      src/docker/tail.rs
  5. 4 3
      src/docker/tests.rs
  6. 5 5
      src/doctor.rs
  7. 22 9
      src/pre_serve.rs
  8. 5 5
      src/tmux/commands.rs
  9. 1 1
      src/tmux/serve.rs

+ 1 - 1
src/before_serve.rs

@@ -245,7 +245,7 @@ mod tests {
             env: HashMap::new(),
             cluster: None,
             entrypoint: vec![],
-            distributed_backend: "mp".to_string(),
+            distributed_backend: crate::config::DistributedBackend::Mp,
             before_serve: vec![],
             mods: Default::default(),
             shm_size: String::new(),

+ 8 - 9
src/config.rs

@@ -1,4 +1,6 @@
-pub use crate::pre_serve::{BeforeServeTable, ModsTable, OnFailure, PreServeCommand};
+pub use crate::pre_serve::{
+    BeforeServeTable, DistributedBackend, ModsTable, OnFailure, PreServeCommand,
+};
 use anyhow::{Context, Result};
 use serde::Deserialize;
 use std::collections::HashMap;
@@ -112,8 +114,8 @@ pub struct Recipe {
     pub cluster: Option<String>,
     #[serde(default)]
     pub entrypoint: Vec<String>,
-    #[serde(default = "default_distributed_backend")]
-    pub distributed_backend: String,
+    #[serde(default)]
+    pub distributed_backend: DistributedBackend,
     /// Shared memory size for the container (e.g. "64gb"). Omitted when empty.
     #[serde(default = "default_shm_size")]
     pub shm_size: String,
@@ -152,9 +154,6 @@ fn resolve_cluster_name<'a>(clusters: &'a [Cluster], name: &str) -> Result<&'a C
 fn default_engine() -> String {
     "vllm".to_string()
 }
-fn default_distributed_backend() -> String {
-    "mp".to_string()
-}
 fn default_shm_size() -> String {
     String::new()
 }
@@ -309,7 +308,7 @@ mod tests {
         assert_eq!(recipe.port, 30000);
         assert!(recipe.args.is_empty());
         assert!(recipe.env.is_empty());
-        assert_eq!(recipe.distributed_backend, "mp");
+        assert_eq!(recipe.distributed_backend, DistributedBackend::Mp);
         assert_eq!(recipe.tensor_parallel_size, 1);
     }
     #[test]
@@ -337,7 +336,7 @@ OTHER = "123"
         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.distributed_backend, DistributedBackend::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");
@@ -483,7 +482,7 @@ ib_disable = 0
     #[test]
     fn test_default_distributed_backend_is_mp() {
         let recipe: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"").unwrap();
-        assert_eq!(recipe.distributed_backend, "mp");
+        assert_eq!(recipe.distributed_backend, DistributedBackend::Mp);
     }
     #[test]
     fn default_config_dir_appends_sm121_launcher() {

+ 2 - 3
src/docker/mod.rs

@@ -23,10 +23,9 @@ pub async fn serve(
     let recipe = load_recipe(recipe_name).await?;
     let rn = recipe.name.clone();
     let re = recipe.engine.clone();
-    let rb = recipe.distributed_backend.clone();
     info!(
         "{}",
-        fluent_i18n::t!("recipe-starting", { "name" => rn, "engine" => re, "backend" => rb })
+        fluent_i18n::t!("recipe-starting", { "name" => rn, "engine" => re, "backend" => recipe.distributed_backend.to_string() })
     );
     let (cluster, nccl): (Cluster, NcclConfig) = if let Some(nn) = node_name {
         let (solo_c, solo_n) = config
@@ -402,7 +401,7 @@ fn build_docker_command(ctx: &DockerCommandCtx, node: &Node, is_head: bool) -> V
     cmd.push(format!("--label sm121.recipe={}", ctx.recipe_name));
     cmd.push(format!("--label sm121.rank={}", node.rank));
     cmd.push(format!("--label sm121.cluster={}", ctx.cluster.name));
-    if ctx.recipe.distributed_backend == "ray" {
+    if ctx.recipe.distributed_backend == crate::config::DistributedBackend::Ray {
         env.push((
             "RAY_ADDRESS",
             format!("{}:6379", ctx.cluster.master_ib_ip).into(),

+ 10 - 11
src/docker/tail.rs

@@ -1,7 +1,7 @@
 //! Builds the "tail" of the docker command — everything after the image
 //! name (entrypoint, model, args, auto-generated flags).
 
-use crate::config::{Node, Recipe};
+use crate::config::{DistributedBackend, Node, Recipe};
 
 use super::shq;
 
@@ -18,11 +18,13 @@ pub(crate) struct TailCtx<'a> {
 /// Build the docker command tail — everything after the image name.
 pub fn build_tail(ctx: &TailCtx<'_>, node: &Node) -> Vec<String> {
     let is_sglang = ctx.recipe.engine == "sglang";
+    use DistributedBackend::{Mp, Ray};
+    let is_distributed = matches!(ctx.recipe.distributed_backend, Mp | Ray);
+
     let mut tail = Vec::new();
 
     if !ctx.recipe.entrypoint.is_empty() {
         let mut ep = ctx.recipe.entrypoint.clone();
-        // Non-mods path: inner shq for sh -c entrypoints
         if ctx.recipe.mods.commands.is_empty() && ep.len() >= 3 && ep[ep.len() - 2] == "-c" {
             let last = ep.len() - 1;
             ep[last] = shq(&ep[last]);
@@ -36,8 +38,8 @@ pub fn build_tail(ctx: &TailCtx<'_>, node: &Node) -> Vec<String> {
     tail.push(ctx.recipe.model.clone());
     tail.extend(ctx.recipe.args.clone());
 
-    let mut flags: Vec<String> = match ctx.recipe.distributed_backend.as_str() {
-        "ray" | "mp" if is_sglang => vec![
+    let mut flags: Vec<String> = match ctx.recipe.distributed_backend {
+        Ray | Mp if is_sglang => vec![
             "--tp".into(),
             ctx.tp_size.to_string(),
             "--nnodes".into(),
@@ -47,7 +49,7 @@ pub fn build_tail(ctx: &TailCtx<'_>, node: &Node) -> Vec<String> {
             "--dist-init-addr".into(),
             format!("{}:{}", ctx.master_ib_ip, ctx.init_port),
         ],
-        "ray" | "mp" => {
+        Ray | Mp => {
             let mut fp = vec![
                 "--tensor-parallel-size".into(),
                 ctx.tp_size.to_string(),
@@ -62,18 +64,15 @@ pub fn build_tail(ctx: &TailCtx<'_>, node: &Node) -> Vec<String> {
                 "--master-port".into(),
                 ctx.init_port.to_string(),
             ];
-            if ctx.recipe.distributed_backend == "ray" {
+            if ctx.recipe.distributed_backend == Ray {
                 fp.push("--distributed-executor-backend".into());
-                fp.push(ctx.recipe.distributed_backend.clone());
+                fp.push("ray".into());
             }
             fp
         }
         _ => vec![],
     };
-    if !is_sglang
-        && node.rank > 0
-        && matches!(ctx.recipe.distributed_backend.as_str(), "mp" | "ray")
-    {
+    if !is_sglang && node.rank > 0 && is_distributed {
         flags.push("--headless".into());
     }
     tail.extend(flags);

+ 4 - 3
src/docker/tests.rs

@@ -1,4 +1,5 @@
 use super::*;
+use crate::config::DistributedBackend;
 use std::collections::HashMap;
 
 fn tcls() -> Cluster {
@@ -50,7 +51,7 @@ fn trec() -> Recipe {
         env: HashMap::from([("K".into(), "V".into())]),
         cluster: None,
         entrypoint: vec![],
-        distributed_backend: "mp".into(),
+        distributed_backend: DistributedBackend::Mp,
         before_serve: vec![],
         mods: Default::default(),
         shm_size: String::new(),
@@ -177,7 +178,7 @@ fn worker_no_port() {
 #[test]
 fn ray_backend() {
     let mut r = trec();
-    r.distributed_backend = "ray".into();
+    r.distributed_backend = DistributedBackend::Ray;
     let cmd = run(&r, &tcls().nodes[0], true);
     assert!(cmd.contains("-e RAY_ADDRESS=1:6379"));
     assert!(cmd.contains("-e RAY_PORT=6379"));
@@ -231,7 +232,7 @@ fn sglang_recipe() -> Recipe {
         env: HashMap::from([("TORCH_NCCL_ASYNC_ERROR_HANDLING".into(), "1".into())]),
         cluster: None,
         entrypoint: vec!["sglang".into(), "serve".into()],
-        distributed_backend: "mp".into(),
+        distributed_backend: DistributedBackend::Mp,
         before_serve: vec![],
         mods: Default::default(),
         shm_size: String::new(),

+ 5 - 5
src/doctor.rs

@@ -44,7 +44,7 @@ pub async fn doctor(
     let recipe = load_recipe(recipe_name).await?;
     let dn = recipe.name.clone();
     let de = recipe.engine.clone();
-    let db = recipe.distributed_backend.clone();
+    let db = recipe.distributed_backend.to_string();
     println!(
         "{}",
         fluent_i18n::t!("doctor-header", {
@@ -64,7 +64,7 @@ pub async fn doctor(
         let config = config.clone();
         let node = node.clone();
         let engine = recipe.engine.clone();
-        let backend = recipe.distributed_backend.clone();
+        let backend = recipe.distributed_backend.to_string();
         let image = recipe.docker_image.clone();
         handles.push(tokio::spawn(async move {
             check_node(&config, &node, &engine, &backend, &image).await
@@ -281,7 +281,7 @@ async fn run_cluster_repair(
 fn generate_repair_dockerfile(base_image: &str, recipe: &Recipe) -> String {
     let mut lines = vec![format!("FROM {}", base_image)];
 
-    if recipe.distributed_backend == "ray" {
+    if recipe.distributed_backend == crate::config::DistributedBackend::Ray {
         lines.push("RUN pip install ray[default] || pip install ray".to_string());
     }
 
@@ -339,7 +339,7 @@ mod tests {
             env: HashMap::new(),
             cluster: None,
             entrypoint: vec![],
-            distributed_backend: "mp".to_string(),
+            distributed_backend: crate::config::DistributedBackend::Mp,
             before_serve: vec![],
             mods: Default::default(),
             shm_size: String::new(),
@@ -360,7 +360,7 @@ mod tests {
     #[test]
     fn test_generate_repair_dockerfile_ray() {
         let mut recipe = test_recipe();
-        recipe.distributed_backend = "ray".to_string();
+        recipe.distributed_backend = crate::config::DistributedBackend::Ray;
         let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);
 
         assert!(dockerfile.contains("FROM base/image:tag"));

+ 22 - 9
src/pre_serve.rs

@@ -9,17 +9,13 @@ use serde::Deserialize;
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Default, Deserialize)]
 #[serde(rename_all = "lowercase")]
 pub enum OnFailure {
-    /// Abort the serve (default).
     #[default]
     Abort,
-    /// Log a warning but continue.
     Warn,
-    /// Continue silently.
     Continue,
 }
 
-/// One entry in a `before_serve` block.  Either a bare command string
-/// (defaults to `on_failure = abort`) or a table with explicit `on_failure`.
+/// One entry in a `before_serve` block.
 #[derive(Clone, Debug, Deserialize)]
 #[serde(untagged)]
 pub enum PreServeCommand {
@@ -47,18 +43,35 @@ impl PreServeCommand {
     }
 }
 
-/// TOML sub-table under `[before_serve."<cluster>"]`.  Currently only
-/// carries `commands`; reserved for future per-cluster tuning.
 #[derive(Clone, Debug, Default, Deserialize)]
 pub struct BeforeServeTable {
     #[serde(default)]
     pub commands: Vec<PreServeCommand>,
 }
 
-/// TOML sub-table under `[mods]` in a recipe.  Mirrors [`BeforeServeTable`].
 #[derive(Clone, Debug, Default, Deserialize)]
 pub struct ModsTable {
-    /// Commands to run inside the container before the engine starts.
     #[serde(default)]
     pub commands: Vec<PreServeCommand>,
 }
+
+/// Distributed execution backend: `mp` (default), `ray`, or `none`.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum DistributedBackend {
+    #[default]
+    Mp,
+    Ray,
+    #[serde(rename = "none")]
+    None_,
+}
+
+impl std::fmt::Display for DistributedBackend {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            DistributedBackend::Mp => write!(f, "mp"),
+            DistributedBackend::Ray => write!(f, "ray"),
+            DistributedBackend::None_ => write!(f, "none"),
+        }
+    }
+}

+ 5 - 5
src/tmux/commands.rs

@@ -1,4 +1,4 @@
-use crate::config::{Cluster, NcclConfig, Node, Recipe};
+use crate::config::{Cluster, DistributedBackend, NcclConfig, Node, Recipe};
 use crate::tmux::TmuxManager;
 
 impl TmuxManager {
@@ -45,8 +45,8 @@ impl TmuxManager {
                 cluster.master_ib_ip.clone(),
             ]);
         } else {
-            match recipe.distributed_backend.as_str() {
-                "ray" => {
+            match recipe.distributed_backend {
+                DistributedBackend::Ray => {
                     parts.extend(vec![
                         "--tensor-parallel-size".to_string(),
                         "1".to_string(),
@@ -75,7 +75,7 @@ impl TmuxManager {
                         "--master-addr".to_string(),
                         cluster.master_ib_ip.clone(),
                         "--distributed-executor-backend".to_string(),
-                        recipe.distributed_backend.clone(),
+                        "mp".to_string(),
                     ]);
                 }
             }
@@ -113,7 +113,7 @@ impl TmuxManager {
         for export in self.build_host_ip_export(&recipe.engine, node) {
             cmds.push(format!("tmux send-keys -t {}:0 '{}' Enter", sess, export));
         }
-        if recipe.distributed_backend == "ray" && node.rank > 0 {
+        if recipe.distributed_backend == DistributedBackend::Ray && node.rank > 0 {
             cmds.push(format!(
                 "tmux send-keys -t {}:0 'export RAY_ADDRESS={}:6379' Enter",
                 sess, cluster.master_ib_ip

+ 1 - 1
src/tmux/serve.rs

@@ -51,7 +51,7 @@ pub async fn tmux_serve(
         session_prefix: recipe_name.to_string(),
     };
 
-    let is_ray = recipe.distributed_backend == "ray";
+    let is_ray = recipe.distributed_backend == crate::config::DistributedBackend::Ray;
 
     if is_ray {
         info!("🐍 Ray backend detected — starting head node first...");