|
|
@@ -35,6 +35,26 @@ pub struct Config {
|
|
|
pub clusters: Vec<Cluster>,
|
|
|
#[serde(default)]
|
|
|
pub nccl: HashMap<String, NcclConfig>,
|
|
|
+ #[serde(default)]
|
|
|
+ pub hf_cache_dir: Option<String>,
|
|
|
+}
|
|
|
+
|
|
|
+impl Config {
|
|
|
+ /// Returns the HF cache directory path for volume mounting.
|
|
|
+ /// If set in config, uses that value; otherwise derives from
|
|
|
+ /// `ssh_user` to produce `/home/<user>/.cache/huggingface`
|
|
|
+ /// (or `/root/.cache/huggingface` for root).
|
|
|
+ pub fn resolved_hf_cache_dir(&self) -> String {
|
|
|
+ if let Some(ref dir) = self.hf_cache_dir {
|
|
|
+ return dir.clone();
|
|
|
+ }
|
|
|
+ let home = if self.ssh_user == "root" {
|
|
|
+ "/root".to_string()
|
|
|
+ } else {
|
|
|
+ format!("/home/{}", self.ssh_user)
|
|
|
+ };
|
|
|
+ format!("{}/.cache/huggingface", home)
|
|
|
+ }
|
|
|
}
|
|
|
|
|
|
#[derive(Clone, Deserialize)]
|
|
|
@@ -310,6 +330,7 @@ nodes = [
|
|
|
ssh_key_path: "/path/to/key".to_string(),
|
|
|
clusters: vec![test_cluster()],
|
|
|
nccl: HashMap::new(),
|
|
|
+ hf_cache_dir: None,
|
|
|
};
|
|
|
config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
|
|
|
|
@@ -325,6 +346,7 @@ nodes = [
|
|
|
ssh_key_path: "/path/to/key".to_string(),
|
|
|
clusters: vec![test_cluster()],
|
|
|
nccl: HashMap::new(),
|
|
|
+ hf_cache_dir: None,
|
|
|
};
|
|
|
config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
|
|
|
|
@@ -340,6 +362,7 @@ nodes = [
|
|
|
ssh_key_path: "/path/to/key".to_string(),
|
|
|
clusters: vec![test_cluster()],
|
|
|
nccl: HashMap::new(),
|
|
|
+ hf_cache_dir: None,
|
|
|
};
|
|
|
let result = select_cluster(&config, "auto");
|
|
|
assert!(result.is_err());
|
|
|
@@ -352,6 +375,7 @@ nodes = [
|
|
|
ssh_key_path: "/path/to/key".to_string(),
|
|
|
clusters: vec![],
|
|
|
nccl: HashMap::new(),
|
|
|
+ hf_cache_dir: None,
|
|
|
};
|
|
|
let result = select_cluster(&config, "nonexistent");
|
|
|
assert!(result.is_err());
|
|
|
@@ -366,6 +390,7 @@ nodes = [
|
|
|
ssh_key_path: "/k".to_string(),
|
|
|
clusters: vec![test_cluster()],
|
|
|
nccl: nccl_map,
|
|
|
+ hf_cache_dir: None,
|
|
|
};
|
|
|
let cluster = select_cluster_by_name(&config, "test-cluster").unwrap();
|
|
|
assert_eq!(cluster.name, "test-cluster");
|
|
|
@@ -379,6 +404,7 @@ nodes = [
|
|
|
ssh_key_path: "/k".to_string(),
|
|
|
clusters: vec![test_cluster()],
|
|
|
nccl: HashMap::new(),
|
|
|
+ hf_cache_dir: None,
|
|
|
};
|
|
|
let result = select_cluster_by_name(&config, "no-such-cluster");
|
|
|
assert!(result.is_err());
|