Add extra_srun_args & scripts in SlurmContainerTestDef (#531)

Author: lilyw97

src/cloudai/workloads/slurm_container/slurm_command_gen_strategy.py

@@ -36,7 +36,7 @@ def gen_srun_prefix(self, slurm_args: dict[str, Any], tr: TestRun, use_pretest_e
         tdef: SlurmContainerTestDefinition = cast(SlurmContainerTestDefinition, tr.test.test_definition)
         slurm_args["image_path"] = tdef.docker_image.installed_path
         cmd = super().gen_srun_prefix(slurm_args, tr)
-        return [*cmd, "--no-container-mount-home"]
+        return [*cmd, *tdef.extra_srun_args, "--no-container-mount-home"]
 
     def generate_test_command(
         self, env_vars: Dict[str, Union[str, List[str]]], cmd_args: Dict[str, Union[str, List[str]]], tr: TestRun

src/cloudai/workloads/slurm_container/slurm_container.py

@@ -16,7 +16,9 @@
 
 from typing import Optional
 
-from cloudai import DockerImage, Installable
+from pydantic import Field
+
+from cloudai import DockerImage, File, Installable
 
 from ...models.workload import CmdArgs, TestDefinition
 
@@ -32,7 +34,8 @@ class SlurmContainerTestDefinition(TestDefinition):
     """Test definition for a generic Slurm container test."""
 
     cmd_args: SlurmContainerCmdArgs
-
+    extra_srun_args: list[str] = Field(default_factory=list)
+    scripts: list[File] = Field(default_factory=list)
     _docker_image: Optional[DockerImage] = None
 
     @property
@@ -43,7 +46,7 @@ def docker_image(self) -> DockerImage:
 
     @property
     def installables(self) -> list[Installable]:
-        return [self.docker_image, *self.git_repos]
+        return [self.docker_image, *self.git_repos, *self.scripts]
 
     @property
     def extra_args_str(self) -> str:

tests/slurm_command_gen_strategy/test_slurm_container_slurm_command_gen_strategy.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 
 from pathlib import Path
+from typing import cast
 
 import pytest
 
@@ -72,3 +73,23 @@ def test_with_nsys(slurm_system: SlurmSystem, test_run: TestRun) -> None:
     )
 
     assert cmd == f'{srun_part} bash -c "{" ".join(nsys.cmd_args)} cmd"'
+
+
+def test_with_extra_srun_args(slurm_system: SlurmSystem, test_run: TestRun) -> None:
+    extra_args = ["--ntasks=1", "--ntasks-per-node=1"]
+    tdef = cast(SlurmContainerTestDefinition, test_run.test.test_definition)
+    tdef.extra_srun_args = extra_args
+
+    cgs = SlurmContainerCommandGenStrategy(slurm_system, {})
+    cmd = cgs.gen_srun_command(test_run)
+
+    srun_part = (
+        f"srun --export=ALL --mpi={slurm_system.mpi} "
+        f"--container-image={test_run.test.test_definition.cmd_args.docker_image_url} "
+        f"--container-mounts={Path.cwd().absolute()}:/cloudai_run_results,"
+        f"{slurm_system.install_path.absolute()}:/cloudai_install "
+        f"{' '.join(extra_args)} "
+        f"--no-container-mount-home"
+    )
+
+    assert cmd == f'{srun_part} bash -c "cmd"'

tests/test_test_definitions.py

@@ -21,7 +21,7 @@
 import toml
 from pydantic import ValidationError
 
-from cloudai import NsysConfiguration, Parser, Registry, TestConfigParsingError, TestDefinition, TestParser
+from cloudai import File, NsysConfiguration, Parser, Registry, TestConfigParsingError, TestDefinition, TestParser
 from cloudai.workloads.chakra_replay import ChakraReplayCmdArgs, ChakraReplayTestDefinition
 from cloudai.workloads.jax_toolbox import (
     GPTCmdArgs,
@@ -35,6 +35,7 @@
 from cloudai.workloads.nccl_test import NCCLCmdArgs, NCCLTestDefinition
 from cloudai.workloads.nemo_launcher import NeMoLauncherCmdArgs, NeMoLauncherTestDefinition
 from cloudai.workloads.nemo_run import NeMoRunCmdArgs, NeMoRunTestDefinition
+from cloudai.workloads.slurm_container import SlurmContainerCmdArgs, SlurmContainerTestDefinition
 from cloudai.workloads.ucc_test import UCCCmdArgs, UCCTestDefinition
 
 TOML_FILES = list(Path("conf").glob("**/*.toml"))
@@ -126,6 +127,12 @@ def test_chakra_docker_image_is_required():
             test_template_name="chakra",
             cmd_args=ChakraReplayCmdArgs(docker_image_url="fake://url/chakra"),
         ),
+        SlurmContainerTestDefinition(
+            name="sc",
+            description="desc",
+            test_template_name="sc",
+            cmd_args=SlurmContainerCmdArgs(docker_image_url="fake://url/sc", cmd="cmd"),
+        ),
     ],
 )
 def test_docker_installable_persists(
@@ -137,6 +144,7 @@ def test_docker_installable_persists(
         NeMoLauncherTestDefinition,
         NemotronTestDefinition,
         UCCTestDefinition,
+        SlurmContainerTestDefinition,
     ],
     tmp_path: Path,
 ):
@@ -159,6 +167,25 @@ def test_python_executable_installable_persists(test: NeMoLauncherTestDefinition
     assert test.python_executable.venv_path == tmp_path
 
 
+@pytest.mark.parametrize(
+    "test",
+    [
+        SlurmContainerTestDefinition(
+            name="sc",
+            description="desc",
+            test_template_name="sc",
+            cmd_args=SlurmContainerCmdArgs(docker_image_url="fake://url/sc", cmd="cmd"),
+            scripts=[File(src=Path("./script1")), File(src=Path("./script2"))],
+        )
+    ],
+)
+def test_slurm_container_installables(test: SlurmContainerTestDefinition):
+    assert len(test.installables) >= 3
+    assert test.docker_image in test.installables
+    assert File(src=Path("./script1")) in test.installables
+    assert File(src=Path("./script2")) in test.installables
+
+
 class TestNsysConfiguration:
     def test_default(self):
         nsys = NsysConfiguration()