|
|
@@ -13,6 +13,10 @@
|
|
|
//! **not** include the entrypoint, `--host`, or `--port` — those are
|
|
|
//! deployment-specific and added by each path.
|
|
|
|
|
|
+mod drivers;
|
|
|
+
|
|
|
+pub use drivers::{DgppDriver, Gr8Driver, SglangDriver, TensorFoldDriver, VllmDriver};
|
|
|
+
|
|
|
use crate::config::{DistributedBackend, Node, Recipe};
|
|
|
|
|
|
/// Context for [`build_engine_args`] — the subset of `DockerCommandCtx` /
|
|
|
@@ -79,154 +83,6 @@ pub trait EngineDriver: Send + Sync {
|
|
|
}
|
|
|
}
|
|
|
|
|
|
-/// 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"
|
|
|
- }
|
|
|
-}
|
|
|
-
|
|
|
-/// gr8 (`kindlingai/gr8`): a from-scratch CUDA engine for full GLM-5.3,
|
|
|
-/// tensor-parallel across 3/4/6 Sparks. Per-node launch (every box runs the
|
|
|
-/// same `gr8 glm53 up` command; each rank finds itself in `--peers`). Model
|
|
|
-/// via `--model`, rendezvous via `--tp`/`--peers` (mgmt-LAN IPs), no headless,
|
|
|
-/// no bind-IP env, fixed API port :8100 (no `--host`/`--port`).
|
|
|
-pub struct Gr8Driver;
|
|
|
-
|
|
|
-impl EngineDriver for Gr8Driver {
|
|
|
- fn model_args(&self, recipe: &Recipe) -> Vec<String> {
|
|
|
- if recipe.model.is_empty() {
|
|
|
- Vec::new()
|
|
|
- } else {
|
|
|
- vec!["--model".into(), recipe.model.clone()]
|
|
|
- }
|
|
|
- }
|
|
|
-
|
|
|
- fn distributed_args(&self, ctx: &EngineFlagsCtx<'_>, _node: &Node) -> Vec<String> {
|
|
|
- vec![
|
|
|
- "--tp".into(),
|
|
|
- ctx.tp_size.to_string(),
|
|
|
- "--peers".into(),
|
|
|
- ctx.peers.join(","),
|
|
|
- ]
|
|
|
- }
|
|
|
-
|
|
|
- fn takes_host_port(&self) -> bool {
|
|
|
- false
|
|
|
- }
|
|
|
-
|
|
|
- fn entrypoint_default(&self) -> Vec<String> {
|
|
|
- vec!["/opt/gr8/gr8".into(), "glm53".into(), "up".into()]
|
|
|
- }
|
|
|
-}
|
|
|
-
|
|
|
-/// DGPP (`HawkBearPig/dgpp`): a C++/CUDA engine, config-JSON driven and
|
|
|
-/// head-coordinated — rank 0 runs `dgpp-cluster up`, which brings up the
|
|
|
-/// peers itself. No CLI model/distributed flags (the cluster JSON carries
|
|
|
-/// model + world size + engine knobs), no bind-IP env, fixed API port
|
|
|
-/// (:18080). `deploy = native`.
|
|
|
-pub struct DgppDriver;
|
|
|
-
|
|
|
-impl EngineDriver for DgppDriver {
|
|
|
- // The cluster JSON carries the model; no CLI model argument.
|
|
|
- fn model_args(&self, _recipe: &Recipe) -> Vec<String> {
|
|
|
- Vec::new()
|
|
|
- }
|
|
|
-
|
|
|
- // World size + rendezvous live in the cluster JSON, not CLI flags.
|
|
|
- fn distributed_args(&self, _ctx: &EngineFlagsCtx<'_>, _node: &Node) -> Vec<String> {
|
|
|
- Vec::new()
|
|
|
- }
|
|
|
-
|
|
|
- fn takes_host_port(&self) -> bool {
|
|
|
- false
|
|
|
- }
|
|
|
-
|
|
|
- fn entrypoint_default(&self) -> Vec<String> {
|
|
|
- vec!["scripts/dgpp-cluster".into(), "up".into()]
|
|
|
- }
|
|
|
-}
|
|
|
-
|
|
|
/// 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).
|
|
|
@@ -235,6 +91,7 @@ pub fn driver_for(engine: &str) -> &'static dyn EngineDriver {
|
|
|
"sglang" => &SglangDriver,
|
|
|
"gr8" => &Gr8Driver,
|
|
|
"dgpp" => &DgppDriver,
|
|
|
+ "tensorfold" => &TensorFoldDriver,
|
|
|
"vllm" => &VllmDriver,
|
|
|
_ => &VllmDriver,
|
|
|
}
|
|
|
@@ -446,4 +303,31 @@ pub(crate) mod tests {
|
|
|
// gr8 does not take --host/--port (fixed API port).
|
|
|
assert!(!driver_for("gr8").takes_host_port());
|
|
|
}
|
|
|
+
|
|
|
+ #[test]
|
|
|
+ fn tensorfold_uses_rank_master_rendezvous() {
|
|
|
+ let mut r = base_recipe();
|
|
|
+ r.engine = "tensorfold".into();
|
|
|
+ r.model = "/model".into();
|
|
|
+ r.entrypoint = vec![
|
|
|
+ "python".into(),
|
|
|
+ "-m".into(),
|
|
|
+ "tensorfold".into(),
|
|
|
+ "serve".into(),
|
|
|
+ ];
|
|
|
+ let c = ctx_peers(&r, 2, 2, &[]);
|
|
|
+ let head = build_engine_args(&c, &tnode(0));
|
|
|
+ assert_eq!(head.first().map(String::as_str), Some("/model"));
|
|
|
+ assert!(head.contains(&"--tp".into()));
|
|
|
+ assert!(head.contains(&"--rank".into()));
|
|
|
+ assert!(head.contains(&"--master".into()));
|
|
|
+ assert!(head.contains(&"--master-port".into()));
|
|
|
+ // no headless flag on workers, no engine bind-IP env
|
|
|
+ assert!(!driver_for("tensorfold").headless_workers());
|
|
|
+ assert!(driver_for("tensorfold").engine_ip_env().is_none());
|
|
|
+ let worker = build_engine_args(&c, &tnode(1));
|
|
|
+ assert!(!worker.contains(&"--headless".into()));
|
|
|
+ // the API port stays configurable (--host/--port on the tail)
|
|
|
+ assert!(driver_for("tensorfold").takes_host_port());
|
|
|
+ }
|
|
|
}
|