[Slurm] pre-api refactor

src/cloudai/systems/slurm/single_sbatch_runner.py

@@ -23,8 +23,8 @@
 
 from cloudai.configurator import CloudAIGymEnv
 from cloudai.configurator.env_params import EnvParams
-from cloudai.core import BaseJob, JobIdRetrievalError, Registry, System, TestRun, TestScenario
-from cloudai.util import CommandShell, format_time_limit, parse_time_limit
+from cloudai.core import BaseJob, Registry, System, TestRun, TestScenario
+from cloudai.util import format_time_limit, parse_time_limit
 
 from .slurm_command_gen_strategy import SlurmCommandGenStrategy
 from .slurm_metadata import SlurmJobMetadata, SlurmStepMetadata
@@ -37,7 +37,6 @@ class SingleSbatchRunner(SlurmRunner):
 
     def __init__(self, mode: str, system: System, test_scenario: TestScenario, output_path: Path) -> None:
         super().__init__(mode, system, test_scenario, output_path)
-        self.cmd_shell = CommandShell()
         self.system = cast(SlurmSystem, system)
         self.job_name = "cloudai-single-sbatch"
 
@@ -249,17 +248,8 @@ def _submit_test(self, tr: TestRun) -> SlurmJob:
 
         job_id = 0
         if self.mode == "run":
-            exec_cmd = f"sbatch {self.scenario_root / 'cloudai_sbatch_script.sh'}"
-            stdout, stderr = self.cmd_shell.execute(exec_cmd).communicate()
-            job_id = self.get_job_id(stdout, stderr)
-            if job_id is None:
-                raise JobIdRetrievalError(
-                    test_name=tr.name,
-                    command=exec_cmd,
-                    stdout=stdout,
-                    stderr=stderr,
-                    message="Failed to retrieve job ID.",
-                )
+            script_path = self.scenario_root / "cloudai_sbatch_script.sh"
+            job_id = self.system.submit_sbatch(script_path, tr.name)
         logging.info(f"Submitted slurm job: {job_id}")
         return SlurmJob(tr, id=job_id)
 

src/cloudai/systems/slurm/slurm_installer.py

@@ -15,7 +15,6 @@
 # limitations under the License.
 
 import logging
-import subprocess
 from pathlib import Path
 
 from cloudai.core import BaseInstaller, DockerImage, Installable, InstallStatusResult
@@ -27,15 +26,6 @@
 class SlurmInstaller(BaseInstaller):
     """Installer for Slurm systems."""
 
-    PREREQUISITES = ("git", "sbatch", "sinfo", "squeue", "srun", "scancel", "sacct")
-    REQUIRED_SRUN_OPTIONS = (
-        "--mpi",
-        "--gpus-per-node",
-        "--ntasks-per-node",
-        "--container-image",
-        "--container-mounts",
-    )
-
     def __init__(self, system: SlurmSystem):
         super().__init__(system)
         self.system = system
@@ -47,34 +37,11 @@ def _check_prerequisites(self) -> InstallStatusResult:
             return InstallStatusResult(False, base_prerequisites_result.message)
 
         try:
-            self._check_required_binaries()
-            self._check_srun_options()
+            self.system.validate_install_environment()
             return InstallStatusResult(True)
         except EnvironmentError as e:
             return InstallStatusResult(False, str(e))
 
-    def _check_required_binaries(self) -> None:
-        for binary in self.PREREQUISITES:
-            if not self._is_binary_installed(binary):
-                raise EnvironmentError(f"Required binary '{binary}' is not installed.")
-
-    def _check_srun_options(self) -> None:
-        """
-        Check for the presence of specific srun options.
-
-        Calls `srun --help` and verifying the options. Raises an exception if any required options are missing.
-        """
-        try:
-            result = subprocess.run(["srun", "--help"], text=True, capture_output=True, check=True)
-            help_output = result.stdout
-        except subprocess.CalledProcessError as e:
-            raise EnvironmentError(f"Failed to execute 'srun --help': {e}") from e
-
-        missing_options = [option for option in self.REQUIRED_SRUN_OPTIONS if option not in help_output]
-        if missing_options:
-            missing_options_str = ", ".join(missing_options)
-            raise EnvironmentError(f"Required srun options missing: {missing_options_str}")
-
     def install_one(self, item: Installable) -> InstallStatusResult:
         logging.debug(f"Attempt to install {item}")
         if isinstance(item, DockerImage):

