Add HF home dir property inside System model

By default, install_dir/huggingface is used. Users can override it by
specifying custom value.

AI Dynamo uses this value for mapping this folder into a container.
Going further, CloudAI can take care about model download using HF CLI.

Author: amaslenn

src/cloudai/_core/base_installer.py

@@ -142,6 +142,8 @@ def is_installed(self, items: Iterable[Installable]) -> InstallStatusResult:
         """
         if not prepare_output_dir(self.system.install_path):
             return InstallStatusResult(False, f"Error preparing install dir '{self.system.install_path.absolute()}'")
+        elif not prepare_output_dir(self.system.hf_home_path):
+            return InstallStatusResult(False, f"Error preparing hf home dir '{self.system.hf_home_path.absolute()}'")
 
         install_results: dict[Installable, InstallStatusResult] = {}
         for item in self.all_items(items):
@@ -180,6 +182,8 @@ def install(self, items: Iterable[Installable]) -> InstallStatusResult:
 
         if not prepare_output_dir(self.system.install_path):
             return InstallStatusResult(False, f"Error preparing install dir '{self.system.install_path.absolute()}'")
+        elif not prepare_output_dir(self.system.hf_home_path):
+            return InstallStatusResult(False, f"Error preparing hf home dir '{self.system.hf_home_path.absolute()}'")
 
         logging.debug(f"Going to install {len(set(items))} uniq item(s) (total is {len(list(items))})")
 

src/cloudai/_core/system.py

@@ -37,6 +37,7 @@ class System(ABC, BaseModel):
     scheduler: str
     install_path: Path
     output_path: Path
+    hf_home_path: Path = Field(default_factory=lambda data: data["install_path"] / "huggingface")
     global_env_vars: dict[str, Any] = Field(default_factory=dict)
     monitor_interval: int = 1
 

src/cloudai/workloads/ai_dynamo/ai_dynamo.py

@@ -64,7 +64,6 @@ class AIDynamoCmdArgs(CmdArgs):
     """Arguments for AI Dynamo."""
 
     docker_image_url: str
-    huggingface_home_host_path: Path = Path.home() / ".cache/huggingface"
     huggingface_home_container_path: Path = Path("/root/.cache/huggingface")
     dynamo: AIDynamoArgs
     genai_perf: GenAIPerfArgs
@@ -103,13 +102,6 @@ def docker_image(self) -> DockerImage:
     def installables(self) -> list[Installable]:
         return [self.docker_image, self.script, self.dynamo_repo, self.python_executable]
 
-    @property
-    def huggingface_home_host_path(self) -> Path:
-        path = Path(self.cmd_args.huggingface_home_host_path)
-        if not path.is_dir():
-            raise FileNotFoundError(f"HuggingFace home path not found at {path}")
-        return path
-
     @property
     def python_executable(self) -> PythonExecutable:
         if not self._python_executable:

src/cloudai/workloads/ai_dynamo/slurm_command_gen_strategy.py

@@ -35,7 +35,7 @@ def _container_mounts(self) -> list[str]:
 
         mounts = [
             f"{dynamo_repo_path}:{dynamo_repo_path}",
-            f"{td.cmd_args.huggingface_home_host_path}:{td.cmd_args.huggingface_home_container_path}",
+            f"{self.system.hf_home_path.absolute()}:{td.cmd_args.huggingface_home_container_path}",
             f"{td.script.installed_path.absolute()!s}:{td.script.installed_path.absolute()!s}",
         ]
 

tests/ref_data/ai-dynamo.sbatch

@@ -10,9 +10,9 @@
 
 export SLURM_JOB_MASTER_NODE=$(scontrol show hostname $SLURM_JOB_NODELIST | head -n 1)
 
-srun --export=ALL --mpi=pmix -N2 --container-image=nvcr.io/nvidia/ai-dynamo:24.09 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output,__INSTALL_DIR__:__INSTALL_DIR__,__OUTPUT_DIR__/output/hf_home:/root/.cache/huggingface,__CLOUDAI_DIR__/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh:__CLOUDAI_DIR__/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh --output=__OUTPUT_DIR__/output/mapping-stdout.txt --error=__OUTPUT_DIR__/output/mapping-stderr.txt bash -c "echo \$(date): \$(hostname):node \${SLURM_NODEID}:rank \${SLURM_PROCID}."
+srun --export=ALL --mpi=pmix -N2 --container-image=nvcr.io/nvidia/ai-dynamo:24.09 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output,__INSTALL_DIR__:__INSTALL_DIR__,__INSTALL_DIR__/huggingface:/root/.cache/huggingface,__CLOUDAI_DIR__/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh:__CLOUDAI_DIR__/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh --output=__OUTPUT_DIR__/output/mapping-stdout.txt --error=__OUTPUT_DIR__/output/mapping-stderr.txt bash -c "echo \$(date): \$(hostname):node \${SLURM_NODEID}:rank \${SLURM_PROCID}."
 
-srun --export=ALL --mpi=pmix -N2 --container-image=nvcr.io/nvidia/ai-dynamo:24.09 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output,__INSTALL_DIR__:__INSTALL_DIR__,__OUTPUT_DIR__/output/hf_home:/root/.cache/huggingface,__CLOUDAI_DIR__/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh:__CLOUDAI_DIR__/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash /cloudai_install/slurm-metadata.sh
+srun --export=ALL --mpi=pmix -N2 --container-image=nvcr.io/nvidia/ai-dynamo:24.09 --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output,__INSTALL_DIR__:__INSTALL_DIR__,__INSTALL_DIR__/huggingface:/root/.cache/huggingface,__CLOUDAI_DIR__/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh:__CLOUDAI_DIR__/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash /cloudai_install/slurm-metadata.sh
 
 num_retries=${DYNAMO_NUM_RETRY_ON_FAILURE:-0}
 for try in $(seq 0 $num_retries); do
@@ -23,7 +23,7 @@ for try in $(seq 0 $num_retries); do
   --mpi=pmix \
   -N2 \
   --container-image=nvcr.io/nvidia/ai-dynamo:24.09 \
-  --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output,__INSTALL_DIR__:__INSTALL_DIR__,__OUTPUT_DIR__/output/hf_home:/root/.cache/huggingface,__CLOUDAI_DIR__/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh:__CLOUDAI_DIR__/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh \
+  --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output,__INSTALL_DIR__:__INSTALL_DIR__,__INSTALL_DIR__/huggingface:/root/.cache/huggingface,__CLOUDAI_DIR__/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh:__CLOUDAI_DIR__/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh \
   --nodes=2 \
   --ntasks=2 \
   --ntasks-per-node=1 \

tests/slurm_command_gen_strategy/test_ai_dynamo_slurm_command_gen_strategy.py

@@ -36,7 +36,6 @@
 def cmd_args() -> AIDynamoCmdArgs:
     return AIDynamoCmdArgs(
         docker_image_url="url",
-        huggingface_home_host_path=Path.home() / ".cache/huggingface",
         huggingface_home_container_path=Path("/root/.cache/huggingface"),
         dynamo=AIDynamoArgs(
             prefill_worker=PrefillWorkerArgs(
@@ -76,10 +75,6 @@ def cmd_args() -> AIDynamoCmdArgs:
 
 @pytest.fixture
 def test_run(tmp_path: Path, cmd_args: AIDynamoCmdArgs) -> TestRun:
-    hf_home = tmp_path / "huggingface"
-    hf_home.mkdir()
-    cmd_args.huggingface_home_host_path = hf_home
-
     dynamo_repo_path = tmp_path / "dynamo_repo"
     dynamo_repo_path.mkdir()
 
@@ -99,20 +94,6 @@ def strategy(slurm_system: SlurmSystem, test_run: TestRun) -> AIDynamoSlurmComma
     return AIDynamoSlurmCommandGenStrategy(slurm_system, test_run)
 
 
-def test_hugging_face_home_path_valid(test_run: TestRun) -> None:
-    td = cast(AIDynamoTestDefinition, test_run.test)
-    path = td.huggingface_home_host_path
-    assert path.exists()
-    assert path.is_dir()
-
-
-def test_hugging_face_home_path_missing(test_run: TestRun) -> None:
-    td = cast(AIDynamoTestDefinition, test_run.test)
-    td.cmd_args.huggingface_home_host_path = Path("/nonexistent")
-    with pytest.raises(FileNotFoundError):
-        _ = td.huggingface_home_host_path
-
-
 def test_container_mounts(strategy: AIDynamoSlurmCommandGenStrategy, test_run: TestRun) -> None:
     mounts = strategy._container_mounts()
     td = cast(AIDynamoTestDefinition, test_run.test)
@@ -121,7 +102,7 @@ def test_container_mounts(strategy: AIDynamoSlurmCommandGenStrategy, test_run: T
 
     assert mounts == [
         f"{dynamo_repo_path!s}:{dynamo_repo_path!s}",
-        f"{td.cmd_args.huggingface_home_host_path!s}:{td.cmd_args.huggingface_home_container_path!s}",
+        f"{strategy.system.hf_home_path.absolute()!s}:{td.cmd_args.huggingface_home_container_path!s}",
         f"{td.script.installed_path.absolute()!s}:{td.script.installed_path.absolute()!s}",
     ]
 

tests/test_acceptance.py

@@ -431,7 +431,6 @@ def test_req(request, slurm_system: SlurmSystem, partial_tr: partial[TestRun]) -
                 ),
                 cmd_args=AIDynamoCmdArgs(
                     docker_image_url="nvcr.io/nvidia/ai-dynamo:24.09",
-                    huggingface_home_host_path=Path.home() / ".cache/huggingface",
                     dynamo=AIDynamoArgs(
                         backend="vllm",
                         prefill_worker=PrefillWorkerArgs(

tests/test_slurm_system.py

@@ -657,3 +657,23 @@ def test_update(
 
         all_nodes_set = set([node for p in slurm_system.partitions for node in p.slurm_nodes])
         assert all_nodes_set == set(expected_nodes)
+
+
+class TestHfHomePath:
+    @pytest.fixture
+    def system_args(self, tmp_path: Path) -> dict:
+        return {
+            "name": "test_system",
+            "install_path": tmp_path / "install",
+            "output_path": tmp_path / "output",
+            "partitions": [],
+            "default_partition": "main",
+        }
+
+    def test_default(self, system_args: dict):
+        system = SlurmSystem(**system_args)
+        assert system.hf_home_path == system_args["install_path"] / "huggingface"
+
+    def test_custom(self, system_args: dict):
+        system = SlurmSystem(**system_args, hf_home_path=system_args["output_path"] / "custom")
+        assert system.hf_home_path == system_args["output_path"] / "custom"