Browse Source

feat: mount HF cache directory in Docker containers

- Add hf_cache_dir to Config (optional, auto-derives from ssh_user:
  /home/<user>/.cache/huggingface or /root/ for root)
- Add -v <cache_dir>:/cache/huggingface and -e HF_HOME=/cache/huggingface
  to Docker run command for all containers
- tmux runtime unchanged (runs natively, no mount needed)
- 25 tests pass, clippy clean, fmt clean
Alexander Korolev 4 tháng trước cách đây
mục cha
commit
a726b83bcf
3 tập tin đã thay đổi với 69 bổ sung và 4 xóa
  1. 26 0
      src/config.rs
  2. 17 2
      src/docker/mod.rs
  3. 26 2
      src/docker/tests.rs

+ 26 - 0
src/config.rs

@@ -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());

+ 17 - 2
src/docker/mod.rs

@@ -39,7 +39,17 @@ pub async fn serve(
         recipe.docker_image = image.clone();
     }
 
-    start_cluster(config, &cluster, nccl, recipe_name, &recipe, dry_run).await?;
+    let cache_dir = config.resolved_hf_cache_dir();
+    start_cluster(
+        config,
+        &cluster,
+        nccl,
+        recipe_name,
+        &recipe,
+        cache_dir,
+        dry_run,
+    )
+    .await?;
     if !dry_run {
         info!("\n📜 Последние логи контейнеров:");
         show_logs(config, &cluster, recipe_name, false, dry_run).await?;
@@ -208,6 +218,7 @@ async fn start_cluster(
     nccl: &NcclConfig,
     recipe_name: &str,
     recipe: &Recipe,
+    cache_dir: String,
     dry_run: bool,
 ) -> Result<()> {
     let nnodes = cluster.nodes.len() as u32;
@@ -233,12 +244,13 @@ async fn start_cluster(
             recipe_name, node.rank
         );
 
+        let cd = cache_dir.clone();
         handles.push(tokio::spawn(async move {
             if !dry {
                 let _ = run_ssh(&config, &node, cleanup, dry).await;
             }
             let cmd = build_docker_command(
-                &cluster, &nccl, &recipe, &rn, &node, is_head, pp_size, nnodes,
+                &cluster, &nccl, &recipe, &rn, &node, is_head, pp_size, nnodes, &cd,
             )
             .join(" ");
             let result = run_ssh_with_retry(&config, &node, cmd, dry, 2).await;
@@ -285,6 +297,7 @@ fn build_docker_command(
     is_head: bool,
     pp_size: u32,
     nnodes: u32,
+    cache_dir: &str,
 ) -> Vec<String> {
     let is_sglang = recipe.engine == "sglang";
 
@@ -294,6 +307,8 @@ fn build_docker_command(
         format!("-e NCCL_IB_HCA={}", nccl.ib_hca),
         format!("-e NCCL_SOCKET_IFNAME={}", nccl.socket_ifname),
         format!("-e NCCL_IB_GID_INDEX={}", nccl.ib_gid_index),
+        format!("-e HF_HOME=/cache/huggingface"),
+        format!("-v {}:/cache/huggingface", cache_dir),
     ];
 
     // Engine-specific host IP env var

+ 26 - 2
src/docker/tests.rs

@@ -52,7 +52,18 @@ fn trec() -> Recipe {
 }
 
 fn run(recipe: &Recipe, node: &Node, head: bool) -> String {
-    build_docker_command(&tcls(), &tnccl(), recipe, "r", node, head, 2, 2).join(" ")
+    build_docker_command(
+        &tcls(),
+        &tnccl(),
+        recipe,
+        "r",
+        node,
+        head,
+        2,
+        2,
+        "/home/user/.cache/huggingface",
+    )
+    .join(" ")
 }
 
 #[test]
@@ -137,7 +148,18 @@ fn sglang_recipe() -> Recipe {
 }
 
 fn srun(recipe: &Recipe, node: &Node, head: bool) -> String {
-    build_docker_command(&tcls(), &tnccl(), recipe, "sg", node, head, 4, 4).join(" ")
+    build_docker_command(
+        &tcls(),
+        &tnccl(),
+        recipe,
+        "sg",
+        node,
+        head,
+        4,
+        4,
+        "/home/user/.cache/huggingface",
+    )
+    .join(" ")
 }
 
 #[test]
@@ -162,6 +184,8 @@ fn sglang_head() {
     assert!(cmd.contains("-p 8000:8000"));
     assert!(cmd.contains("-e TORCH_NCCL_ASYNC_ERROR_HANDLING=1"));
     assert!(cmd.contains("--port 8000"));
+    assert!(cmd.contains("-e HF_HOME=/cache/huggingface"));
+    assert!(cmd.contains("-v /home/user/.cache/huggingface:/cache/huggingface"));
 }
 
 #[test]