Merge pull request #394 from NVIDIA/am/no-cache-check

Allow skipping cache validation

Author: amaslenn

USER_GUIDE.md

@@ -438,6 +438,14 @@ CloudAI runs all slurm jobs using containers. To simplify file system related ta
     These mounts are not verified for validity and do not override default mounts.
 1. Test specific mounts can be mounted in-code.
 
+### Head node without shared storage available on compute nodes
+When compute nodes don't share file system with head node, `--enable-cache-without-check` for `run` and `dry-run` skips real check for cache existence, but still builds all paths correctly. The flow is like this:
+
+1. _[on the head node]_ run `cloudai install`
+1. _[on the head node]_ copy cache to compute nodes
+1. Modify system.toml to set compute nodes' installation root
+1. Run `cloudai run --enable-cache-without-check ...`
+
 #### Dev details
 `SlurmCommandGenStrategy` defines abstract method `_container_mounts(tr: TestRun)` that must be implemented by every subclass. This method is used in `SlurmCommandGenStrategy.container_mounts(tr: TestRun)` (defined as `@final`) where mounts like `/cloudai_run_results` (default mount), `TestDefinition.extra_container_mounts` (from Test TOML) and test specific mounts (defined in-code) are added.
 

src/cloudai/_core/base_installer.py

@@ -187,6 +187,23 @@ def uninstall(self, items: Iterable[Installable]) -> InstallStatusResult:
         nfailed = len([result for result in uninstall_results.values() if result != "Success"])
         return InstallStatusResult(False, f"{nfailed} item(s) failed to uninstall.", uninstall_results)
 
+    @final
+    def mark_as_installed(self, items: Iterable[Installable]) -> InstallStatusResult:
+        """
+        Mark the installable items as installed.
+
+        Args:
+            items (Iterable[Installable]): Items to mark as installed.
+
+        Returns:
+            InstallStatusResult: Result containing the status and error message if any.
+        """
+        install_results = {}
+        for item in items:
+            self.mark_as_installed_one(item)
+
+        return InstallStatusResult(True, "All items marked as installed successfully.", install_results)
+
     @abstractmethod
     def install_one(self, item: Installable) -> InstallStatusResult: ...
 
@@ -195,3 +212,6 @@ def uninstall_one(self, item: Installable) -> InstallStatusResult: ...
 
     @abstractmethod
     def is_installed_one(self, item: Installable) -> InstallStatusResult: ...
+
+    @abstractmethod
+    def mark_as_installed_one(self, item: Installable) -> InstallStatusResult: ...

src/cloudai/cli/cli.py

@@ -124,7 +124,7 @@ def add_run_and_dry_run(self):
             desc = "Execute the test scenarios."
             if mode == "dry-run":
                 desc = "Perform a dry-run of the test scenarios without executing them."
-            self.add_command(
+            p = self.add_command(
                 mode,
                 desc,
                 handle_dry_run_and_run,
@@ -133,6 +133,12 @@ def add_run_and_dry_run(self):
                 test_scenario=True,
                 output_dir=False,
             )
+            p.add_argument(
+                "--enable-cache-without-check",
+                action="store_true",
+                help="Enable cache without checking.",
+                default=False,
+            )
 
     def add_install_and_uninstall(self):
         for mode in {"install", "uninstall"}:

src/cloudai/cli/handlers.py

@@ -191,7 +191,10 @@ def handle_dry_run_and_run(args: argparse.Namespace) -> int:
         raise NotImplementedError(f"No installer available for scheduler: {system.scheduler}")
     installer = installer_class(system)
 
-    result = installer.is_installed(installables)
+    if args.enable_cache_without_check:
+        result = installer.mark_as_installed(installables)
+    else:
+        result = installer.is_installed(installables)
 
     if args.mode == "run" and not result.success:
         logging.error("CloudAI has not been installed. Please run install mode first.")

src/cloudai/installer/slurm_installer.py

@@ -177,6 +177,23 @@ def is_installed_one(self, item: Installable) -> InstallStatusResult:
 
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
 
+    def mark_as_installed_one(self, item: Installable) -> InstallStatusResult:
+        if isinstance(item, DockerImage):
+            item.installed_path = self.system.install_path / item.cache_filename
+            return InstallStatusResult(True)
+        elif isinstance(item, GitRepo):
+            item.installed_path = self.system.install_path / item.repo_name
+            return InstallStatusResult(True)
+        elif isinstance(item, PythonExecutable):
+            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, File):
+            item.installed_path = self.system.install_path / item.src.name
+            return InstallStatusResult(True)
+
+        return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
+
     def _install_docker_image(self, item: DockerImage) -> DockerImageCacheResult:
         res = self.docker_image_cache_manager.ensure_docker_image(item.url, item.cache_filename)
         if res.success and res.docker_image_path:

