Merge remote-tracking branch 'origin/main' into disable-dset-check

Author: srivatsankrishnan

src/cloudai/schema/test_template/nemo_launcher/slurm_command_gen_strategy.py

@@ -78,6 +78,20 @@ def gen_exec_command(
         del self.final_cmd_args["repository_commit_hash"]
         del self.final_cmd_args["docker_image_url"]
 
+        if self.slurm_system.account is not None:
+            self.final_cmd_args["cluster.account"] = self.slurm_system.account
+            self.final_cmd_args["cluster.job_name_prefix"] = f"{self.slurm_system.account}-cloudai.nemo:"
+        self.final_cmd_args["cluster.gpus_per_node"] = (
+            self.slurm_system.gpus_per_node if self.slurm_system.gpus_per_node is not None else "null"
+        )
+
+        if ("data_dir" in self.final_cmd_args) and (self.final_cmd_args["data_dir"] == "DATA_DIR"):
+            raise ValueError(
+                "The 'data_dir' field of the NeMo launcher test contains the placeholder 'DATA_DIR'. "
+                "Please update the test schema TOML file with a valid path to the dataset. "
+                "The 'data_dir' field must point to an actual dataset location, not a placeholder."
+            )
+
         cmd_args_str = self._generate_cmd_args_str(self.final_cmd_args, nodes)
 
         full_cmd = f"python {launcher_path}/launcher_scripts/main.py {cmd_args_str}"

src/cloudai/util/docker_image_cache_manager.py

@@ -82,7 +82,7 @@ def __init__(self, success: bool, docker_image_path: Optional[Path] = None, mess
             message (str): A message providing additional information about the result.
         """
         self.success = success
-        self._docker_image_path = docker_image_path or Path()
+        self.docker_image_path = docker_image_path
         self.message = message
 
     def __bool__(self):
@@ -103,26 +103,6 @@ def __str__(self):
         """
         return self.message
 
-    @property
-    def docker_image_path(self) -> Path:
-        """
-        Get the path to the Docker image.
-
-        Returns
-            Path: Absolute path to the Docker image.
-        """
-        return self._docker_image_path.absolute()
-
-    @docker_image_path.setter
-    def docker_image_path(self, value: Path) -> None:
-        """
-        Set the path to the Docker image.
-
-        Args:
-            value (Path): The path to the Docker image.
-        """
-        self._docker_image_path = value
-
 
 class DockerImageCacheManager:
     """
@@ -189,7 +169,7 @@ def check_docker_image_exists(
         docker_image_path = Path(docker_image_url)
         if docker_image_path.is_file() and docker_image_path.exists():
             return DockerImageCacheResult(
-                True, docker_image_path, f"Docker image file path is valid: {docker_image_url}."
+                True, docker_image_path.absolute(), f"Docker image file path is valid: {docker_image_url}."
             )
 
         # Check if the cache file exists
@@ -208,7 +188,7 @@ def check_docker_image_exists(
         if docker_image_path.is_file() and docker_image_path.exists():
             message = f"Cached Docker image already exists at {docker_image_path}."
             logging.debug(message)
-            return DockerImageCacheResult(True, docker_image_path, message)
+            return DockerImageCacheResult(True, docker_image_path.absolute(), message)
 
         message = f"Docker image does not exist at the specified path: {docker_image_path}."
         logging.debug(message)
@@ -234,7 +214,7 @@ def cache_docker_image(
         if docker_image_path.is_file():
             success_message = f"Cached Docker image already exists at {docker_image_path}."
             logging.info(success_message)
-            return DockerImageCacheResult(True, docker_image_path, success_message)
+            return DockerImageCacheResult(True, docker_image_path.absolute(), success_message)
 
         prerequisite_check = self._check_prerequisites(docker_image_url)
         if not prerequisite_check:
@@ -281,7 +261,7 @@ def cache_docker_image(
             success_message = f"Docker image cached successfully at {docker_image_path}."
             logging.debug(success_message)
             logging.debug(f"Command used: {enroot_import_cmd}, stdout: {p.stdout}, stderr: {p.stderr}")
-            return DockerImageCacheResult(True, docker_image_path, success_message)
+            return DockerImageCacheResult(True, docker_image_path.absolute(), success_message)
         except subprocess.CalledProcessError as e:
             error_message = (
                 f"Failed to import Docker image from {docker_image_url}. Command: {enroot_import_cmd}. Error: {e}"
@@ -416,14 +396,14 @@ def uninstall_cached_image(self, subdir_name: str, docker_image_filename: str) -
                     subdir_path.rmdir()
                     success_message = f"Subdirectory removed successfully: {subdir_path}."
                     logging.info(success_message)
-                    return DockerImageCacheResult(True, subdir_path, success_message)
+                    return DockerImageCacheResult(True, subdir_path.absolute(), success_message)
             except OSError as e:
                 error_message = f"Failed to remove subdirectory {subdir_path}. Error: {e}"
                 logging.error(error_message)
                 return DockerImageCacheResult(False, subdir_path, error_message)
         success_message = f"Cached Docker image uninstalled successfully from {subdir_path}."
         logging.info(success_message)
-        return DockerImageCacheResult(True, subdir_path, success_message)
+        return DockerImageCacheResult(True, subdir_path.absolute(), success_message)
 
     def remove_cached_image(self, subdir_name: str, docker_image_filename: str) -> DockerImageCacheResult:
         """
@@ -442,11 +422,11 @@ def remove_cached_image(self, subdir_name: str, docker_image_filename: str) -> D
                 docker_image_path.unlink()
                 success_message = f"Cached Docker image removed successfully from {docker_image_path}."
                 logging.info(success_message)
-                return DockerImageCacheResult(True, docker_image_path, success_message)
+                return DockerImageCacheResult(True, docker_image_path.absolute(), success_message)
             except OSError as e:
                 error_message = f"Failed to remove cached Docker image at {docker_image_path}. Error: {e}"
                 logging.error(error_message)
                 return DockerImageCacheResult(False, docker_image_path, error_message)
         success_message = f"No cached Docker image found to remove at {docker_image_path}."
         logging.info(success_message)
-        return DockerImageCacheResult(True, docker_image_path, success_message)
+        return DockerImageCacheResult(True, docker_image_path.absolute(), success_message)

tests/test_docker_image_cache_manager.py

@@ -71,6 +71,7 @@ def test_cache_docker_image(mock_check_prerequisites, mock_run, mock_access, moc
     mock_is_file.return_value = True
     result = manager.cache_docker_image("docker.io/hello-world", "subdir", "image.tar.gz")
     assert result.success
+    assert result.docker_image_path == Path("/fake/install/path/subdir/image.tar.gz")
     assert result.message == "Cached Docker image already exists at /fake/install/path/subdir/image.tar.gz."
 
     # Test creating subdirectory when it doesn't exist
@@ -301,9 +302,18 @@ def test_check_docker_image_accessibility_with_enroot(mock_tempdir, mock_which,
     assert "docker://docker.io/hello-world" in command
 
 
-def test_docker_image_path_is_always_absolute():
-    r = DockerImageCacheResult(True, Path("relative/path"), "message")
-    assert r.docker_image_path.is_absolute()
+@patch("pathlib.Path.is_file")
+@patch("pathlib.Path.exists")
+def test_ensure_docker_image_no_local_cache(mock_exists, mock_is_file):
+    manager = DockerImageCacheManager(Path("/fake/install/path"), False, "default")
 
-    r.docker_image_path = Path("another/relative/path")
-    assert r.docker_image_path.is_absolute()
+    mock_is_file.return_value = False
+    mock_exists.return_value = True
+
+    result = manager.ensure_docker_image("docker.io/hello-world", "subdir", "docker_image.sqsh")
+
+    assert result.success
+    assert result.docker_image_path is not None
+    assert str(result.docker_image_path) == "docker.io/hello-world"
+    assert not result.docker_image_path.is_absolute()
+    assert result.message == ""

tests/test_slurm_command_gen_strategy.py

@@ -298,6 +298,101 @@ def test_invalid_tokenizer_path(self, nemo_cmd_gen: NeMoLauncherSlurmCommandGenS
                 nodes=[],
             )
 
+    def test_account_in_command(self, nemo_cmd_gen: NeMoLauncherSlurmCommandGenStrategy):
+        extra_env_vars = {"TEST_VAR_1": "value1"}
+        cmd_args = {
+            "docker_image_url": "fake",
+            "repository_url": "fake",
+            "repository_commit_hash": "fake",
+        }
+
+        nemo_cmd_gen.slurm_system.account = "test_account"
+        cmd = nemo_cmd_gen.gen_exec_command(
+            cmd_args=cmd_args,
+            extra_env_vars=extra_env_vars,
+            extra_cmd_args="",
+            output_path=Path(""),
+            num_nodes=1,
+            nodes=[],
+        )
+
+        assert "cluster.account=test_account" in cmd
+        assert "cluster.job_name_prefix=test_account-cloudai.nemo:" in cmd
+
+    def test_no_account_in_command(self, nemo_cmd_gen: NeMoLauncherSlurmCommandGenStrategy):
+        extra_env_vars = {"TEST_VAR_1": "value1"}
+        cmd_args = {
+            "docker_image_url": "fake",
+            "repository_url": "fake",
+            "repository_commit_hash": "fake",
+        }
+
+        nemo_cmd_gen.slurm_system.account = None
+        cmd = nemo_cmd_gen.gen_exec_command(
+            cmd_args=cmd_args,
+            extra_env_vars=extra_env_vars,
+            extra_cmd_args="",
+            output_path=Path(""),
+            num_nodes=1,
+            nodes=[],
+        )
+
+        assert "cluster.account" not in cmd
+        assert "cluster.job_name_prefix" not in cmd
+
+    def test_gpus_per_node_value(self, nemo_cmd_gen: NeMoLauncherSlurmCommandGenStrategy):
+        extra_env_vars = {"TEST_VAR_1": "value1"}
+        cmd_args = {
+            "docker_image_url": "fake",
+            "repository_url": "fake",
+            "repository_commit_hash": "fake",
+        }
+
+        nemo_cmd_gen.slurm_system.gpus_per_node = 4
+        cmd = nemo_cmd_gen.gen_exec_command(
+            cmd_args=cmd_args,
+            extra_env_vars=extra_env_vars,
+            extra_cmd_args="",
+            output_path=Path(""),
+            num_nodes=1,
+            nodes=[],
+        )
+
+        assert "cluster.gpus_per_node=4" in cmd
+
+        nemo_cmd_gen.slurm_system.gpus_per_node = None
+        cmd = nemo_cmd_gen.gen_exec_command(
+            cmd_args=cmd_args,
+            extra_env_vars=extra_env_vars,
+            extra_cmd_args="",
+            output_path=Path(""),
+            num_nodes=1,
+            nodes=[],
+        )
+
+        assert "cluster.gpus_per_node=null" in cmd
+
+    def test_data_dir_validation(self, nemo_cmd_gen: NeMoLauncherSlurmCommandGenStrategy):
+        extra_env_vars = {"TEST_VAR_1": "value1"}
+        cmd_args = {
+            "docker_image_url": "fake",
+            "repository_url": "fake",
+            "repository_commit_hash": "fake",
+            "data_dir": "DATA_DIR",  # Invalid placeholder
+        }
+
+        with pytest.raises(
+            ValueError, match="The 'data_dir' field of the NeMo launcher test contains the placeholder 'DATA_DIR'."
+        ):
+            nemo_cmd_gen.gen_exec_command(
+                cmd_args=cmd_args,
+                extra_env_vars=extra_env_vars,
+                extra_cmd_args="",
+                output_path=Path(""),
+                num_nodes=1,
+                nodes=[],
+            )
+
 
 class TestWriteSbatchScript:
     MANDATORY_ARGS = {