|
|
@@ -1,8 +1,18 @@
|
|
|
-use crate::config::{Cluster, DistributedBackend, NcclConfig, Node, Recipe};
|
|
|
+use crate::config::{Cluster, Node, Recipe};
|
|
|
+use crate::engine_flags::{EngineFlagsCtx, build_engine_args};
|
|
|
use crate::tmux::TmuxManager;
|
|
|
|
|
|
impl TmuxManager {
|
|
|
- /// Builds the full model launch command for a given node.
|
|
|
+ /// Build the engine launch command for a node.
|
|
|
+ ///
|
|
|
+ /// Uses `recipe.entrypoint` when present (e.g. `["vllm", "serve"]` →
|
|
|
+ /// `vllm serve`), falling back to the hardcoded Python module path
|
|
|
+ /// when the entrypoint is empty. This lets the same recipe work
|
|
|
+ /// under both Docker and tmux deployment.
|
|
|
+ ///
|
|
|
+ /// Appends shared engine args (model, recipe args, distributed flags)
|
|
|
+ /// from [`crate::engine_flags::build_engine_args`], plus the port
|
|
|
+ /// flag (head node only, or all nodes for SGLang).
|
|
|
pub fn build_launch_command(&self, recipe: &Recipe, node: &Node, cluster: &Cluster) -> String {
|
|
|
let nnodes = cluster.nodes.len() as u32;
|
|
|
let pp_size = recipe.pipeline_parallel_size.unwrap_or(nnodes);
|
|
|
@@ -10,78 +20,35 @@ impl TmuxManager {
|
|
|
|
|
|
let mut parts: Vec<String> = Vec::new();
|
|
|
|
|
|
- // Engine entrypoint
|
|
|
- match recipe.engine.as_str() {
|
|
|
- "vllm" => {
|
|
|
- parts.push("python -m vllm.entrypoints.openai.api_server".to_string());
|
|
|
- }
|
|
|
- "sglang" => {
|
|
|
- parts.push("python -m sglang.launch_server".to_string());
|
|
|
- }
|
|
|
- other => {
|
|
|
- parts.push(format!("python -m {}", other));
|
|
|
- }
|
|
|
- }
|
|
|
-
|
|
|
- // Model — engine-specific
|
|
|
- if is_sglang {
|
|
|
- parts.push("--model-path".to_string());
|
|
|
- }
|
|
|
- parts.push(recipe.model.clone());
|
|
|
-
|
|
|
- // Recipe-specific args
|
|
|
- parts.extend(recipe.args.to_vec());
|
|
|
-
|
|
|
- // Distributed flags — engine-specific
|
|
|
- if is_sglang {
|
|
|
- parts.extend(vec![
|
|
|
- "--tp".to_string(),
|
|
|
- recipe.tensor_parallel_size.to_string(),
|
|
|
- "--nnodes".to_string(),
|
|
|
- nnodes.to_string(),
|
|
|
- "--node-rank".to_string(),
|
|
|
- node.rank.to_string(),
|
|
|
- "--master-addr".to_string(),
|
|
|
- cluster.master_ib_ip.clone(),
|
|
|
- ]);
|
|
|
+ // Entrypoint — from recipe when present, else hardcoded Python path.
|
|
|
+ if !recipe.entrypoint.is_empty() {
|
|
|
+ parts.extend(recipe.entrypoint.clone());
|
|
|
} else {
|
|
|
- match recipe.distributed_backend {
|
|
|
- Some(DistributedBackend::Ray) => {
|
|
|
- parts.extend(vec![
|
|
|
- "--tensor-parallel-size".to_string(),
|
|
|
- "1".to_string(),
|
|
|
- "--pipeline-parallel-size".to_string(),
|
|
|
- pp_size.to_string(),
|
|
|
- "--nnodes".to_string(),
|
|
|
- nnodes.to_string(),
|
|
|
- "--node-rank".to_string(),
|
|
|
- node.rank.to_string(),
|
|
|
- "--master-addr".to_string(),
|
|
|
- cluster.master_ib_ip.clone(),
|
|
|
- "--distributed-executor-backend".to_string(),
|
|
|
- "ray".to_string(),
|
|
|
- ]);
|
|
|
+ match recipe.engine.as_str() {
|
|
|
+ "vllm" => {
|
|
|
+ parts.push("python -m vllm.entrypoints.openai.api_server".to_string());
|
|
|
}
|
|
|
- _ => {
|
|
|
- parts.extend(vec![
|
|
|
- "--tensor-parallel-size".to_string(),
|
|
|
- "1".to_string(),
|
|
|
- "--pipeline-parallel-size".to_string(),
|
|
|
- pp_size.to_string(),
|
|
|
- "--nnodes".to_string(),
|
|
|
- nnodes.to_string(),
|
|
|
- "--node-rank".to_string(),
|
|
|
- node.rank.to_string(),
|
|
|
- "--master-addr".to_string(),
|
|
|
- cluster.master_ib_ip.clone(),
|
|
|
- "--distributed-executor-backend".to_string(),
|
|
|
- "mp".to_string(),
|
|
|
- ]);
|
|
|
+ "sglang" => {
|
|
|
+ parts.push("python -m sglang.launch_server".to_string());
|
|
|
+ }
|
|
|
+ other => {
|
|
|
+ parts.push(format!("python -m {}", other));
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
- // Port — SGLang needs it on all nodes, vLLM only on head node
|
|
|
+ // Shared engine args: model, recipe args, distributed flags, headless.
|
|
|
+ let ectx = EngineFlagsCtx {
|
|
|
+ recipe,
|
|
|
+ tp_size: recipe.tensor_parallel_size,
|
|
|
+ pp_size,
|
|
|
+ nnodes,
|
|
|
+ init_port: 25000,
|
|
|
+ master_ib_ip: &cluster.master_ib_ip,
|
|
|
+ };
|
|
|
+ parts.extend(build_engine_args(&ectx, node));
|
|
|
+
|
|
|
+ // Port — SGLang needs it on all nodes, vLLM only on head node.
|
|
|
if is_sglang || node.rank == 0 {
|
|
|
parts.extend(vec!["--port".to_string(), recipe.port.to_string()]);
|
|
|
}
|
|
|
@@ -94,7 +61,7 @@ impl TmuxManager {
|
|
|
&self,
|
|
|
node: &Node,
|
|
|
recipe: &Recipe,
|
|
|
- nccl: &NcclConfig,
|
|
|
+ nccl: &crate::config::NcclConfig,
|
|
|
cluster: &Cluster,
|
|
|
) -> Vec<String> {
|
|
|
let sess = self.session_name(recipe, node.rank);
|
|
|
@@ -113,7 +80,9 @@ 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 == Some(DistributedBackend::Ray) && node.rank > 0 {
|
|
|
+ if recipe.distributed_backend == Some(crate::config::DistributedBackend::Ray)
|
|
|
+ && node.rank > 0
|
|
|
+ {
|
|
|
cmds.push(format!(
|
|
|
"tmux send-keys -t {}:0 'export RAY_ADDRESS={}:6379' Enter",
|
|
|
sess, cluster.master_ib_ip
|
|
|
@@ -150,3 +119,108 @@ impl TmuxManager {
|
|
|
cmds
|
|
|
}
|
|
|
}
|
|
|
+
|
|
|
+#[cfg(test)]
|
|
|
+mod tests {
|
|
|
+ use super::*;
|
|
|
+ use crate::config::{DeployTarget, DistributedBackend, RecipeArgs};
|
|
|
+ use std::collections::HashMap;
|
|
|
+
|
|
|
+ fn trecipe() -> Recipe {
|
|
|
+ Recipe {
|
|
|
+ engine: "vllm".into(),
|
|
|
+ name: "T".into(),
|
|
|
+ docker_image: "img".into(),
|
|
|
+ model: "org/model".into(),
|
|
|
+ port: 8000,
|
|
|
+ required_nodes: None,
|
|
|
+ pipeline_parallel_size: Some(1),
|
|
|
+ tensor_parallel_size: 1,
|
|
|
+ args: RecipeArgs::List(vec!["--x".into(), "4096".into()]),
|
|
|
+ env: HashMap::new(),
|
|
|
+ cluster: None,
|
|
|
+ entrypoint: vec!["vllm".into(), "serve".into()],
|
|
|
+ distributed_backend: Some(DistributedBackend::Mp),
|
|
|
+ before_serve: vec![],
|
|
|
+ after_stop: vec![],
|
|
|
+ solo: None,
|
|
|
+ deploy: DeployTarget::Docker,
|
|
|
+ mods: Default::default(),
|
|
|
+ shm_size: String::new(),
|
|
|
+ sync_image: true,
|
|
|
+ docker_image_dir: None,
|
|
|
+ extra_mounts: vec![],
|
|
|
+ build: Default::default(),
|
|
|
+ composed_model: None,
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ fn tcluster() -> Cluster {
|
|
|
+ Cluster {
|
|
|
+ name: "c".into(),
|
|
|
+ master_ib_ip: "10.0.0.1".into(),
|
|
|
+ nodes: vec![Node {
|
|
|
+ name: "n0".into(),
|
|
|
+ ssh_ip: "10.0.0.1".into(),
|
|
|
+ ib_ip: "10.0.0.1".into(),
|
|
|
+ rank: 0,
|
|
|
+ host_key: None,
|
|
|
+ entity: None,
|
|
|
+ }],
|
|
|
+ extra_entities: vec![],
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn entrypoint_used_when_present() {
|
|
|
+ let r = trecipe();
|
|
|
+ let c = tcluster();
|
|
|
+ let mgr = TmuxManager;
|
|
|
+ let cmd = mgr.build_launch_command(&r, &c.nodes[0], &c);
|
|
|
+ // Should start with "vllm serve" (from entrypoint), not "python -m"
|
|
|
+ assert!(cmd.starts_with("vllm serve "), "cmd: {cmd}");
|
|
|
+ // Should include recipe args
|
|
|
+ assert!(cmd.contains("--x 4096"));
|
|
|
+ }
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn entrypoint_fallback_when_empty() {
|
|
|
+ let mut r = trecipe();
|
|
|
+ r.entrypoint = vec![];
|
|
|
+ let c = tcluster();
|
|
|
+ let mgr = TmuxManager;
|
|
|
+ let cmd = mgr.build_launch_command(&r, &c.nodes[0], &c);
|
|
|
+ // Should fall back to hardcoded python module path
|
|
|
+ assert!(
|
|
|
+ cmd.starts_with("python -m vllm.entrypoints.openai.api_server "),
|
|
|
+ "cmd: {cmd}"
|
|
|
+ );
|
|
|
+ }
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn sglang_entrypoint_fallback() {
|
|
|
+ let mut r = trecipe();
|
|
|
+ r.entrypoint = vec![];
|
|
|
+ r.engine = "sglang".into();
|
|
|
+ let c = tcluster();
|
|
|
+ let mgr = TmuxManager;
|
|
|
+ let cmd = mgr.build_launch_command(&r, &c.nodes[0], &c);
|
|
|
+ assert!(
|
|
|
+ cmd.starts_with("python -m sglang.launch_server "),
|
|
|
+ "cmd: {cmd}"
|
|
|
+ );
|
|
|
+ }
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn single_node_no_distributed_flags() {
|
|
|
+ let r = trecipe();
|
|
|
+ let c = tcluster();
|
|
|
+ let mgr = TmuxManager;
|
|
|
+ let cmd = mgr.build_launch_command(&r, &c.nodes[0], &c);
|
|
|
+ // Single node → no --nnodes, --master-addr, etc.
|
|
|
+ assert!(!cmd.contains("--nnodes"));
|
|
|
+ assert!(!cmd.contains("--master-addr"));
|
|
|
+ // But port should be present (head node)
|
|
|
+ assert!(cmd.contains("--port 8000"));
|
|
|
+ }
|
|
|
+}
|