src/cloudai/installer/standalone_installer.py

@@ -48,3 +48,6 @@ def uninstall_one(self, item: Installable) -> InstallStatusResult:
 
     def is_installed_one(self, item: Installable) -> InstallStatusResult:
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
+
+    def mark_as_installed_one(self, item: Installable) -> InstallStatusResult:
+        return InstallStatusResult(False, f"Unsupported item type: {type(item)}")

tests/test_acceptance.py

@@ -79,6 +79,7 @@ def test_slurm(tmp_path: Path, scenario: Dict):
         hook_dir=Path("conf/common/hook"),
         test_scenario=test_scenario_path,
         output_dir=tmp_path,
+        enable_cache_without_check=False,
     )
     with (
         patch("asyncio.sleep", return_value=None),

tests/test_base_installer.py

@@ -43,6 +43,9 @@ def uninstall_one(self, item: Installable) -> InstallStatusResult:
     def is_installed_one(self, item: Installable) -> InstallStatusResult:
         return InstallStatusResult(success=True)
 
+    def mark_as_installed_one(self, item: Installable) -> InstallStatusResult:
+        return InstallStatusResult(success=True)
+
 
 @pytest.fixture
 def docker_image() -> DockerImage:

tests/test_cli.py

@@ -290,6 +290,7 @@ def test_run_dry_run_modes(self, cli: CloudAICLI):
                 tests_dir=Path("tests_dir"),
                 test_scenario=Path("test_scenario"),
                 output_dir=None,
+                enable_cache_without_check=False,
             )
 
     @pytest.mark.parametrize(

tests/test_slurm_installer.py

@@ -16,6 +16,7 @@
 
 from pathlib import Path
 from subprocess import CompletedProcess
+from typing import cast
 from unittest.mock import Mock, patch
 
 import pytest
@@ -24,6 +25,7 @@
 from cloudai.installer.slurm_installer import SlurmInstaller
 from cloudai.systems.slurm.slurm_system import SlurmSystem
 from cloudai.util.docker_image_cache_manager import DockerImageCacheResult
+from cloudai.workloads.nemo_launcher import NeMoLauncherCmdArgs, NeMoLauncherTestDefinition
 
 
 @pytest.fixture
@@ -348,7 +350,7 @@ def test_check_supported(slurm_system: SlurmSystem):
     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")
+    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)
@@ -360,6 +362,9 @@ def test_check_supported(slurm_system: SlurmSystem):
         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
@@ -368,7 +373,35 @@ def __hash__(self) -> int:
             return hash("MyInstallable")
 
     unsupported = MyInstallable()
-    for func in [installer.install_one, installer.uninstall_one, installer.is_installed_one]:
+    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_git_repo():
+    git = GitRepo(url="./git_url", commit="commit_hash")
+    assert git.container_mount == f"/git/{git.repo_name}"
+
+    git.mount_as = "/my_mount"
+    assert git.container_mount == git.mount_as
+
+
+def test_mark_as_installed(slurm_system: SlurmSystem):
+    tdef = NeMoLauncherTestDefinition(
+        name="name", description="desc", test_template_name="tt", cmd_args=NeMoLauncherCmdArgs()
+    )
+    docker = cast(DockerImage, tdef.installables[0])
+    py_script = cast(PythonExecutable, tdef.installables[1])
+
+    installer = SlurmInstaller(slurm_system)
+    res = installer.mark_as_installed(tdef.installables)
+
+    assert res.success
+    assert docker.installed_path == slurm_system.install_path / docker.cache_filename
+    assert py_script.git_repo.installed_path == slurm_system.install_path / py_script.git_repo.repo_name