Переглянути джерело

feat(recipe): add Gemma 4 12B IT (MTP) recipe + --node flag for single-node serve

New recipe recipes/gemma4-12b/recipe.toml:
  - vllm/vllm-openai:gemma4-unified image
  - google/gemma-4-12B-it model
  - BF16, FP8 KV cache, MTP speculative decoding (5 tokens)
  - 196K context, tool/reasoning parsers, multimodal limits
  - port 8000 (matching spark-arena.com recipe)

New --node flag on serve command:
  - Filter cluster nodes to a single named node
  - Enables single-node inference on multi-node clusters
  - Error: bail if node name not found in cluster
  - Useful for solo-only recipes like Gemma 4 12B
Alexander Korolev 4 місяців тому
батько
коміт
71855e9f45
4 змінених файлів з 71 додано та 4 видалено
  1. 38 0
      recipes/gemma4-12b/recipe.toml
  2. 4 0
      src/cli.rs
  3. 18 3
      src/docker/mod.rs
  4. 11 1
      src/main.rs

+ 38 - 0
recipes/gemma4-12b/recipe.toml

@@ -0,0 +1,38 @@
+engine = "vllm"
+name = "Gemma 4 12B IT (MTP)"
+docker_image = "vllm/vllm-openai:gemma4-unified"
+model = "google/gemma-4-12B-it"
+port = 8000
+tensor_parallel_size = 1
+pipeline_parallel_size = 1
+distributed_backend = "mp"
+
+# Single-node recipe: Gemma 4 12B with BF16, FP8 KV cache,
+# MTP speculative decoding (5 tokens), 196K context.
+
+entrypoint = ["vllm", "serve"]
+
+args = [
+    "--dtype", "bfloat16",
+    "--kv-cache-dtype", "fp8",
+    "--gpu-memory-utilization", "0.75",
+    "--max-model-len", "196608",
+    "--max-num-seqs", "2",
+    "--max-num-batched-tokens", "8192",
+    "--enable-prefix-caching",
+    "--enable-chunked-prefill",
+    "--enable-auto-tool-choice",
+    "--tool-call-parser", "gemma4",
+    "--reasoning-parser", "gemma4",
+    "--generation-config", "vllm",
+    "--safetensors-load-strategy", "prefetch",
+    "--speculative-config", """
+'{"method":"mtp","model":"google/gemma-4-12B-it-assistant","num_speculative_tokens":5}'""",
+    "--limit-mm-per-prompt", """
+'{"image":4,"audio":1,"video":1}'""",
+    "--trust-remote-code",
+    "--served-model-name", "google/gemma-4-12B-it",
+]
+
+[env]
+TRITON_CACHE_DIR = "/root/.cache/triton"

+ 4 - 0
src/cli.rs

@@ -23,6 +23,10 @@ pub enum Commands {
         #[arg(long, default_value = "auto")]
         cluster: String,
 
+        /// Target a specific node within the cluster (default: all).
+        #[arg(long)]
+        node: Option<String>,
+
         /// Keep containers after stop (default: auto-remove on stop)
         #[arg(long, default_value = "false")]
         keep: bool,

+ 18 - 3
src/docker/mod.rs

@@ -9,6 +9,7 @@ pub async fn serve(
     config: &Config,
     recipe_name: &str,
     cluster_name: &str,
+    node_name: Option<&str>,
     keep: bool,
     skip_before_serve: bool,
     dry_run: bool,
@@ -49,6 +50,7 @@ pub async fn serve(
         &recipe,
         cache_dir,
         keep,
+        node_name,
         dry_run,
     )
     .await?;
@@ -263,17 +265,30 @@ async fn start_cluster(
     recipe: &Recipe,
     cache_dir: std::path::PathBuf,
     keep: bool,
+    node_name: Option<&str>,
     dry_run: bool,
 ) -> Result<()> {
-    let nnodes = cluster.nodes.len() as u32;
+    let nodes: Vec<&Node> = if let Some(nn) = node_name {
+        cluster.nodes.iter().filter(|n| n.name == nn).collect()
+    } else {
+        cluster.nodes.iter().collect()
+    };
+    if nodes.is_empty() {
+        anyhow::bail!(
+            "Node '{:?}' not found in cluster '{}'",
+            node_name,
+            cluster.name
+        );
+    }
+    let nnodes = nodes.len() as u32;
     let pp_size = if recipe.pipeline_parallel_size.unwrap_or(0) == 0 {
         nnodes
     } else {
         recipe.pipeline_parallel_size.unwrap()
     };
     let mut handles = Vec::new();
-    for node in &cluster.nodes {
-        let node = node.clone();
+    for node in &nodes {
+        let node = (*node).clone();
         let cluster = cluster.clone();
         let nccl = nccl.clone();
         let recipe = recipe.clone();

+ 11 - 1
src/main.rs

@@ -80,10 +80,20 @@ async fn main() -> Result<()> {
         Commands::Serve {
             recipe,
             cluster,
+            node,
             keep,
             skip_before_serve,
         } => {
-            docker::serve(&config, &recipe, &cluster, keep, skip_before_serve, dry_run).await?;
+            docker::serve(
+                &config,
+                &recipe,
+                &cluster,
+                node.as_deref(),
+                keep,
+                skip_before_serve,
+                dry_run,
+            )
+            .await?;
         }
         Commands::Logs {
             recipe: recipe_name,