Merge pull request #210 from TaekyungHeo/nemo-bug-fix-slurm-args

Add Support for Cluster Account and gpus_per_node in Command Generation

Author: TaekyungHeo

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}"

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 = {