src/cloudai/systems/slurm/slurm_runner.py

@@ -15,14 +15,12 @@
 # limitations under the License.
 
 import logging
-import re
 from pathlib import Path
 from typing import cast
 
 import toml
 
-from cloudai.core import BaseJob, BaseRunner, JobIdRetrievalError, System, TestRun, TestScenario
-from cloudai.util import CommandShell
+from cloudai.core import BaseJob, BaseRunner, System, TestRun, TestScenario
 
 from .slurm_command_gen_strategy import SlurmCommandGenStrategy
 from .slurm_job import SlurmJob
@@ -31,17 +29,11 @@
 
 
 class SlurmRunner(BaseRunner):
-    """
-    Implementation of the Runner for a system using Slurm.
-
-    Attributes
-        cmd_shell (CommandShell): An instance of CommandShell for executing system commands.
-    """
+    """Implementation of the Runner for a system using Slurm."""
 
     def __init__(self, mode: str, system: System, test_scenario: TestScenario, output_path: Path) -> None:
         super().__init__(mode, system, test_scenario, output_path)
         self.system = cast(SlurmSystem, system)
-        self.cmd_shell = CommandShell()
         self.pinned_nodes: dict[str, list[str]] = {}
 
     def submit_test(self, tr: TestRun) -> None:
@@ -50,33 +42,13 @@ def submit_test(self, tr: TestRun) -> None:
             logging.info("Forcing test case '%s' to use pinned nodes: %s", tr.name, ",".join(tr.nodes))
         super().submit_test(tr)
 
-    def get_job_id(self, stdout: str, stderr: str) -> int | None:
-        match = re.search(r"Submitted batch job (\d+)", stdout)
-        if match:
-            return int(match.group(1))
-
-        match = re.search(r"submitted with Job ID (\d+)", stdout)  # NemoLauncher specific
-        if match:
-            return int(match.group(1))
-
-        return None
-
     def _submit_test(self, tr: TestRun) -> SlurmJob:
         logging.info(f"Running test: {tr.name}")
         exec_cmd = self.get_cmd_gen_strategy(self.system, tr).gen_exec_command()
         logging.debug(f"Executing command for test {tr.name}: {exec_cmd}")
         job_id = 0
         if self.mode == "run":
-            stdout, stderr = self.cmd_shell.execute(exec_cmd).communicate()
-            job_id = self.get_job_id(stdout, stderr)
-            if job_id is None:
-                raise JobIdRetrievalError(
-                    test_name=str(tr.name),
-                    command=exec_cmd,
-                    stdout=stdout,
-                    stderr=stderr,
-                    message="Failed to retrieve job ID.",
-                )
+            job_id = self.system.submit_job(exec_cmd, str(tr.name))
         logging.info(f"Submitted slurm job: {job_id}")
         return SlurmJob(tr, id=job_id)
 

src/cloudai/systems/slurm/slurm_system.py

@@ -18,14 +18,17 @@
 
 import logging
 import re
+import shlex
+import shutil
+import subprocess
 import time
 from copy import copy
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, ClassVar, Dict, Iterable, List, Optional, Tuple, Union
 
 from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator
 
-from cloudai.core import BaseJob, File, Installable, System
+from cloudai.core import BaseJob, File, Installable, JobIdRetrievalError, System
 from cloudai.models.scenario import ReportConfig, parse_reports_spec
 from cloudai.util import CommandShell
 
@@ -97,6 +100,12 @@ class SlurmPartition(BaseModel):
 class SlurmSystem(System):
     """Represents a Slurm system."""
 
+    def submit_sbatch(self, script_path: Path, operation_name: str, *, wait: bool = False) -> int:
+        """Submit an sbatch script without exposing the CLI transport to callers."""
+        wait_arg = " --wait" if wait else ""
+        command = f"sbatch{wait_arg} {shlex.quote(str(script_path))}"
+        return self.submit_job(command, operation_name)
+
     default_partition: str
     partitions: List[SlurmPartition]
     account: Optional[str] = None
@@ -120,6 +129,23 @@ class SlurmSystem(System):
 
     group_allocated: set[SlurmNode] = Field(default_factory=set, exclude=True)
 
