Bläddra i källkod

refactor: DistributedBackend as Option — remove None_ variant

- Recipe.distributed_backend is now Option<DistributedBackend>
- None means single-node (absent from TOML)
- Enum only has Mp and Ray
- Updated all match/comparison sites to use Some(...)
- Removed distributed_backend = "none" from gemma4-12b and test-mods recipes
Alexander Korolev 4 månader sedan
förälder
incheckning
5931869d69

+ 0 - 1
recipes/gemma4-12b/recipe.toml

@@ -6,7 +6,6 @@ model = "google/gemma-4-12B-it"
 # model = "RedHatAI/gemma-4-12B-it-NVFP4"
 tensor_parallel_size = 1
 pipeline_parallel_size = 1
-distributed_backend = "none"
 
 # Single-node recipe: Gemma 4 12B with BF16, FP8 KV cache,
 # MTP speculative decoding (5 tokens), 196K context.

+ 0 - 1
recipes/test-mods/recipe.toml

@@ -6,7 +6,6 @@ port = 30000
 required_nodes = 1
 tensor_parallel_size = 1
 pipeline_parallel_size = 1
-distributed_backend = "none"
 
 # Entrypoint must be set when mods are present.
 entrypoint = ["sh", "-c", "i=0; while true; do echo \"[$(hostname)] heartbeat $i\"; i=$((i+1)); sleep 2; done"]

+ 1 - 1
src/before_serve.rs

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

+ 4 - 5
src/config.rs

