Add --fit flag to serve: scale recipe to N nodes or whole cluster

--fit=N scales required_nodes + tensor_parallel_size to N (auto-selects
the smallest cluster with >= N free nodes). Bare --fit scales to the
whole selected cluster (TP = cluster node count); pair with
--cluster <name> to target a specific cluster. Incompatible with --solo.
pipeline_parallel_size is left untouched.

TP tracks the node count via the existing required_nodes == Some(tp_size)
auto-track branch in start_cluster — no new TP-computation code.

AGENTS.md

@@ -34,6 +34,15 @@ Use to run a lightweight test recipe alongside a live serve, or to force onto
 a busy node. Dangerous for GPU recipes (OOM, port conflicts); the operator is
 responsible for coexistence.
 
+`serve --fit` (optional value) scales a recipe to a node count and lets
+`tensor_parallel_size` track it (overrides the recipe's `required_nodes` and
+`tensor_parallel_size`; `pipeline_parallel_size` is left untouched).
+`--fit=N` scales to N nodes; bare `--fit` uses the whole selected cluster
+(TP = cluster node count). With `--cluster auto`, `--fit=N` auto-picks the
+smallest cluster with ≥N free nodes; bare `--fit` + `--cluster auto` picks
+the smallest cluster, so pass `--cluster <name> --fit` to scale onto a
+specific larger cluster. Incompatible with `--solo`.
+
 `serve --mirror` (optional value) mounts a **per-node local mirror** of the HF
 cache instead of the shared one at `/root/.cache/huggingface`. Bare `--mirror`
 uses `~/.cache/hf-models`; `--mirror <path>` uses that exact path (on each

src/cli.rs

@@ -79,6 +79,14 @@ pub enum Commands {
         /// uses ~/.cache/hf-models; `--mirror <path>` uses that exact path.
         #[arg(long, num_args = 0..=1, default_missing_value = "")]
         mirror: Option<String>,
+
+        /// Scale the recipe to a node count and let tensor-parallel track it
+        /// (overrides `required_nodes` and `tensor_parallel_size`).
+        /// `--fit=N` scales to N nodes; bare `--fit` uses the whole selected
+        /// cluster. Incompatible with `--solo`. `pipeline_parallel_size` is
+        /// left untouched.
+        #[arg(long, num_args = 0..=1, default_missing_value = "")]
+        fit: Option<String>,
     },
 
     /// Show container logs

src/docker/mod.rs

@@ -40,6 +40,10 @@ pub struct ServeOpts {
     /// of the shared HF cache. `Some("")` = bare `--mirror` → default
     /// `~/.cache/hf-models`; `Some(path)` = that exact path; `None` = off.
     pub mirror: Option<String>,
+    /// CLI override for recipe scaling: `Some("")` = bare `--fit` → scale to
+    /// the whole selected cluster; `Some("N")` = scale to N nodes;
+    /// `None` = unset. See `docker::serve`.
+    pub fit: Option<String>,
 }
 /// Shell-quote a string for safe interpolation into a remote `sh -c`
 /// command (single-quote wrapping, embedded quotes escaped). Re-exported
@@ -53,13 +57,41 @@ pub async fn serve(
     node_name: Option<&str>,
     opts: ServeOpts,
 ) -> Result<Cluster> {
-    let recipe = load_recipe(recipe_name).await?;
+    let mut recipe = load_recipe(recipe_name).await?;
     let rn = recipe.name.clone();
     let re = recipe.engine.clone();
     info!(
         "{}",
         fluent_i18n::t!("recipe-starting", { "name" => rn, "engine" => re, "backend" => recipe.distributed_backend.map_or("none".to_string(), |b| b.to_string()) })
     );
+    // `--fit`: override required_nodes + tensor_parallel_size so TP tracks
+    // the node count (the recipe's `required_nodes == Some(tp_size)` path in
+    // `start_cluster` makes tp_size follow nnodes). `pipeline_parallel_size`
+    // is left untouched. Bare `--fit` (whole cluster) is handled after
+    // cluster selection, when nnodes is known; for now just clear the cap so
+    // selection/trim keeps every cluster node.
+    //
+    // Mutually exclusive with `--solo` (solo = 1 node, no dist flags).
+    let fit: Option<u32> = match opts.fit.as_deref() {
+        None => None,
+        Some("") => Some(0),
+        Some(s) => Some(
+            s.parse()
+                .context(format!("--fit expects a node count, got '{}'", s))?,
+        ),
+    };
+    if fit.is_some() && opts.solo.is_some() {
+        anyhow::bail!("--fit and --solo are mutually exclusive");
+    }
+    match fit {
+        Some(0) => recipe.required_nodes = None,
+        Some(n) => {
+            recipe.required_nodes = Some(n);
+            recipe.tensor_parallel_size = n;
+            info!("--fit: scaling to {} nodes (tp = {})", n, n);
+        }
+        None => {}
+    }
     // Prepare cluster: probe, auto-power-on if requested, select.
     use crate::ends::{SshNodePinger, SshProber, ha_controller_from_config};
     use crate::orchestrate::{BootTiming, PrepareRequest, prepare_cluster};
@@ -123,6 +155,19 @@ pub async fn serve(
         );
         cluster.nodes = vec![solo_node];
     }
+    // Bare `--fit` (whole cluster): now that `cluster.nodes` is resolved, fix
+    // tp to the node count so `start_cluster`'s `tp_size` auto-track fires
+    // (required_nodes == Some(tp_size) → tp_size = nnodes). required_nodes is
+    // already None (no trim), so every node is kept.
+    if fit == Some(0) {
+        let nn = cluster.nodes.len() as u32;
+        recipe.required_nodes = Some(nn);
+        recipe.tensor_parallel_size = nn;
+        info!(
+            "--fit: scaling to whole cluster '{}' ({} nodes, tp = {})",
+            cluster.name, nn, nn
+        );
+    }
     crate::before_serve::run(
         config,
         &cluster,
@@ -170,7 +215,6 @@ pub async fn serve(
     )
     .await?;
     let image = try_repaired_image(config, &cluster, &recipe).await?;
-    let mut recipe = recipe;
     if image != recipe.docker_image {
         recipe.docker_image = image.clone();
     }

src/main.rs

@@ -105,6 +105,7 @@ async fn main() -> Result<()> {
             force,
             solo,
             mirror,
+            fit,
         } => {
             handle_serve(
                 &config,
@@ -123,6 +124,7 @@ async fn main() -> Result<()> {
                     force,
                     solo,
                     mirror,
+                    fit,
                 },
                 follow,
                 dry_run,
@@ -202,10 +204,6 @@ async fn main() -> Result<()> {
     Ok(())
 }
 
-// ---------------------------------------------------------------------------
-// Per-command handler functions
-// ---------------------------------------------------------------------------
-
 /// Handle `serve` command: launch inference containers on a cluster.
 ///
 /// Dispatches to `docker::serve` or `tmux::serve` based on the recipe's
@@ -240,6 +238,7 @@ async fn handle_serve(
                 (opts.keep, "--keep"),
                 (opts.solo.is_some(), "--solo/--no-solo"),
                 (opts.mirror.is_some(), "--mirror"),
+                (opts.fit.is_some(), "--fit"),
             ]
             .iter()
             .filter_map(|(set, flag)| if *set { Some(*flag) } else { None })
@@ -295,6 +294,7 @@ async fn handle_build_cmd(
         force: false,
         solo: None,
         mirror: None,
+        fit: None,
     };
     docker::build::handle_build(config, &cluster, recipe, &recipe_obj, &opts, build_node).await?;
     Ok(())