Merge pull request #737 from NVIDIA/am/hf-model-k8s

amaslenn · web-flow · commit 5e1a1cf60ab9 · 2025-12-04T13:28:19.000+01:00
Add support for HF model in K8s
diff --git a/src/cloudai/systems/kubernetes/kubernetes_installer.py b/src/cloudai/systems/kubernetes/kubernetes_installer.py
@@ -22,13 +22,28 @@
 from pathlib import Path
 from shutil import rmtree
 
-from cloudai.core import BaseInstaller, DockerImage, File, GitRepo, Installable, InstallStatusResult, PythonExecutable
+from cloudai.core import (
+    BaseInstaller,
+    DockerImage,
+    File,
+    GitRepo,
+    HFModel,
+    Installable,
+    InstallStatusResult,
+    PythonExecutable,
+    System,
+)
+from cloudai.util.hf_model_manager import HFModelManager
 from cloudai.util.lazy_imports import lazy
 
 
 class KubernetesInstaller(BaseInstaller):
     """Installer for Kubernetes systems."""
 
+    def __init__(self, system: System) -> None:
+        super().__init__(system)
+        self.hf_model_manager = HFModelManager(system.hf_home_path)
+
     def _check_prerequisites(self) -> InstallStatusResult:
         """
         Check for the presence of required binaries and Kubernetes configurations.
@@ -70,6 +85,8 @@ def install_one(self, item: Installable) -> InstallStatusResult:
             item.installed_path = self.system.install_path / item.src.name
             shutil.copyfile(item.src, item.installed_path, follow_symlinks=False)
             return InstallStatusResult(True)
+        elif isinstance(item, HFModel):
+            return self.hf_model_manager.download_model(item)
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
 
     def uninstall_one(self, item: Installable) -> InstallStatusResult:
@@ -86,6 +103,8 @@ def uninstall_one(self, item: Installable) -> InstallStatusResult:
                 return InstallStatusResult(True)
             logging.debug(f"File {item.installed_path} does not exist.")
             return InstallStatusResult(True)
+        elif isinstance(item, HFModel):
+            return self.hf_model_manager.remove_model(item)
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
 
     def is_installed_one(self, item: Installable) -> InstallStatusResult:
@@ -99,6 +118,8 @@ def is_installed_one(self, item: Installable) -> InstallStatusResult:
             return InstallStatusResult(False, f"Git repository {item.url} not cloned")
         elif isinstance(item, PythonExecutable):
             return self._is_python_executable_installed(item)
+        elif isinstance(item, HFModel):
+            return self.hf_model_manager.is_model_downloaded(item)
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
 
     def mark_as_installed_one(self, item: Installable) -> InstallStatusResult:
@@ -111,6 +132,9 @@ def mark_as_installed_one(self, item: Installable) -> InstallStatusResult:
             item.git_repo.installed_path = self.system.install_path / item.git_repo.repo_name
             item.venv_path = self.system.install_path / item.venv_name
             return InstallStatusResult(True)
+        elif isinstance(item, HFModel):
+            item.installed_path = self.system.hf_home_path  # fake path is OK here as the whole HF home will be mounted
+            return InstallStatusResult(True)
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
 
     def _install_one_git_repo(self, item: GitRepo) -> InstallStatusResult:
diff --git a/src/cloudai/systems/slurm/slurm_installer.py b/src/cloudai/systems/slurm/slurm_installer.py
@@ -52,7 +52,7 @@ def __init__(self, system: SlurmSystem):
         super().__init__(system)
         self.system = system
         self.docker_image_cache_manager = DockerImageCacheManager(system)
-        self.hf_model_downloader = HFModelManager(system.hf_home_path)
+        self.hf_model_manager = HFModelManager(system.hf_home_path)
 
     def _check_prerequisites(self) -> InstallStatusResult:
         base_prerequisites_result = super()._check_prerequisites()
@@ -102,7 +102,7 @@ def install_one(self, item: Installable) -> InstallStatusResult:
             shutil.copyfile(item.src, item.installed_path, follow_symlinks=False)
             return InstallStatusResult(True)
         elif isinstance(item, HFModel):
-            return self.hf_model_downloader.download_model(item)
+            return self.hf_model_manager.download_model(item)
 
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
 
@@ -123,7 +123,7 @@ def uninstall_one(self, item: Installable) -> InstallStatusResult:
             logging.debug(f"File {item.installed_path} does not exist.")
             return InstallStatusResult(True)
         elif isinstance(item, HFModel):
-            return self.hf_model_downloader.remove_model(item)
+            return self.hf_model_manager.remove_model(item)
 
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
 
@@ -149,7 +149,7 @@ def is_installed_one(self, item: Installable) -> InstallStatusResult:
                 return InstallStatusResult(True)
             return InstallStatusResult(False, f"File {item.installed_path} does not exist")
         elif isinstance(item, HFModel):
-            return self.hf_model_downloader.is_model_downloaded(item)
+            return self.hf_model_manager.is_model_downloaded(item)
 
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
 
diff --git a/tests/test_base_installer.py b/tests/test_base_installer.py
@@ -17,13 +17,25 @@
 import shutil
 from concurrent.futures import Future
 from pathlib import Path
-from typing import Generator, cast
+from typing import Any, Generator, cast
 from unittest.mock import Mock, patch
 
 import pytest
 
-from cloudai.core import BaseInstaller, DockerImage, File, GitRepo, Installable, InstallStatusResult
+from cloudai.core import (
+    BaseInstaller,
+    DockerImage,
+    File,
+    GitRepo,
+    HFModel,
+    Installable,
+    InstallStatusResult,
+    PythonExecutable,
+)
+from cloudai.systems.kubernetes.kubernetes_installer import KubernetesInstaller
+from cloudai.systems.kubernetes.kubernetes_system import KubernetesSystem
 from cloudai.systems.slurm import SlurmInstaller, SlurmSystem
+from cloudai.systems.slurm.docker_image_cache_manager import DockerImageCacheResult
 from cloudai.util import prepare_output_dir
 
 
@@ -280,3 +292,63 @@ def test_order_of_items_does_not_matter(self, installer: SlurmInstaller):
             "First ('self', present in the statuses) file was not marked as installed"
         )
         assert f2._installed_path is not None, "Second file was not marked as installed"
+
+
+@pytest.fixture(params=["k8s", "slurm"])
+def installer(
+    request: Any, kubernetes_system: KubernetesSystem, slurm_system: SlurmSystem
+) -> KubernetesInstaller | SlurmInstaller:
+    installer = KubernetesInstaller(kubernetes_system) if request.param == "k8s" else SlurmInstaller(slurm_system)
+
+    installer.system.install_path.mkdir(parents=True)
+    installer._check_low_thread_environment = lambda threshold=None: False
+    return installer
+
+
+def test_check_supported(installer: KubernetesInstaller | SlurmInstaller):
+    if isinstance(installer, SlurmInstaller):
+        installer._install_docker_image = lambda item: DockerImageCacheResult(True)
+        installer._uninstall_docker_image = lambda item: DockerImageCacheResult(True)
+        installer.docker_image_cache_manager.check_docker_image_exists = Mock(return_value=DockerImageCacheResult(True))
+    installer._install_python_executable = lambda item: InstallStatusResult(True)
+    installer._uninstall_python_executable = lambda item: InstallStatusResult(True)
+    installer._is_python_executable_installed = lambda item: InstallStatusResult(True)
+    installer.hf_model_manager = Mock()
+
+    git = GitRepo(url="./git_url", commit="commit_hash")
+    items = [DockerImage("fake_url/img"), PythonExecutable(git), HFModel("model_name")]
+    if isinstance(installer, SlurmInstaller):
+        items.append(File(Path(__file__)))
+    for item in items:
+        res = installer.install_one(item)
+        assert res.success, f"Failed to install {item} for {installer.__class__.__name__=} {res.message=}"
+
+        res = installer.is_installed_one(item)
+        assert res.success, f"Failed to check installation of {item} for {installer.__class__.__name__=} {res.message=}"
+
+        res = installer.uninstall_one(item)
+        assert res.success, f"Failed to uninstall {item} for {installer.__class__.__name__=} {res.message=}"
+
+        res = installer.mark_as_installed_one(item)
+        assert res.success, f"Failed to mark as installed {item} for {installer.__class__.__name__=} {res.message=}"
+
+
+class MyInstallable(Installable):
+    def __eq__(self, other: object) -> bool:
+        return True
+
+    def __hash__(self) -> int:
+        return hash("MyInstallable")
+
+
+def test_check_unsupported(installer: KubernetesInstaller | SlurmInstaller):
+    unsupported = MyInstallable()
+    for func in [
+        installer.install_one,
+        installer.uninstall_one,
+        installer.is_installed_one,
+        installer.mark_as_installed_one,
+    ]:
+        res = func(unsupported)
+        assert not res.success
+        assert res.message == f"Unsupported item type: {type(unsupported)}"
diff --git a/tests/test_slurm_installer.py b/tests/test_slurm_installer.py
@@ -21,8 +21,7 @@
 
 import pytest
 
-from cloudai.core import DockerImage, File, GitRepo, Installable, InstallStatusResult, PythonExecutable
-from cloudai.systems.slurm.docker_image_cache_manager import DockerImageCacheResult
+from cloudai.core import DockerImage, File, GitRepo, InstallStatusResult, PythonExecutable
 from cloudai.systems.slurm.slurm_installer import SlurmInstaller
 from cloudai.systems.slurm.slurm_system import SlurmSystem
 from cloudai.workloads.nemo_launcher import NeMoLauncherCmdArgs, NeMoLauncherTestDefinition
@@ -323,50 +322,6 @@ def test_is_installed_checks_content(self, installer: SlurmInstaller, f: File):
         assert not res.success
 
 
-def test_check_supported(slurm_system: SlurmSystem):
-    slurm_system.install_path.mkdir()
-    installer = SlurmInstaller(slurm_system)
-    installer._install_docker_image = lambda item: DockerImageCacheResult(True)
-    installer._install_python_executable = lambda item: InstallStatusResult(True)
-    installer._uninstall_docker_image = lambda item: DockerImageCacheResult(True)
-    installer._uninstall_python_executable = lambda item: InstallStatusResult(True)
-    installer._is_python_executable_installed = lambda item: InstallStatusResult(True)
-    installer.docker_image_cache_manager.check_docker_image_exists = Mock(return_value=DockerImageCacheResult(True))
-
-    git = GitRepo(url="./git_url", commit="commit_hash")
-    items = [DockerImage("fake_url/img"), PythonExecutable(git), File(Path(__file__))]
-    for item in items:
-        res = installer.install_one(item)
-        assert res.success
-
-        res = installer.is_installed_one(item)
-        assert res.success
-
-        res = installer.uninstall_one(item)
-        assert res.success
-
-        res = installer.mark_as_installed_one(item)
-        assert res.success
-
-    class MyInstallable(Installable):
-        def __eq__(self, other: object) -> bool:
-            return True
-
-        def __hash__(self) -> int:
-            return hash("MyInstallable")
-
-    unsupported = MyInstallable()
-    for func in [
-        installer.install_one,
-        installer.uninstall_one,
-        installer.is_installed_one,
-        installer.mark_as_installed_one,
-    ]:
-        res = func(unsupported)
-        assert not res.success
-        assert res.message == f"Unsupported item type: {type(unsupported)}"
-
-
 def test_mark_as_installed(slurm_system: SlurmSystem):
     tdef = NeMoLauncherTestDefinition(
         name="name", description="desc", test_template_name="tt", cmd_args=NeMoLauncherCmdArgs()