@@ -114,8 +114,7 @@ pub struct Recipe {
     pub cluster: Option<String>,
     #[serde(default)]
     pub entrypoint: Vec<String>,
-    #[serde(default)]
-    pub distributed_backend: DistributedBackend,
+    pub distributed_backend: Option<DistributedBackend>,
     /// Shared memory size for the container (e.g. "64gb"). Omitted when empty.
     #[serde(default = "default_shm_size")]
     pub shm_size: String,
@@ -308,7 +307,7 @@ mod tests {
         assert_eq!(recipe.port, 30000);
         assert!(recipe.args.is_empty());
         assert!(recipe.env.is_empty());
-        assert_eq!(recipe.distributed_backend, DistributedBackend::None_);
+        assert_eq!(recipe.distributed_backend, None);
         assert_eq!(recipe.tensor_parallel_size, 1);
     }
     #[test]
@@ -336,7 +335,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, DistributedBackend::Ray);
+        assert_eq!(recipe.distributed_backend, Some(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");
@@ -482,7 +481,7 @@ ib_disable = 0
     #[test]
     fn test_default_distributed_backend_is_none() {
         let recipe: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"").unwrap();
-        assert_eq!(recipe.distributed_backend, DistributedBackend::None_);
+        assert_eq!(recipe.distributed_backend, None);
     }
     #[test]
     fn default_config_dir_appends_sm121_launcher() {

+ 2 - 2
src/docker/mod.rs

@@ -25,7 +25,7 @@ pub async fn serve(
     let re = recipe.engine.clone();
     info!(
         "{}",
-        fluent_i18n::t!("recipe-starting", { "name" => rn, "engine" => re, "backend" => recipe.distributed_backend.to_string() })
+        fluent_i18n::t!("recipe-starting", { "name" => rn, "engine" => re, "backend" => recipe.distributed_backend.map_or("none".to_string(), |b| b.to_string()) })
     );
     let (cluster, nccl): (Cluster, NcclConfig) = if let Some(nn) = node_name {
         let (solo_c, solo_n) = config
@@ -401,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 == crate::config::DistributedBackend::Ray {
+    if ctx.recipe.distributed_backend == Some(crate::config::DistributedBackend::Ray) {
         env.push((
             "RAY_ADDRESS",
             format!("{}:6379", ctx.cluster.master_ib_ip).into(),

+ 7 - 5
src/docker/tail.rs

@@ -18,8 +18,10 @@ 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 is_distributed = matches!(
+        ctx.recipe.distributed_backend,
+        Some(DistributedBackend::Mp | DistributedBackend::Ray)
+    );
 
     let mut tail = Vec::new();
 
@@ -39,7 +41,7 @@ pub fn build_tail(ctx: &TailCtx<'_>, node: &Node) -> Vec<String> {
     tail.extend(ctx.recipe.args.clone());
 
     let mut flags: Vec<String> = match ctx.recipe.distributed_backend {
-        Ray | Mp if is_sglang => vec![
+        Some(DistributedBackend::Ray | DistributedBackend::Mp) if is_sglang => vec![
             "--tp".into(),
             ctx.tp_size.to_string(),
             "--nnodes".into(),
@@ -49,7 +51,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 => {
+        Some(DistributedBackend::Ray | DistributedBackend::Mp) => {
             let mut fp = vec![
                 "--tensor-parallel-size".into(),
                 ctx.tp_size.to_string(),
@@ -64,7 +66,7 @@ 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 == Some(DistributedBackend::Ray) {
                 fp.push("--distributed-executor-backend".into());
                 fp.push("ray".into());
             }

+ 3 - 3
src/docker/tests.rs

@@ -51,7 +51,7 @@ fn trec() -> Recipe {
         env: HashMap::from([("K".into(), "V".into())]),
         cluster: None,
         entrypoint: vec![],
-        distributed_backend: DistributedBackend::Mp,
+        distributed_backend: Some(DistributedBackend::Mp),
         before_serve: vec![],
         mods: Default::default(),
         shm_size: String::new(),
@@ -178,7 +178,7 @@ fn worker_no_port() {
 #[test]
 fn ray_backend() {
     let mut r = trec();
-    r.distributed_backend = DistributedBackend::Ray;
+    r.distributed_backend = Some(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"));
@@ -232,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: DistributedBackend::Mp,
+        distributed_backend: Some(DistributedBackend::Mp),
         before_serve: vec![],
         mods: Default::default(),
         shm_size: String::new(),

+ 9 - 5
src/doctor.rs

@@ -44,7 +44,9 @@ 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.to_string();
+    let db = recipe
+        .distributed_backend
+        .map_or("none".to_string(), |b| b.to_string());
     println!(
         "{}",
         fluent_i18n::t!("doctor-header", {
@@ -64,7 +66,9 @@ pub async fn doctor(
         let config = config.clone();
         let node = node.clone();
         let engine = recipe.engine.clone();
-        let backend = recipe.distributed_backend.to_string();
+        let backend = recipe
+            .distributed_backend
+            .map_or("none".to_string(), |b| b.to_string());
         let image = recipe.docker_image.clone();
         handles.push(tokio::spawn(async move {
             check_node(&config, &node, &engine, &backend, &image).await
@@ -281,7 +285,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 == crate::config::DistributedBackend::Ray {
+    if recipe.distributed_backend == Some(crate::config::DistributedBackend::Ray) {
         lines.push("RUN pip install ray[default] || pip install ray".to_string());
     }
 
@@ -339,7 +343,7 @@ mod tests {
             env: HashMap::new(),
             cluster: None,
             entrypoint: vec![],
-            distributed_backend: crate::config::DistributedBackend::Mp,
+            distributed_backend: Some(crate::config::DistributedBackend::Mp),
             before_serve: vec![],
             mods: Default::default(),
             shm_size: String::new(),
@@ -360,7 +364,7 @@ mod tests {
     #[test]
     fn test_generate_repair_dockerfile_ray() {
         let mut recipe = test_recipe();
-        recipe.distributed_backend = crate::config::DistributedBackend::Ray;
+        recipe.distributed_backend = Some(crate::config::DistributedBackend::Ray);
         let dockerfile = generate_repair_dockerfile("base/image:tag", &recipe);
 
         assert!(dockerfile.contains("FROM base/image:tag"));

+ 2 - 7
src/pre_serve.rs

@@ -55,16 +55,12 @@ pub struct ModsTable {
     pub commands: Vec<PreServeCommand>,
 }
 
-/// Distributed execution backend — `mp`, `ray`, or absent (single-node).
-/// Defaults to [`None_`] when not specified in the recipe.
-#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Deserialize)]
+/// Distributed execution backend — `mp` or `ray`.  Absent means single-node.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
 #[serde(rename_all = "lowercase")]
 pub enum DistributedBackend {
     Mp,
     Ray,
-    #[default]
-    #[serde(rename = "none")]
-    None_,
 }
 
 impl std::fmt::Display for DistributedBackend {
@@ -72,7 +68,6 @@ impl std::fmt::Display for DistributedBackend {
         match self {
             DistributedBackend::Mp => write!(f, "mp"),
             DistributedBackend::Ray => write!(f, "ray"),
-            DistributedBackend::None_ => write!(f, "none"),
         }
     }
 }

+ 2 - 2
src/tmux/commands.rs

@@ -46,7 +46,7 @@ impl TmuxManager {
             ]);
         } else {
             match recipe.distributed_backend {
-                DistributedBackend::Ray => {
+                Some(DistributedBackend::Ray) => {
                     parts.extend(vec![
                         "--tensor-parallel-size".to_string(),
                         "1".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 == DistributedBackend::Ray && node.rank > 0 {
+        if recipe.distributed_backend == Some(DistributedBackend::Ray) && node.rank > 0 {
             cmds.push(format!(
                 "tmux send-keys -t {}:0 'export RAY_ADDRESS={}:6379' Enter",
                 sess, cluster.master_ib_ip

+ 7 - 2
src/tmux/serve.rs

@@ -20,7 +20,12 @@ pub async fn tmux_serve(
 
     info!(
         "🚀 TMUX launch: {} ({}) [backend: {}] [engine: {}]",
-        recipe.name, recipe_name, recipe.distributed_backend, recipe.engine
+        recipe.name,
+        recipe_name,
+        recipe
+            .distributed_backend
+            .map_or("none".into(), |b| b.to_string()),
+        recipe.engine
     );
 
     let choice = if cluster_name != "auto" {
@@ -51,7 +56,7 @@ pub async fn tmux_serve(
         session_prefix: recipe_name.to_string(),
     };
 
-    let is_ray = recipe.distributed_backend == crate::config::DistributedBackend::Ray;
+    let is_ray = recipe.distributed_backend == Some(crate::config::DistributedBackend::Ray);
 
     if is_ray {
         info!("🐍 Ray backend detected — starting head node first...");