From ebd965b69f9d598939d313d1904451e92e51bd2a Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Thu, 27 Aug 2026 13:37:30 +0200 Subject: [PATCH 1/8] pre-api slurm refactor --- .../slurm/docker_image_cache_manager.py | 59 ++---------- .../systems/slurm/single_sbatch_runner.py | 16 +--- src/cloudai/systems/slurm/slurm_installer.py | 35 +------ src/cloudai/systems/slurm/slurm_runner.py | 34 +------ src/cloudai/systems/slurm/slurm_system.py | 91 ++++++++++++++++++- tests/systems/slurm/test_system.py | 22 ++++- tests/test_docker_image_cache_manager.py | 46 ++++------ tests/test_get_job_id.py | 29 +++--- 8 files changed, 155 insertions(+), 177 deletions(-) diff --git a/src/cloudai/systems/slurm/docker_image_cache_manager.py b/src/cloudai/systems/slurm/docker_image_cache_manager.py index 3f1effa73..737da93fa 100644 --- a/src/cloudai/systems/slurm/docker_image_cache_manager.py +++ b/src/cloudai/systems/slurm/docker_image_cache_manager.py @@ -19,8 +19,6 @@ import logging import os import shutil -import subprocess -from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Optional @@ -186,42 +184,6 @@ def check_docker_image_exists(self, docker_image_url: str, docker_image_filename logging.debug(message) return DockerImageCacheResult(False, Path(), message) - def _import_docker_image( - self, srun_prefix: str, docker_image_url: str, docker_image_path: Path - ) -> DockerImageCacheResult: - job_name = "CloudAI_install_docker_image" - if self.system.account: - job_name = f"{self.system.account}-{job_name}.{datetime.now().strftime('%Y%m%d_%H%M%S')}" - else: - job_name = f"{job_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}" - - # Use -N1 --ntasks=1 to ensure only one compute node downloads the image - enroot_import_cmd = f"{srun_prefix} -N1 --ntasks=1 --job-name={job_name} enroot import -o {docker_image_path} docker://{docker_image_url}" - logging.debug(f"Importing Docker image: {enroot_import_cmd}") - try: - p = subprocess.run(enroot_import_cmd, shell=True, check=True, capture_output=True, text=True) - - if "Disk quota exceeded" in p.stderr or "Write error" in p.stderr: - error_message = ( - f"Failed to cache Docker image {docker_image_url}. Command: {enroot_import_cmd}. " - f"Error: '{p.stderr}'\n\n" - "This error indicates a disk-related issue. Please check if the disk is full or not usable. " - "If the disk is full, consider using a different disk or removing unnecessary files." - ) - logging.error(error_message) - return DockerImageCacheResult(False, Path(), error_message) - - 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.absolute(), success_message) - except subprocess.CalledProcessError as e: - error_message = ( - f"Failed to import Docker image {docker_image_url}. Command: {enroot_import_cmd}. Error: {e.stderr}" - ) - logging.debug(error_message) - return DockerImageCacheResult(False, message=error_message) - def cache_docker_image(self, docker_image_url: str, docker_image_filename: str) -> DockerImageCacheResult: """ Cache the Docker image locally using enroot import. @@ -245,25 +207,20 @@ def cache_docker_image(self, docker_image_url: str, docker_image_filename: str) logging.error(error_message) return DockerImageCacheResult(False, Path(), error_message) - prerequisite_check = self._check_prerequisites() - if not prerequisite_check: - logging.error(f"Prerequisite check failed: {prerequisite_check.message}") - return DockerImageCacheResult(False, Path(), prerequisite_check.message) - if not os.access(self.system.install_path, os.W_OK): error_message = f"No permission to write in install path {self.system.install_path}." logging.error(error_message) return DockerImageCacheResult(False, Path(), error_message) - srun_prefix = f"srun --export=ALL --partition={self.system.default_partition}" - if self.system.account: - srun_prefix += f" --account={self.system.account}" - if self.system.supports_gpu_directives: - srun_prefix += " --gres=gpu:1" - if self.system.extra_srun_args: - srun_prefix += f" {self.system.extra_srun_args}" + try: + self.system.import_docker_image(docker_image_url, docker_image_path) + except RuntimeError as exc: + logging.debug(str(exc)) + return DockerImageCacheResult(False, message=str(exc)) - return self._import_docker_image(srun_prefix, docker_image_url, docker_image_path) + success_message = f"Docker image cached successfully at {docker_image_path}." + logging.debug(success_message) + return DockerImageCacheResult(True, docker_image_path.absolute(), success_message) def _check_prerequisites(self) -> PrerequisiteCheckResult: """ diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py index 3a49041f4..1aaf353bf 100644 --- a/src/cloudai/systems/slurm/single_sbatch_runner.py +++ b/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" @@ -250,16 +249,7 @@ 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.", - ) + job_id = self.system.submit_job(exec_cmd, tr.name) logging.info(f"Submitted slurm job: {job_id}") return SlurmJob(tr, id=job_id) diff --git a/src/cloudai/systems/slurm/slurm_installer.py b/src/cloudai/systems/slurm/slurm_installer.py index 58aa7a5d5..641d9e203 100644 --- a/src/cloudai/systems/slurm/slurm_installer.py +++ b/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): diff --git a/src/cloudai/systems/slurm/slurm_runner.py b/src/cloudai/systems/slurm/slurm_runner.py index 1980211ee..9c7013b49 100644 --- a/src/cloudai/systems/slurm/slurm_runner.py +++ b/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) diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index edef2989b..ee8853808 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -18,14 +18,17 @@ import logging import re +import shutil +import subprocess import time from copy import copy +from datetime import datetime 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 @@ -120,6 +123,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 +276,73 @@ 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 import_docker_image(self, docker_image_url: str, docker_image_path: Path) -> None: + """Import a Docker image into the shared Enroot cache through Slurm.""" + job_name = "CloudAI_install_docker_image" + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + job_name = f"{self.account}-{job_name}.{timestamp}" if self.account else f"{job_name}_{timestamp}" + + command = f"srun --export=ALL --partition={self.default_partition}" + if self.account: + command += f" --account={self.account}" + if self.supports_gpu_directives: + command += " --gres=gpu:1" + if self.extra_srun_args: + command += f" {self.extra_srun_args}" + command += ( + f" -N1 --ntasks=1 --job-name={job_name} enroot import -o {docker_image_path} docker://{docker_image_url}" + ) + + logging.debug("Importing Docker image: %s", command) + try: + result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True) + except subprocess.CalledProcessError as exc: + raise RuntimeError(f"Failed to import Docker image {docker_image_url}: {exc.stderr}") from exc + + if "Disk quota exceeded" in result.stderr or "Write error" in result.stderr: + raise RuntimeError( + f"Failed to cache Docker image {docker_image_url}: {result.stderr}. " + "Please check whether the target disk is full or unusable." + ) + 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. diff --git a/tests/systems/slurm/test_system.py b/tests/systems/slurm/test_system.py index 0f1b69ebf..da05ccf22 100644 --- a/tests/systems/slurm/test_system.py +++ b/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", [ diff --git a/tests/test_docker_image_cache_manager.py b/tests/test_docker_image_cache_manager.py index 3d183b6f3..5a9e409ff 100644 --- a/tests/test_docker_image_cache_manager.py +++ b/tests/test_docker_image_cache_manager.py @@ -23,7 +23,6 @@ from cloudai.systems.slurm.docker_image_cache_manager import ( DockerImageCacheManager, DockerImageCacheResult, - PrerequisiteCheckResult, ) from cloudai.systems.slurm.slurm_system import SlurmSystem @@ -67,11 +66,8 @@ def test_ensure_docker_image_url_cache_enabled(mock_access, mock_exists, mock_is @patch("pathlib.Path.is_file") @patch("pathlib.Path.exists") @patch("os.access") -@patch("subprocess.run") -@patch("cloudai.systems.slurm.docker_image_cache_manager.DockerImageCacheManager._check_prerequisites") -def test_cache_docker_image( - mock_check_prerequisites, mock_run, mock_access, mock_exists, mock_is_file, slurm_system: SlurmSystem -): +@patch("cloudai.systems.slurm.slurm_system.subprocess.run") +def test_cache_docker_image(mock_run, mock_access, mock_exists, mock_is_file, slurm_system: SlurmSystem): manager = DockerImageCacheManager(slurm_system) # Test when cached file already exists @@ -90,9 +86,6 @@ def test_cache_docker_image( ] # install_path exists, subdir_path does not, install_path again result = manager.cache_docker_image("docker.io/hello-world", "image.tar.gz") - # Ensure prerequisites are always met for the following tests - mock_check_prerequisites.return_value = PrerequisiteCheckResult(True, "All prerequisites are met.") - # Reset the mock calls mock_run.reset_mock() mock_exists.side_effect = None @@ -169,22 +162,20 @@ def test_uninstall_cached_image(mock_is_file, mock_unlink, slurm_system: SlurmSy @patch("shutil.which") -def test_check_prerequisites(mock_which, slurm_system: SlurmSystem): - manager = DockerImageCacheManager(slurm_system) - - # Ensure enroot and srun are installed - mock_which.side_effect = lambda x: x in ["enroot", "srun"] - - # Test all prerequisites met - result = manager._check_prerequisites() - assert result.success - assert result.message == "All prerequisites are met." +@patch("cloudai.systems.slurm.slurm_system.subprocess.run") +def test_check_prerequisites(mock_run, mock_which, slurm_system: SlurmSystem): + mock_which.return_value = "/usr/bin/tool" + mock_run.return_value = subprocess.CompletedProcess( + args=["srun", "--help"], + returncode=0, + stdout=" ".join(slurm_system._REQUIRED_SRUN_OPTIONS), + stderr="", + ) + slurm_system.validate_install_environment() - # Test srun not installed mock_which.side_effect = lambda x: x != "srun" - result = manager._check_prerequisites() - assert not result.success - assert result.message == "srun are required for caching Docker images but are not installed." + with pytest.raises(EnvironmentError, match="Required binary 'srun' is not installed"): + slurm_system.validate_install_environment() def test_ensure_docker_image_no_local_cache(slurm_system: SlurmSystem): @@ -209,9 +200,9 @@ def test_docker_import_with_extra_system_config( slurm_system.install_path.mkdir(parents=True, exist_ok=True) manager = DockerImageCacheManager(slurm_system) - manager._check_prerequisites = lambda: PrerequisiteCheckResult(True, "All prerequisites are met.") - with patch("subprocess.run") as mock_run: + with patch("cloudai.systems.slurm.slurm_system.subprocess.run") as mock_run: + mock_run.return_value = subprocess.CompletedProcess(args=["cmd"], returncode=0, stdout="", stderr="") res = manager.cache_docker_image("docker.io/hello-world", "docker_image.sqsh") assert res.success @@ -245,10 +236,9 @@ def test_docker_import_with_extra_srun_args(slurm_system: SlurmSystem): slurm_system.install_path.mkdir(parents=True, exist_ok=True) manager = DockerImageCacheManager(slurm_system) - manager._check_prerequisites = lambda: PrerequisiteCheckResult(True, "All prerequisites are met.") - with patch("subprocess.run") as mock_run: - mock_run.return_value = subprocess.CompletedProcess(args=["cmd"], returncode=0, stderr="") + with patch("cloudai.systems.slurm.slurm_system.subprocess.run") as mock_run: + mock_run.return_value = subprocess.CompletedProcess(args=["cmd"], returncode=0, stdout="", stderr="") res = manager.cache_docker_image("docker.io/hello-world", "docker_image.sqsh") assert res.success diff --git a/tests/test_get_job_id.py b/tests/test_get_job_id.py index e57ee984e..8c306481c 100644 --- a/tests/test_get_job_id.py +++ b/tests/test_get_job_id.py @@ -14,7 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import subprocess from pathlib import Path from unittest.mock import Mock, patch @@ -24,20 +23,9 @@ from cloudai.systems.lsf.lsf_runner import LSFRunner from cloudai.systems.lsf.lsf_system import LSFSystem from cloudai.systems.slurm import SlurmJob, SlurmRunner, SlurmSystem -from cloudai.util import CommandShell from cloudai.workloads.sleep.sleep import SleepCmdArgs, SleepTestDefinition -class MockCommandShell(CommandShell): - def execute(self, command): - mock_popen = Mock(spec=subprocess.Popen) - mock_popen.communicate.return_value = ( - "", - "sbatch: error: Batch job submission failed: Requested node configuration is not available", - ) - return mock_popen - - @pytest.fixture def test_scenario(slurm_system: SlurmSystem) -> TestScenario: test_scenario = TestScenario( @@ -58,15 +46,22 @@ def test_scenario(slurm_system: SlurmSystem) -> TestScenario: @pytest.fixture def slurm_runner(slurm_system: SlurmSystem, test_scenario: TestScenario) -> SlurmRunner: - runner = SlurmRunner( + return SlurmRunner( mode="run", system=slurm_system, test_scenario=test_scenario, output_path=slurm_system.output_path ) - runner.cmd_shell = MockCommandShell() - return runner def test_job_id_retrieval_error(slurm_runner: SlurmRunner): tr = slurm_runner.test_scenario.test_runs[0] + slurm_runner.system.submit_job = Mock( + side_effect=JobIdRetrievalError( + test_name=str(tr.name), + command="sbatch script.sh", + stdout="", + stderr="sbatch: error: Batch job submission failed: Requested node configuration is not available", + message="Failed to retrieve job ID.", + ) + ) with pytest.raises(JobIdRetrievalError) as excinfo: slurm_runner._submit_test(tr) assert "Failed to retrieve job ID." in str(excinfo.value) @@ -83,8 +78,8 @@ def test_job_id_retrieval_error(slurm_runner: SlurmRunner): ("", "sbatch: error: Batch job submission failed:...", None), ], ) -def test_slurm_get_job_id(slurm_runner: SlurmRunner, stdout: str, stderr: str, expected_job_id: int | None): - res = slurm_runner.get_job_id(stdout, stderr) +def test_slurm_get_job_id(stdout: str, stderr: str, expected_job_id: int | None): + res = SlurmSystem._parse_submitted_job_id(stdout) assert res == expected_job_id From 673ed2a98ec6c37792708654aaaccd54c16a99da Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Fri, 28 Aug 2026 14:33:14 +0200 Subject: [PATCH 2/8] make slurm system class responsibilities clear --- .../slurm/docker_image_cache_manager.py | 278 +++++++----------- .../systems/slurm/single_sbatch_runner.py | 4 +- src/cloudai/systems/slurm/slurm_system.py | 9 +- tests/test_docker_image_cache_manager.py | 266 ++++------------- tests/test_toml_files.py | 2 +- 5 files changed, 171 insertions(+), 388 deletions(-) diff --git a/src/cloudai/systems/slurm/docker_image_cache_manager.py b/src/cloudai/systems/slurm/docker_image_cache_manager.py index 737da93fa..8101e4d91 100644 --- a/src/cloudai/systems/slurm/docker_image_cache_manager.py +++ b/src/cloudai/systems/slurm/docker_image_cache_manager.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -18,147 +18,56 @@ import logging import os -import shutil +import shlex +from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Optional +from cloudai.core import JobIdRetrievalError + if TYPE_CHECKING: from cloudai.systems.slurm import SlurmSystem -class PrerequisiteCheckResult: - """ - Class representing the result of a prerequisite check. - - Attributes - success (bool): Indicates whether the prerequisite check was successful. - message (str): A message providing additional information about the result. - """ - - def __init__(self, success: bool, message: str = "") -> None: - """ - Initialize the PrerequisiteCheckResult. - - Args: - success (bool): Indicates whether the prerequisite check was successful. - message (str): A message providing additional information about the result. - """ - self.success = success - self.message = message - - def __bool__(self): - """ - Return the success status as a boolean. - - Returns - bool: True if the check was successful, False otherwise. - """ - return self.success - - def __str__(self): - """ - Return the message as a string. - - Returns - str: The message providing additional information about the result. - """ - return self.message - - class DockerImageCacheResult: - """ - Class representing the result of a Docker image caching operation. - - Attributes - success (bool): Indicates whether the operation was successful. - docker_image_path (Path): The path to the Docker image. - message (str): A message providing additional information about the result. - """ - - def __init__(self, success: bool, docker_image_path: Optional[Path] = None, message: str = "") -> None: - """ - Initialize the DockerImageCacheResult. - - Args: - success (bool): Indicates whether the operation was successful. - docker_image_path (Path): The path to the Docker image. - message (str): A message providing additional information about the result. - """ + """Result of a Docker image caching operation.""" + + def __init__( + self, + success: bool, + docker_image_path: Optional[Path] = None, + message: str = "", + ) -> None: self.success = success self.docker_image_path = docker_image_path self.message = message def __bool__(self): - """ - Return the success status as a boolean. - - Returns - bool: True if the operation was successful, False otherwise. - """ + """Return whether the cache operation succeeded.""" return self.success def __str__(self): - """ - Return the message as a string. - - Returns - str: The message providing additional information about the result. - """ + """Return the result message.""" return self.message class DockerImageCacheManager: - """ - Manages the caching of Docker images for installation strategies. - - Attributes - system (SlurmSystem): The Slurm system configuration. - """ + """Generate and interpret jobs which cache Docker images on a Slurm system.""" def __init__(self, system: SlurmSystem) -> None: self.system = system def ensure_docker_image(self, docker_image_url: str, docker_image_filename: str) -> DockerImageCacheResult: - """ - Ensure the Docker image exists by checking and optionally caching it. - - Args: - docker_image_url (str): URL or file path of the Docker image. - docker_image_filename (str): Docker image filename. - - Returns: - DockerImageCacheResult: Result of ensuring the Docker image exists. - """ - image_check_result = self.check_docker_image_exists(docker_image_url, docker_image_filename) - if image_check_result.success: - return image_check_result - + result = self.check_docker_image_exists(docker_image_url, docker_image_filename) + if result.success: + return result if self.system.cache_docker_images_locally: return self.cache_docker_image(docker_image_url, docker_image_filename) - - return image_check_result + return result def check_docker_image_exists(self, docker_image_url: str, docker_image_filename: str) -> DockerImageCacheResult: - """ - Check if the Docker image exists without caching it. - - Args: - docker_image_url (str): URL or file path of the Docker image. - docker_image_filename (str): Docker image filename. - - Returns: - DockerImageCacheResult: Result of the Docker image existence check. - """ - logging.debug( - f"Checking if Docker image exists: docker_image_url={docker_image_url}, " - f"subdir_name={self.system.install_path}, " - f"docker_image_filename={docker_image_filename}, " - f"cache_docker_images_locally={self.system.cache_docker_images_locally}" - ) - - # If not caching locally, return True. Defer checking URL accessibility to srun. if not self.system.cache_docker_images_locally: - return DockerImageCacheResult(True, None, "") + return DockerImageCacheResult(True) docker_image_path = Path(docker_image_url) if docker_image_path.is_file() and docker_image_path.exists(): @@ -168,7 +77,6 @@ def check_docker_image_exists(self, docker_image_url: str, docker_image_filename f"Docker image file path is valid: {docker_image_url}.", ) - # Check if the cache file exists if not self.system.install_path.exists(): message = f"Install path {self.system.install_path.absolute()} does not exist." logging.debug(message) @@ -184,85 +92,99 @@ def check_docker_image_exists(self, docker_image_url: str, docker_image_filename logging.debug(message) return DockerImageCacheResult(False, Path(), message) - def cache_docker_image(self, docker_image_url: str, docker_image_filename: str) -> DockerImageCacheResult: - """ - Cache the Docker image locally using enroot import. - - Args: - docker_image_url (str): URL of the Docker image. - docker_image_filename (str): Docker image filename. + def _job_name(self) -> str: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + if self.system.account: + return f"{self.system.account}-CloudAI_install_docker_image.{timestamp}" + return f"CloudAI_install_docker_image_{timestamp}" + + def _write_import_script(self, docker_image_url: str, docker_image_path: Path) -> tuple[Path, Path]: + job_name = self._job_name() + script_path = self.system.install_path / f".{job_name}.sh" + stdout_path = self.system.install_path / f".{job_name}.out" + stderr_path = self.system.install_path / f".{job_name}.err" + directives = [ + "#!/bin/bash", + f"#SBATCH --job-name={job_name}", + f"#SBATCH --output={stdout_path}", + f"#SBATCH --error={stderr_path}", + f"#SBATCH --partition={self.system.default_partition}", + "#SBATCH -N1", + "#SBATCH --ntasks=1", + ] + if self.system.account: + directives.append(f"#SBATCH --account={self.system.account}") + if self.system.supports_gpu_directives: + directives.append("#SBATCH --gres=gpu:1") + directives.extend(f"#SBATCH {arg}" for arg in self.system.extra_sbatch_args) + + srun = ["srun", "--export=ALL", "--ntasks=1"] + if self.system.extra_srun_args: + srun.append(self.system.extra_srun_args) + srun.extend( + [ + "enroot import -o", + shlex.quote(str(docker_image_path)), + shlex.quote(f"docker://{docker_image_url}"), + ] + ) + script_path.write_text("\n".join([*directives, "", " ".join(srun), ""]), encoding="utf-8") + return script_path, stderr_path - Returns: - DockerImageCacheResult: Result of the Docker image caching operation. - """ + def cache_docker_image(self, docker_image_url: str, docker_image_filename: str) -> DockerImageCacheResult: docker_image_path = self.system.install_path / docker_image_filename - 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.absolute(), success_message) + message = f"Cached Docker image already exists at {docker_image_path}." + logging.info(message) + return DockerImageCacheResult(True, docker_image_path.absolute(), message) if not self.system.install_path.exists(): - error_message = f"Install path {self.system.install_path.absolute()} does not exist." - logging.error(error_message) - return DockerImageCacheResult(False, Path(), error_message) - + message = f"Install path {self.system.install_path.absolute()} does not exist." + logging.error(message) + return DockerImageCacheResult(False, Path(), message) if not os.access(self.system.install_path, os.W_OK): - error_message = f"No permission to write in install path {self.system.install_path}." - logging.error(error_message) - return DockerImageCacheResult(False, Path(), error_message) + message = f"No permission to write in install path {self.system.install_path}." + logging.error(message) + return DockerImageCacheResult(False, Path(), message) + script_path, stderr_path = self._write_import_script(docker_image_url, docker_image_path) try: - self.system.import_docker_image(docker_image_url, docker_image_path) - except RuntimeError as exc: - logging.debug(str(exc)) - return DockerImageCacheResult(False, message=str(exc)) - - success_message = f"Docker image cached successfully at {docker_image_path}." - logging.debug(success_message) - return DockerImageCacheResult(True, docker_image_path.absolute(), success_message) - - def _check_prerequisites(self) -> PrerequisiteCheckResult: - """ - Check prerequisites for caching Docker image. + self.system.submit_sbatch(script_path, "Docker image import", wait=True) + except JobIdRetrievalError as error: + message = f"Failed to import Docker image {docker_image_url}: {error}" + logging.error(message) + return DockerImageCacheResult(False, message=message) - Returns: - PrerequisiteCheckResult: Result of the prerequisite check. - """ - required_binaries = ["srun"] - missing_binaries = [binary for binary in required_binaries if not shutil.which(binary)] + stderr = stderr_path.read_text(encoding="utf-8") if stderr_path.is_file() else "" + if docker_image_path.is_file(): + message = f"Docker image cached successfully at {docker_image_path}." + logging.debug(message) + return DockerImageCacheResult(True, docker_image_path.absolute(), message) - if missing_binaries: - missing_binaries_str = ", ".join(missing_binaries) - logging.error(f"{missing_binaries_str} are required for caching Docker images but are not installed.") - return PrerequisiteCheckResult( - False, - f"{missing_binaries_str} are required for caching Docker images but are not installed.", + if "Disk quota exceeded" in stderr or "Write error" in stderr: + message = ( + f"Failed to cache Docker image {docker_image_url}. Error: '{stderr}'\n\n" + "This error indicates a disk-related issue. Please check if the disk is full or not usable. " + "If the disk is full, consider using a different disk or removing unnecessary files." ) - - return PrerequisiteCheckResult(True, "All prerequisites are met.") + else: + message = f"Failed to import Docker image {docker_image_url}. Error: {stderr or 'image was not created'}" + logging.error(message) + return DockerImageCacheResult(False, message=message) def uninstall_cached_image(self, docker_image_filename: str) -> DockerImageCacheResult: - """ - Remove an existing cached Docker image. - - Args: - docker_image_filename (str): Docker image filename. - - Returns: - DockerImageCacheResult: Result of the removal operation. - """ docker_image_path = self.system.install_path / docker_image_filename if docker_image_path.is_file(): try: 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.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.warning(success_message) - return DockerImageCacheResult(True, docker_image_path.absolute(), success_message) + message = f"Cached Docker image removed successfully from {docker_image_path}." + logging.info(message) + return DockerImageCacheResult(True, docker_image_path.absolute(), message) + except OSError as error: + message = f"Failed to remove cached Docker image at {docker_image_path}. Error: {error}" + logging.error(message) + return DockerImageCacheResult(False, docker_image_path, message) + + message = f"No cached Docker image found to remove at {docker_image_path}." + logging.warning(message) + return DockerImageCacheResult(True, docker_image_path.absolute(), message) diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py index 1aaf353bf..0476e2f92 100644 --- a/src/cloudai/systems/slurm/single_sbatch_runner.py +++ b/src/cloudai/systems/slurm/single_sbatch_runner.py @@ -248,8 +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'}" - job_id = self.system.submit_job(exec_cmd, tr.name) + 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) diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index ee8853808..1baa525ff 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -18,6 +18,7 @@ import logging import re +import shlex import shutil import subprocess import time @@ -100,6 +101,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 @@ -314,7 +321,7 @@ def validate_install_environment(self) -> None: if missing_options: raise EnvironmentError(f"Required srun options missing: {', '.join(missing_options)}") - def import_docker_image(self, docker_image_url: str, docker_image_path: Path) -> None: + def _deprecated_import_docker_image(self, docker_image_url: str, docker_image_path: Path) -> None: """Import a Docker image into the shared Enroot cache through Slurm.""" job_name = "CloudAI_install_docker_image" timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") diff --git a/tests/test_docker_image_cache_manager.py b/tests/test_docker_image_cache_manager.py index 5a9e409ff..26ca24e48 100644 --- a/tests/test_docker_image_cache_manager.py +++ b/tests/test_docker_image_cache_manager.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,242 +14,96 @@ # See the License for the specific language governing permissions and # limitations under the License. -import subprocess from pathlib import Path from unittest.mock import patch import pytest -from cloudai.systems.slurm.docker_image_cache_manager import ( - DockerImageCacheManager, - DockerImageCacheResult, -) +from cloudai.core import JobIdRetrievalError +from cloudai.systems.slurm.docker_image_cache_manager import DockerImageCacheManager from cloudai.systems.slurm.slurm_system import SlurmSystem -@patch("pathlib.Path.is_file") -@patch("pathlib.Path.exists") -@patch("os.access") -def test_ensure_docker_image_file_exists(mock_access, mock_exists, mock_is_file, slurm_system: SlurmSystem): - manager = DockerImageCacheManager(slurm_system) - mock_is_file.return_value = True - mock_exists.return_value = True - result = manager.ensure_docker_image("/tmp/existing_file.sqsh", "docker_image.sqsh") - assert result.success - assert result.docker_image_path == Path("/tmp/existing_file.sqsh") - assert result.message == "Docker image file path is valid: /tmp/existing_file.sqsh." - - -@patch("pathlib.Path.is_file") -@patch("pathlib.Path.exists") -@patch("os.access") -def test_ensure_docker_image_url_cache_enabled(mock_access, mock_exists, mock_is_file, slurm_system: SlurmSystem): - manager = DockerImageCacheManager(slurm_system) - mock_is_file.return_value = False - mock_exists.return_value = True - mock_access.return_value = True - with patch.object( - manager, - "cache_docker_image", - return_value=DockerImageCacheResult( - True, - Path("/fake/install/path/subdir/docker_image.sqsh"), - "Docker image cached successfully.", - ), - ): - result = manager.ensure_docker_image("docker.io/hello-world", "docker_image.sqsh") - assert result.success - assert result.docker_image_path == Path("/fake/install/path/subdir/docker_image.sqsh") - assert result.message == "Docker image cached successfully." - - -@patch("pathlib.Path.is_file") -@patch("pathlib.Path.exists") -@patch("os.access") -@patch("cloudai.systems.slurm.slurm_system.subprocess.run") -def test_cache_docker_image(mock_run, mock_access, mock_exists, mock_is_file, slurm_system: SlurmSystem): - manager = DockerImageCacheManager(slurm_system) - - # Test when cached file already exists - mock_is_file.return_value = True - result = manager.cache_docker_image("docker.io/hello-world", "image.tar.gz") - assert result.success - assert result.docker_image_path == slurm_system.install_path / "image.tar.gz" - assert result.message == f"Cached Docker image already exists at {slurm_system.install_path}/image.tar.gz." - - # Test creating subdirectory when it doesn't exist - mock_is_file.return_value = False - mock_exists.side_effect = [ - True, - False, - True, - ] # install_path exists, subdir_path does not, install_path again - result = manager.cache_docker_image("docker.io/hello-world", "image.tar.gz") - - # Reset the mock calls - mock_run.reset_mock() - mock_exists.side_effect = None - - # Test caching success with subprocess command (removal of default partition keyword) - mock_is_file.return_value = False - mock_exists.side_effect = [ - True, - True, - True, - True, - True, - ] # Ensure all path checks return True - mock_run.return_value = subprocess.CompletedProcess(args=["cmd"], returncode=0, stderr="") - result = manager.cache_docker_image("docker.io/hello-world", "image.tar.gz") - - assert mock_run.call_count == 1 - actual_command = mock_run.call_args[0][0] - assert f"srun --export=ALL --partition={slurm_system.default_partition}" in actual_command - assert "--ntasks=1" in actual_command - assert "-N1" in actual_command - assert "--job-name=CloudAI_install_docker_image_" in actual_command - assert f"enroot import -o {slurm_system.install_path}/image.tar.gz docker://docker.io/hello-world" in actual_command - assert mock_run.call_args[1] == {"shell": True, "check": True, "capture_output": True, "text": True} - - assert result.success - assert result.message == f"Docker image cached successfully at {slurm_system.install_path}/image.tar.gz." - - # Test caching failure due to subprocess error - mock_is_file.return_value = False - mock_run.side_effect = subprocess.CalledProcessError(1, "cmd") - result = manager.cache_docker_image("docker.io/hello-world", "image.tar.gz") - assert not result.success - - # Test caching failure due to disk-related errors - mock_is_file.return_value = False - mock_run.side_effect = None - mock_run.return_value = subprocess.CompletedProcess(args=["cmd"], returncode=1, stderr="Disk quota exceeded\n") - mock_exists.side_effect = [True, True, True, True, True] - result = manager.cache_docker_image("docker.io/hello-world", "image.tar.gz") - assert not result.success - assert "Disk quota exceeded" in result.message +def test_ensure_existing_image_file(slurm_system: SlurmSystem, tmp_path: Path): + image = tmp_path / "existing.sqsh" + image.touch() + slurm_system.cache_docker_images_locally = True - mock_run.return_value = subprocess.CompletedProcess(args=["cmd"], returncode=1, stderr="Write error\n") - result = manager.cache_docker_image("docker.io/hello-world", "image.tar.gz") - assert not result.success - assert "Write error" in result.message - - -@patch("pathlib.Path.unlink") -@patch("pathlib.Path.is_file") -def test_uninstall_cached_image(mock_is_file, mock_unlink, slurm_system: SlurmSystem): - # Mock setup - manager = DockerImageCacheManager(slurm_system) - - # Test successful removal - mock_is_file.return_value = True - result = manager.uninstall_cached_image("image.tar.gz") - assert result.success - assert result.message == f"Cached Docker image removed successfully from {slurm_system.install_path}/image.tar.gz." - mock_unlink.assert_called_once() - - # Test failed removal due to OSError - mock_unlink.side_effect = OSError("Mocked OSError") - result = manager.uninstall_cached_image("image.tar.gz") - assert not result.success - assert "Failed to remove cached Docker image" in result.message + result = DockerImageCacheManager(slurm_system).ensure_docker_image(str(image), "cached.sqsh") - # Test no file to remove - mock_is_file.return_value = False - result = manager.uninstall_cached_image("image.tar.gz") assert result.success - assert result.message == f"No cached Docker image found to remove at {slurm_system.install_path}/image.tar.gz." - - -@patch("shutil.which") -@patch("cloudai.systems.slurm.slurm_system.subprocess.run") -def test_check_prerequisites(mock_run, mock_which, slurm_system: SlurmSystem): - mock_which.return_value = "/usr/bin/tool" - mock_run.return_value = subprocess.CompletedProcess( - args=["srun", "--help"], - returncode=0, - stdout=" ".join(slurm_system._REQUIRED_SRUN_OPTIONS), - stderr="", - ) - slurm_system.validate_install_environment() - - mock_which.side_effect = lambda x: x != "srun" - with pytest.raises(EnvironmentError, match="Required binary 'srun' is not installed"): - slurm_system.validate_install_environment() + assert result.docker_image_path == image def test_ensure_docker_image_no_local_cache(slurm_system: SlurmSystem): slurm_system.cache_docker_images_locally = False - manager = DockerImageCacheManager(slurm_system) - result = manager.ensure_docker_image("docker.io/hello-world", "docker_image.sqsh") + result = DockerImageCacheManager(slurm_system).ensure_docker_image("docker.io/hello-world", "image.sqsh") assert result.success assert result.docker_image_path is None - assert result.message == "" -@pytest.mark.parametrize( - "account, supports_gpu_directives", [(None, False), ("test_account", True), (None, False), ("test_account", True)] -) -def test_docker_import_with_extra_system_config( - slurm_system: SlurmSystem, account: str | None, supports_gpu_directives: bool | None +@pytest.mark.parametrize("account,supports_gpu", [(None, False), ("test-account", True)]) +def test_cache_docker_image_submits_one_node_sbatch_job( + slurm_system: SlurmSystem, account: str | None, supports_gpu: bool ): + slurm_system.cache_docker_images_locally = True slurm_system.account = account - slurm_system.supports_gpu_directives_cache = supports_gpu_directives + slurm_system.supports_gpu_directives_cache = supports_gpu + slurm_system.extra_srun_args = "--reservation test-reservation" slurm_system.install_path.mkdir(parents=True, exist_ok=True) + image_path = slurm_system.install_path / "image.sqsh" + + def submit(script_path: Path, operation_name: str, *, wait: bool) -> int: + content = script_path.read_text(encoding="utf-8") + assert "#SBATCH --partition=" + slurm_system.default_partition in content + assert "#SBATCH -N1" in content + assert "#SBATCH --ntasks=1" in content + assert "srun --export=ALL --ntasks=1 --reservation test-reservation enroot import -o" in content + assert "docker://docker.io/hello-world" in content + if account: + assert f"#SBATCH --account={account}" in content + if supports_gpu: + assert "#SBATCH --gres=gpu:1" in content + assert operation_name == "Docker image import" + assert wait is True + image_path.touch() + return 123 + + with patch.object(SlurmSystem, "submit_sbatch", side_effect=submit) as submit_sbatch: + result = DockerImageCacheManager(slurm_system).cache_docker_image("docker.io/hello-world", "image.sqsh") - manager = DockerImageCacheManager(slurm_system) - - with patch("cloudai.systems.slurm.slurm_system.subprocess.run") as mock_run: - mock_run.return_value = subprocess.CompletedProcess(args=["cmd"], returncode=0, stdout="", stderr="") - res = manager.cache_docker_image("docker.io/hello-world", "docker_image.sqsh") - assert res.success - - assert mock_run.call_count == 1 - - actual_command = mock_run.call_args[0][0] - - expected_prefix = f"srun --export=ALL --partition={slurm_system.default_partition}" - assert expected_prefix in actual_command - assert "-N1" in actual_command - - if account: - assert f"--account={account}" in actual_command - assert f"--job-name={account}-CloudAI_install_docker_image." in actual_command - else: - assert "--job-name=CloudAI_install_docker_image_" in actual_command + assert result.success + assert result.docker_image_path == image_path + submit_sbatch.assert_called_once() - if supports_gpu_directives: - assert "--gres=gpu:1" in actual_command - assert ( - f"enroot import -o {slurm_system.install_path}/docker_image.sqsh docker://docker.io/hello-world" - in actual_command +def test_cache_docker_image_reports_submission_failure(slurm_system: SlurmSystem): + slurm_system.cache_docker_images_locally = True + slurm_system.supports_gpu_directives_cache = False + slurm_system.install_path.mkdir(parents=True, exist_ok=True) + error = JobIdRetrievalError( + test_name="Docker image import", + command="sbatch image.sh", + stdout="", + stderr="submission failed", + message="Failed to retrieve job ID.", ) - assert mock_run.call_args[1] == {"shell": True, "check": True, "capture_output": True, "text": True} - - -def test_docker_import_with_extra_srun_args(slurm_system: SlurmSystem): - slurm_system.extra_srun_args = "--reservation test_reservation --w hgx-isr1-pre-09" - slurm_system.install_path.mkdir(parents=True, exist_ok=True) + with patch.object(SlurmSystem, "submit_sbatch", side_effect=error): + result = DockerImageCacheManager(slurm_system).cache_docker_image("docker.io/hello-world", "image.sqsh") - manager = DockerImageCacheManager(slurm_system) + assert not result.success + assert "Failed to import Docker image" in result.message - with patch("cloudai.systems.slurm.slurm_system.subprocess.run") as mock_run: - mock_run.return_value = subprocess.CompletedProcess(args=["cmd"], returncode=0, stdout="", stderr="") - res = manager.cache_docker_image("docker.io/hello-world", "docker_image.sqsh") - assert res.success - assert mock_run.call_count == 1 - actual_command = mock_run.call_args[0][0] +def test_uninstall_cached_image(slurm_system: SlurmSystem): + slurm_system.install_path.mkdir(parents=True, exist_ok=True) + image_path = slurm_system.install_path / "image.sqsh" + image_path.touch() - # Check that extra_srun_args are included - assert "--reservation test_reservation" in actual_command - assert "--w hgx-isr1-pre-09" in actual_command + result = DockerImageCacheManager(slurm_system).uninstall_cached_image("image.sqsh") - # Check that the command still has the basic structure - expected_prefix = f"srun --export=ALL --partition={slurm_system.default_partition}" - assert expected_prefix in actual_command - assert "enroot import -o" in actual_command + assert result.success + assert not image_path.exists() diff --git a/tests/test_toml_files.py b/tests/test_toml_files.py index 819353312..1fedb72db 100644 --- a/tests/test_toml_files.py +++ b/tests/test_toml_files.py @@ -59,7 +59,7 @@ def test_toml_files(toml_file: Path): @pytest.mark.parametrize("system_file", ALL_SYSTEMS, ids=lambda x: str(x)) @patch("kubernetes.config.load_kube_config") @patch("pathlib.Path.exists", return_value=True) -def test_systems(mock_exists, mock_load_kube_config, system_file: Path): +def test_systems(mock_load_kube_config, system_file: Path): """ Validate the syntax of a system configuration file. From 36b8c075fb2c31025f2c0e067e40643f212127c6 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 7 Sep 2026 14:01:47 +0200 Subject: [PATCH 3/8] fix tests --- .../systems/slurm/docker_image_cache_manager.py | 10 ++++++---- tests/test_docker_image_cache_manager.py | 3 +++ tests/test_get_job_id.py | 16 +++++++--------- tests/test_toml_files.py | 2 +- 4 files changed, 17 insertions(+), 14 deletions(-) diff --git a/src/cloudai/systems/slurm/docker_image_cache_manager.py b/src/cloudai/systems/slurm/docker_image_cache_manager.py index 8101e4d91..af2e444d8 100644 --- a/src/cloudai/systems/slurm/docker_image_cache_manager.py +++ b/src/cloudai/systems/slurm/docker_image_cache_manager.py @@ -20,6 +20,7 @@ import os import shlex from datetime import datetime +from hashlib import sha256 from pathlib import Path from typing import TYPE_CHECKING, Optional @@ -92,14 +93,15 @@ def check_docker_image_exists(self, docker_image_url: str, docker_image_filename logging.debug(message) return DockerImageCacheResult(False, Path(), message) - def _job_name(self) -> str: + def _job_name(self, docker_image_url: str) -> str: timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + image_hash = sha256(docker_image_url.encode()).hexdigest()[:8] if self.system.account: - return f"{self.system.account}-CloudAI_install_docker_image.{timestamp}" - return f"CloudAI_install_docker_image_{timestamp}" + return f"{self.system.account}-CloudAI_install_docker_image.{image_hash}.{timestamp}" + return f"CloudAI_install_docker_image_{image_hash}_{timestamp}" def _write_import_script(self, docker_image_url: str, docker_image_path: Path) -> tuple[Path, Path]: - job_name = self._job_name() + job_name = self._job_name(docker_image_url) script_path = self.system.install_path / f".{job_name}.sh" stdout_path = self.system.install_path / f".{job_name}.out" stderr_path = self.system.install_path / f".{job_name}.err" diff --git a/tests/test_docker_image_cache_manager.py b/tests/test_docker_image_cache_manager.py index 26ca24e48..b2e9c8166 100644 --- a/tests/test_docker_image_cache_manager.py +++ b/tests/test_docker_image_cache_manager.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from hashlib import sha256 from pathlib import Path from unittest.mock import patch @@ -57,6 +58,8 @@ def test_cache_docker_image_submits_one_node_sbatch_job( def submit(script_path: Path, operation_name: str, *, wait: bool) -> int: content = script_path.read_text(encoding="utf-8") + image_hash = sha256(b"docker.io/hello-world").hexdigest()[:8] + assert image_hash in script_path.name assert "#SBATCH --partition=" + slurm_system.default_partition in content assert "#SBATCH -N1" in content assert "#SBATCH --ntasks=1" in content diff --git a/tests/test_get_job_id.py b/tests/test_get_job_id.py index 8c306481c..0d27b6570 100644 --- a/tests/test_get_job_id.py +++ b/tests/test_get_job_id.py @@ -53,16 +53,14 @@ def slurm_runner(slurm_system: SlurmSystem, test_scenario: TestScenario) -> Slur def test_job_id_retrieval_error(slurm_runner: SlurmRunner): tr = slurm_runner.test_scenario.test_runs[0] - slurm_runner.system.submit_job = Mock( - side_effect=JobIdRetrievalError( - test_name=str(tr.name), - command="sbatch script.sh", - stdout="", - stderr="sbatch: error: Batch job submission failed: Requested node configuration is not available", - message="Failed to retrieve job ID.", - ) + error = JobIdRetrievalError( + test_name=str(tr.name), + command="sbatch script.sh", + stdout="", + stderr="sbatch: error: Batch job submission failed: Requested node configuration is not available", + message="Failed to retrieve job ID.", ) - with pytest.raises(JobIdRetrievalError) as excinfo: + with patch.object(SlurmSystem, "submit_job", side_effect=error), pytest.raises(JobIdRetrievalError) as excinfo: slurm_runner._submit_test(tr) assert "Failed to retrieve job ID." in str(excinfo.value) assert "sbatch: error: Batch job submission failed: Requested node configuration is not available" in str( diff --git a/tests/test_toml_files.py b/tests/test_toml_files.py index 1fedb72db..762b4c7f4 100644 --- a/tests/test_toml_files.py +++ b/tests/test_toml_files.py @@ -59,7 +59,7 @@ def test_toml_files(toml_file: Path): @pytest.mark.parametrize("system_file", ALL_SYSTEMS, ids=lambda x: str(x)) @patch("kubernetes.config.load_kube_config") @patch("pathlib.Path.exists", return_value=True) -def test_systems(mock_load_kube_config, system_file: Path): +def test_systems(_mock_exists, mock_load_kube_config, system_file: Path): """ Validate the syntax of a system configuration file. From c7e1d494ca25336c975e011ea4ee62da8b0a9740 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 7 Sep 2026 14:10:50 +0200 Subject: [PATCH 4/8] remove unused method --- src/cloudai/systems/slurm/slurm_system.py | 30 ----------------------- 1 file changed, 30 deletions(-) diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index 1baa525ff..f27367e12 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -23,7 +23,6 @@ import subprocess import time from copy import copy -from datetime import datetime from pathlib import Path from typing import Any, ClassVar, Dict, Iterable, List, Optional, Tuple, Union @@ -321,35 +320,6 @@ def validate_install_environment(self) -> None: if missing_options: raise EnvironmentError(f"Required srun options missing: {', '.join(missing_options)}") - def _deprecated_import_docker_image(self, docker_image_url: str, docker_image_path: Path) -> None: - """Import a Docker image into the shared Enroot cache through Slurm.""" - job_name = "CloudAI_install_docker_image" - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - job_name = f"{self.account}-{job_name}.{timestamp}" if self.account else f"{job_name}_{timestamp}" - - command = f"srun --export=ALL --partition={self.default_partition}" - if self.account: - command += f" --account={self.account}" - if self.supports_gpu_directives: - command += " --gres=gpu:1" - if self.extra_srun_args: - command += f" {self.extra_srun_args}" - command += ( - f" -N1 --ntasks=1 --job-name={job_name} enroot import -o {docker_image_path} docker://{docker_image_url}" - ) - - logging.debug("Importing Docker image: %s", command) - try: - result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True) - except subprocess.CalledProcessError as exc: - raise RuntimeError(f"Failed to import Docker image {docker_image_url}: {exc.stderr}") from exc - - if "Disk quota exceeded" in result.stderr or "Write error" in result.stderr: - raise RuntimeError( - f"Failed to cache Docker image {docker_image_url}: {result.stderr}. " - "Please check whether the target disk is full or unusable." - ) - 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. From b2d50098ece6c4d20cc26d2400d2d0c6ae77bc85 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 14 Sep 2026 19:34:59 +0200 Subject: [PATCH 5/8] Clean successful image import artifacts and cover failures --- .../slurm/docker_image_cache_manager.py | 5 +++ tests/test_docker_image_cache_manager.py | 45 +++++++++++++++++-- 2 files changed, 47 insertions(+), 3 deletions(-) diff --git a/src/cloudai/systems/slurm/docker_image_cache_manager.py b/src/cloudai/systems/slurm/docker_image_cache_manager.py index af2e444d8..91a3b067b 100644 --- a/src/cloudai/systems/slurm/docker_image_cache_manager.py +++ b/src/cloudai/systems/slurm/docker_image_cache_manager.py @@ -159,6 +159,11 @@ def cache_docker_image(self, docker_image_url: str, docker_image_filename: str) stderr = stderr_path.read_text(encoding="utf-8") if stderr_path.is_file() else "" if docker_image_path.is_file(): + for artifact in (script_path, script_path.with_suffix(".out"), stderr_path): + try: + artifact.unlink(missing_ok=True) + except OSError as error: + logging.warning(f"Failed to remove Docker image import artifact {artifact}: {error}") message = f"Docker image cached successfully at {docker_image_path}." logging.debug(message) return DockerImageCacheResult(True, docker_image_path.absolute(), message) diff --git a/tests/test_docker_image_cache_manager.py b/tests/test_docker_image_cache_manager.py index b2e9c8166..6d4853612 100644 --- a/tests/test_docker_image_cache_manager.py +++ b/tests/test_docker_image_cache_manager.py @@ -16,7 +16,7 @@ from hashlib import sha256 from pathlib import Path -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest @@ -45,9 +45,15 @@ def test_ensure_docker_image_no_local_cache(slurm_system: SlurmSystem): assert result.docker_image_path is None -@pytest.mark.parametrize("account,supports_gpu", [(None, False), ("test-account", True)]) +@pytest.mark.parametrize( + "account,supports_gpu,cleanup_error", [(None, False, False), ("test-account", True, False), (None, False, True)] +) def test_cache_docker_image_submits_one_node_sbatch_job( - slurm_system: SlurmSystem, account: str | None, supports_gpu: bool + slurm_system: SlurmSystem, + monkeypatch: pytest.MonkeyPatch, + account: str | None, + supports_gpu: bool, + cleanup_error: bool, ): slurm_system.cache_docker_images_locally = True slurm_system.account = account @@ -55,6 +61,8 @@ def test_cache_docker_image_submits_one_node_sbatch_job( slurm_system.extra_srun_args = "--reservation test-reservation" slurm_system.install_path.mkdir(parents=True, exist_ok=True) image_path = slurm_system.install_path / "image.sqsh" + if cleanup_error: + monkeypatch.setattr(Path, "unlink", Mock(side_effect=PermissionError("cleanup failed"))) def submit(script_path: Path, operation_name: str, *, wait: bool) -> int: content = script_path.read_text(encoding="utf-8") @@ -71,6 +79,8 @@ def submit(script_path: Path, operation_name: str, *, wait: bool) -> int: assert "#SBATCH --gres=gpu:1" in content assert operation_name == "Docker image import" assert wait is True + script_path.with_suffix(".out").touch() + script_path.with_suffix(".err").touch() image_path.touch() return 123 @@ -79,6 +89,9 @@ def submit(script_path: Path, operation_name: str, *, wait: bool) -> int: assert result.success assert result.docker_image_path == image_path + assert image_path.is_file() + expected_suffixes = {".sh", ".out", ".err", ".sqsh"} if cleanup_error else {".sqsh"} + assert {path.suffix for path in slurm_system.install_path.iterdir()} == expected_suffixes submit_sbatch.assert_called_once() @@ -101,6 +114,32 @@ def test_cache_docker_image_reports_submission_failure(slurm_system: SlurmSystem assert "Failed to import Docker image" in result.message +@pytest.mark.parametrize( + "stderr,expected_message", + [ + ("Disk quota exceeded", "disk-related issue"), + ("Write error", "disk-related issue"), + ("import failed", "Error: import failed"), + ("", "image was not created"), + ], +) +def test_cache_docker_image_reports_missing_image(slurm_system: SlurmSystem, stderr: str, expected_message: str): + slurm_system.supports_gpu_directives_cache = False + slurm_system.install_path.mkdir(parents=True, exist_ok=True) + + def submit(script_path: Path, operation_name: str, *, wait: bool) -> int: + script_path.with_suffix(".out").touch() + script_path.with_suffix(".err").write_text(stderr, encoding="utf-8") + return 123 + + with patch.object(SlurmSystem, "submit_sbatch", side_effect=submit): + result = DockerImageCacheManager(slurm_system).cache_docker_image("docker.io/hello-world", "image.sqsh") + + assert not result.success + assert expected_message in result.message + assert {path.suffix for path in slurm_system.install_path.iterdir()} == {".sh", ".out", ".err"} + + def test_uninstall_cached_image(slurm_system: SlurmSystem): slurm_system.install_path.mkdir(parents=True, exist_ok=True) image_path = slurm_system.install_path / "image.sqsh" From 5793499c70659a680e491229e55617b7c87e6850 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 14 Sep 2026 20:00:35 +0200 Subject: [PATCH 6/8] Revert unrelated mock fixture rename --- tests/test_toml_files.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_toml_files.py b/tests/test_toml_files.py index 762b4c7f4..819353312 100644 --- a/tests/test_toml_files.py +++ b/tests/test_toml_files.py @@ -59,7 +59,7 @@ def test_toml_files(toml_file: Path): @pytest.mark.parametrize("system_file", ALL_SYSTEMS, ids=lambda x: str(x)) @patch("kubernetes.config.load_kube_config") @patch("pathlib.Path.exists", return_value=True) -def test_systems(_mock_exists, mock_load_kube_config, system_file: Path): +def test_systems(mock_exists, mock_load_kube_config, system_file: Path): """ Validate the syntax of a system configuration file. From 2368de94a882873979175a1c45c432ed91e36bc8 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 14 Sep 2026 20:40:42 +0200 Subject: [PATCH 7/8] Fix unused mock fixture lint --- tests/test_toml_files.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_toml_files.py b/tests/test_toml_files.py index 819353312..762b4c7f4 100644 --- a/tests/test_toml_files.py +++ b/tests/test_toml_files.py @@ -59,7 +59,7 @@ def test_toml_files(toml_file: Path): @pytest.mark.parametrize("system_file", ALL_SYSTEMS, ids=lambda x: str(x)) @patch("kubernetes.config.load_kube_config") @patch("pathlib.Path.exists", return_value=True) -def test_systems(mock_exists, mock_load_kube_config, system_file: Path): +def test_systems(_mock_exists, mock_load_kube_config, system_file: Path): """ Validate the syntax of a system configuration file. From 10dcbc14165b518138eb88db09c6a2a39d8610f0 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Fri, 18 Sep 2026 17:07:20 +0200 Subject: [PATCH 8/8] Handle failures from waited Slurm jobs --- src/cloudai/core.py | 2 ++ .../slurm/docker_image_cache_manager.py | 4 ++-- src/cloudai/systems/slurm/slurm_system.py | 15 +++++++++---- tests/test_docker_image_cache_manager.py | 22 ++++++++++++++++++- tests/test_get_job_id.py | 13 ++++++++++- 5 files changed, 48 insertions(+), 8 deletions(-) diff --git a/src/cloudai/core.py b/src/cloudai/core.py index 7dffb5493..2ebd13c8c 100644 --- a/src/cloudai/core.py +++ b/src/cloudai/core.py @@ -23,6 +23,7 @@ from ._core.base_system_parser import BaseSystemParser from ._core.command_gen_strategy import CommandGenStrategy from ._core.exceptions import ( + JobFailureError, JobIdRetrievalError, MissingTestError, SystemConfigParsingError, @@ -98,6 +99,7 @@ "InstallStatusResult", "Installable", "JUnitReporter", + "JobFailureError", "JobIdRetrievalError", "JobStatusResult", "JsonGenStrategy", diff --git a/src/cloudai/systems/slurm/docker_image_cache_manager.py b/src/cloudai/systems/slurm/docker_image_cache_manager.py index 91a3b067b..6cd9f0eb7 100644 --- a/src/cloudai/systems/slurm/docker_image_cache_manager.py +++ b/src/cloudai/systems/slurm/docker_image_cache_manager.py @@ -24,7 +24,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Optional -from cloudai.core import JobIdRetrievalError +from cloudai.core import JobFailureError, JobIdRetrievalError if TYPE_CHECKING: from cloudai.systems.slurm import SlurmSystem @@ -152,7 +152,7 @@ def cache_docker_image(self, docker_image_url: str, docker_image_filename: str) script_path, stderr_path = self._write_import_script(docker_image_url, docker_image_path) try: self.system.submit_sbatch(script_path, "Docker image import", wait=True) - except JobIdRetrievalError as error: + except (JobFailureError, JobIdRetrievalError) as error: message = f"Failed to import Docker image {docker_image_url}: {error}" logging.error(message) return DockerImageCacheResult(False, message=message) diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index f27367e12..2f7a21b7a 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -28,7 +28,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator -from cloudai.core import BaseJob, File, Installable, JobIdRetrievalError, System +from cloudai.core import BaseJob, File, Installable, JobFailureError, JobIdRetrievalError, System from cloudai.models.scenario import ReportConfig, parse_reports_spec from cloudai.util import CommandShell @@ -104,7 +104,7 @@ def submit_sbatch(self, script_path: Path, operation_name: str, *, wait: bool = """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) + return self.submit_job(command, operation_name, check_return_code=wait) default_partition: str partitions: List[SlurmPartition] @@ -292,9 +292,10 @@ def _parse_submitted_job_id(stdout: str) -> int | None: 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: + def submit_job(self, submission_command: str, test_name: str, *, check_return_code: bool = False) -> int: """Submit a generated Slurm workload and return its job ID.""" - stdout, stderr = self.cmd_shell.execute(submission_command).communicate() + process = self.cmd_shell.execute(submission_command) + stdout, stderr = process.communicate() job_id = self._parse_submitted_job_id(stdout) if job_id is None: raise JobIdRetrievalError( @@ -304,6 +305,12 @@ def submit_job(self, submission_command: str, test_name: str) -> int: stderr=stderr, message="Failed to retrieve job ID.", ) + if check_return_code and process.returncode != 0: + raise JobFailureError( + test_name=test_name, + message=f"Slurm job {job_id} failed.", + details=stderr, + ) return job_id def validate_install_environment(self) -> None: diff --git a/tests/test_docker_image_cache_manager.py b/tests/test_docker_image_cache_manager.py index 6d4853612..daa3917e8 100644 --- a/tests/test_docker_image_cache_manager.py +++ b/tests/test_docker_image_cache_manager.py @@ -20,7 +20,7 @@ import pytest -from cloudai.core import JobIdRetrievalError +from cloudai.core import JobFailureError, JobIdRetrievalError from cloudai.systems.slurm.docker_image_cache_manager import DockerImageCacheManager from cloudai.systems.slurm.slurm_system import SlurmSystem @@ -114,6 +114,26 @@ def test_cache_docker_image_reports_submission_failure(slurm_system: SlurmSystem assert "Failed to import Docker image" in result.message +def test_cache_docker_image_reports_failed_job_with_partial_image(slurm_system: SlurmSystem): + slurm_system.cache_docker_images_locally = True + slurm_system.supports_gpu_directives_cache = False + slurm_system.install_path.mkdir(parents=True, exist_ok=True) + + def fail_after_creating_image(*_args, **_kwargs): + (slurm_system.install_path / "image.sqsh").touch() + raise JobFailureError( + test_name="Docker image import", + message="Slurm job 123 failed.", + details="import failed", + ) + + with patch.object(SlurmSystem, "submit_sbatch", side_effect=fail_after_creating_image): + result = DockerImageCacheManager(slurm_system).cache_docker_image("docker.io/hello-world", "image.sqsh") + + assert not result.success + assert "Slurm job 123 failed" in result.message + + @pytest.mark.parametrize( "stderr,expected_message", [ diff --git a/tests/test_get_job_id.py b/tests/test_get_job_id.py index 6482b121f..9743fbff2 100644 --- a/tests/test_get_job_id.py +++ b/tests/test_get_job_id.py @@ -19,7 +19,7 @@ import pytest -from cloudai.core import JobIdRetrievalError, TestRun, TestScenario +from cloudai.core import JobFailureError, JobIdRetrievalError, TestRun, TestScenario from cloudai.systems.lsf.lsf_runner import LSFRunner from cloudai.systems.lsf.lsf_system import LSFSystem from cloudai.systems.slurm import SlurmJob, SlurmRunner, SlurmSystem @@ -99,6 +99,17 @@ def test_slurm_get_job_id(stdout: str, stderr: str, expected_job_id: int | None) assert res == expected_job_id +def test_submit_sbatch_wait_propagates_job_failure(slurm_system: SlurmSystem, tmp_path: Path): + process = Mock(returncode=1) + process.communicate.return_value = ("Submitted batch job 123", "job failed") + slurm_system.cmd_shell.execute = Mock(return_value=process) + + with pytest.raises(JobFailureError, match="Slurm job 123 failed"): + slurm_system.submit_sbatch(tmp_path / "job.sh", "Docker image import", wait=True) + + slurm_system.cmd_shell.execute.assert_called_once_with(f"sbatch --wait {tmp_path / 'job.sh'}") + + def test_slurm_runner_on_job_completion_calls_cleanup(slurm_runner: SlurmRunner): tr = slurm_runner.test_scenario.test_runs[0] job = SlurmJob(tr, id=1)