|
@@ -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'"));
|
|
|
|
|
+ }
|
|
|
|
|
+}
|