Merge pull request #548 from NVIDIA/am/mark-docker-as-installed

Make sure mark_as_installed respects system config

Author: amaslenn

src/cloudai/_core/base_installer.py

@@ -150,7 +150,7 @@ def is_installed(self, items: Iterable[Installable]) -> InstallStatusResult:
             logging.debug(f"Installation check for {item!r}: {result.success}, {result.message}")
             install_results[item] = result
 
-        self._populate_sucessful_install(items, install_results)
+        self._populate_successful_install(items, install_results)
 
         nfailed = len([result for result in install_results.values() if not result.success])
         if nfailed:
@@ -204,7 +204,7 @@ def install(self, items: Iterable[Installable]) -> InstallStatusResult:
                     logging.error(f"{done}/{total} Installation failed for {item!r}: {e}")
                     install_results[item] = InstallStatusResult(False, str(e))
 
-        self._populate_sucessful_install(items, install_results)
+        self._populate_successful_install(items, install_results)
 
         all_success = all(result.success for result in install_results.values())
         if all_success:
@@ -213,22 +213,13 @@ def install(self, items: Iterable[Installable]) -> InstallStatusResult:
         nfailed = len([result for result in install_results.values() if not result.success])
         return InstallStatusResult(False, f"{nfailed} item(s) failed to install.", install_results)
 
-    def _populate_sucessful_install(
+    def _populate_successful_install(
         self, items: Iterable[Installable], install_results: dict[Installable, InstallStatusResult]
     ):
-        dups: dict[Installable, list[Installable]] = {}
         for item in self.all_items(items, with_duplicates=True):
-            if item in dups:
-                dups[item].append(item)
-            else:
-                dups[item] = [item]
-
-        for item, res in install_results.items():
-            if not res.success:
+            if item not in install_results or not install_results[item].success:
                 continue
-
-            for dup in dups[item]:
-                self.mark_as_installed_one(dup)
+            self.mark_as_installed_one(item)
 
     @final
     def uninstall(self, items: Iterable[Installable]) -> InstallStatusResult:

src/cloudai/installer/slurm_installer.py

@@ -145,7 +145,8 @@ def is_installed_one(self, item: Installable) -> InstallStatusResult:
 
     def mark_as_installed_one(self, item: Installable) -> InstallStatusResult:
         if isinstance(item, DockerImage):
-            item.installed_path = self.system.install_path / item.cache_filename
+            if self.system.cache_docker_images_locally:
+                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

tests/test_base_installer.py

@@ -208,7 +208,7 @@ def test_system_installables_are_used(slurm_system: SlurmSystem):
     installer.uninstall_one = Mock(return_value=InstallStatusResult(True))
     installer.is_installed_one = Mock(return_value=InstallStatusResult(True))
     installer.mark_as_installed_one = Mock(return_value=InstallStatusResult(True))
-    installer._populate_sucessful_install = Mock()
+    installer._populate_successful_install = Mock()
 
     installer.install([])
     assert installer.install_one.call_count == len(slurm_system.system_installables())
@@ -253,3 +253,26 @@ def test_both_is_installed(self, installer: SlurmInstaller):
         assert f1._installed_path is not None, "First file is not installed"
         assert f2._installed_path is not None, "Second file is not installed"
         assert f1._installed_path == f2._installed_path, "Files are installed to different paths"
+
+    def test_order_of_items_does_not_matter(self, installer: SlurmInstaller):
+        f = installer.system.install_path / "file"
+        f1, f2 = File(src=f), File(src=f)
+        assert f1._installed_path is None, "First file is installed before testing"
+        assert f2._installed_path is None, "Second file is installed before testing"
+
+        installer._populate_successful_install([f1, f2], {})
+        assert f1._installed_path is None, "First file was marked as installed, but should not be"
+        assert f2._installed_path is None, "Second file was marked as installed, but should not be"
+
+        installer._populate_successful_install([f1, f2], {f1: InstallStatusResult(success=True)})
+        assert f1._installed_path is not None, (
+            "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"
+
+        f1._installed_path, f2._installed_path = None, None
+        installer._populate_successful_install([f2, f1], {f1: InstallStatusResult(success=True)})
+        assert f1._installed_path is not None, (
+            "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"

tests/test_slurm_installer.py

@@ -486,3 +486,16 @@ def test_mark_as_installed(slurm_system: SlurmSystem):
     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
+
+
+@pytest.mark.parametrize("cache", [True, False])
+def test_mark_as_installed_docker_image_system_is_respected(slurm_system: SlurmSystem, cache: bool):
+    slurm_system.cache_docker_images_locally = cache
+    docker = DockerImage(url="fake_url/img")
+    installer = SlurmInstaller(slurm_system)
+    res = installer.mark_as_installed([docker])
+    assert res.success
+    if cache:
+        assert docker.installed_path == slurm_system.install_path / docker.cache_filename
+    else:
+        assert docker.installed_path == docker.url