+    _REQUIRED_BINARIES: ClassVar[tuple[str, ...]] = (
+        "git",
+        "sbatch",
+        "sinfo",
+        "squeue",
+        "srun",
+        "scancel",
+        "sacct",
+    )
+    _REQUIRED_SRUN_OPTIONS: ClassVar[tuple[str, ...]] = (
+        "--mpi",
+        "--gpus-per-node",
+        "--ntasks-per-node",
+        "--container-image",
+        "--container-mounts",
+    )
+
     @field_validator("reports", mode="before")
     @classmethod
     def parse_reports(cls, value: dict[str, Any] | None) -> dict[str, ReportConfig] | None:
@@ -256,6 +282,44 @@ def _is_transient_status_error(self, stderr: str) -> bool:
         ]
         return any(p in stderr for p in patterns)
 
+    @staticmethod
+    def _parse_submitted_job_id(stdout: str) -> int | None:
+        match = re.search(r"Submitted batch job (\d+)", stdout)
+        if match:
+            return int(match.group(1))
+
+        # Some launchers submit Slurm jobs themselves and use this output format.
+        match = re.search(r"submitted with Job ID (\d+)", stdout)
+        return int(match.group(1)) if match else None
+
+    def submit_job(self, submission_command: str, test_name: str) -> int:
+        """Submit a generated Slurm workload and return its job ID."""
+        stdout, stderr = self.cmd_shell.execute(submission_command).communicate()
+        job_id = self._parse_submitted_job_id(stdout)
+        if job_id is None:
+            raise JobIdRetrievalError(
+                test_name=test_name,
+                command=submission_command,
+                stdout=stdout,
+                stderr=stderr,
+                message="Failed to retrieve job ID.",
+            )
+        return job_id
+
+    def validate_install_environment(self) -> None:
+        """Validate that the configured Slurm environment can run CloudAI workloads."""
+        for binary in self._REQUIRED_BINARIES:
+            if shutil.which(binary) is None:
+                raise EnvironmentError(f"Required binary '{binary}' is not installed.")
+
+        try:
+            result = subprocess.run(["srun", "--help"], text=True, capture_output=True, check=True)
+        except subprocess.CalledProcessError as exc:
+            raise EnvironmentError(f"Failed to execute 'srun --help': {exc}") from exc
+        missing_options = [option for option in self._REQUIRED_SRUN_OPTIONS if option not in result.stdout]
+        if missing_options:
+            raise EnvironmentError(f"Required srun options missing: {', '.join(missing_options)}")
+
     def is_job_running(self, job: BaseJob, retry_threshold: int = 3) -> bool:
         """
         Determine if a specified Slurm job is currently running by checking its presence and state in the job queue.

tests/systems/slurm/test_system.py

@@ -22,7 +22,7 @@
 import toml
 from pydantic import ValidationError
 
-from cloudai.core import BaseJob, TestRun
+from cloudai.core import BaseJob, JobIdRetrievalError, TestRun
 from cloudai.models.scenario import ReportConfig
 from cloudai.systems.slurm import (
     SlurmCommandGenStrategy,
@@ -35,6 +35,26 @@
 from cloudai.workloads.nccl_test import NCCLCmdArgs, NCCLTestDefinition
 
 
+@patch("cloudai.systems.slurm.slurm_system.CommandShell.execute")
+def test_submit_job_returns_parsed_job_id(mock_execute: Mock, slurm_system: SlurmSystem):
+    process = Mock()
+    process.communicate.return_value = ("Submitted batch job 123", "")
+    mock_execute.return_value = process
+
+    assert slurm_system.submit_job("sbatch script.sh", "test") == 123
+    mock_execute.assert_called_once_with("sbatch script.sh")
+
+
+@patch("cloudai.systems.slurm.slurm_system.CommandShell.execute")
+def test_submit_job_raises_semantic_error(mock_execute: Mock, slurm_system: SlurmSystem):
+    process = Mock()
+    process.communicate.return_value = ("", "submission failed")
+    mock_execute.return_value = process
+
+    with pytest.raises(JobIdRetrievalError, match="Failed to retrieve job ID"):
+        slurm_system.submit_job("sbatch script.sh", "test")
+
+
 @pytest.mark.parametrize(
     "squeue_output,expected_nodes",
     [