Jelajahi Sumber

feat(recipe): add venv field for declarative tmux venv activation

Alexander Korolev 2 bulan lalu
induk
melakukan
d9b922e01b

+ 1 - 0
AGENTS.md

@@ -51,6 +51,7 @@ port = 30000                        # default
 required_nodes = 2                  # optional; also caps the launch set (N nodes, not all)
 solo = false                       # optional; single-node on first free node, no dist flags
 deploy = "docker"                   # optional; "docker" (default) or "tmux"
+venv = "~/venvs/vllm025"            # optional; Python venv for tmux (activates before entrypoint)
 tensor_parallel_size = 2
 pipeline_parallel_size = 1          # vLLM only, 0 = auto
 distributed_backend = "mp"          # or "ray"

+ 3 - 2
recipes/laguna-s2.1-nvfp4/recipe.toml

@@ -30,8 +30,8 @@
 #   hf download poolside/Laguna-S-2.1-NVFP4
 #   hf download poolside/Laguna-S-2.1-DFlash-NVFP4
 #
-# Then ensure `vllm` is on PATH (activate the venv or add its bin/ to PATH).
-# The entrypoint below (`["vllm", "serve"]`) is used directly by tmux.
+# The `venv` field below is used by tmux to activate the venv before
+# running the entrypoint.
 
 engine = "vllm"
 name = "Laguna S 2.1 NVFP4 (vLLM)"
@@ -40,6 +40,7 @@ model = "poolside/Laguna-S-2.1-NVFP4"
 tensor_parallel_size = 1
 pipeline_parallel_size = 1
 deploy = "tmux"
+venv = "~/venvs/vllm025"
 
 entrypoint = ["vllm", "serve"]
 

+ 1 - 0
src/after_stop.rs

