|
|
@@ -0,0 +1,338 @@
|
|
|
+//! First-class inference-engine drivers.
|
|
|
+//!
|
|
|
+//! A recipe's `engine` string selects an [`EngineDriver`] that encapsulates
|
|
|
+//! everything the launcher needs to know about that engine: how the model is
|
|
|
+//! passed, the distributed-rendezvous flags, whether workers run headless, the
|
|
|
+//! bind-IP env var, port placement, the tmux entrypoint fallback, and the
|
|
|
+//! doctor install hint. This replaces the `recipe.engine == "sglang"` branches
|
|
|
+//! that were previously duplicated across `engine_flags`, the docker env
|
|
|
+//! builder, the tmux command builder, doctor, and publish-render.
|
|
|
+//!
|
|
|
+//! `build_engine_args` is the shared argument builder used by both the Docker
|
|
|
+//! (`docker/tail.rs`) and tmux (`tmux/commands.rs`) deployment paths. It does
|
|
|
+//! **not** include the entrypoint, `--host`, or `--port` — those are
|
|
|
+//! deployment-specific and added by each path.
|
|
|
+
|
|
|
+use crate::config::{DistributedBackend, Node, Recipe};
|
|
|
+
|
|
|
+/// Context for [`build_engine_args`] — the subset of `DockerCommandCtx` /
|
|
|
+/// `TmuxManager` state that the shared flag builder needs.
|
|
|
+pub struct EngineFlagsCtx<'a> {
|
|
|
+ pub recipe: &'a Recipe,
|
|
|
+ pub tp_size: u32,
|
|
|
+ pub pp_size: u32,
|
|
|
+ pub nnodes: u32,
|
|
|
+ pub init_port: u16,
|
|
|
+ pub master_ib_ip: &'a str,
|
|
|
+}
|
|
|
+
|
|
|
+/// Per-engine launch semantics. See the module docs.
|
|
|
+pub trait EngineDriver: Send + Sync {
|
|
|
+ /// The model argument: positional (vLLM), `--model-path` (SGLang),
|
|
|
+ /// `--model` (gr8), or none (DGPP — the config file carries it). An empty
|
|
|
+ /// recipe model yields no model argument (entrypoint-driven images).
|
|
|
+ fn model_args(&self, recipe: &Recipe) -> Vec<String> {
|
|
|
+ if recipe.model.is_empty() {
|
|
|
+ Vec::new()
|
|
|
+ } else {
|
|
|
+ vec![recipe.model.clone()]
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ /// Distributed-rendezvous flags. Only called for multi-node distributed
|
|
|
+ /// runs (`nnodes > 1` with an Mp/Ray backend).
|
|
|
+ fn distributed_args(&self, ctx: &EngineFlagsCtx<'_>, node: &Node) -> Vec<String>;
|
|
|
+
|
|
|
+ /// Non-head workers run headless (vLLM yes; SGLang/gr8/DGPP no).
|
|
|
+ fn headless_workers(&self) -> bool {
|
|
|
+ false
|
|
|
+ }
|
|
|
+
|
|
|
+ /// Env var carrying the node's bind IP (`VLLM_HOST_IP` / `SGLANG_WORKER_IP`
|
|
|
+ /// / none).
|
|
|
+ fn engine_ip_env(&self) -> Option<&'static str> {
|
|
|
+ None
|
|
|
+ }
|
|
|
+
|
|
|
+ /// Emit the `--port` flag on every node (SGLang) vs the head only
|
|
|
+ /// (vLLM/gr8/DGPP).
|
|
|
+ fn port_on_all_nodes(&self) -> bool {
|
|
|
+ false
|
|
|
+ }
|
|
|
+
|
|
|
+ /// tmux entrypoint fallback when the recipe sets none.
|
|
|
+ fn entrypoint_default(&self) -> Vec<String>;
|
|
|
+
|
|
|
+ /// doctor install hint (empty = nothing to install, e.g. prebuilt images).
|
|
|
+ fn install_hint(&self) -> &'static str {
|
|
|
+ ""
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+/// vLLM: positional model, `--tensor-parallel-size`/`--nnodes`/`--master-addr`
|
|
|
+/// rendezvous, `--headless` workers, `VLLM_HOST_IP`.
|
|
|
+pub struct VllmDriver;
|
|
|
+
|
|
|
+impl EngineDriver for VllmDriver {
|
|
|
+ fn distributed_args(&self, ctx: &EngineFlagsCtx<'_>, node: &Node) -> Vec<String> {
|
|
|
+ let mut fp = vec![
|
|
|
+ "--tensor-parallel-size".into(),
|
|
|
+ ctx.tp_size.to_string(),
|
|
|
+ "--pipeline-parallel-size".into(),
|
|
|
+ ctx.pp_size.to_string(),
|
|
|
+ "--nnodes".into(),
|
|
|
+ ctx.nnodes.to_string(),
|
|
|
+ "--node-rank".into(),
|
|
|
+ node.rank.to_string(),
|
|
|
+ "--master-addr".into(),
|
|
|
+ ctx.master_ib_ip.to_string(),
|
|
|
+ "--master-port".into(),
|
|
|
+ ctx.init_port.to_string(),
|
|
|
+ ];
|
|
|
+ if ctx.recipe.distributed_backend == Some(DistributedBackend::Ray) {
|
|
|
+ fp.push("--distributed-executor-backend".into());
|
|
|
+ fp.push("ray".into());
|
|
|
+ }
|
|
|
+ fp
|
|
|
+ }
|
|
|
+
|
|
|
+ fn headless_workers(&self) -> bool {
|
|
|
+ true
|
|
|
+ }
|
|
|
+
|
|
|
+ fn engine_ip_env(&self) -> Option<&'static str> {
|
|
|
+ Some("VLLM_HOST_IP")
|
|
|
+ }
|
|
|
+
|
|
|
+ fn entrypoint_default(&self) -> Vec<String> {
|
|
|
+ vec!["python -m vllm.entrypoints.openai.api_server".into()]
|
|
|
+ }
|
|
|
+
|
|
|
+ fn install_hint(&self) -> &'static str {
|
|
|
+ "pip install vllm"
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+/// SGLang: `--model-path`, `--tp`/`--nnodes`/`--dist-init-addr` rendezvous,
|
|
|
+/// no headless (port on every node), `SGLANG_WORKER_IP`.
|
|
|
+pub struct SglangDriver;
|
|
|
+
|
|
|
+impl EngineDriver for SglangDriver {
|
|
|
+ fn model_args(&self, recipe: &Recipe) -> Vec<String> {
|
|
|
+ if recipe.model.is_empty() {
|
|
|
+ Vec::new()
|
|
|
+ } else {
|
|
|
+ vec!["--model-path".into(), recipe.model.clone()]
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ fn distributed_args(&self, ctx: &EngineFlagsCtx<'_>, node: &Node) -> Vec<String> {
|
|
|
+ vec![
|
|
|
+ "--tp".into(),
|
|
|
+ ctx.tp_size.to_string(),
|
|
|
+ "--nnodes".into(),
|
|
|
+ ctx.nnodes.to_string(),
|
|
|
+ "--node-rank".into(),
|
|
|
+ node.rank.to_string(),
|
|
|
+ "--dist-init-addr".into(),
|
|
|
+ format!("{}:{}", ctx.master_ib_ip, ctx.init_port),
|
|
|
+ ]
|
|
|
+ }
|
|
|
+
|
|
|
+ fn engine_ip_env(&self) -> Option<&'static str> {
|
|
|
+ Some("SGLANG_WORKER_IP")
|
|
|
+ }
|
|
|
+
|
|
|
+ fn port_on_all_nodes(&self) -> bool {
|
|
|
+ true
|
|
|
+ }
|
|
|
+
|
|
|
+ fn entrypoint_default(&self) -> Vec<String> {
|
|
|
+ vec!["python -m sglang.launch_server".into()]
|
|
|
+ }
|
|
|
+
|
|
|
+ fn install_hint(&self) -> &'static str {
|
|
|
+ "pip install sglang"
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+/// Resolve the driver for a recipe's `engine` string. Unknown engines fall
|
|
|
+/// back to the vLLM driver (matching the historical default: an engine with no
|
|
|
+/// dedicated driver still gets the vLLM-style tail).
|
|
|
+pub fn driver_for(engine: &str) -> &'static dyn EngineDriver {
|
|
|
+ match engine {
|
|
|
+ "sglang" => &SglangDriver,
|
|
|
+ "vllm" => &VllmDriver,
|
|
|
+ _ => &VllmDriver,
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+/// Build the engine-specific arguments for a recipe: model path, recipe args,
|
|
|
+/// and distributed-rendezvous flags. Delegates the engine-specific pieces to
|
|
|
+/// the [`EngineDriver`].
|
|
|
+///
|
|
|
+/// Does **not** include the entrypoint, `--host`, or `--port` — those are
|
|
|
+/// deployment-specific and added by each path:
|
|
|
+/// - Docker (`tail.rs`): always adds `--host 0.0.0.0 --port <port>`.
|
|
|
+/// - tmux (`commands.rs`): adds `--port <port>` on the head node (or all
|
|
|
+/// nodes for engines whose driver sets `port_on_all_nodes`).
|
|
|
+///
|
|
|
+/// Single-node runs emit **no** distributed flags — the engine defaults to
|
|
|
+/// single-node. Multi-node runs get the full rendezvous flag set.
|
|
|
+pub fn build_engine_args(ctx: &EngineFlagsCtx<'_>, node: &Node) -> Vec<String> {
|
|
|
+ let driver = driver_for(&ctx.recipe.engine);
|
|
|
+ let is_distributed = matches!(
|
|
|
+ ctx.recipe.distributed_backend,
|
|
|
+ Some(DistributedBackend::Mp | DistributedBackend::Ray)
|
|
|
+ );
|
|
|
+
|
|
|
+ let mut args = driver.model_args(ctx.recipe);
|
|
|
+
|
|
|
+ // Recipe-specific args (e.g. --kv-cache-dtype, --max-model-len).
|
|
|
+ args.extend(ctx.recipe.args.to_vec());
|
|
|
+
|
|
|
+ // Distributed flags — only for multi-node distributed runs.
|
|
|
+ if ctx.nnodes > 1 && is_distributed {
|
|
|
+ args.extend(driver.distributed_args(ctx, node));
|
|
|
+ }
|
|
|
+
|
|
|
+ // Headless workers (vLLM) in distributed mode.
|
|
|
+ if ctx.nnodes > 1 && driver.headless_workers() && node.rank > 0 && is_distributed {
|
|
|
+ args.push("--headless".into());
|
|
|
+ }
|
|
|
+
|
|
|
+ args
|
|
|
+}
|
|
|
+
|
|
|
+#[cfg(test)]
|
|
|
+mod tests {
|
|
|
+ use super::*;
|
|
|
+ use crate::config::DeployTarget;
|
|
|
+ use crate::pre_serve::{RecipeArgs, VenvSetup};
|
|
|
+ use std::collections::HashMap;
|
|
|
+
|
|
|
+ pub(crate) fn base_recipe() -> Recipe {
|
|
|
+ Recipe {
|
|
|
+ build_profile: None,
|
|
|
+ parent_model: None,
|
|
|
+ archived: false,
|
|
|
+ engine: "vllm".into(),
|
|
|
+ name: "T".into(),
|
|
|
+ docker_image: "img".into(),
|
|
|
+ model: "mdl".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,
|
|
|
+ venv: None,
|
|
|
+ venv_setup: VenvSetup::default(),
|
|
|
+ venv_setup_script: None,
|
|
|
+ mods: Default::default(),
|
|
|
+ shm_size: String::new(),
|
|
|
+ devices: vec![],
|
|
|
+ memory_limit: None,
|
|
|
+ cpus: None,
|
|
|
+ sync_image: true,
|
|
|
+ docker_image_dir: None,
|
|
|
+ extra_mounts: vec![],
|
|
|
+ extra_mounts_manifest: None,
|
|
|
+ build: Default::default(),
|
|
|
+ composed_model: None,
|
|
|
+ staging: vec![],
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ pub(crate) fn tnode(rank: u32) -> Node {
|
|
|
+ Node {
|
|
|
+ name: format!("n{}", rank),
|
|
|
+ ssh_ip: "10.0.0.1".into(),
|
|
|
+ ib_ip: Some("10.0.0.1".into()),
|
|
|
+ rank,
|
|
|
+ host_key: None,
|
|
|
+ entity: None,
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ fn ctx<'a>(r: &'a Recipe, tp: u32, nnodes: u32) -> EngineFlagsCtx<'a> {
|
|
|
+ EngineFlagsCtx {
|
|
|
+ recipe: r,
|
|
|
+ tp_size: tp,
|
|
|
+ pp_size: 1,
|
|
|
+ nnodes,
|
|
|
+ init_port: 25000,
|
|
|
+ master_ib_ip: "10.0.0.1",
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn single_node_no_distributed_flags() {
|
|
|
+ let r = base_recipe();
|
|
|
+ let args = build_engine_args(&ctx(&r, 1, 1), &tnode(0));
|
|
|
+ assert_eq!(args, vec!["mdl", "--x", "4096"]);
|
|
|
+ }
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn multi_node_vllm_mp_flags() {
|
|
|
+ let r = base_recipe();
|
|
|
+ let c = ctx(&r, 2, 2);
|
|
|
+ let head_args = build_engine_args(&c, &tnode(0));
|
|
|
+ assert!(head_args.contains(&"--tensor-parallel-size".into()));
|
|
|
+ assert!(head_args.contains(&"2".into()));
|
|
|
+ assert!(head_args.contains(&"--nnodes".into()));
|
|
|
+ assert!(head_args.contains(&"--master-addr".into()));
|
|
|
+ assert!(head_args.contains(&"10.0.0.1".into()));
|
|
|
+ assert!(head_args.contains(&"--master-port".into()));
|
|
|
+ assert!(head_args.contains(&"25000".into()));
|
|
|
+ assert!(!head_args.contains(&"--headless".into()));
|
|
|
+
|
|
|
+ let worker_args = build_engine_args(&c, &tnode(1));
|
|
|
+ assert!(worker_args.contains(&"--headless".into()));
|
|
|
+ }
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn multi_node_sglang_flags() {
|
|
|
+ let mut r = base_recipe();
|
|
|
+ r.engine = "sglang".into();
|
|
|
+ let c = ctx(&r, 2, 2);
|
|
|
+ let args = build_engine_args(&c, &tnode(0));
|
|
|
+ assert!(args.contains(&"--model-path".into()));
|
|
|
+ assert!(args.contains(&"--tp".into()));
|
|
|
+ assert!(args.contains(&"--dist-init-addr".into()));
|
|
|
+ assert!(args.contains(&"10.0.0.1:25000".into()));
|
|
|
+ assert!(!args.contains(&"--headless".into()));
|
|
|
+ }
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn ray_backend_adds_executor() {
|
|
|
+ let mut r = base_recipe();
|
|
|
+ r.distributed_backend = Some(DistributedBackend::Ray);
|
|
|
+ let c = ctx(&r, 1, 2);
|
|
|
+ let args = build_engine_args(&c, &tnode(0));
|
|
|
+ assert!(args.contains(&"--distributed-executor-backend".into()));
|
|
|
+ assert!(args.contains(&"ray".into()));
|
|
|
+ }
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn driver_registry_maps_engines() {
|
|
|
+ // vLLM: headless workers, no port-on-all-nodes.
|
|
|
+ assert!(driver_for("vllm").headless_workers());
|
|
|
+ assert!(!driver_for("vllm").port_on_all_nodes());
|
|
|
+ // SGLang: port on every node, --model-path model arg.
|
|
|
+ assert!(driver_for("sglang").port_on_all_nodes());
|
|
|
+ assert_eq!(
|
|
|
+ driver_for("sglang").model_args(&base_recipe()),
|
|
|
+ vec!["--model-path".to_string(), "mdl".to_string()]
|
|
|
+ );
|
|
|
+ // Unknown engines fall back to the vLLM driver (historical default).
|
|
|
+ assert!(driver_for("something-new").headless_workers());
|
|
|
+ }
|
|
|
+}
|