Sfoglia il codice sorgente

feat(tmux): support inline venv_setup + venv_setup_script

Add VenvSetup untagged enum (like RecipeArgs) accepting either an inline
multi-line string (one command per line; blank lines and # comments
skipped) or a Vec<String> list. Add venv_setup_script field for
referencing a shell script file in the recipe directory (read and run
on each node if venv doesn't exist; takes precedence over venv_setup).
The laguna recipe now uses inline format. AGENTS.md documents both
forms. 149 tests pass, make check green.
Alexander Korolev 2 mesi fa
parent
commit
afa6999213

+ 31 - 10
AGENTS.md

@@ -53,6 +53,7 @@ solo = false                       # optional; single-node on first free node, n
 deploy = "docker"                   # optional; "docker" (default) or "tmux"
 venv = "~/venvs/vllm025"            # optional; Python venv for tmux (activates before entrypoint)
 venv_setup = [...]                 # optional; commands to create venv if it doesn't exist (tmux only)
+venv_setup_script = "venv.sh"      # optional; script file (relative to recipe dir) to create venv (tmux only)
 tensor_parallel_size = 2
 pipeline_parallel_size = 1          # vLLM only, 0 = auto
 distributed_backend = "mp"          # or "ray"
@@ -127,29 +128,49 @@ served via either `docker` or `tmux` without changes.
 
 When `deploy = "tmux"`, the `venv` field specifies a Python virtualenv path that is
 activated before the entrypoint (`source <venv>/bin/activate && ...`). The `venv_setup`
-field is an optional list of shell commands to create and initialize the venv **if it
+field is an optional set of shell commands to create and initialize the venv **if it
 does not already exist** on a node. Before serve, the launcher checks
-`test -f <venv>/bin/activate` on each node; if absent, it runs the `venv_setup` commands
+`test -f <venv>/bin/activate` on each node; if absent, it runs the setup commands
 in order, then re-checks. A failure aborts the serve.
 
+`venv_setup` accepts either an inline multi-line string (one command per line;
+blank lines and `#` comments are skipped) or a `Vec<String>` list:
+
 ```toml
 deploy = "tmux"
 venv = "~/venvs/vllm025"
-venv_setup = [
-    "uv venv ~/venvs/vllm025 -p 3.12",
-    "uv pip install vllm==0.25.1 --torch-backend=cu130",
-    "uv pip install flashinfer-python==0.6.15.dev20260712 \
+venv_setup = """
+    uv venv ~/venvs/vllm025 -p 3.12
+    uv pip install vllm==0.25.1 --torch-backend=cu130
+    uv pip install flashinfer-python==0.6.15.dev20260712 \
       flashinfer-cubin==0.6.15.dev20260712 \
       flashinfer-jit-cache==0.6.15.dev20260712 \
       --extra-index-url https://flashinfer.ai/whl/nightly/ \
-      --extra-index-url https://flashinfer.ai/whl/nightly/cu130/",
-    "hf download poolside/Laguna-S-2.1-NVFP4",
-    "hf download poolside/Laguna-S-2.1-DFlash-NVFP4",
+      --extra-index-url https://flashinfer.ai/whl/nightly/cu130/
+    hf download poolside/Laguna-S-2.1-NVFP4
+    hf download poolside/Laguna-S-2.1-DFlash-NVFP4
+"""
+```
+
+Or as a list:
+
+```toml
+venv_setup = [
+    "uv venv ~/venvs/vllm025 -p 3.12",
+    "uv pip install vllm==0.25.1 --torch-backend=cu130",
 ]
 ```
 
+Alternatively, `venv_setup_script` references a shell script file (relative to the
+recipe directory) that is read and run on each node if the venv doesn't exist.
+Script file takes precedence over `venv_setup` if both are set:
+
+```toml
+venv_setup_script = "venv_setup.sh"
+```
+
 If the venv already exists on a node, the setup commands are skipped. `venv_setup`
-is ignored in Docker mode (the container has its own Python).
+and `venv_setup_script` are ignored in Docker mode (the container has its own Python).
 
 ## Env vars (auto-set)
 

+ 7 - 7
recipes/laguna-s2.1-nvfp4/recipe.toml

@@ -41,13 +41,13 @@ tensor_parallel_size = 1
 pipeline_parallel_size = 1
 deploy = "tmux"
 venv = "~/venvs/vllm025"
-venv_setup = [
-    "uv venv ~/venvs/vllm025 -p 3.12",
-    "uv pip install vllm==0.25.1 --torch-backend=cu130",
-    "uv pip install flashinfer-python==0.6.15.dev20260712 flashinfer-cubin==0.6.15.dev20260712 flashinfer-jit-cache==0.6.15.dev20260712 --extra-index-url https://flashinfer.ai/whl/nightly/ --extra-index-url https://flashinfer.ai/whl/nightly/cu130/",
-    "hf download poolside/Laguna-S-2.1-NVFP4",
-    "hf download poolside/Laguna-S-2.1-DFlash-NVFP4",
-]
+venv_setup = """
+    uv venv ~/venvs/vllm025 -p 3.12
+    uv pip install vllm==0.25.1 --torch-backend=cu130
+    uv pip install flashinfer-python==0.6.15.dev20260712 flashinfer-cubin==0.6.15.dev20260712 flashinfer-jit-cache==0.6.15.dev20260712 --extra-index-url https://flashinfer.ai/whl/nightly/ --extra-index-url https://flashinfer.ai/whl/nightly/cu130/
+    hf download poolside/Laguna-S-2.1-NVFP4
+    hf download poolside/Laguna-S-2.1-DFlash-NVFP4
+"""
 
 entrypoint = ["vllm", "serve"]
 

+ 3 - 3
specs/state.yaml

@@ -23,9 +23,9 @@ session:
 
 git:
   branch: main
-  head: d9b922e
-  clean: true                    # venv field + deploy field + structured-output fix landed + pushed
+  head: 5f64696
+  clean: true                    # venv_setup + venv + deploy + structured-output fix landed + pushed
 
 handoff:
   next_skill: survey-context    # venv_setup feature shipped; survey to pick next task
-  note: "venv_setup field added to Recipe; prepare_venv creates venvs on tmux nodes if absent. All 147 tests pass, make check green."
+  note: "venv_setup field added to Recipe; prepare_venv creates venvs on tmux nodes if absent. All 147 tests pass, make check green. Solar recipe change (--served-model-name) remains uncommitted."

+ 3 - 2
src/after_stop.rs

@@ -145,7 +145,7 @@ fn truncate(s: &str, max: usize) -> String {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::config::{Node, PreServeCommand};
+    use crate::config::{Node, PreServeCommand, VenvSetup};
 
     fn test_cluster() -> Cluster {
         Cluster {
@@ -200,7 +200,8 @@ mod tests {
             solo: None,
             deploy: crate::config::DeployTarget::Docker,
             venv: None,
-            venv_setup: vec![],
+            venv_setup: VenvSetup::default(),
+            venv_setup_script: None,
             mods: Default::default(),
             sync_image: true,
             docker_image_dir: None,

+ 3 - 2
src/before_serve.rs

@@ -177,7 +177,7 @@ fn truncate(s: &str, max: usize) -> String {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::config::{Node, PreServeCommand};
+    use crate::config::{Node, PreServeCommand, VenvSetup};
     use std::collections::HashMap;
 
     fn test_cluster() -> Cluster {
@@ -310,7 +310,8 @@ mod tests {
             solo: None,
             deploy: crate::config::DeployTarget::Docker,
             venv: None,
-            venv_setup: vec![],
+            venv_setup: VenvSetup::default(),
+            venv_setup_script: None,
             mods: Default::default(),
             sync_image: true,
             docker_image_dir: None,

+ 52 - 5
src/config/tests.rs

@@ -227,7 +227,7 @@ fn test_venv_defaults_to_none() {
 }
 
 #[test]
-fn test_venv_setup_deserializes() {
+fn test_venv_setup_list_deserializes() {
     let r: Recipe = toml::from_str(
         r#"name = "x"
 model = "m"
@@ -241,18 +241,65 @@ venv_setup = [
     )
     .unwrap();
     assert_eq!(r.venv.as_deref(), Some("~/venvs/vllm025"));
-    assert_eq!(r.venv_setup.len(), 2);
-    assert_eq!(r.venv_setup[0], "uv venv ~/venvs/vllm025 -p 3.12");
+    let setup = r.venv_setup.to_vec();
+    assert_eq!(setup.len(), 2);
+    assert_eq!(setup[0], "uv venv ~/venvs/vllm025 -p 3.12");
     assert_eq!(
-        r.venv_setup[1],
+        setup[1],
         "uv pip install vllm==0.25.1 --torch-backend=cu130"
     );
 }
 
+#[test]
+fn test_venv_setup_inline_deserializes() {
+    let r: Recipe = toml::from_str(
+        r#"name = "x"
+model = "m"
+deploy = "tmux"
+venv = "~/venvs/vllm025"
+venv_setup = """
+    uv venv ~/venvs/vllm025 -p 3.12
+    uv pip install vllm==0.25.1 --torch-backend=cu130
+"""
+"#,
+    )
+    .unwrap();
+    let setup = r.venv_setup.to_vec();
+    assert_eq!(setup.len(), 2);
+    assert_eq!(setup[0], "uv venv ~/venvs/vllm025 -p 3.12");
+    assert_eq!(
+        setup[1],
+        "uv pip install vllm==0.25.1 --torch-backend=cu130"
+    );
+}
+
+#[test]
+fn test_venv_setup_inline_skips_comments_and_blanks() {
+    let r: Recipe = toml::from_str(
+        r#"name = "x"
+model = "m"
+deploy = "tmux"
+venv = "~/venvs/vllm025"
+venv_setup = """
+    # create the venv
+    uv venv ~/venvs/vllm025 -p 3.12
+
+    # install deps
+    uv pip install vllm==0.25.1
+"""
+"#,
+    )
+    .unwrap();
+    let setup = r.venv_setup.to_vec();
+    assert_eq!(setup.len(), 2);
+    assert_eq!(setup[0], "uv venv ~/venvs/vllm025 -p 3.12");
+    assert_eq!(setup[1], "uv pip install vllm==0.25.1");
+}
+
 #[test]
 fn test_venv_setup_defaults_to_empty() {
     let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
-    assert!(r.venv_setup.is_empty());
+    assert!(r.venv_setup.to_vec().is_empty());
 }
 
 #[test]

+ 24 - 4
src/config/types.rs

@@ -3,6 +3,7 @@ use std::collections::HashMap;
 
 pub use crate::pre_serve::{
     BeforeServeTable, DistributedBackend, ModsTable, OnFailure, PreServeCommand, RecipeArgs,
+    VenvSetup,
 };
 
 #[derive(Clone, Debug, Deserialize)]
@@ -169,14 +170,24 @@ pub struct Recipe {
     #[serde(default)]
     pub venv: Option<String>,
     /// Shell commands to create and initialize the venv at `venv` if it
-    /// does not already exist on a node. Each command runs as-is over SSH
-    /// (shell expands `~` and `$VAR`). Only used in tmux deployment mode.
+    /// does not already exist on a node. Either an inline multi-line string
+    /// (one command per line; blank lines and `#` comments are skipped) or
+    /// a `Vec<String>` list. Only used in tmux deployment mode.
     ///
     /// The launcher checks `test -f <venv>/bin/activate` on each node
     /// before serve; if absent, it runs these commands in order, then
     /// re-checks. A failure aborts the serve.
     ///
-    /// Example:
+    /// Example (inline):
+    /// ```toml
+    /// venv = "~/venvs/vllm025"
+    /// venv_setup = """
+    ///     uv venv ~/venvs/vllm025 -p 3.12
+    ///     uv pip install vllm==0.25.1 --torch-backend=cu130
+    /// """
+    /// ```
+    ///
+    /// Example (list):
     /// ```toml
     /// venv = "~/venvs/vllm025"
     /// venv_setup = [
@@ -185,7 +196,16 @@ pub struct Recipe {
     /// ]
     /// ```
     #[serde(default)]
-    pub venv_setup: Vec<String>,
+    pub venv_setup: VenvSetup,
+    /// Path to a shell script (relative to the recipe directory) that creates
+    /// and initializes the venv. When set, the launcher reads the file from
+    /// `recipes/<recipe_name>/<venv_setup_script>` and runs it on each node
+    /// if the venv doesn't exist. Mutually exclusive with `venv_setup` — if
+    /// both are set, `venv_setup_script` takes precedence.
+    ///
+    /// Example: `venv_setup_script = "venv_setup.sh"`
+    #[serde(default)]
+    pub venv_setup_script: Option<String>,
     #[serde(default)]
     pub mods: ModsTable,
     #[serde(default = "default_sync_image")]

+ 3 - 1
src/docker/command_tests.rs

@@ -1,6 +1,7 @@
 use super::*;
 use crate::config::{
     Cluster, DeployTarget, DistributedBackend, ExtraMount, NcclConfig, Node, Recipe, RecipeArgs,
+    VenvSetup,
 };
 use std::collections::HashMap;
 
@@ -90,7 +91,8 @@ fn base_recipe() -> Recipe {
         solo: None,
         deploy: DeployTarget::Docker,
         venv: None,
-        venv_setup: vec![],
+        venv_setup: VenvSetup::default(),
+        venv_setup_script: None,
         mods: Default::default(),
         shm_size: String::new(),
         sync_image: true,

+ 5 - 3
src/docker/tests.rs

@@ -1,6 +1,6 @@
 use super::command::{DockerCommandCtx, build_docker_command};
 use super::*;
-use crate::config::{DeployTarget, DistributedBackend, RecipeArgs};
+use crate::config::{DeployTarget, DistributedBackend, RecipeArgs, VenvSetup};
 use std::collections::HashMap;
 
 fn tcls() -> Cluster {
@@ -63,7 +63,8 @@ fn trec() -> Recipe {
         solo: None,
         deploy: DeployTarget::Docker,
         venv: None,
-        venv_setup: vec![],
+        venv_setup: VenvSetup::default(),
+        venv_setup_script: None,
         mods: Default::default(),
         shm_size: String::new(),
         sync_image: true,
@@ -255,7 +256,8 @@ fn sglang_recipe() -> Recipe {
         solo: None,
         deploy: DeployTarget::Docker,
         venv: None,
-        venv_setup: vec![],
+        venv_setup: VenvSetup::default(),
+        venv_setup_script: None,
         mods: Default::default(),
         shm_size: String::new(),
         sync_image: true,

+ 3 - 1
src/doctor.rs

@@ -348,6 +348,7 @@ pub async fn try_repaired_image(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::config::VenvSetup;
     use std::collections::HashMap;
 
     fn test_recipe() -> Recipe {
@@ -373,7 +374,8 @@ mod tests {
             solo: None,
             deploy: crate::config::DeployTarget::Docker,
             venv: None,
-            venv_setup: vec![],
+            venv_setup: VenvSetup::default(),
+            venv_setup_script: None,
             mods: Default::default(),
             sync_image: true,
             docker_image_dir: None,

+ 3 - 2
src/engine_flags.rs

@@ -99,7 +99,7 @@ pub fn build_engine_args(ctx: &EngineFlagsCtx<'_>, node: &Node) -> Vec<String> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::config::{DeployTarget, RecipeArgs};
+    use crate::config::{DeployTarget, RecipeArgs, VenvSetup};
     use std::collections::HashMap;
 
     fn base_recipe() -> Recipe {
@@ -122,7 +122,8 @@ mod tests {
             solo: None,
             deploy: DeployTarget::Docker,
             venv: None,
-            venv_setup: vec![],
+            venv_setup: VenvSetup::default(),
+            venv_setup_script: None,
             mods: Default::default(),
             shm_size: String::new(),
             sync_image: true,

+ 47 - 0
src/pre_serve.rs

@@ -132,6 +132,53 @@ impl RecipeArgs {
     }
 }
 
+/// Venv setup commands — either an inline multi-line TOML string (one
+/// command per line) or a `Vec<String>` list. Accepted in TOML as either:
+///
+/// ```toml
+/// venv_setup = """
+///   uv venv ~/venvs/vllm025 -p 3.12
+///   uv pip install vllm==0.25.1 --torch-backend=cu130
+/// """
+/// ```
+///
+/// ```toml
+/// venv_setup = [
+///     "uv venv ~/venvs/vllm025 -p 3.12",
+///     "uv pip install vllm==0.25.1 --torch-backend=cu130",
+/// ]
+/// ```
+///
+/// Blank lines and lines starting with `#` are skipped (shell comments).
+#[derive(Clone, Debug, Deserialize)]
+#[serde(untagged)]
+pub enum VenvSetup {
+    List(Vec<String>),
+    Inline(String),
+}
+
+impl Default for VenvSetup {
+    fn default() -> Self {
+        VenvSetup::List(Vec::new())
+    }
+}
+
+impl VenvSetup {
+    /// Expand to a list of commands. For `Inline`, each non-blank,
+    /// non-comment line becomes one command. For `List`, each entry is
+    /// a command as-is.
+    pub fn to_vec(&self) -> Vec<String> {
+        match self {
+            VenvSetup::List(v) => v.clone(),
+            VenvSetup::Inline(s) => s
+                .lines()
+                .map(|l| l.trim().to_string())
+                .filter(|l| !l.is_empty() && !l.starts_with('#'))
+                .collect(),
+        }
+    }
+}
+
 impl std::fmt::Display for DistributedBackend {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {

+ 3 - 2
src/tmux/commands.rs

@@ -132,7 +132,7 @@ impl TmuxManager {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::config::{DeployTarget, DistributedBackend, RecipeArgs};
+    use crate::config::{DeployTarget, DistributedBackend, RecipeArgs, VenvSetup};
     use std::collections::HashMap;
 
     fn trecipe() -> Recipe {
@@ -155,7 +155,8 @@ mod tests {
             solo: None,
             deploy: DeployTarget::Docker,
             venv: None,
-            venv_setup: vec![],
+            venv_setup: VenvSetup::default(),
+            venv_setup_script: None,
             mods: Default::default(),
             shm_size: String::new(),
             sync_image: true,

+ 33 - 6
src/tmux/serve.rs

@@ -69,7 +69,7 @@ pub async fn tmux_serve(
     crate::before_serve::run(config, &cluster, &recipe, skip_before_serve, dry_run).await?;
 
     // Prepare the Python venv on each node if it doesn't exist yet.
-    prepare_venv(config, &recipe, &cluster, dry_run).await?;
+    prepare_venv(config, &recipe, &cluster, recipe_name, dry_run).await?;
 
     let manager = TmuxManager;
 
@@ -230,19 +230,46 @@ async fn wait_for_head_ready(
 /// Prepares the Python venv on each node if it doesn't exist.
 ///
 /// Checks `test -f <venv>/bin/activate` on each node; if absent, runs
-/// `recipe.venv_setup` commands in order, then re-checks. A failure on
-/// any command aborts the serve. Nodes are prepared in parallel.
+/// the setup commands in order, then re-checks. The setup commands come
+/// from either `recipe.venv_setup_script` (a script file read from the
+/// recipe directory) or `recipe.venv_setup` (inline list or multi-line
+/// string). A failure on any command aborts the serve. Nodes are
+/// prepared in parallel.
 async fn prepare_venv(
     config: &Config,
     recipe: &Recipe,
     cluster: &crate::config::Cluster,
+    recipe_name: &str,
     dry_run: bool,
 ) -> Result<()> {
     let venv = match &recipe.venv {
-        Some(v) if !recipe.venv_setup.is_empty() => v,
-        _ => return Ok(()),
+        Some(v) => v,
+        None => return Ok(()),
+    };
+
+    // Resolve setup commands: script file takes precedence over inline.
+    let setup: Vec<String> = if let Some(script_path) = &recipe.venv_setup_script {
+        let full_path = format!("recipes/{}/{}", recipe_name, script_path);
+        if dry_run {
+            info!("  read venv_setup_script: {} (dry-run)", full_path);
+        } else {
+            info!("  read venv_setup_script: {}", full_path);
+        }
+        let body = tokio::fs::read_to_string(&full_path)
+            .await
+            .map_err(|e| anyhow::anyhow!("Cannot read venv_setup_script '{}': {}", full_path, e))?;
+        body.lines()
+            .map(|l| l.trim().to_string())
+            .filter(|l| !l.is_empty() && !l.starts_with('#'))
+            .collect()
+    } else {
+        recipe.venv_setup.to_vec()
     };
 
+    if setup.is_empty() {
+        return Ok(());
+    }
+
     info!("{}", fluent_i18n::t!("venv-preparing", { "venv" => venv }));
 
     let mut handles = Vec::with_capacity(cluster.nodes.len());
@@ -250,7 +277,7 @@ async fn prepare_venv(
         let config = config.clone();
         let node = node.clone();
         let venv = venv.clone();
-        let setup = recipe.venv_setup.clone();
+        let setup = setup.clone();
         let dry = dry_run;
 
         handles.push(crate::locale::spawn_with_locale(async move {