@@ -199,6 +199,7 @@ mod tests {
             after_stop,
             solo: None,
             deploy: crate::config::DeployTarget::Docker,
+            venv: None,
             mods: Default::default(),
             sync_image: true,
             docker_image_dir: None,

+ 1 - 0
src/before_serve.rs

@@ -309,6 +309,7 @@ mod tests {
             after_stop: vec![],
             solo: None,
             deploy: crate::config::DeployTarget::Docker,
+            venv: None,
             mods: Default::default(),
             sync_image: true,
             docker_image_dir: None,

+ 16 - 0
src/config/tests.rs

@@ -210,6 +210,22 @@ fn test_deploy_explicit_docker() {
     assert_eq!(r.deploy, DeployTarget::Docker);
 }
 
+#[test]
+fn test_venv_field() {
+    let r: Recipe = toml::from_str(
+        "name = \"x\"\nmodel = \"m\"\ndeploy = \"tmux\"\nvenv = \"~/venvs/vllm025\"\n",
+    )
+    .unwrap();
+    assert_eq!(r.deploy, DeployTarget::Tmux);
+    assert_eq!(r.venv.as_deref(), Some("~/venvs/vllm025"));
+}
+
+#[test]
+fn test_venv_defaults_to_none() {
+    let r: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"\n").unwrap();
+    assert_eq!(r.venv, None);
+}
+
 #[test]
 fn default_config_dir_appends_sm121() {
     let Some(d) = default_config_dir() else {

+ 9 - 0
src/config/types.rs

@@ -159,6 +159,15 @@ pub struct Recipe {
     /// sync) are rejected in tmux mode.
     #[serde(default)]
     pub deploy: DeployTarget,
+    /// Python virtualenv path for tmux deployment. When set, the tmux
+    /// launch command prepends `source <venv>/bin/activate &&` before
+    /// the entrypoint, so the engine runs from the specified venv instead
+    /// of the system Python. Docker-only (ignored in Docker mode — the
+    /// container has its own Python).
+    ///
+    /// Example: `venv = "~/venvs/vllm025"`
+    #[serde(default)]
+    pub venv: Option<String>,
     #[serde(default)]
     pub mods: ModsTable,
     #[serde(default = "default_sync_image")]

+ 2 - 2
src/docker/command.rs

@@ -4,12 +4,11 @@
 //! [`build_docker_command`] is the single entry point; everything else here is
 //! private to this module.
 
+use super::{mods, tail};
 use crate::config::{Cluster, Config, NcclConfig, Node, Recipe};
 use anyhow::Result;
 use std::borrow::Cow;
 
-use super::{mods, tail};
-
 /// Shared context for [`build_docker_command`] — constructed once, cloned per node.
 #[derive(Clone)]
 pub(crate) struct DockerCommandCtx {
@@ -386,6 +385,7 @@ mod tests {
             after_stop: vec![],
             solo: None,
             deploy: DeployTarget::Docker,
+            venv: None,
             mods: Default::default(),
             shm_size: String::new(),
             sync_image: true,

+ 2 - 0
src/docker/tests.rs

@@ -62,6 +62,7 @@ fn trec() -> Recipe {
         after_stop: vec![],
         solo: None,
         deploy: DeployTarget::Docker,
+        venv: None,
         mods: Default::default(),
         shm_size: String::new(),
         sync_image: true,
@@ -252,6 +253,7 @@ fn sglang_recipe() -> Recipe {
         after_stop: vec![],
         solo: None,
         deploy: DeployTarget::Docker,
+        venv: None,
         mods: Default::default(),
         shm_size: String::new(),
         sync_image: true,

+ 1 - 0
src/doctor.rs

@@ -372,6 +372,7 @@ mod tests {
             after_stop: vec![],
             solo: None,
             deploy: crate::config::DeployTarget::Docker,
+            venv: None,
             mods: Default::default(),
             sync_image: true,
             docker_image_dir: None,

+ 1 - 0
src/engine_flags.rs

@@ -121,6 +121,7 @@ mod tests {
             after_stop: vec![],
             solo: None,
             deploy: DeployTarget::Docker,
+            venv: None,
             mods: Default::default(),
             shm_size: String::new(),
             sync_image: true,

+ 34 - 1
src/tmux/commands.rs

@@ -20,6 +20,15 @@ impl TmuxManager {
 
         let mut parts: Vec<String> = Vec::new();
 
+        // Venv activation — prepend `source <venv>/bin/activate &&` when
+        // recipe.venv is set. The venv path is expanded for ~ by the
+        // caller (before_serve runs on the host via SSH, so ~ works).
+        let venv_prefix = if let Some(ref venv) = recipe.venv {
+            format!("source {}/bin/activate && ", venv)
+        } else {
+            String::new()
+        };
+
         // Entrypoint — from recipe when present, else hardcoded Python path.
         if !recipe.entrypoint.is_empty() {
             parts.extend(recipe.entrypoint.clone());
@@ -53,7 +62,7 @@ impl TmuxManager {
             parts.extend(vec!["--port".to_string(), recipe.port.to_string()]);
         }
 
-        parts.join(" ")
+        format!("{}{}", venv_prefix, parts.join(" "))
     }
 
     /// Builds the full sequence of tmux commands for a session, with multi-window layout.
@@ -145,6 +154,7 @@ mod tests {
             after_stop: vec![],
             solo: None,
             deploy: DeployTarget::Docker,
+            venv: None,
             mods: Default::default(),
             shm_size: String::new(),
             sync_image: true,
@@ -223,4 +233,27 @@ mod tests {
         // But port should be present (head node)
         assert!(cmd.contains("--port 8000"));
     }
+
+    #[test]
+    fn venv_prefix_prepended() {
+        let mut r = trecipe();
+        r.venv = Some("~/venvs/vllm025".into());
+        let c = tcluster();
+        let mgr = TmuxManager;
+        let cmd = mgr.build_launch_command(&r, &c.nodes[0], &c);
+        assert!(
+            cmd.starts_with("source ~/venvs/vllm025/bin/activate && vllm serve "),
+            "cmd: {cmd}"
+        );
+    }
+
+    #[test]
+    fn venv_prefix_absent_when_unset() {
+        let r = trecipe();
+        let c = tcluster();
+        let mgr = TmuxManager;
+        let cmd = mgr.build_launch_command(&r, &c.nodes[0], &c);
+        // No venv → command starts directly with entrypoint
+        assert!(cmd.starts_with("vllm serve "), "cmd: {cmd}");
+    }
 }