Преглед на файлове

feat(recipe): model-subfolder mounts resolved per-node from HF cache

Extends [[extra_mounts]] with a repo+subfolder variant (mutually exclusive
with alias/host) that resolves the host path at serve time, per node, from
the node's HF cache — eliminating the fragile manual symlink.

  [[extra_mounts]]
  repo = "aidendle94/GLM-5.2-Hybrid-FP8-MXFP4"
  subfolder = "mtp-draft"          # optional; omit = whole snapshot dir
  container = "/models/mtp-draft"   # fixed path referenced in args
  read_only = true

The launcher reads <hf_cache>/hub/models--<repo-dashed>/refs/main for the
current commit hash, then bind-mounts snapshots/<hash>/<subfolder> at the
fixed container path. Since refs/main always points at the latest commit,
the mount survives `huggingface-cli download` updates — no symlink to break.

Resolution is per-node (the snapshot hash can differ across nodes); if the
repo or subfolder is absent on a node, serve aborts with a clear error
pointing at `huggingface-cli download <repo>` (the launcher never
downloads). Dry-run skips SSH and emits a descriptive
<hf_cache>/hub/.../snapshots/<refs-main>/<subfolder> placeholder.

resolve_extra_mounts now partitions into static mounts (alias/host,
resolved once) and model specs (repo/subfolder, deferred to per-node
resolution in start_cluster's spawn loop). Validation enforces that repo is
mutually exclusive with alias/host, and subfolder requires repo. Pure helper
model_cache_dir_name(repo) maps org/name to models--org--name.

Applied to glm52-hybrid-fp8-mxfp4: replaced the fragile
/root/.cache/huggingface/links/mtp-draft symlink reference in
--speculative-config with the stable /models/mtp-draft mount path.

3 new tests (model_cache_dir_name mapping, model-subfolder partition +
placeholder, existing static-mount tests updated).
Alexander Korolev преди 3 месеца
родител
ревизия
6cd245f9a8
променени са 5 файла, в които са добавени 275 реда и са изтрити 37 реда
  1. 11 0
      AGENTS.md
  2. 22 11
      recipes/glm52-hybrid-fp8-mxfp4/recipe.toml
  3. 47 7
      src/config/types.rs
  4. 129 17
      src/docker/command.rs
  5. 66 2
      src/docker/mod.rs

+ 11 - 0
AGENTS.md

@@ -45,6 +45,17 @@ container = "/root/.cache/jit"     # container-side target (required)
 read_only = false                 # optional, default false (adds :ro)
 env = "XDG_CACHE_HOME"            # optional: injects -e <name>=<container path>
 
+# Model subfolder mount (mutually exclusive with alias/host): resolves at
+# serve time, per node, from the HF cache — refs/main → latest snapshot, so
+# it survives `huggingface-cli download` updates with no fragile symlink.
+# `container` is the fixed stable path the recipe references in args. Absent
+# subfolder → serve aborts with a hint to `huggingface-cli download <repo>`.
+[[extra_mounts]]
+repo = "aidendle94/GLM-5.2-Hybrid-FP8-MXFP4"
+subfolder = "mtp-draft"           # optional; omit = mount whole snapshot dir
+container = "/models/mtp-draft"
+read_only = true
+
 [mods]
 commands = [                         # run *inside* container before engine
     "pip install flash-attn",

+ 22 - 11
recipes/glm52-hybrid-fp8-mxfp4/recipe.toml

@@ -35,16 +35,6 @@ distributed_backend = "mp"
 shm_size = "10gb"
 docker_image_dir = "glm52-int4-int8"
 
-# Persistent JIT/torch-inductor cache (avoids ~20 min recompilation on cold boot).
-# `alias = "jit_cache"` resolves via [mount_aliases] in config.toml; if absent,
-# falls back to host "~/.cache/jit" (expanded to /home/<ssh_user>/.cache/jit).
-# `env` injects -e XDG_CACHE_HOME=/root/.cache/jit so vLLM/torch write there.
-[[extra_mounts]]
-alias = "jit_cache"
-host = "~/.cache/jit"
-container = "/root/.cache/jit"
-env = "XDG_CACHE_HOME"
-
 entrypoint = ["vllm", "serve"]
 
 args = [
@@ -74,7 +64,7 @@ args = [
     "--gpu-memory-utilization", "0.89",
     # Speculative decode — MTP k=4 with adaptive depths
     "--speculative-config", """
-'{"model":"/root/.cache/huggingface/links/mtp-draft","method":"mtp","num_speculative_tokens":4,"draft_attention_backend":"B12X_MLA_SPARSE","adaptive_speculative_tokens_window":32}'""",
+'{"model":"/models/mtp-draft","method":"mtp","num_speculative_tokens":4,"draft_attention_backend":"B12X_MLA_SPARSE","adaptive_speculative_tokens_window":32}'""",
     # Compilation — FULL_AND_PIECEWISE with all custom ops, 8 capture sizes
     "--compilation-config", """
 '{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"],"cudagraph_capture_sizes":[1,2,3,4,5,6,8]}'""",
@@ -83,6 +73,27 @@ args = [
     "--cpu-distributed-timeout-seconds", "3600",
 ]
 
+# Persistent JIT/torch-inductor cache (avoids ~20 min recompilation on cold boot).
+# `alias = "jit_cache"` resolves via [mount_aliases] in config.toml; if absent,
+# falls back to host "~/.cache/jit" (expanded to /home/<ssh_user>/.cache/jit).
+# `env` injects -e XDG_CACHE_HOME=/root/.cache/jit so vLLM/torch write there.
+[[extra_mounts]]
+alias = "jit_cache"
+host = "~/.cache/jit"
+container = "/root/.cache/jit"
+env = "XDG_CACHE_HOME"
+
+# MTP draft model — a subfolder of the main GLM-5.2 repo. Resolved at serve
+# time, per node, from the HF cache (refs/main → latest snapshot), so the
+# mount survives `huggingface-cli download` updates without a fragile symlink.
+# The recipe references the fixed /models/mtp-draft path in --speculative-config.
+[[extra_mounts]]
+repo = "aidendle94/GLM-5.2-Hybrid-FP8-MXFP4"
+subfolder = "mtp-draft"
+container = "/models/mtp-draft"
+read_only = true
+
+
 [env]
 # ── b12x suite ────────────────────────────────────────────────────────────
 VLLM_USE_V2_MODEL_RUNNER = "1"

+ 47 - 7
src/config/types.rs

@@ -169,20 +169,42 @@ impl Config {
 
 /// A recipe-level extra bind-mount for `docker run`.
 ///
-/// `alias` resolves against `Config.mount_aliases` and, when present there,
-/// overrides `host` (so an operator can repoint a mount site-wide without
-/// editing recipes). If the alias is absent from config, `host` is used as the
-/// default. `~` (and `~/`) at the start of the resolved path is expanded to the
-/// remote SSH user's home. At least one of `alias` or `host` must be set.
+/// Two mutually exclusive kinds:
+/// - **Static** (`alias`/`host`): host path is known at recipe-write time.
+///   `alias` resolves against `Config.mount_aliases` and, when present there,
+///   overrides `host`. If the alias is absent from config, `host` is used as
+///   the default. `~` (and `~/`) at the start of the resolved path is expanded
+///   to the remote SSH user's home. At least one of `alias` or `host` must be
+///   set.
+/// - **Model subfolder** (`repo` + optional `subfolder`): the host path is
+///   resolved **at serve time, per node**, from the node's HF cache —
+///   `<hf_cache>/hub/<model_cache_dir_name(repo)>/snapshots/<refs/main>/<subfolder>`.
+///   This pins the current snapshot without a fragile symlink: `refs/main`
+///   always points at the latest commit, so the mount survives `huggingface-cli
+///   download` updates. The recipe references the fixed `container` path in
+///   args (e.g. `--speculative-config`), so a vLLM JSON arg stays stable.
+///   If the subfolder is absent on a node, serve aborts with a hint to
+///   `huggingface-cli download <repo>` (the launcher never downloads).
 #[derive(Clone, Debug, Default, Deserialize)]
 pub struct ExtraMount {
     /// Name resolved via `Config.mount_aliases`; overrides `host` when present.
+    /// Mutually exclusive with `repo`.
     #[serde(default)]
     pub alias: Option<String>,
     /// Fallback host path (used when the alias is not in `mount_aliases`).
+    /// Mutually exclusive with `repo`.
     #[serde(default)]
     pub host: Option<String>,
-    /// Container-side mount target.
+    /// HF Hub repo id, e.g. `aidendle94/GLM-5.2-Hybrid-FP8-MXFP4`. When set,
+    /// the host path is resolved at serve time from the node's HF cache
+    /// (so `alias`/`host` must be empty).
+    #[serde(default)]
+    pub repo: Option<String>,
+    /// Subfolder within the repo snapshot to mount (e.g. `mtp-draft`).
+    /// Empty/omitted = mount the whole snapshot dir.
+    #[serde(default)]
+    pub subfolder: Option<String>,
+    /// Container-side mount target (always the stable path referenced in args).
     pub container: String,
     /// Mount read-only (`:ro`).
     #[serde(default)]
@@ -193,6 +215,14 @@ pub struct ExtraMount {
     pub env: Option<String>,
 }
 
+/// Directory name a repo id maps to inside `<hf_cache>/hub/`.
+/// `org/name` -> `models--org--name` (HF convention: prefix `models--`,
+/// replace `/` with `--`). Pure & deterministic; used by the model-subfolder
+/// mount resolver.
+pub fn model_cache_dir_name(repo: &str) -> String {
+    format!("models--{}", repo.replace('/', "--"))
+}
+
 /// Expand a leading `~` / `~/` to the remote SSH user's home directory.
 /// Paths without a leading `~` are returned unchanged. Used for host paths
 /// before they are sent to `docker run -v`.
@@ -240,9 +270,19 @@ pub fn resolve_mount_host(
 
 #[cfg(test)]
 mod tests {
-    use super::{expand_tilde, resolve_mount_host};
+    use super::{expand_tilde, model_cache_dir_name, resolve_mount_host};
     use std::collections::HashMap;
 
+    #[test]
+    fn model_cache_dir_name_maps_repo_id() {
+        assert_eq!(
+            model_cache_dir_name("aidendle94/GLM-5.2-Hybrid-FP8-MXFP4"),
+            "models--aidendle94--GLM-5.2-Hybrid-FP8-MXFP4"
+        );
+        assert_eq!(model_cache_dir_name("org/name"), "models--org--name");
+        assert_eq!(model_cache_dir_name("bare"), "models--bare");
+    }
+
     #[test]
     fn expand_tilde_user_and_root() {
         assert_eq!(expand_tilde("~/.cache", "kilork"), "/home/kilork/.cache");

+ 129 - 17
src/docker/command.rs

@@ -24,13 +24,17 @@ pub(crate) struct DockerCommandCtx {
     pub keep: bool,
     pub init_port: u16,
     pub mods_remote_path: String,
-    /// Recipe `extra_mounts` with host paths already resolved (alias lookup +
-    /// `~` expansion) against `Config.mount_aliases` / `ssh_user`.
+    /// Static `extra_mounts` (alias/host) with host paths already resolved
+    /// once against `Config.mount_aliases` / `ssh_user`. Model-subfolder mounts
+    /// are NOT here — `start_cluster` resolves those per node and merges the
+    /// results into this `Vec` (as `ResolvedMount`s) before calling
+    /// [`build_docker_command`].
     pub extra_mounts: Vec<ResolvedMount>,
 }
 
 /// A recipe `ExtraMount` after host-path resolution (alias + tilde), ready to
-/// turn into a `-v` flag / `-e` env pair.
+/// turn into a `-v` flag / `-e` env pair. Also the final form a model-subfolder
+/// mount takes after per-node resolution.
 #[derive(Clone)]
 pub(crate) struct ResolvedMount {
     pub host: String,
@@ -39,15 +43,51 @@ pub(crate) struct ResolvedMount {
     pub env: Option<String>,
 }
 
-/// Resolve every `extra_mount` in a recipe to a concrete host path using the
-/// config alias registry and the SSH user (for `~` expansion). Done once per
-/// serve; the resulting `Vec<ResolvedMount>` is cloned into each node's
-/// [`DockerCommandCtx`].
-pub(crate) fn resolve_extra_mounts(config: &Config, recipe: &Recipe) -> Result<Vec<ResolvedMount>> {
-    recipe
-        .extra_mounts
-        .iter()
-        .map(|m| {
+/// A model-subfolder mount spec (`repo`+`subfolder`) before per-node
+/// resolution. The host path is resolved at serve time from each node's HF
+/// cache (so it can differ across nodes); `container` is the fixed path the
+/// recipe references in args.
+#[derive(Clone)]
+pub(crate) struct ModelMountSpec {
+    pub repo: String,
+    pub subfolder: Option<String>,
+    pub container: String,
+    pub read_only: bool,
+    pub env: Option<String>,
+}
+
+/// Partition a recipe's `extra_mounts` into static (resolved now) and
+/// model-subfolder (deferred to per-node resolution). Validates that a mount
+/// is one or the other — `repo` set means `alias`/`host` must be empty, and
+/// `subfolder` without `repo` is an error.
+pub(crate) fn resolve_extra_mounts(
+    config: &Config,
+    recipe: &Recipe,
+) -> Result<ResolvedExtraMounts> {
+    let mut static_mounts = Vec::new();
+    let mut model_mounts = Vec::new();
+    for m in &recipe.extra_mounts {
+        if m.repo.as_deref().is_some_and(|r| !r.trim().is_empty()) {
+            if m.alias.is_some() || m.host.is_some() {
+                anyhow::bail!(
+                    "extra_mount '{}': `repo` is mutually exclusive with `alias`/`host`",
+                    m.container
+                );
+            }
+            model_mounts.push(ModelMountSpec {
+                repo: m.repo.as_ref().unwrap().trim().to_string(),
+                subfolder: m
+                    .subfolder
+                    .as_deref()
+                    .filter(|s| !s.trim().is_empty())
+                    .map(|s| s.trim().to_string()),
+                container: m.container.clone(),
+                read_only: m.read_only,
+                env: m.env.clone(),
+            });
+        } else if m.subfolder.as_deref().is_some_and(|s| !s.trim().is_empty()) {
+            anyhow::bail!("extra_mount '{}': `subfolder` requires `repo`", m.container);
+        } else {
             let host = crate::config::resolve_mount_host(
                 m.alias.as_deref(),
                 m.host.as_deref(),
@@ -55,14 +95,43 @@ pub(crate) fn resolve_extra_mounts(config: &Config, recipe: &Recipe) -> Result<V
                 &config.ssh_user,
                 &m.container,
             )?;
-            Ok(ResolvedMount {
+            static_mounts.push(ResolvedMount {
                 host,
                 container: m.container.clone(),
                 read_only: m.read_only,
                 env: m.env.clone(),
-            })
-        })
-        .collect()
+            });
+        }
+    }
+    Ok(ResolvedExtraMounts {
+        static_mounts,
+        model_mounts,
+    })
+}
+
+/// Result of [`resolve_extra_mounts`]: static mounts ready to clone per node,
+/// model-subfolder specs deferred to per-node resolution.
+pub(crate) struct ResolvedExtraMounts {
+    pub static_mounts: Vec<ResolvedMount>,
+    pub model_mounts: Vec<ModelMountSpec>,
+}
+
+/// Dry-run placeholder for a model-subfolder mount: nodes may be off, so no
+/// HF cache to read. Emits a descriptive path so the printed `docker run`
+/// command shows *what* would be mounted, not a broken lookup.
+pub(crate) fn model_mount_placeholder(spec: &ModelMountSpec) -> ResolvedMount {
+    let sub = spec.subfolder.as_deref().unwrap_or("");
+    let host = format!(
+        "<hf_cache>/hub/{}/snapshots/<refs/main>/{}",
+        crate::config::model_cache_dir_name(&spec.repo),
+        sub
+    );
+    ResolvedMount {
+        host,
+        container: spec.container.clone(),
+        read_only: spec.read_only,
+        env: spec.env.clone(),
+    }
 }
 
 /// Collect NCCL cluster env vars for multi-node deployments.
@@ -274,6 +343,7 @@ mod tests {
     }
 
     fn run_with(recipe: &Recipe) -> String {
+        let resolved = resolve_extra_mounts(&minimal_config(), recipe).unwrap();
         let ctx = DockerCommandCtx {
             cluster: tcls(),
             nccl: tnccl(),
@@ -286,7 +356,7 @@ mod tests {
             keep: false,
             init_port: 25000,
             mods_remote_path: String::new(),
-            extra_mounts: resolve_extra_mounts(&minimal_config(), recipe).unwrap_or_default(),
+            extra_mounts: resolved.static_mounts,
         };
         build_docker_command(&ctx, &tcls().nodes[0], true).join(" ")
     }
@@ -321,6 +391,8 @@ mod tests {
         r.extra_mounts.push(ExtraMount {
             alias: Some("models".into()),
             host: Some("~/.cache/models".into()),
+            repo: None,
+            subfolder: None,
             container: "/models:ro".into(),
             read_only: true,
             env: Some("MODEL_DIR".into()),
@@ -339,6 +411,8 @@ mod tests {
         r.extra_mounts.push(ExtraMount {
             alias: Some("no_such_alias".into()),
             host: Some("~/.cache/jit".into()),
+            repo: None,
+            subfolder: None,
             container: "/root/.cache/jit".into(),
             read_only: false,
             env: None,
@@ -348,4 +422,42 @@ mod tests {
         assert!(cmd.contains("-v /home/u/.cache/jit:/root/.cache/jit"));
         assert!(!cmd.contains(":ro"));
     }
+
+    #[test]
+    fn model_subfolder_mount_uses_placeholder_in_dry_run() {
+        // In build_docker_command (no per-node SSH resolution here), a model
+        // mount shows the descriptive placeholder host so the printed command
+        // documents what would be mounted. start_cluster replaces this with
+        // the real resolved path at serve time (or keeps it under dry-run).
+        let mut r = base_recipe();
+        r.extra_mounts.push(ExtraMount {
+            alias: None,
+            host: None,
+            repo: Some("aidendle94/GLM-5.2-Hybrid-FP8-MXFP4".into()),
+            subfolder: Some("mtp-draft".into()),
+            container: "/models/mtp-draft".into(),
+            read_only: true,
+            env: Some("MTP_DRAFT_PATH".into()),
+        });
+        let cmd = run_with(&r);
+        // model_mounts are NOT in extra_mounts here (run_with doesn't resolve
+        // them per-node), so the -v flag for the model mount is absent — only
+        // start_cluster emits it. Verify the spec is parsed + the env is NOT
+        // injected yet (env injection follows extra_mounts, which is empty of
+        // the model mount here). This test guards the partition.
+        assert!(!cmd.contains("-v /models/mtp-draft"));
+        assert!(!cmd.contains("MTP_DRAFT_PATH"));
+        // But the recipe parsed into a model_mount spec:
+        let resolved = resolve_extra_mounts(&minimal_config(), &r).unwrap();
+        assert_eq!(resolved.model_mounts.len(), 1);
+        assert_eq!(resolved.static_mounts.len(), 0);
+        let ph = model_mount_placeholder(&resolved.model_mounts[0]);
+        assert_eq!(ph.container, "/models/mtp-draft");
+        assert!(
+            ph.host
+                .contains("models--aidendle94--GLM-5.2-Hybrid-FP8-MXFP4")
+        );
+        assert!(ph.host.contains("mtp-draft"));
+        assert!(ph.read_only);
+    }
 }

+ 66 - 2
src/docker/mod.rs

@@ -278,6 +278,10 @@ async fn start_cluster(
     } else {
         recipe.tensor_parallel_size
     };
+    // Resolve static mounts once (alias/host + tilde). Model-subfolder mounts
+    // (`repo`+`subfolder`) are kept as specs and resolved per node inside the
+    // spawn loop, since the HF snapshot hash can differ across nodes.
+    let resolved = command::resolve_extra_mounts(config, recipe)?;
     let mut handles = Vec::new();
     for node in &nodes {
         let node = (*node).clone();
@@ -302,13 +306,32 @@ async fn start_cluster(
             keep: opts.keep,
             init_port: config.init_port,
             mods_remote_path: opts.mods_remote_path.clone(),
-            extra_mounts: command::resolve_extra_mounts(&config, recipe)?,
+            extra_mounts: resolved.static_mounts.clone(),
         };
+        let hf_cache = config.resolved_hf_cache_dir();
+        let model_specs = resolved.model_mounts.clone();
         handles.push(tokio::spawn(async move {
             if !dry {
                 let _ = run_ssh(&config, &node, cleanup, dry).await;
             }
-            let cmd = command::build_docker_command(&dm_ctx, &node, is_head).join(" ");
+            // Resolve model-subfolder mounts per node from that node's HF cache
+            // (the snapshot hash can differ across nodes). Dry-run skips SSH
+            // and emits a descriptive placeholder path instead.
+            let mut ctx = dm_ctx.clone();
+            for spec in &model_specs {
+                let resolved_mount = if dry {
+                    command::model_mount_placeholder(spec)
+                } else {
+                    match resolve_model_mount_on_node(&config, &node, spec, &hf_cache).await {
+                        Ok(m) => m,
+                        Err(e) => {
+                            return (node.name, node.rank, engine, Err(e));
+                        }
+                    }
+                };
+                ctx.extra_mounts.push(resolved_mount);
+            }
+            let cmd = command::build_docker_command(&ctx, &node, is_head).join(" ");
             let result = run_ssh_with_retry(&config, &node, cmd, dry, 2).await;
             (node.name, node.rank, engine, result)
         }));
@@ -336,5 +359,46 @@ async fn start_cluster(
     Ok(())
 }
 
+/// Resolve a model-subfolder mount's host path on a specific node by reading
+/// its HF cache: `<hf_cache>/hub/<model_cache_dir_name(repo)>/snapshots/<refs/main>[/<subfolder>]`.
+/// `refs/main` always points at the latest commit, so the mount survives
+/// `huggingface-cli download` updates (no fragile symlink). If the repo or
+/// subfolder is absent, returns an error pointing at `huggingface-cli download`
+/// (the launcher never downloads).
+///
+/// Docker requires the bind-mount source to exist, so a `test -d` guards it.
+async fn resolve_model_mount_on_node(
+    config: &Config,
+    node: &Node,
+    spec: &command::ModelMountSpec,
+    hf_cache: &std::path::Path,
+) -> Result<command::ResolvedMount> {
+    let model_dir = crate::config::model_cache_dir_name(&spec.repo);
+    let sub = spec.subfolder.as_deref().unwrap_or("");
+    let hub_dir = format!("{}/hub/{}", hf_cache.display(), model_dir);
+    // `cat refs/main` prints the current commit hash; trim trailing newline.
+    let cmd = format!(
+        "ref=\"$(cat '{}/refs/main' 2>/dev/null)\" && test -n \"$ref\" && \
+         dir=\"{}/snapshots/$ref/{}\" && test -d \"$dir\" && printf '%s' \"$dir\"",
+        hub_dir, hub_dir, sub
+    );
+    let resolved_host = run_ssh(config, node, cmd, false).await?;
+    if resolved_host.trim().is_empty() {
+        let hint = format!(
+            "Model subfolder '{}' not found in HF cache on '{}': \
+             repo '{}' is not downloaded there or 'refs/main' is missing. \
+             Run `huggingface-cli download {}` on that node.",
+            sub, node.name, spec.repo, spec.repo
+        );
+        anyhow::bail!(hint);
+    }
+    Ok(command::ResolvedMount {
+        host: resolved_host,
+        container: spec.container.clone(),
+        read_only: spec.read_only,
+        env: spec.env.clone(),
+    })
+}
+
 #[cfg(test)]
 mod tests;