Browse Source

feat: serve --hf-cache <path> overrides the host-side HF cache dir per serve (config hf_cache_dir stays the global knob; mutual exclusion with --mirror)

Alexander Korolev 4 tuần trước cách đây
mục cha
commit
51b7404aa8
8 tập tin đã thay đổi với 57 bổ sung và 11 xóa
  1. 7 0
      AGENTS.md
  2. 1 0
      README.md
  3. 7 0
      src/cli.rs
  4. 1 0
      src/docker/build.rs
  5. 27 10
      src/docker/mirror.rs
  6. 11 1
      src/docker/mod.rs
  7. 2 0
      src/main.rs
  8. 1 0
      src/web/api.rs

+ 7 - 0
AGENTS.md

@@ -117,6 +117,13 @@ 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.
 
+`serve --hf-cache <path>` overrides the **host-side HF cache dir** for that
+serve (config `hf_cache_dir`, else `~/.cache/huggingface`); `~` expands to
+the node's SSH-user home (quote it — the shell would expand it locally first).
+The mount target stays `/root/.cache/huggingface`, so model path refs and
+repo mounts resolve as usual. Mutually exclusive with `--mirror` (both pick
+the HF mount source — the error says so).
+
 **Cluster health is cluster-level** (`health`, web dashboard): a cluster is
 UP only when *all* its nodes are reachable AND its **blocking**
 `extra_entities` report `on`. HA entities are probed via HA REST; a

+ 1 - 0
README.md

@@ -58,6 +58,7 @@ sm121 serve --recipe R --solo                    # first free node, single-node
 sm121 serve --recipe R --force                   # skip occupancy check (co-serve)
 sm121 serve --recipe R --fit=2                   # scale to N nodes, TP tracks
 sm121 serve --recipe R --mirror                  # per-node local HF mirror (no NAS at load)
+sm121 serve --recipe R --hf-cache ~/hf-local     # override the host-side HF cache dir
 sm121 serve --recipe R --clear                   # rm -rf host extra_mounts dirs first
 sm121 serve --recipe R --build-node dgx1         # build the image first (if stale)
 sm121 serve --recipe R --sync-image              # push the image to launch nodes first

+ 7 - 0
src/cli.rs

@@ -83,9 +83,16 @@ pub enum Commands {
 
         /// Mount a per-node local HF mirror (bare = ~/.cache/hf-models;
         /// `--mirror <path>` = that path; first serve copies per node).
+        /// Mutually exclusive with `--hf-cache` (both pick the HF mount source).
         #[arg(long, num_args = 0..=1, default_missing_value = "")]
         mirror: Option<String>,
 
+        /// Override the host-side HF cache dir for this serve (config
+        /// `hf_cache_dir`, else `~/.cache/huggingface`). `~` expands to the
+        /// node's SSH-user home. Mutually exclusive with `--mirror`.
+        #[arg(long)]
+        hf_cache: Option<String>,
+
         /// Scale the recipe to N nodes (`--fit=N`; bare = whole cluster, TP
         /// tracks it; incompatible with --solo).
         #[arg(long, num_args = 0..=1, default_missing_value = "")]

+ 1 - 0
src/docker/build.rs

@@ -347,6 +347,7 @@ pub async fn handle_build_cmd(
         force: false,
         solo: None,
         mirror: None,
+        hf_cache: None,
         fit: None,
         clear: false,
     };

+ 27 - 10
src/docker/mirror.rs

@@ -30,13 +30,19 @@ pub fn resolve_mirror_dir(config: &Config, arg: &str) -> PathBuf {
 }
 
 /// 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());
+/// the cache dir to mount at `/root/.cache/huggingface`: mirror when set,
+/// else the `--hf-cache` override (tilde-expanded per node), else the shared
+/// HF cache from config.
+pub fn resolve(
+    config: &Config,
+    mirror: Option<&str>,
+    hf_cache: Option<&str>,
+) -> (Option<PathBuf>, PathBuf) {
+    let mirror = mirror.map(|m| resolve_mirror_dir(config, m));
+    let cache_dir = mirror.clone().unwrap_or_else(|| match hf_cache {
+        Some(p) => PathBuf::from(crate::config::expand_tilde(p, &config.ssh_user)),
+        None => config.resolved_hf_cache_dir(),
+    });
     (mirror, cache_dir)
 }
 
@@ -233,17 +239,28 @@ mod tests {
     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);
+        let (mirror, cache) = resolve(&c, None, 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(""));
+        let (mirror, cache) = resolve(&c, Some(""), None);
         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"));
+        let (mirror, cache) = resolve(&c, Some("/abs/mirror"), None);
         assert_eq!(mirror, Some(PathBuf::from("/abs/mirror")));
         assert_eq!(cache, PathBuf::from("/abs/mirror"));
+        // `--hf-cache` override (tilde expands to the node's SSH-user home).
+        let (mirror, cache) = resolve(&c, None, Some("~/hf-local"));
+        assert!(mirror.is_none());
+        assert_eq!(cache, PathBuf::from("/home/u/hf-local"));
+        // An absolute override passes through.
+        let (_, cache) = resolve(&c, None, Some("/mnt/nas/ai/huggingface"));
+        assert_eq!(cache, PathBuf::from("/mnt/nas/ai/huggingface"));
+        // A mirror wins over the shared/override source (the flag conflict
+        // is rejected upstream in serve()).
+        let (_, cache) = resolve(&c, Some(""), Some("~/hf-local"));
+        assert_eq!(cache, PathBuf::from("/home/u/.cache/hf-models"));
     }
 
     #[test]

+ 11 - 1
src/docker/mod.rs

@@ -47,6 +47,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 the host-side HF cache dir (`--hf-cache`); wins over
+    /// the config `hf_cache_dir` default. Mutually exclusive with `mirror`
+    /// (both pick the HF mount source).
+    pub hf_cache: 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`.
@@ -255,7 +259,13 @@ pub async fn serve(
     // `--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 opts.mirror.is_some() && opts.hf_cache.is_some() {
+        anyhow::bail!(
+            "--hf-cache and --mirror are mutually exclusive (both pick the HF mount source)"
+        );
+    }
+    let (mirror, cache_dir) =
+        mirror::resolve(config, opts.mirror.as_deref(), opts.hf_cache.as_deref());
     if let Some(ref m) = mirror {
         info!(
             "serving from local HF mirror {} (models copied per node on first use)",

+ 2 - 0
src/main.rs

@@ -90,6 +90,7 @@ async fn main() -> Result<()> {
             force,
             solo,
             mirror,
+            hf_cache,
             fit,
             clear,
             no_push,
@@ -113,6 +114,7 @@ async fn main() -> Result<()> {
                     force,
                     solo,
                     mirror,
+                    hf_cache,
                     fit,
                     clear,
                 },

+ 1 - 0
src/web/api.rs

@@ -74,6 +74,7 @@ pub(crate) async fn serve(
         force: req.force,
         solo: req.solo,
         mirror: req.mirror,
+        hf_cache: None,
         fit: req.fit,
         clear: req.clear,
     };