|
@@ -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
|
|
/// 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)
|
|
(mirror, cache_dir)
|
|
|
}
|
|
}
|
|
|
|
|
|
|
@@ -233,17 +239,28 @@ mod tests {
|
|
|
fn resolve_returns_mirror_and_falls_back_to_shared() {
|
|
fn resolve_returns_mirror_and_falls_back_to_shared() {
|
|
|
let c = cfg();
|
|
let c = cfg();
|
|
|
// No flag -> mount the shared HF cache, no mirror.
|
|
// 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!(mirror.is_none());
|
|
|
assert_eq!(cache, PathBuf::from("/home/u/.cache/huggingface"));
|
|
assert_eq!(cache, PathBuf::from("/home/u/.cache/huggingface"));
|
|
|
// Bare `--mirror` -> default mirror becomes the mounted cache dir.
|
|
// 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!(mirror, Some(PathBuf::from("/home/u/.cache/hf-models")));
|
|
|
assert_eq!(cache, PathBuf::from("/home/u/.cache/hf-models"));
|
|
assert_eq!(cache, PathBuf::from("/home/u/.cache/hf-models"));
|
|
|
// Explicit path.
|
|
// 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!(mirror, Some(PathBuf::from("/abs/mirror")));
|
|
|
assert_eq!(cache, 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]
|
|
#[test]
|