launcher/src/docker/mirror.rs

//! `--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,
/// 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)
}

/// 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};
    use crate::pre_serve::{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,
            mikrotik: None,
            global_mods_dir: None,
            drop_caches: true,
            mount_aliases: HashMap::new(),
            recipe_sources: Vec::new(),
            docker_image_sources: Vec::new(),
            cluster_sources: Vec::new(),
            registry: None,
        }
    }

    fn trec() -> Recipe {
        Recipe {
            build_profile: None,
            parent_model: None,
            archived: false,
            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, 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(""), 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"), 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]
    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![],
            local_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'"));
    }
}