Merge pull request #140 from NVIDIA/am/slurm-common

Move parts of srun CLI generation into base class

Author: TaekyungHeo

pyproject.toml

@@ -15,7 +15,7 @@
 
 [project]
 name = "cloudai"
-version = "0.7.12"
+version = "0.7.13"
 dependencies = [
     "bokeh==3.4.1",
     "pandas==2.2.1",

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

@@ -49,7 +49,7 @@ def gen_exec_command(
 
         job_name_prefix = "chakra_replay"
         slurm_args = self._parse_slurm_args(job_name_prefix, final_env_vars, final_cmd_args, num_nodes, nodes)
-        srun_command = self._generate_srun_command(slurm_args, final_env_vars, final_cmd_args, extra_cmd_args)
+        srun_command = self.generate_full_srun_command(slurm_args, final_env_vars, final_cmd_args, extra_cmd_args)
         return self._write_sbatch_script(slurm_args, env_vars_str, srun_command, output_path)
 
     def _parse_slurm_args(
@@ -69,23 +69,15 @@ def _parse_slurm_args(
 
         return base_args
 
-    def _generate_srun_command(
-        self,
-        slurm_args: Dict[str, Any],
-        env_vars: Dict[str, str],
-        cmd_args: Dict[str, str],
-        extra_cmd_args: str,
-    ) -> str:
+    def generate_test_command(
+        self, slurm_args: Dict[str, Any], env_vars: Dict[str, str], cmd_args: Dict[str, str], extra_cmd_args: str
+    ) -> List[str]:
         srun_command_parts = [
-            "srun",
-            f"--mpi={slurm_args['mpi']}",
-            f'--container-image={slurm_args["image_path"]}',
-            f'--container-mounts={slurm_args["container_mounts"]}',
             "python /workspace/param/train/comms/pt/commsTraceReplay.py",
             f'--trace-type {cmd_args["trace_type"]}',
             f'--trace-path {cmd_args["trace_path"]}',
             f'--backend {cmd_args["backend"]}',
             f'--device {cmd_args["device"]}',
             extra_cmd_args,
         ]
-        return " \\\n".join(srun_command_parts)
+        return srun_command_parts

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

@@ -52,7 +52,7 @@ def gen_exec_command(
         env_vars_str = self._format_env_vars(final_env_vars)
 
         slurm_args = self._parse_slurm_args("JaxToolbox", final_env_vars, final_cmd_args, num_nodes, nodes)
-        srun_command = self._generate_srun_command(slurm_args, final_env_vars, final_cmd_args, extra_cmd_args)
+        srun_command = self.generate_full_srun_command(slurm_args, final_env_vars, final_cmd_args, extra_cmd_args)
         return self._write_sbatch_script(slurm_args, env_vars_str, srun_command, output_path)
 
     def _format_xla_flags(self, cmd_args: Dict[str, str]) -> str:
@@ -131,18 +131,14 @@ def _parse_slurm_args(
 
         return base_args
 
-    def _generate_srun_command(
-        self,
-        slurm_args: Dict[str, Any],
-        env_vars: Dict[str, str],
-        cmd_args: Dict[str, str],
-        extra_cmd_args: str,
+    def generate_full_srun_command(
+        self, slurm_args: Dict[str, Any], env_vars: Dict[str, str], cmd_args: Dict[str, str], extra_cmd_args: str
     ) -> str:
         self._create_run_script(slurm_args, env_vars, cmd_args, extra_cmd_args)
 
         srun_command_parts = [
             "srun",
-            f"--mpi={slurm_args['mpi']}",
+            f"--mpi={self.slurm_system.mpi}",
             "--export=ALL",
             f"-o {slurm_args['output']}",
             f"-e {slurm_args['error']}",

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

@@ -45,7 +45,7 @@ def gen_exec_command(
             raise KeyError("Subtest name not specified or unsupported.")
 
         slurm_args = self._parse_slurm_args(subtest_name, final_env_vars, final_cmd_args, num_nodes, nodes)
-        srun_command = self._generate_srun_command(slurm_args, final_env_vars, final_cmd_args, extra_cmd_args)
+        srun_command = self.generate_full_srun_command(slurm_args, final_env_vars, final_cmd_args, extra_cmd_args)
         return self._write_sbatch_script(slurm_args, env_vars_str, srun_command, output_path)
 
     def _parse_slurm_args(
@@ -76,24 +76,10 @@ def _parse_slurm_args(
 
         return base_args
 
-    def _generate_srun_command(
-        self,
-        slurm_args: Dict[str, Any],
-        env_vars: Dict[str, str],
-        cmd_args: Dict[str, str],
-        extra_cmd_args: str,
-    ) -> str:
-        srun_command_parts = [
-            "srun",
-            f"--mpi={slurm_args['mpi']}",
-            f"--container-image={slurm_args['image_path']}",
-        ]
-
-        if slurm_args.get("container_mounts"):
-            srun_command_parts.append(f"--container-mounts={slurm_args['container_mounts']}")
-
-        srun_command_parts.append(f"/usr/local/bin/{cmd_args['subtest_name']}")
-
+    def generate_test_command(
+        self, slurm_args: Dict[str, Any], env_vars: Dict[str, str], cmd_args: Dict[str, str], extra_cmd_args: str
+    ) -> List[str]:
+        srun_command_parts = [f"/usr/local/bin/{cmd_args['subtest_name']}"]
         nccl_test_args = [
             "nthreads",
             "ngpus",
@@ -119,4 +105,4 @@ def _generate_srun_command(
         if extra_cmd_args:
             srun_command_parts.append(extra_cmd_args)
 
-        return " \\\n".join(srun_command_parts)
+        return srun_command_parts

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

@@ -37,20 +37,13 @@ def gen_exec_command(
         env_vars_str = self._format_env_vars(final_env_vars)
 
         slurm_args = self._parse_slurm_args("sleep", final_env_vars, final_cmd_args, num_nodes, nodes)
-        srun_command = self._generate_srun_command(slurm_args, final_env_vars, final_cmd_args, extra_cmd_args)
+        srun_command = self.generate_full_srun_command(slurm_args, final_env_vars, final_cmd_args, extra_cmd_args)
         return self._write_sbatch_script(slurm_args, env_vars_str, srun_command, output_path)
 
-    def _generate_srun_command(
-        self,
-        slurm_args: Dict[str, Any],
-        env_vars: Dict[str, str],
-        cmd_args: Dict[str, str],
-        extra_cmd_args: str,
+    def generate_full_srun_command(
+        self, slurm_args: Dict[str, Any], env_vars: Dict[str, str], cmd_args: Dict[str, str], extra_cmd_args: str
     ) -> str:
-        srun_command_parts = [
-            "srun",
-            f"--mpi={slurm_args['mpi']}",
-        ]
+        srun_command_parts = ["srun", f"--mpi={self.slurm_system.mpi}"]
 
         sec = cmd_args["seconds"]
         srun_command_parts.append(f"sleep {sec}")

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

@@ -44,7 +44,7 @@ def gen_exec_command(
             raise KeyError("Collective name not specified or unsupported.")
 
         slurm_args = self._parse_slurm_args(collective, final_env_vars, final_cmd_args, num_nodes, nodes)
-        srun_command = self._generate_srun_command(slurm_args, final_env_vars, final_cmd_args, extra_cmd_args)
+        srun_command = self.generate_full_srun_command(slurm_args, final_env_vars, final_cmd_args, extra_cmd_args)
         return self._write_sbatch_script(slurm_args, env_vars_str, srun_command, output_path)
 
     def _parse_slurm_args(
@@ -69,19 +69,10 @@ def _parse_slurm_args(
 
         return base_args
 
-    def _generate_srun_command(
-        self,
-        slurm_args: Dict[str, Any],
-        env_vars: Dict[str, str],
-        cmd_args: Dict[str, str],
-        extra_cmd_args: str,
-    ) -> str:
-        srun_command_parts = [
-            "srun",
-            f"--mpi={slurm_args['mpi']}",
-            f"--container-image={slurm_args['image_path']}",
-            "/opt/hpcx/ucc/bin/ucc_perftest",
-        ]
+    def generate_test_command(
+        self, slurm_args: Dict[str, Any], env_vars: Dict[str, str], cmd_args: Dict[str, str], extra_cmd_args: str
+    ) -> List[str]:
+        srun_command_parts = ["/opt/hpcx/ucc/bin/ucc_perftest"]
 
         # Add collective, minimum bytes, and maximum bytes options if available
         if "collective" in cmd_args:
@@ -99,4 +90,4 @@ def _generate_srun_command(
         if extra_cmd_args:
             srun_command_parts.append(extra_cmd_args)
 
-        return " \\\n".join(srun_command_parts)
+        return srun_command_parts

src/cloudai/systems/slurm/strategy/slurm_command_gen_strategy.py

@@ -31,12 +31,7 @@ class SlurmCommandGenStrategy(CommandGenStrategy):
             properties and methods.
     """
 
-    def __init__(
-        self,
-        system: SlurmSystem,
-        env_vars: Dict[str, Any],
-        cmd_args: Dict[str, Any],
-    ) -> None:
+    def __init__(self, system: SlurmSystem, env_vars: Dict[str, Any], cmd_args: Dict[str, Any]) -> None:
         """
         Initialize a new SlurmCommandGenStrategy instance.
 
@@ -125,8 +120,6 @@ def _parse_slurm_args(
             slurm_args["account"] = self.slurm_system.account
         if self.slurm_system.distribution:
             slurm_args["distribution"] = self.slurm_system.distribution
-        if self.slurm_system.mpi:
-            slurm_args["mpi"] = self.slurm_system.mpi
         if self.slurm_system.gpus_per_node:
             slurm_args["gpus_per_node"] = self.slurm_system.gpus_per_node
         if self.slurm_system.ntasks_per_node:
@@ -136,27 +129,28 @@ def _parse_slurm_args(
 
         return slurm_args
 
-    def _generate_srun_command(
-        self,
-        slurm_args: Dict[str, Any],
-        env_vars: Dict[str, str],
-        cmd_args: Dict[str, str],
-        extra_cmd_args: str,
+    def generate_full_srun_command(
+        self, slurm_args: Dict[str, Any], env_vars: Dict[str, str], cmd_args: Dict[str, str], extra_cmd_args: str
     ) -> str:
-        """
-        Generate the srun command string for executing the test.
-
-        Args:
-            slurm_args (Dict[str, Any]): Arguments containing Slurm job settings including image path and container
-                mounts.
-            env_vars (Dict[str, str]): Environment variables.
-            cmd_args (Dict[str, str]): Command-line arguments.
-            extra_cmd_args (str): Additional command-line arguments to be included in the srun command.
-
-        Returns:
-            str: The complete srun command to execute the test.
-        """
-        return ""
+        srun_command_parts = self.generate_srun_command(slurm_args, env_vars, cmd_args, extra_cmd_args)
+        test_command_parts = self.generate_test_command(slurm_args, env_vars, cmd_args, extra_cmd_args)
+        return " \\\n".join(srun_command_parts + test_command_parts)
+
+    def generate_srun_command(
+        self, slurm_args: Dict[str, Any], env_vars: Dict[str, str], cmd_args: Dict[str, str], extra_cmd_args: str
+    ) -> List[str]:
+        srun_command_parts = ["srun", f"--mpi={self.slurm_system.mpi}"]
+        if slurm_args.get("image_path"):
+            srun_command_parts.append(f'--container-image={slurm_args["image_path"]}')
+            if slurm_args.get("container_mounts"):
+                srun_command_parts.append(f'--container-mounts={slurm_args["container_mounts"]}')
+
+        return srun_command_parts
+
+    def generate_test_command(
+        self, slurm_args: Dict[str, Any], env_vars: Dict[str, str], cmd_args: Dict[str, str], extra_cmd_args: str
+    ) -> List[str]:
+        return []
 
     def _write_sbatch_script(self, args: Dict[str, Any], env_vars_str: str, srun_command: str, output_path: str) -> str:
         """

tests/test_slurm_command_gen_strategy.py

@@ -16,6 +16,7 @@
 from pathlib import Path
 
 import pytest
+from cloudai.schema.test_template.nccl_test.slurm_command_gen_strategy import NcclTestSlurmCommandGenStrategy
 from cloudai.schema.test_template.nemo_launcher.slurm_command_gen_strategy import (
     NeMoLauncherSlurmCommandGenStrategy,
 )
@@ -39,6 +40,7 @@ def slurm_system(tmp_path: Path) -> SlurmSystem:
                 SlurmNode(name="node4", partition="main", state=SlurmNodeState.IDLE),
             ]
         },
+        mpi="fake-mpi",
     )
     Path(slurm_system.install_path).mkdir()
     Path(slurm_system.output_path).mkdir()
@@ -112,6 +114,51 @@ def test_only_nodes(strategy_fixture: SlurmCommandGenStrategy):
     assert slurm_args["num_nodes"] == len(nodes)
 
 
+class TestGenerateSrunCommand__CmdGeneration:
+    def test_generate_test_command(self, strategy_fixture: SlurmCommandGenStrategy):
+        test_command = strategy_fixture.generate_test_command({}, {}, {}, "")
+        assert test_command == []
+
+    def test_generate_srun_command(self, strategy_fixture: SlurmCommandGenStrategy):
+        srun_command = strategy_fixture.generate_srun_command({}, {}, {}, "")
+        assert srun_command == ["srun", f"--mpi={strategy_fixture.slurm_system.mpi}"]
+
+    def test_generate_srun_command_with_container_image(self, strategy_fixture: SlurmCommandGenStrategy):
+        slurm_args = {"image_path": "fake_image_path"}
+        srun_command = strategy_fixture.generate_srun_command(slurm_args, {}, {}, "")
+        assert srun_command == [
+            "srun",
+            f"--mpi={strategy_fixture.slurm_system.mpi}",
+            "--container-image=fake_image_path",
+        ]
+
+    def test_generate_srun_command_with_container_image_and_mounts(self, strategy_fixture: SlurmCommandGenStrategy):
+        slurm_args = {"image_path": "fake_image_path", "container_mounts": "fake_mounts"}
+        srun_command = strategy_fixture.generate_srun_command(slurm_args, {}, {}, "")
+        assert srun_command == [
+            "srun",
+            f"--mpi={strategy_fixture.slurm_system.mpi}",
+            "--container-image=fake_image_path",
+            "--container-mounts=fake_mounts",
+        ]
+
+    def test_generate_srun_empty_str(self, strategy_fixture: SlurmCommandGenStrategy):
+        slurm_args = {"image_path": "", "container_mounts": ""}
+        srun_command = strategy_fixture.generate_srun_command(slurm_args, {}, {}, "")
+        assert srun_command == ["srun", f"--mpi={strategy_fixture.slurm_system.mpi}"]
+
+        slurm_args = {"image_path": "fake", "container_mounts": ""}
+        srun_command = strategy_fixture.generate_srun_command(slurm_args, {}, {}, "")
+        assert srun_command == ["srun", f"--mpi={strategy_fixture.slurm_system.mpi}", "--container-image=fake"]
+
+    def test_generate_full_srun_command(self, strategy_fixture: SlurmCommandGenStrategy):
+        strategy_fixture.generate_srun_command = lambda *_, **__: ["srun", "--test", "test_arg"]
+        strategy_fixture.generate_test_command = lambda *_, **__: ["test_command"]
+
+        full_srun_command = strategy_fixture.generate_full_srun_command({}, {}, {}, "")
+        assert full_srun_command == " \\\n".join(["srun", "--test", "test_arg", "test_command"])
+
+
 class TestNeMoLauncherSlurmCommandGenStrategy__GenExecCommand:
     @pytest.fixture
     def nemo_cmd_gen(self, slurm_system: SlurmSystem) -> NeMoLauncherSlurmCommandGenStrategy:
@@ -305,3 +352,37 @@ def test_disable_output_and_error(self, add_arg: str, strategy_fixture: SlurmCom
 
         self.assert_positional_lines(file_contents.splitlines())
         assert f"--{add_arg}=" not in file_contents
+
+
+class TestNCCLSlurmCommandGen:
+    def get_cmd(self, slurm_system: SlurmSystem, slurm_args: dict, cmd_args: dict) -> str:
+        return NcclTestSlurmCommandGenStrategy(slurm_system, {}, {}).generate_full_srun_command(
+            slurm_args, {}, cmd_args, ""
+        )
+
+    def test_only_mandatory(self, slurm_system: SlurmSystem) -> None:
+        slurm_args = {"image_path": "fake_image_path"}
+        cmd_args = {"subtest_name": "fake_subtest_name"}
+        cmd = self.get_cmd(slurm_system, slurm_args, cmd_args)
+        assert cmd == " \\\n".join(
+            [
+                "srun",
+                f"--mpi={slurm_system.mpi}",
+                f"--container-image={slurm_args['image_path']}",
+                f"/usr/local/bin/{cmd_args['subtest_name']}",
+            ]
+        )
+
+    def test_with_container_mounts(self, slurm_system: SlurmSystem) -> None:
+        slurm_args = {"image_path": "fake_image_path", "container_mounts": "fake_mounts"}
+        cmd_args = {"subtest_name": "fake_subtest_name"}
+        cmd = self.get_cmd(slurm_system, slurm_args, cmd_args)
+        assert cmd == " \\\n".join(
+            [
+                "srun",
+                f"--mpi={slurm_system.mpi}",
+                f"--container-image={slurm_args['image_path']}",
+                f"--container-mounts={slurm_args['container_mounts']}",
+                f"/usr/local/bin/{cmd_args['subtest_name']}",
+            ]
+        )