diff --git a/src/cloudai/systems/slurm/docker_image_cache_manager.py b/src/cloudai/systems/slurm/docker_image_cache_manager.py index 3f1effa73..af2e444d8 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,149 +18,57 @@ import logging import os -import shutil -import subprocess +import shlex from datetime import datetime +from hashlib import sha256 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(): @@ -170,7 +78,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) @@ -186,126 +93,100 @@ 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" + 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: - 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) + 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(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" + 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) - 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) + 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 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. - - Returns: - DockerImageCacheResult: Result of the Docker image caching operation. - """ 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) - - 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) - + 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) - - 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}" - - return self._import_docker_image(srun_prefix, docker_image_url, docker_image_path) + message = f"No permission to write in install path {self.system.install_path}." + logging.error(message) + return DockerImageCacheResult(False, Path(), message) - def _check_prerequisites(self) -> PrerequisiteCheckResult: - """ - Check prerequisites for caching Docker image. + 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: + 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 3a49041f4..0476e2f92 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" @@ -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) 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..f27367e12 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 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. 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..b2e9c8166 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,252 +14,99 @@ # See the License for the specific language governing permissions and # limitations under the License. -import subprocess +from hashlib import sha256 from pathlib import Path from unittest.mock import patch import pytest -from cloudai.systems.slurm.docker_image_cache_manager import ( - DockerImageCacheManager, - DockerImageCacheResult, - PrerequisiteCheckResult, -) +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("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 -): - 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") - - # 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 - - # 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 - - 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 - - # 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." +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 + result = DockerImageCacheManager(slurm_system).ensure_docker_image(str(image), "cached.sqsh") -@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." - - # 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." + 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") + 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 + 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) - manager._check_prerequisites = lambda: PrerequisiteCheckResult(True, "All prerequisites are met.") - - with patch("subprocess.run") as mock_run: - 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} - + with patch.object(SlurmSystem, "submit_sbatch", side_effect=error): + result = DockerImageCacheManager(slurm_system).cache_docker_image("docker.io/hello-world", "image.sqsh") -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) - - manager = DockerImageCacheManager(slurm_system) - manager._check_prerequisites = lambda: PrerequisiteCheckResult(True, "All prerequisites are met.") + assert not result.success + assert "Failed to import Docker image" in result.message - with patch("subprocess.run") as mock_run: - mock_run.return_value = subprocess.CompletedProcess(args=["cmd"], returncode=0, 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_get_job_id.py b/tests/test_get_job_id.py index e57ee984e..0d27b6570 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,16 +46,21 @@ 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] - with pytest.raises(JobIdRetrievalError) as excinfo: + 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 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( @@ -83,8 +76,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 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.