Jelajahi Sumber

feat(serve): add --mirror — serve from a per-node local HF mirror

--mirror [PATH]: mount a per-node mirror of the HF cache at
/root/.cache/huggingface instead of the shared cache. Bare --mirror uses
~/.cache/hf-models; --mirror <path> uses that exact path (per node).

Before launch, each serving node checks the mirror for every repo the recipe
needs (main model + repo-based {{path}} refs + composed-model sources) and,
when missing, copies it from the shared cache (cp -a, one-time per-node cost).
A repo absent from both locations aborts with a huggingface-cli download hint.
The mirror keeps the HF hub layout, so refs/main path refs, {{repo}} mounts,
and composed-model symlinks resolve unchanged.

Ships in a new src/docker/mirror.rs: resolve() (flag->dir+cache-dir),
repos_for_recipe(), mirror_copy_command() (MIRROR_PRESENT/COPIED/MISSING
markers, replace-broken-dir), ensure_mirrors() (parallel per node, uses
run_ssh_long_running for the big copy, dry-run prints only). Wired into
serve()/start_cluster(); hf_cache for path-ref/composed resolution now follows
the mounted cache dir. --mirror is ignored with a warning for tmux deploys.
6 new unit tests; make check green.
Alexander Korolev 2 bulan lalu
induk
melakukan
08d1e9fe03
6 mengubah file dengan 357 tambahan dan 5 penghapusan
  1. 11 0
      AGENTS.md
  2. 5 3
      recipes/deepseek-v4-flash-eugr/recipe.toml
  3. 7 0
      src/cli.rs
  4. 297 0
      src/docker/mirror.rs
  5. 33 2
      src/docker/mod.rs
  6. 4 0
      src/main.rs

+ 11 - 0
AGENTS.md

@@ -34,6 +34,17 @@ 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
 a busy node. Dangerous for GPU recipes (OOM, port conflicts); the operator is
 responsible for coexistence.
 responsible for coexistence.
 
 
