Browse Source

feat: RecipeArgs enum — inline string or list for recipe args

- RecipeArgs supports both TOML multiline string (split by whitespace)
  and the existing Vec<String> array format
- Default is empty List variant
- to_vec() converts either variant to Vec<String>
- All call sites updated to use .to_vec() instead of .clone()
- Added deserialize_args_inline test
Alexander Korolev 4 months ago
parent
commit
4936530710
7 changed files with 49 additions and 11 deletions
  1. 1 1
      src/before_serve.rs
  2. 4 4
      src/config.rs
  3. 1 1
      src/docker/tail.rs
  4. 3 3
      src/docker/tests.rs
  5. 4 1
      src/doctor.rs
  6. 35 0
      src/pre_serve.rs
  7. 1 1
      src/tmux/commands.rs

+ 1 - 1
src/before_serve.rs

@@ -241,7 +241,7 @@ mod tests {
             required_nodes: None,
             pipeline_parallel_size: None,
             tensor_parallel_size: 1,
-            args: vec![],
+            args: crate::config::RecipeArgs::default(),
             env: HashMap::new(),
             cluster: None,
             entrypoint: vec![],

+ 4 - 4
src/config.rs

@@ -1,5 +1,5 @@
 pub use crate::pre_serve::{
-    BeforeServeTable, DistributedBackend, ModsTable, OnFailure, PreServeCommand,
+    BeforeServeTable, DistributedBackend, ModsTable, OnFailure, PreServeCommand, RecipeArgs,
 };
 use anyhow::{Context, Result};
 use serde::Deserialize;
@@ -107,7 +107,7 @@ pub struct Recipe {
     #[serde(default = "default_tensor_parallel_size")]
     pub tensor_parallel_size: u32,
     #[serde(default)]
-    pub args: Vec<String>,
+    pub args: RecipeArgs,
     #[serde(default)]
     pub env: HashMap<String, String>,
     #[serde(default)]
@@ -305,7 +305,7 @@ mod tests {
         assert_eq!(recipe.model, "test/model");
         assert_eq!(recipe.engine, "vllm");
         assert_eq!(recipe.port, 30000);
-        assert!(recipe.args.is_empty());
+        assert!(recipe.args.to_vec().is_empty());
         assert!(recipe.env.is_empty());
         assert_eq!(recipe.distributed_backend, None);
         assert_eq!(recipe.tensor_parallel_size, 1);
@@ -337,7 +337,7 @@ OTHER = "123"
         assert_eq!(recipe.pipeline_parallel_size, Some(2));
         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.args.to_vec(), vec!["--arg1", "val1", "--arg2"]);
         assert_eq!(recipe.env.get("KEY").unwrap(), "value");
         assert_eq!(recipe.env.get("OTHER").unwrap(), "123");
     }

+ 1 - 1
src/docker/tail.rs

@@ -38,7 +38,7 @@ pub fn build_tail(ctx: &TailCtx<'_>, node: &Node) -> Vec<String> {
         tail.push("--model-path".to_string());
     }
     tail.push(ctx.recipe.model.clone());
-    tail.extend(ctx.recipe.args.clone());
+    tail.extend(ctx.recipe.args.to_vec());
 
     let mut flags: Vec<String> = match ctx.recipe.distributed_backend {
         Some(DistributedBackend::Ray | DistributedBackend::Mp) if is_sglang => vec![

+ 3 - 3
src/docker/tests.rs

@@ -1,5 +1,5 @@
 use super::*;
-use crate::config::DistributedBackend;
+use crate::config::{DistributedBackend, RecipeArgs};
 use std::collections::HashMap;
 
 fn tcls() -> Cluster {
@@ -47,7 +47,7 @@ fn trec() -> Recipe {
         required_nodes: Some(2),
         pipeline_parallel_size: Some(2),
         tensor_parallel_size: 1,
-        args: vec!["--x".into(), "4096".into()],
+        args: RecipeArgs::List(vec!["--x".into(), "4096".into()]),
         env: HashMap::from([("K".into(), "V".into())]),
         cluster: None,
         entrypoint: vec![],
@@ -228,7 +228,7 @@ fn sglang_recipe() -> Recipe {
         required_nodes: Some(4),
         pipeline_parallel_size: None,
         tensor_parallel_size: 4,
-        args: vec!["--mem-fraction-static".into(), "0.8".into()],
+        args: RecipeArgs::List(vec!["--mem-fraction-static".into(), "0.8".into()]),
         env: HashMap::from([("TORCH_NCCL_ASYNC_ERROR_HANDLING".into(), "1".into())]),
         cluster: None,
         entrypoint: vec!["sglang".into(), "serve".into()],

+ 4 - 1
src/doctor.rs

@@ -339,7 +339,10 @@ mod tests {
             required_nodes: Some(2),
             pipeline_parallel_size: Some(2),
             tensor_parallel_size: 1,
-            args: vec!["--max-model-len".to_string(), "4096".to_string()],
+            args: crate::config::RecipeArgs::List(vec![
+                "--max-model-len".to_string(),
+                "4096".to_string(),
+            ]),
             env: HashMap::new(),
             cluster: None,
             entrypoint: vec![],

+ 35 - 0
src/pre_serve.rs

@@ -63,6 +63,41 @@ pub enum DistributedBackend {
     Ray,
 }
 
+/// Recipe arguments — either an inline TOML string (split by whitespace) or
+/// a `Vec<String>` list.  Accepted in TOML as either:
+///
+/// ```toml
+/// args = """
+///   --gpu-memory-utilization 0.90
+///   --max-model-len 262144
+/// """
+/// ```
+///
+/// ```toml
+/// args = ["--gpu-memory-utilization", "0.90"]
+/// ```
+#[derive(Clone, Debug, Deserialize)]
+#[serde(untagged)]
+pub enum RecipeArgs {
+    List(Vec<String>),
+    Inline(String),
+}
+
+impl Default for RecipeArgs {
+    fn default() -> Self {
+        RecipeArgs::List(Vec::new())
+    }
+}
+
+impl RecipeArgs {
+    pub fn to_vec(&self) -> Vec<String> {
+        match self {
+            RecipeArgs::Inline(s) => s.split_whitespace().map(String::from).collect(),
+            RecipeArgs::List(v) => v.clone(),
+        }
+    }
+}
+
 impl std::fmt::Display for DistributedBackend {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {

+ 1 - 1
src/tmux/commands.rs

@@ -30,7 +30,7 @@ impl TmuxManager {
         parts.push(recipe.model.clone());
 
         // Recipe-specific args
-        parts.extend(recipe.args.clone());
+        parts.extend(recipe.args.to_vec());
 
         // Distributed flags — engine-specific
         if is_sglang {