+`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
+node). The mirror keeps the HF hub layout (`<mirror>/hub/models--org--model/...`),
+so `refs/main` path refs, `{{repo}}` mounts, and composed-model symlinks work
+unchanged. Before launch each node checks its mirror and, if a required repo
+(main model + repo-based path refs + composed sources) is missing, copies it
+from the shared cache (`cp -a`, one-time per-node cost); a repo absent from
+BOTH locations aborts with a `huggingface-cli download` hint. Later serves read
+the model locally — no streaming weights over slow shared storage at load time.
+
 > ⚠️ **Cluster may be serving live traffic.** A model (often the one serving this
 > ⚠️ **Cluster may be serving live traffic.** A model (often the one serving this
 > very agent) may be running on the cluster at any time. Before running `stop`,
 > very agent) may be running on the cluster at any time. Before running `stop`,
 > `power off`, or `power switch`, ALWAYS first run `health` and confirm with the
 > `power off`, or `power switch`, ALWAYS first run `health` and confirm with the

+ 5 - 3
recipes/deepseek-v4-flash-eugr/recipe.toml

@@ -49,10 +49,10 @@ entrypoint = ["vllm", "serve"]
 
 
 args = [
 args = [
     "--kv-cache-dtype", "fp8",
     "--kv-cache-dtype", "fp8",
-    "--kv-cache-memory-bytes", "20000000000",
+    "--kv-cache-memory-bytes", "21474836480",
     "--block-size", "256",
     "--block-size", "256",
     "--load-format", "instanttensor",
     "--load-format", "instanttensor",
-    "--max-model-len", "1048576",
+    "--max-model-len", "524288",
     "--max-num-seqs", "8",
     "--max-num-seqs", "8",
     "--max-num-batched-tokens", "8192",
     "--max-num-batched-tokens", "8192",
     "--max-cudagraph-capture-size", "64",
     "--max-cudagraph-capture-size", "64",
@@ -89,7 +89,7 @@ args = [
     "--enable-auto-tool-choice",
     "--enable-auto-tool-choice",
     "--generation-config", "vllm",
     "--generation-config", "vllm",
     "--default-chat-template-kwargs.thinking=true",
     "--default-chat-template-kwargs.thinking=true",
-    "--default-chat-template-kwargs.reasoning_effort=high",
+    "--default-chat-template-kwargs.reasoning_effort=max",
     "--disable-custom-all-reduce",
     "--disable-custom-all-reduce",
     "--trust-remote-code",
     "--trust-remote-code",
 ]
 ]
@@ -142,3 +142,5 @@ NCCL_SOCKET_FAMILY = "inet"
 # PCIe allreduce not needed on GB10 single-NIC
 # PCIe allreduce not needed on GB10 single-NIC
 VLLM_ENABLE_PCIE_ALLREDUCE = "0"
 VLLM_ENABLE_PCIE_ALLREDUCE = "0"
 VLLM_PCIE_ALLREDUCE_BACKEND = "cpp"
 VLLM_PCIE_ALLREDUCE_BACKEND = "cpp"
+
+VLLM_DSV4_REASONING_CONTENT = "1"

+ 7 - 0
src/cli.rs

@@ -72,6 +72,13 @@ pub enum Commands {
         /// sets `solo = true`.
         /// sets `solo = true`.
         #[arg(long)]
         #[arg(long)]
         solo: Option<bool>,
         solo: Option<bool>,
+
+        /// Mirror the model into a per-node local folder and mount that
+        /// instead of the shared HF cache (first serve copies the model from
+        /// the shared cache; later serves load it locally). Without a value
+        /// uses ~/.cache/hf-models; `--mirror <path>` uses that exact path.
+        #[arg(long, num_args = 0..=1, default_missing_value = "")]
+        mirror: Option<String>,
     },
     },
 
 
     /// Show container logs
     /// Show container logs

+ 297 - 0
src/docker/mirror.rs

@@ -0,0 +1,297 @@
+//! `--mirror`: serve from a per-node local mirror of the HF cache instead of
+//! the shared one. On first serve each node copies the required HF repos
+//! (main model + `{{repo}}` path refs + composed-model sources) from the
+//! shared cache into the mirror; later serves mount the mirror directly.
+//! Local reads skip streaming the weights over slow shared storage at load
+//! time — the copy is a one-time, per-node cost.
+//!
+//! The mirror keeps the HF hub layout (`<mirror>/hub/models--org--model/...`),
+//! so the container still sees the canonical `/root/.cache/huggingface`
+//! structure: `refs/main`, path-ref resolution, and composed-model symlinks
+//! all keep working unchanged.
+
+use crate::config::{Config, Node, Recipe, model_cache_dir_name};
+use crate::docker::shq;
+use crate::ssh::run_ssh_long_running;
+use anyhow::{Result, bail};
+use std::path::{Path, PathBuf};
+use tracing::{info, warn};
+
+/// Default mirror root when `--mirror` is given without a value.
+pub const DEFAULT_MIRROR: &str = "~/.cache/hf-models";
+
+/// Resolve the `--mirror` argument to the host-side mirror dir that is
+/// mounted at `/root/.cache/huggingface`. A bare `--mirror` (empty value)
+/// uses the default `~/.cache/hf-models`; `--mirror <path>` uses that path.
+/// `~` / `~/` expands to the SSH user's home (mirrors are per-node dirs).
+pub fn resolve_mirror_dir(config: &Config, arg: &str) -> PathBuf {
+    let p = if arg.is_empty() { DEFAULT_MIRROR } else { arg };
+    PathBuf::from(crate::config::expand_tilde(p, &config.ssh_user))
+}
+
+/// Resolve the `--mirror` flag for a serve into the mirror dir (if any) and
+/// the cache dir to mount at `/root/.cache/huggingface` (mirror when set,
+/// otherwise the shared HF cache).
+pub fn resolve(config: &Config, arg: Option<&str>) -> (Option<PathBuf>, PathBuf) {
+    let mirror = arg.map(|m| resolve_mirror_dir(config, m));
+    let cache_dir = mirror
+        .clone()
+        .unwrap_or_else(|| config.resolved_hf_cache_dir());
+    (mirror, cache_dir)
+}
+
+/// All HF repos a recipe needs at serve time, deduped and sorted (stable
+/// order): the main model, every repo-based path ref in `extra_mounts`, and
+/// every composed-model source repo.
+pub fn repos_for_recipe(recipe: &Recipe) -> Vec<String> {
+    let mut repos = vec![recipe.model.clone()];
+    for m in &recipe.extra_mounts {
+        if let Some(r) = &m.repo {
+            repos.push(r.clone());
+        }
+    }
+    if let Some(cm) = &recipe.composed_model {
+        repos.extend(cm.sources.iter().map(|s| s.repo.clone()));
+    }
+    repos.sort();
+    repos.dedup();
+    repos
+}
+
+/// Remote command that ensures one repo's HF hub dir exists in the mirror,
+/// copying it from the shared cache when missing. Prints exactly one marker
+/// line: `MIRROR_PRESENT`, `MIRROR_COPIED`, or `MIRROR_MISSING_SOURCE`.
+/// A broken/partial mirror dir (`snapshots`/`refs/main` missing) is replaced.
+pub fn mirror_copy_command(source: &Path, mirror: &Path, repo: &str) -> String {
+    let dir = model_cache_dir_name(repo);
+    let src_dir = format!("{}/hub/{}", source.display(), dir);
+    let dst_dir = format!("{}/hub/{}", mirror.display(), dir);
+    let hub = format!("{}/hub", mirror.display());
+    format!(
+        "if test -d {d} && test -f {r}; then echo MIRROR_PRESENT; \
+         elif test -d {s}; then rm -rf {d}; mkdir -p {hub}; cp -a {s} {hub}/. && echo MIRROR_COPIED; \
+         else echo MIRROR_MISSING_SOURCE; fi",
+        d = shq(&dst_dir),
+        r = shq(&format!("{dst_dir}/refs/main")),
+        s = shq(&src_dir),
+        hub = shq(&hub),
+    )
+}
+
+/// Per node (in parallel): ensure every repo the recipe needs exists in the
+/// mirror, copying missing ones from the shared cache. Aborts when a repo is
+/// absent from both locations (hint: `huggingface-cli download`). Dry-run
+/// only prints the intended copies — no SSH.
+pub async fn ensure_mirrors(
+    config: &Config,
+    nodes: &[&Node],
+    recipe: &Recipe,
+    source: &Path,
+    mirror: &Path,
+    dry: bool,
+) -> Result<()> {
+    let repos = repos_for_recipe(recipe);
+    if dry {
+        for node in nodes {
+            for repo in &repos {
+                info!(
+                    "dry-run: would mirror {} on {} from {}",
+                    repo,
+                    node.name,
+                    source.display()
+                );
+            }
+        }
+        return Ok(());
+    }
+    let mut handles = Vec::new();
+    for node in nodes {
+        let node = (*node).clone();
+        let config = config.clone();
+        let repos = repos.clone();
+        let source = source.to_path_buf();
+        let mirror = mirror.to_path_buf();
+        handles.push(crate::locale::spawn_with_locale(async move {
+            let mut missing = Vec::new();
+            for repo in &repos {
+                let cmd = mirror_copy_command(&source, &mirror, repo);
+                let out = run_ssh_long_running(&config, &node, cmd, false).await?;
+                match out.trim() {
+                    "MIRROR_PRESENT" => info!("{}: {} already mirrored", node.name, repo),
+                    "MIRROR_COPIED" => {
+                        info!("{}: copied {} into mirror {}", node.name, repo, mirror.display())
+                    }
+                    "MIRROR_MISSING_SOURCE" => missing.push(repo.clone()),
+                    other => warn!("{}: mirror status for {}: {:?}", node.name, repo, other),
+                }
+            }
+            if !missing.is_empty() {
+                bail!(
+                    "{}: model repos absent from shared HF cache: {} — run `huggingface-cli download <repo>` on that node first",
+                    node.name,
+                    missing.join(", ")
+                );
+            }
+            Ok(())
+        }));
+    }
+    for r in futures::future::join_all(handles).await {
+        r??;
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::config::{DeployTarget, DistributedBackend, RecipeArgs, VenvSetup};
+    use std::collections::HashMap;
+
+    fn cfg() -> Config {
+        Config {
+            ssh_user: "u".into(),
+            ssh_key_path: "/tmp/k".into(),
+            clusters: vec![],
+            nccl: HashMap::new(),
+            hf_cache_dir: None,
+            init_port: 25000,
+            before_serve: HashMap::new(),
+            exec_aliases: HashMap::new(),
+            home_assistant: None,
+            global_mods_dir: None,
+            drop_caches: true,
+            mount_aliases: HashMap::new(),
+        }
+    }
+
+    fn trec() -> Recipe {
+        Recipe {
+            engine: "vllm".into(),
+            name: "T".into(),
+            docker_image: "img".into(),
+            model: "org/main".into(),
+            port: 8000,
+            required_nodes: Some(1),
+            pipeline_parallel_size: Some(1),
+            tensor_parallel_size: 1,
+            args: RecipeArgs::List(vec![]),
+            env: HashMap::new(),
+            cluster: None,
+            entrypoint: vec![],
+            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(),
+            sync_image: true,
+            docker_image_dir: None,
+            extra_mounts: vec![],
+            build: Default::default(),
+            composed_model: None,
+        }
+    }
+
+    #[test]
+    fn mirror_dir_default_tilde_and_explicit() {
+        let c = cfg();
+        // Bare `--mirror` -> default under the SSH user's home.
+        assert_eq!(
+            resolve_mirror_dir(&c, ""),
+            PathBuf::from("/home/u/.cache/hf-models")
+        );
+        // `~/...` expands to the SSH user's home.
+        assert_eq!(resolve_mirror_dir(&c, "~/m"), PathBuf::from("/home/u/m"));
+        // Absolute paths pass through unchanged.
+        assert_eq!(
+            resolve_mirror_dir(&c, "/abs/path"),
+            PathBuf::from("/abs/path")
+        );
+        // root user -> /root.
+        let mut root = cfg();
+        root.ssh_user = "root".into();
+        assert_eq!(
+            resolve_mirror_dir(&root, ""),
+            PathBuf::from("/root/.cache/hf-models")
+        );
+    }
+
+    #[test]
+    fn resolve_returns_mirror_and_falls_back_to_shared() {
+        let c = cfg();
+        // No flag -> mount the shared HF cache, no mirror.
+        let (mirror, cache) = resolve(&c, None);
+        assert!(mirror.is_none());
+        assert_eq!(cache, PathBuf::from("/home/u/.cache/huggingface"));
+        // Bare `--mirror` -> default mirror becomes the mounted cache dir.
+        let (mirror, cache) = resolve(&c, Some(""));
+        assert_eq!(mirror, Some(PathBuf::from("/home/u/.cache/hf-models")));
+        assert_eq!(cache, PathBuf::from("/home/u/.cache/hf-models"));
+        // Explicit path.
+        let (mirror, cache) = resolve(&c, Some("/abs/mirror"));
+        assert_eq!(mirror, Some(PathBuf::from("/abs/mirror")));
+        assert_eq!(cache, PathBuf::from("/abs/mirror"));
+    }
+
+    #[test]
+    fn repos_gather_model_pathrefs_and_composed() {
+        let mut r = trec();
+        r.extra_mounts = vec![
+            crate::config::ExtraMount {
+                repo: Some("org/draft".into()),
+                container: "draft".into(),
+                ..Default::default()
+            },
+            crate::config::ExtraMount {
+                host: Some("/x".into()),
+                container: "/y".into(),
+                ..Default::default()
+            },
+        ];
+        r.composed_model = Some(crate::config::ComposedModel {
+            mount: "/root/model".into(),
+            sources: vec![
+                crate::config::ComposedSource {
+                    repo: "org/base".into(),
+                    glob: None,
+                    exclude: None,
+                },
+                crate::config::ComposedSource {
+                    repo: "org/main".into(), // duplicate of model — dedup
+                    glob: None,
+                    exclude: None,
+                },
+            ],
+            files: vec![],
+        });
+        assert_eq!(
+            repos_for_recipe(&r),
+            vec!["org/base", "org/draft", "org/main"]
+        );
+    }
+
+    #[test]
+    fn copy_command_branches_and_quoting() {
+        let c = mirror_copy_command(
+            Path::new("/home/u/.cache/huggingface"),
+            Path::new("/home/u/.cache/hf-models"),
+            "org/main",
+        );
+        // All three markers present, in order.
+        assert!(c.contains("MIRROR_PRESENT"));
+        assert!(c.contains("MIRROR_COPIED"));
+        assert!(c.contains("MIRROR_MISSING_SOURCE"));
+        // Hub layout + model dir name.
+        assert!(c.contains("models--org--main"));
+        assert!(c.contains("/home/u/.cache/huggingface/hub/"));
+        assert!(c.contains("/home/u/.cache/hf-models/hub/"));
+        // refs/main presence gate.
+        assert!(c.contains("refs/main"));
+        // Paths are shell-quoted.
+        assert!(c.contains("'/home/u/.cache/hf-models/hub/models--org--main'"));
+    }
+}

+ 33 - 2
src/docker/mod.rs

@@ -4,6 +4,7 @@ pub mod compose;
 pub mod deploy;
 pub mod deploy;
 pub mod health;
 pub mod health;
 pub mod logs;
 pub mod logs;
+pub mod mirror;
 pub mod mods;
 pub mod mods;
 pub(crate) mod path_ref;
 pub(crate) mod path_ref;
 pub use logs::show_logs;
 pub use logs::show_logs;
@@ -34,6 +35,10 @@ pub struct ServeOpts {
     /// CLI override for recipe `solo`: `Some(true)` (--solo), `Some(false)`
     /// CLI override for recipe `solo`: `Some(true)` (--solo), `Some(false)`
     /// (--no-solo), `None` (unset → use recipe `solo` field).
     /// (--no-solo), `None` (unset → use recipe `solo` field).
     pub solo: Option<bool>,
     pub solo: Option<bool>,
+    /// Mirror the model into a per-node local folder and mount that instead
+    /// of the shared HF cache. `Some("")` = bare `--mirror` → default
+    /// `~/.cache/hf-models`; `Some(path)` = that exact path; `None` = off.
+    pub mirror: Option<String>,
 }
 }
 /// Shell-quote a string for safe interpolation into a remote `sh -c`
 /// Shell-quote a string for safe interpolation into a remote `sh -c`
 /// command (single-quote wrapping, embedded quotes escaped). Re-exported
 /// command (single-quote wrapping, embedded quotes escaped). Re-exported
@@ -199,7 +204,16 @@ pub async fn serve(
         }
         }
     }
     }
 
 
-    let cache_dir = config.resolved_hf_cache_dir();
+    // `--mirror`: mount a per-node local mirror of the HF cache instead of
+    // the shared one (first serve copies the model per node; later serves
+    // load it locally). See mirror.rs.
+    let (mirror, cache_dir) = mirror::resolve(config, opts.mirror.as_deref());
+    if let Some(ref m) = mirror {
+        info!(
+            "serving from local HF mirror {} (models copied per node on first use)",
+            m.display()
+        );
+    }
     start_cluster(
     start_cluster(
         config,
         config,
         &cluster,
         &cluster,
@@ -208,6 +222,7 @@ pub async fn serve(
         &recipe,
         &recipe,
         StartClusterOpts {
         StartClusterOpts {
             cache_dir,
             cache_dir,
+            mirror,
             keep: opts.keep,
             keep: opts.keep,
             node_name: node_name.map(String::from),
             node_name: node_name.map(String::from),
             mods_remote_path,
             mods_remote_path,
@@ -296,6 +311,9 @@ pub async fn resolve_cluster_by_container_label<'c>(
 /// Options for [start_cluster].
 /// Options for [start_cluster].
 struct StartClusterOpts {
 struct StartClusterOpts {
     cache_dir: std::path::PathBuf,
     cache_dir: std::path::PathBuf,
+    /// When set, the per-node mirror dir mounted at /root/.cache/huggingface
+    /// (repos are copied into it from the shared cache before launch).
+    mirror: Option<std::path::PathBuf>,
     keep: bool,
     keep: bool,
     node_name: Option<String>,
     node_name: Option<String>,
     mods_remote_path: String,
     mods_remote_path: String,
@@ -330,6 +348,19 @@ async fn start_cluster(
         }
         }
     };
     };
     let nnodes = nodes.len() as u32;
     let nnodes = nodes.len() as u32;
+    // `--mirror`: ensure every repo the recipe needs exists in the per-node
+    // mirror before launch (copies from the shared HF cache when missing).
+    if let Some(ref mirror) = opts.mirror {
+        mirror::ensure_mirrors(
+            config,
+            &nodes,
+            recipe,
+            &config.resolved_hf_cache_dir(),
+            mirror,
+            opts.dry_run,
+        )
+        .await?;
+    }
     let pp_size = if recipe.pipeline_parallel_size.unwrap_or(0) == 0
     let pp_size = if recipe.pipeline_parallel_size.unwrap_or(0) == 0
         || recipe.required_nodes == recipe.pipeline_parallel_size
         || recipe.required_nodes == recipe.pipeline_parallel_size
     {
     {
@@ -372,7 +403,7 @@ async fn start_cluster(
             mods_remote_path: opts.mods_remote_path.clone(),
             mods_remote_path: opts.mods_remote_path.clone(),
             extra_mounts: resolved.static_mounts.clone(),
             extra_mounts: resolved.static_mounts.clone(),
         };
         };
-        let hf_cache = config.resolved_hf_cache_dir();
+        let hf_cache = opts.cache_dir.clone();
         let path_refs = resolved.path_refs.clone();
         let path_refs = resolved.path_refs.clone();
         handles.push(crate::locale::spawn_with_locale(async move {
         handles.push(crate::locale::spawn_with_locale(async move {
             if !dry {
             if !dry {

+ 4 - 0
src/main.rs

@@ -103,6 +103,7 @@ async fn main() -> Result<()> {
             force_build,
             force_build,
             force,
             force,
             solo,
             solo,
+            mirror,
         } => {
         } => {
             handle_serve(
             handle_serve(
                 &config,
                 &config,
@@ -120,6 +121,7 @@ async fn main() -> Result<()> {
                     force_build,
                     force_build,
                     force,
                     force,
                     solo,
                     solo,
+                    mirror,
                 },
                 },
                 follow,
                 follow,
                 dry_run,
                 dry_run,
@@ -236,6 +238,7 @@ async fn handle_serve(
                 (opts.force_build, "--force-build"),
                 (opts.force_build, "--force-build"),
                 (opts.keep, "--keep"),
                 (opts.keep, "--keep"),
                 (opts.solo.is_some(), "--solo/--no-solo"),
                 (opts.solo.is_some(), "--solo/--no-solo"),
+                (opts.mirror.is_some(), "--mirror"),
             ]
             ]
             .iter()
             .iter()
             .filter_map(|(set, flag)| if *set { Some(*flag) } else { None })
             .filter_map(|(set, flag)| if *set { Some(*flag) } else { None })
@@ -290,6 +293,7 @@ async fn handle_build_cmd(
         force_build,
         force_build,
         force: false,
         force: false,
         solo: None,
         solo: None,
+        mirror: None,
     };
     };
     docker::build::handle_build(config, &cluster, recipe, &recipe_obj, &opts, build_node).await?;
     docker::build::handle_build(config, &cluster, recipe, &recipe_obj, &opts, build_node).await?;
     Ok(())
     Ok(())