From 8ea541b76f85e890c15e8d71301c31ab16089abd Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 8 Sep 2026 19:07:40 +0200 Subject: [PATCH 1/5] Add Slurm REST API transport --- conf/common/system/example_slurm_cluster.toml | 9 + doc/USER_GUIDE.rst | 21 + pyproject.toml | 1 + src/cloudai/systems/slurm/__init__.py | 3 +- src/cloudai/systems/slurm/slurm_system.py | 453 +++++++++++++++++- tests/systems/slurm/test_system.py | 193 +++++++- uv.lock | 2 + 7 files changed, 677 insertions(+), 5 deletions(-) diff --git a/conf/common/system/example_slurm_cluster.toml b/conf/common/system/example_slurm_cluster.toml index a66815c81..17ca534ac 100644 --- a/conf/common/system/example_slurm_cluster.toml +++ b/conf/common/system/example_slurm_cluster.toml @@ -60,3 +60,12 @@ NCCL_IB_QPS_PER_CONNECTION = "4" # Device Visibility Configuration MELLANOX_VISIBLE_DEVICES = "0,3,4,5,6,9,10,11" CUDA_VISIBLE_DEVICES = "0,1,2,3,4,5,6,7" + +# Uncomment to use Slurm REST API instead of local Slurm CLI tools. +# [slurm_api] +# url = "https://slurm-api.example.com" +# verify_certs = true +# +# [slurm_api.headers] +# X-SLURM-USER-NAME = "${SLURM_USER}" +# X-SLURM-USER-TOKEN = "${SLURM_JWT}" diff --git a/doc/USER_GUIDE.rst b/doc/USER_GUIDE.rst index ea42e516d..9edfc54b8 100644 --- a/doc/USER_GUIDE.rst +++ b/doc/USER_GUIDE.rst @@ -79,6 +79,27 @@ Field Descriptions - Specifies whether CloudAI should cache remote Docker images locally during installation. If set to ``true``, CloudAI will cache the Docker images, enabling local access without needing to download them each time a test is run. This approach saves network bandwidth but requires more disk capacity. If set to ``false``, CloudAI will allow Slurm to download the Docker images as needed when they are not cached locally by Slurm. * - **global_env_vars** - Lists all global environment variables that will be applied globally whenever tests are run. + * - **[optional] slurm_api** + - Uses Slurm REST API instead of local Slurm CLI tools. Set ``url`` and optional ``verify_certs`` and ``headers``. + Header values may reference environment variables using ``${NAME}``. + +Slurm REST API +~~~~~~~~~~~~~~ + +CloudAI uses Slurm REST API v0.0.43 when ``slurm_api`` is configured. Both the ``slurm`` and ``slurmdb`` endpoints +must be enabled by the service. + +.. code-block:: toml + + [slurm_api] + url = "https://slurm-api.example.com" + verify_certs = true + + [slurm_api.headers] + X-SLURM-USER-NAME = "${SLURM_USER}" + X-SLURM-USER-TOKEN = "${SLURM_JWT}" + +Workloads that invoke their own Slurm launcher instead of producing an sbatch script are not supported in REST mode. RunAI Scheduler --------------- diff --git a/pyproject.toml b/pyproject.toml index 4507e5048..3404704e6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "jinja2~=3.1.6", "websockets~=16.0", "rich~=14.3", + "requests~=2.33", "click~=8.3", "huggingface-hub~=1.4", "numpy>=2.4.6; python_version >= '3.14'", diff --git a/src/cloudai/systems/slurm/__init__.py b/src/cloudai/systems/slurm/__init__.py index 1a92a318a..c7391b82c 100644 --- a/src/cloudai/systems/slurm/__init__.py +++ b/src/cloudai/systems/slurm/__init__.py @@ -21,10 +21,11 @@ from .slurm_metadata import SlurmJobMetadata, SlurmStepMetadata, SlurmSystemMetadata from .slurm_node import SlurmNode, SlurmNodeState from .slurm_runner import SlurmRunner -from .slurm_system import SlurmGroup, SlurmPartition, SlurmSystem, parse_node_list +from .slurm_system import SlurmAPIConfig, SlurmGroup, SlurmPartition, SlurmSystem, parse_node_list __all__ = [ "SingleSbatchRunner", + "SlurmAPIConfig", "SlurmCommandGenStrategy", "SlurmGroup", "SlurmInstaller", diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index f27367e12..ae48a9cc9 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -17,20 +17,24 @@ from __future__ import annotations import logging +import math +import os import re import shlex import shutil import subprocess import time from copy import copy +from datetime import datetime, timezone from pathlib import Path from typing import Any, ClassVar, Dict, Iterable, List, Optional, Tuple, Union +import requests from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator from cloudai.core import BaseJob, File, Installable, JobIdRetrievalError, System from cloudai.models.scenario import ReportConfig, parse_reports_spec -from cloudai.util import CommandShell +from cloudai.util import CommandShell, parse_time_limit from .slurm_job import SlurmJob from .slurm_metadata import SlurmStepMetadata @@ -44,6 +48,24 @@ class DataRepositoryConfig(BaseModel): verify_certs: bool = True +class SlurmAPIConfig(BaseModel): + """Connection details for a Slurm REST API endpoint.""" + + model_config = ConfigDict(extra="forbid") + + url: str + headers: dict[str, str] = Field(default_factory=dict) + verify_certs: bool = True + + @field_validator("url") + @classmethod + def _validate_url(cls, value: str) -> str: + value = value.strip().rstrip("/") + if not value: + raise ValueError("slurm_api.url must be non-blank") + return value + + def parse_node_list(node_list: str) -> List[str]: """ Expand a list of node names (with ranges) into a flat list of individual node names, keeping leading zeroes. @@ -102,6 +124,9 @@ class SlurmSystem(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.""" + if self.uses_slurm_api: + return self._submit_sbatch_rest(script_path, operation_name, wait=wait) + wait_arg = " --wait" if wait else "" command = f"sbatch{wait_arg} {shlex.quote(str(script_path))}" return self.submit_job(command, operation_name) @@ -123,6 +148,7 @@ def submit_sbatch(self, script_path: Path, operation_name: str, *, wait: bool = status_retry_pause_seconds: int = Field(default=10, ge=0) supports_gpu_directives_cache: Optional[bool] = Field(default=None, exclude=True) container_mount_home: bool = False + slurm_api: Optional[SlurmAPIConfig] = None data_repository: Optional[DataRepositoryConfig] = None reports: Optional[dict[str, ReportConfig]] = None @@ -145,6 +171,37 @@ def submit_sbatch(self, script_path: Path, operation_name: str, *, wait: bool = "--container-image", "--container-mounts", ) + _REST_API_VERSION: ClassVar[str] = "v0.0.43" + _REST_TIMEOUT_SECONDS: ClassVar[int] = 30 + _TERMINAL_JOB_STATES: ClassVar[frozenset[str]] = frozenset( + { + "BOOT_FAIL", + "CANCELLED", + "COMPLETED", + "DEADLINE", + "FAILED", + "NODE_FAIL", + "OUT_OF_MEMORY", + "PREEMPTED", + "REVOKED", + "SPECIAL_EXIT", + "TIMEOUT", + } + ) + _DIRECTIVE_FIELDS: ClassVar[dict[str, str]] = { + "--job-name": "name", + "-J": "name", + "--output": "standard_output", + "-o": "standard_output", + "--error": "standard_error", + "-e": "standard_error", + "--partition": "partition", + "-p": "partition", + "--account": "account", + "-A": "account", + "--reservation": "reservation", + "--distribution": "distribution", + } @field_validator("reports", mode="before") @classmethod @@ -158,6 +215,296 @@ def _reject_blank_transient_patterns(cls, value: list[str]) -> list[str]: raise ValueError("extra_transient_status_errors entries must be non-blank") return value + @property + def uses_slurm_api(self) -> bool: + """Whether Slurm communication uses slurmrestd instead of local CLI tools.""" + return self.slurm_api is not None + + def _rest_headers(self) -> dict[str, str]: + assert self.slurm_api is not None + headers: dict[str, str] = {} + for name, value in self.slurm_api.headers.items(): + expanded = os.path.expandvars(value) + if re.search(r"\$(?:[A-Za-z_][A-Za-z0-9_]*|\{[^}]+\})", expanded): + raise EnvironmentError(f"Environment variable referenced by Slurm API header '{name}' is not set.") + headers[name] = expanded + return headers + + @staticmethod + def _rest_message(item: object) -> str: + if not isinstance(item, dict): + return str(item) + return str(item.get("error") or item.get("description") or item) + + def _rest_request( + self, + method: str, + service: str, + path: str, + *, + payload: dict[str, object] | None = None, + retry_threshold: int = 1, + ) -> dict[str, Any]: + assert self.slurm_api is not None + url = f"{self.slurm_api.url}/{service}/{self._REST_API_VERSION}/{path.lstrip('/')}" + last_error = "" + + for attempt in range(retry_threshold): + try: + response = requests.request( + method, + url, + headers=self._rest_headers(), + json=payload, + timeout=self._REST_TIMEOUT_SECONDS, + verify=self.slurm_api.verify_certs, + ) + response.raise_for_status() + data = response.json() + if not isinstance(data, dict): + raise RuntimeError(f"Slurm API returned a non-object response from {url}.") + if errors := data.get("errors"): + details = "; ".join(self._rest_message(error) for error in errors) + raise RuntimeError(f"Slurm API request failed: {details}") + for warning in data.get("warnings", []): + logging.warning("Slurm API warning: %s", self._rest_message(warning)) + return data + except (requests.RequestException, ValueError, RuntimeError) as exc: + last_error = str(exc) + if attempt + 1 < retry_threshold: + logging.warning( + "Slurm API request failed; retrying (%d/%d): %s", + attempt + 1, + retry_threshold, + exc, + ) + time.sleep(self.status_retry_pause_seconds) + + raise RuntimeError(f"Slurm API request failed after {retry_threshold} attempt(s): {last_error}") + + @staticmethod + def _directive_value(args: list[str], index: int, option: str) -> tuple[str, int]: + token = args[index] + if "=" in token: + return token.split("=", 1)[1], index + 1 + if option == "--nodes" and token.startswith("-N") and token != "-N": + return token[2:], index + 1 + if index + 1 >= len(args): + raise ValueError(f"SBATCH directive '{option}' requires a value.") + return args[index + 1], index + 2 + + @staticmethod + def _gpu_tres(value: str, *, from_gres: bool) -> str: + if from_gres: + return ",".join(item if item.startswith("gres/") else f"gres/{item}" for item in value.split(",")) + return f"gres/gpu:{value}" + + @staticmethod + def _set_directive(job: dict[str, object], field: str, value: object, option: str) -> None: + existing = job.get(field) + if existing is not None and existing != value: + raise ValueError(f"Conflicting SBATCH directives for '{option}'.") + job[field] = value + + def _apply_sbatch_args(self, job: dict[str, object], args: list[str]) -> None: # noqa: C901 + index = 0 + while index < len(args): + token = args[index] + option = token.split("=", 1)[0] + if token.startswith("-N"): + option = "--nodes" + elif option == "-n": + option = "--ntasks" + elif option == "-D": + option = "--chdir" + value, index = self._directive_value(args, index, option) + + if option in self._DIRECTIVE_FIELDS: + self._set_directive(job, self._DIRECTIVE_FIELDS[option], value, option) + elif option == "--nodes": + self._set_directive(job, "nodes", str(value), option) + elif option == "--nodelist": + self._set_directive(job, "required_nodes", str(value).split(","), option) + elif option == "--exclude": + self._set_directive(job, "excluded_nodes", str(value).split(","), option) + elif option == "--ntasks": + self._set_directive(job, "tasks", int(value), option) + elif option == "--ntasks-per-node": + self._set_directive(job, "tasks_per_node", int(value), option) + elif option == "--time": + minutes = ( + int(value) if str(value).isdigit() else math.ceil(parse_time_limit(str(value)).total_seconds() / 60) + ) + self._set_directive(job, "time_limit", {"set": True, "number": minutes}, option) + elif option in {"--gres", "--gpus-per-node"}: + tres = self._gpu_tres(str(value), from_gres=option == "--gres") + self._set_directive(job, "tres_per_node", tres, option) + elif option in {"--chdir", "-D"}: + self._set_directive(job, "current_working_directory", value, option) + else: + raise ValueError(f"SBATCH directive '{option}' is not supported by CloudAI's Slurm REST transport.") + + def _rest_job_description(self, script: str, script_path: Path) -> dict[str, object]: + job: dict[str, object] = {} + for line in script.splitlines(): + stripped = line.strip() + if not stripped or (stripped.startswith("#") and not stripped.startswith("#SBATCH")): + continue + if not stripped.startswith("#SBATCH"): + break + args = shlex.split(stripped.removeprefix("#SBATCH").strip()) + self._apply_sbatch_args(job, args) + + job.setdefault("current_working_directory", str(script_path.parent.absolute())) + job["environment"] = [f"PATH={os.environ.get('PATH', '/usr/local/bin:/usr/bin:/bin')}"] + return job + + def _submit_sbatch_rest(self, script_path: Path, operation_name: str, *, wait: bool = False) -> int: + try: + script = script_path.read_text(encoding="utf-8") + data = self._rest_request( + "POST", + "slurm", + "job/submit", + payload={"script": script, "job": self._rest_job_description(script, script_path)}, + ) + except (OSError, RuntimeError, ValueError) as exc: + raise JobIdRetrievalError( + test_name=operation_name, + command=f"POST /slurm/{self._REST_API_VERSION}/job/submit", + stdout="", + stderr=str(exc), + message="Failed to submit job through Slurm REST API.", + ) from exc + + job_id = data.get("job_id") + if not isinstance(job_id, int): + raise JobIdRetrievalError( + test_name=operation_name, + command=f"POST /slurm/{self._REST_API_VERSION}/job/submit", + stdout=str(data), + stderr="", + message="Failed to retrieve job ID.", + ) + + if wait: + while not self._is_rest_job_completed(job_id): + time.sleep(self.monitor_interval) + return job_id + + @staticmethod + def _rest_values(value: object) -> list[str]: + if isinstance(value, dict): + value = value.get("current", []) + if isinstance(value, list): + return [str(item) for item in value] + if value is None: + return [] + return [item for item in re.split(r"[,+]", str(value)) if item] + + @classmethod + def _rest_states(cls, record: dict[str, Any]) -> list[str]: + value = record.get("state", record.get("job_state")) + return [state.upper().rstrip("+") for state in cls._rest_values(value)] + + def _rest_node_state(self, node: dict[str, Any]) -> SlurmNodeState: + states = [self.convert_state_to_enum(state) for state in self._rest_states(node)] + ordinary_states = { + SlurmNodeState.ALLOCATED, + SlurmNodeState.ALLOCATED_COMPLETING, + SlurmNodeState.COMPLETING, + SlurmNodeState.IDLE, + SlurmNodeState.MIXED_ALLOCATION, + } + fallback = states[0] if states else SlurmNodeState.UNKNOWN_STATE + return next((state for state in states if state not in ordinary_states), fallback) + + @staticmethod + def _rest_number(value: object) -> int: + if isinstance(value, dict): + if value.get("set") is False: + return 0 + value = value.get("number", 0) + if not isinstance(value, (str, int, float)): + return 0 + try: + return int(value or 0) + except (TypeError, ValueError): + return 0 + + @classmethod + def _rest_time(cls, value: object) -> str: + if isinstance(value, str) and not value.isdigit(): + return value + timestamp = cls._rest_number(value) + if not timestamp: + return "" + return datetime.fromtimestamp(timestamp, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + @classmethod + def _rest_exit_code(cls, record: dict[str, Any]) -> str: + exit_code = record.get("exit_code") + if isinstance(exit_code, str): + return exit_code + if not isinstance(exit_code, dict): + return "0:0" + return_code = cls._rest_number(exit_code.get("return_code")) + signal_value = exit_code.get("signal") + if isinstance(signal_value, dict): + signal_value = signal_value.get("id", signal_value) + signal = cls._rest_number(signal_value) + return f"{return_code}:{signal}" + + def _rest_accounting_job(self, job_id: int, retry_threshold: int = 3) -> dict[str, Any] | None: + data = self._rest_request("GET", "slurmdb", f"job/{job_id}", retry_threshold=retry_threshold) + jobs = data.get("jobs", []) + if not isinstance(jobs, list): + raise RuntimeError("Slurm API returned an invalid jobs response.") + return next( + (item for item in jobs if isinstance(item, dict) and self._rest_number(item.get("job_id")) == job_id), + None, + ) + + def _rest_job_states(self, job_id: int, retry_threshold: int = 3) -> list[str]: + job = self._rest_accounting_job(job_id, retry_threshold) + if job is None: + return [] + states = self._rest_states(job) + for step in job.get("steps", []): + if isinstance(step, dict): + states.extend(self._rest_states(step)) + return states + + def _is_rest_job_completed(self, job_id: int, retry_threshold: int = 3) -> bool: + states = self._rest_job_states(job_id, retry_threshold) + if "RUNNING" in states: + return False + return any(state in self._TERMINAL_JOB_STATES for state in states) + + @classmethod + def _rest_step_metadata(cls, job: dict[str, Any]) -> list[SlurmStepMetadata]: + job_id = cls._rest_number(job.get("job_id")) + records = [job, *(step for step in job.get("steps", []) if isinstance(step, dict))] + metadata: list[SlurmStepMetadata] = [] + for index, record in enumerate(records): + step = record.get("step", {}) if isinstance(record.get("step"), dict) else {} + times = record.get("time", {}) if isinstance(record.get("time"), dict) else {} + states = cls._rest_states(record) + metadata.append( + SlurmStepMetadata( + job_id=job_id, + step_id="" if index == 0 else str(step.get("id", "")), + name=str(record.get("name", step.get("name", ""))), + state=states[0] if states else "", + exit_code=cls._rest_exit_code(record), + start_time=cls._rest_time(times.get("start")), + end_time=cls._rest_time(times.get("end")), + elapsed_time_sec=cls._rest_number(times.get("elapsed")), + submit_line=str(record.get("submit_line", "")), + ) + ) + return metadata + @property def groups(self) -> Dict[str, Dict[str, List[SlurmNode]]]: groups: Dict[str, Dict[str, List[SlurmNode]]] = {} @@ -186,6 +533,22 @@ def supports_gpu_directives(self) -> bool: if self.supports_gpu_directives_cache is not None: return self.supports_gpu_directives_cache + if self.uses_slurm_api: + try: + data = self._rest_request("GET", "slurm", "nodes/") + except RuntimeError as exc: + logging.warning("Error checking GPU support: %s", exc) + self.supports_gpu_directives_cache = True + return True + + self.supports_gpu_directives_cache = any( + "gpu" in str(node.get(field, "")).lower() + for node in data.get("nodes", []) + if isinstance(node, dict) + for field in ("gres", "tres") + ) + return self.supports_gpu_directives_cache + stdout, stderr = self.fetch_command_output("scontrol show config") if stderr: logging.warning(f"Error checking GPU support: {stderr}") @@ -218,6 +581,17 @@ def update(self) -> None: self.update_nodes_state_and_user(self.group_allocated) def nodes_from_sinfo(self) -> list[SlurmNode]: + if self.uses_slurm_api: + data = self._rest_request("GET", "slurm", "nodes/") + nodes: list[SlurmNode] = [] + for node in data.get("nodes", []): + if not isinstance(node, dict) or not node.get("name"): + continue + state = self._rest_node_state(node) + for partition in self._rest_values(node.get("partitions")): + nodes.append(SlurmNode(name=str(node["name"]), partition=partition, state=state)) + return nodes + sinfo_output, _ = self.fetch_command_output("sinfo --noheader -o '%P|%t|%u|%N'") nodes: list[SlurmNode] = [] for line in sinfo_output.split("\n"): @@ -237,6 +611,25 @@ def nodes_from_sinfo(self) -> list[SlurmNode]: return nodes def nodes_from_squeue(self) -> list[SlurmNode]: + if self.uses_slurm_api: + data = self._rest_request("GET", "slurm", "jobs/") + nodes: list[SlurmNode] = [] + for job in data.get("jobs", []): + if not isinstance(job, dict) or not {"RUNNING", "PENDING"}.intersection(self._rest_states(job)): + continue + partition = str(job.get("partition", "")) + user = str(job.get("user_name", job.get("user", "N/A"))) + for node_name in parse_node_list(str(job.get("nodes", ""))): + nodes.append( + SlurmNode( + name=node_name, + partition=partition, + state=SlurmNodeState.ALLOCATED, + user=user, + ) + ) + return nodes + squeue_output, _ = self.fetch_command_output("squeue --states=running,pending --noheader -o '%P|%T|%N|%u'") nodes: list[SlurmNode] = [] for line in squeue_output.split("\n"): @@ -294,6 +687,29 @@ def _parse_submitted_job_id(stdout: str) -> int | None: def submit_job(self, submission_command: str, test_name: str) -> int: """Submit a generated Slurm workload and return its job ID.""" + if self.uses_slurm_api: + args = shlex.split(submission_command) + if not args or Path(args[0]).name != "sbatch": + raise JobIdRetrievalError( + test_name=test_name, + command=submission_command, + stdout="", + stderr="Slurm REST mode only supports submission of an sbatch script.", + message="Failed to submit job through Slurm REST API.", + ) + + wait = "--wait" in args[1:-1] + unsupported = [arg for arg in args[1:-1] if arg != "--wait"] + if unsupported or len(args) < 2: + raise JobIdRetrievalError( + test_name=test_name, + command=submission_command, + stdout="", + stderr=f"Unsupported sbatch command arguments: {' '.join(unsupported)}", + message="Failed to submit job through Slurm REST API.", + ) + return self._submit_sbatch_rest(Path(args[-1]), test_name, wait=wait) + stdout, stderr = self.cmd_shell.execute(submission_command).communicate() job_id = self._parse_submitted_job_id(stdout) if job_id is None: @@ -308,6 +724,16 @@ def submit_job(self, submission_command: str, test_name: str) -> int: def validate_install_environment(self) -> None: """Validate that the configured Slurm environment can run CloudAI workloads.""" + if self.uses_slurm_api: + if shutil.which("git") is None: + raise EnvironmentError("Required binary 'git' is not installed.") + try: + self._rest_request("GET", "slurm", "ping/") + self._rest_request("GET", "slurmdb", "ping/") + except RuntimeError as exc: + raise EnvironmentError(f"Failed to access the Slurm REST API: {exc}") from exc + return + for binary in self._REQUIRED_BINARIES: if shutil.which(binary) is None: raise EnvironmentError(f"Required binary '{binary}' is not installed.") @@ -338,6 +764,10 @@ def is_job_running(self, job: BaseJob, retry_threshold: int = 3) -> bool: RuntimeError: If an error occurs that prevents determination of the job's running status, or if the status cannot be determined after the specified number of retries. """ + if self.uses_slurm_api: + assert isinstance(job.id, int) + return "RUNNING" in self._rest_job_states(job.id, retry_threshold) + retry_count = 0 command = f"sacct -j {job.id} --format=State --noheader" @@ -388,6 +818,10 @@ def is_job_completed(self, job: BaseJob, retry_threshold: int = 3) -> bool: Raises: RuntimeError: If unable to determine job status after retries, or if a non-retryable error is encountered. """ + if self.uses_slurm_api: + assert isinstance(job.id, int) + return self._is_rest_job_completed(job.id, retry_threshold) + retry_count = 0 command = f"sacct -j {job.id} --format=State --noheader" @@ -424,6 +858,11 @@ def is_job_completed(self, job: BaseJob, retry_threshold: int = 3) -> bool: return False def get_job_status(self, job: BaseJob, retry_threshold: int = 3) -> list[SlurmStepMetadata]: + if self.uses_slurm_api: + assert isinstance(job.id, int) + rest_job = self._rest_accounting_job(job.id, retry_threshold) + return self._rest_step_metadata(rest_job) if rest_job else [] + retry_count = 0 command = ( f"sacct -j {job.id} --format=JobID,JobName,State,ExitCode,Start,End,ElapsedRAW,SubmitLine " @@ -706,6 +1145,9 @@ def scancel(self, job_id: int) -> None: Args: job_id (int): The ID of the job to cancel. """ + if self.uses_slurm_api: + self._rest_request("DELETE", "slurm", f"job/{job_id}") + return self.cmd_shell.execute(f"scancel {job_id}") def fetch_command_output(self, command: str) -> Tuple[str, str]: @@ -870,8 +1312,13 @@ def system_installables(self) -> list[Installable]: return [File(Path(__file__).parent.absolute() / "slurm-metadata.sh")] def complete_job(self, job: SlurmJob) -> list[str]: - out, _ = self.fetch_command_output(f"sacct -j {job.id} -p --noheader -X --format=NodeList") - spec = out.splitlines()[0] if out.splitlines() else out + if self.uses_slurm_api: + assert isinstance(job.id, int) + rest_job = self._rest_accounting_job(job.id) + spec = str(rest_job.get("nodes", "")) if rest_job else "" + else: + out, _ = self.fetch_command_output(f"sacct -j {job.id} -p --noheader -X --format=NodeList") + spec = out.splitlines()[0] if out.splitlines() else out nodelist = sorted(set(parse_node_list(spec.strip().replace("|", "")))) to_unlock = [node for node in self.group_allocated if node.name in nodelist] self.group_allocated.difference_update(to_unlock) diff --git a/tests/systems/slurm/test_system.py b/tests/systems/slurm/test_system.py index da05ccf22..1e79f59db 100644 --- a/tests/systems/slurm/test_system.py +++ b/tests/systems/slurm/test_system.py @@ -14,9 +14,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import re from pathlib import Path -from unittest.mock import Mock, patch +from unittest.mock import Mock, call, patch import pytest import toml @@ -25,7 +26,9 @@ from cloudai.core import BaseJob, JobIdRetrievalError, TestRun from cloudai.models.scenario import ReportConfig from cloudai.systems.slurm import ( + SlurmAPIConfig, SlurmCommandGenStrategy, + SlurmJob, SlurmNode, SlurmNodeState, SlurmSystem, @@ -35,6 +38,194 @@ from cloudai.workloads.nccl_test import NCCLCmdArgs, NCCLTestDefinition +@pytest.fixture +def rest_slurm_system(slurm_system: SlurmSystem) -> SlurmSystem: + slurm_system.slurm_api = SlurmAPIConfig( + url="https://slurm.example.com/", + headers={ + "X-SLURM-USER-NAME": "${SLURM_USER}", + "X-SLURM-USER-TOKEN": "${SLURM_JWT}", + }, + verify_certs=False, + ) + return slurm_system + + +def test_slurm_api_request_expands_headers(rest_slurm_system: SlurmSystem, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("SLURM_USER", "cloudai") + monkeypatch.setenv("SLURM_JWT", "secret") + response = Mock() + response.json.return_value = {"pings": [], "errors": [], "warnings": []} + + with patch("cloudai.systems.slurm.slurm_system.requests.request", return_value=response) as request: + rest_slurm_system._rest_request("GET", "slurm", "ping/") + + request.assert_called_once_with( + "GET", + "https://slurm.example.com/slurm/v0.0.43/ping/", + headers={"X-SLURM-USER-NAME": "cloudai", "X-SLURM-USER-TOKEN": "secret"}, + json=None, + timeout=30, + verify=False, + ) + + +def test_submit_job_through_slurm_api(rest_slurm_system: SlurmSystem, tmp_path: Path): + script = """#!/bin/bash +# generated by CloudAI +#SBATCH --job-name=rest-test +#SBATCH --output=/shared/stdout.txt +#SBATCH --error=/shared/stderr.txt +#SBATCH --partition=gpu +#SBATCH --account=cloudai +#SBATCH --reservation=nightly +#SBATCH --distribution=block +#SBATCH -N 2 +#SBATCH --exclude=node03,node04 +#SBATCH --gpus-per-node=8 +#SBATCH --gres=gpu:8 +#SBATCH --ntasks-per-node=8 +#SBATCH --time=00:20:30 + +srun --container-image=image.sqsh benchmark +""" + script_path = tmp_path / "job.sbatch" + script_path.write_text(script) + + with patch.object(rest_slurm_system, "_rest_request", return_value={"job_id": 123}) as request: + job_id = rest_slurm_system.submit_job(f"sbatch {script_path}", "rest-test") + + assert job_id == 123 + payload = request.call_args.kwargs["payload"] + assert payload["script"] == script + assert payload["job"] == { + "name": "rest-test", + "standard_output": "/shared/stdout.txt", + "standard_error": "/shared/stderr.txt", + "partition": "gpu", + "account": "cloudai", + "reservation": "nightly", + "distribution": "block", + "nodes": "2", + "excluded_nodes": ["node03", "node04"], + "tres_per_node": "gres/gpu:8", + "tasks_per_node": 8, + "time_limit": {"set": True, "number": 21}, + "current_working_directory": str(tmp_path), + "environment": [f"PATH={os.environ['PATH']}"], + } + + +def test_slurm_api_rejects_unsupported_sbatch_directive(rest_slurm_system: SlurmSystem, tmp_path: Path): + script_path = tmp_path / "job.sbatch" + script_path.write_text("#!/bin/bash\n#SBATCH --qos=high\nsrun true\n") + + with ( + patch.object(rest_slurm_system, "_rest_request") as request, + pytest.raises(JobIdRetrievalError, match="Failed to submit job through Slurm REST API"), + ): + rest_slurm_system.submit_sbatch(script_path, "rest-test") + + request.assert_not_called() + + +def test_slurm_api_rejects_non_sbatch_launcher(rest_slurm_system: SlurmSystem): + with pytest.raises(JobIdRetrievalError, match="Failed to submit job through Slurm REST API"): + rest_slurm_system.submit_job("python launcher.py", "launcher") + + +def test_slurm_api_job_lifecycle(rest_slurm_system: SlurmSystem): + response = { + "jobs": [ + { + "job_id": 42, + "name": "rest-test", + "state": {"current": ["COMPLETED"]}, + "exit_code": {"return_code": {"number": 0}, "signal": {"id": {"number": 0}}}, + "time": {"start": {"number": 100}, "end": {"number": 120}, "elapsed": {"number": 20}}, + "nodes": "node[01-02]", + "submit_line": "sbatch job.sbatch", + "steps": [ + { + "step": {"id": "batch", "name": "batch"}, + "state": ["COMPLETED"], + "exit_code": {"return_code": {"number": 0}, "signal": {"id": {"number": 0}}}, + "time": { + "start": {"number": 100}, + "end": {"number": 120}, + "elapsed": {"number": 20}, + }, + } + ], + } + ] + } + job = SlurmJob(test_run=Mock(), id=42) + + with patch.object(rest_slurm_system, "_rest_request", return_value=response): + assert rest_slurm_system.is_job_running(job) is False + assert rest_slurm_system.is_job_completed(job) is True + assert rest_slurm_system.complete_job(job) == ["node01", "node02"] + metadata = rest_slurm_system.get_job_status(job) + + assert [(item.step_id, item.name, item.state) for item in metadata] == [ + ("", "rest-test", "COMPLETED"), + ("batch", "batch", "COMPLETED"), + ] + assert metadata[0].exit_code == "0:0" + assert metadata[0].elapsed_time_sec == 20 + + +def test_slurm_api_nodes_cancel_and_validation(rest_slurm_system: SlurmSystem): + nodes_response = { + "nodes": [ + {"name": "node01", "partitions": ["main"], "state": ["IDLE", "DRAIN"], "gres": "gpu:8"}, + {"name": "node02", "partitions": ["main", "backup"], "state": ["ALLOCATED"]}, + ] + } + jobs_response = { + "jobs": [ + { + "job_id": 42, + "partition": "main", + "job_state": ["RUNNING"], + "nodes": "node02", + "user_name": "cloudai", + } + ] + } + + def request(_method: str, service: str, path: str, **_kwargs): + if service == "slurm" and path == "nodes/": + return nodes_response + if service == "slurm" and path == "jobs/": + return jobs_response + return {} + + rest_slurm_system.supports_gpu_directives_cache = None + with ( + patch.object(rest_slurm_system, "_rest_request", side_effect=request) as rest_request, + patch("cloudai.systems.slurm.slurm_system.shutil.which", return_value="/usr/bin/git"), + ): + assert rest_slurm_system.supports_gpu_directives is True + assert [(node.name, node.partition, node.state) for node in rest_slurm_system.nodes_from_sinfo()] == [ + ("node01", "main", SlurmNodeState.DRAINED), + ("node02", "main", SlurmNodeState.ALLOCATED), + ("node02", "backup", SlurmNodeState.ALLOCATED), + ] + assert rest_slurm_system.nodes_from_squeue() == [ + SlurmNode(name="node02", partition="main", state=SlurmNodeState.ALLOCATED, user="cloudai") + ] + rest_slurm_system.scancel(42) + rest_slurm_system.validate_install_environment() + + assert rest_request.call_args_list[-3:] == [ + call("DELETE", "slurm", "job/42"), + call("GET", "slurm", "ping/"), + call("GET", "slurmdb", "ping/"), + ] + + @patch("cloudai.systems.slurm.slurm_system.CommandShell.execute") def test_submit_job_returns_parsed_job_id(mock_execute: Mock, slurm_system: SlurmSystem): process = Mock() diff --git a/uv.lock b/uv.lock index 4405ede12..f8db48cf4 100644 --- a/uv.lock +++ b/uv.lock @@ -278,6 +278,7 @@ dependencies = [ { name = "pandas" }, { name = "pydantic" }, { name = "pyyaml" }, + { name = "requests" }, { name = "rich" }, { name = "tbparse" }, { name = "toml" }, @@ -346,6 +347,7 @@ requires-dist = [ { name = "pytest-cov", marker = "extra == 'dev'", specifier = "~=7.0" }, { name = "pytest-deadfixtures", marker = "extra == 'dev'", specifier = "~=3.1" }, { name = "pyyaml", specifier = "~=6.0" }, + { name = "requests", specifier = "~=2.33" }, { name = "rich", specifier = "~=14.3" }, { name = "ruff", marker = "extra == 'dev'", specifier = "~=0.15" }, { name = "sphinx", marker = "extra == 'docs'", specifier = "~=8.1" }, From 046840136acb2b7a96de690aec19c7df71ee20c8 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Wed, 9 Sep 2026 14:31:39 +0200 Subject: [PATCH 2/5] adjust slurm rest api implementation to target 22.05.8 version --- doc/USER_GUIDE.rst | 4 +-- src/cloudai/systems/slurm/slurm_system.py | 26 +++++++-------- tests/systems/slurm/test_system.py | 39 +++++++++++------------ 3 files changed, 32 insertions(+), 37 deletions(-) diff --git a/doc/USER_GUIDE.rst b/doc/USER_GUIDE.rst index 9edfc54b8..627b6eeee 100644 --- a/doc/USER_GUIDE.rst +++ b/doc/USER_GUIDE.rst @@ -86,8 +86,8 @@ Field Descriptions Slurm REST API ~~~~~~~~~~~~~~ -CloudAI uses Slurm REST API v0.0.43 when ``slurm_api`` is configured. Both the ``slurm`` and ``slurmdb`` endpoints -must be enabled by the service. +CloudAI uses the Slurm 22.05 REST API v0.0.38 when ``slurm_api`` is configured. Both the ``slurm`` and ``slurmdb`` +endpoints must be enabled by the service. .. code-block:: toml diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index ae48a9cc9..295bdb088 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -171,7 +171,7 @@ def submit_sbatch(self, script_path: Path, operation_name: str, *, wait: bool = "--container-image", "--container-mounts", ) - _REST_API_VERSION: ClassVar[str] = "v0.0.43" + _REST_API_VERSION: ClassVar[str] = "v0.0.38" _REST_TIMEOUT_SECONDS: ClassVar[int] = 30 _TERMINAL_JOB_STATES: ClassVar[frozenset[str]] = frozenset( { @@ -294,10 +294,8 @@ def _directive_value(args: list[str], index: int, option: str) -> tuple[str, int return args[index + 1], index + 2 @staticmethod - def _gpu_tres(value: str, *, from_gres: bool) -> str: - if from_gres: - return ",".join(item if item.startswith("gres/") else f"gres/{item}" for item in value.split(",")) - return f"gres/gpu:{value}" + def _gpu_gres(value: str, *, from_gres: bool) -> str: + return value if from_gres else f"gpu:{value}" @staticmethod def _set_directive(job: dict[str, object], field: str, value: object, option: str) -> None: @@ -322,11 +320,11 @@ def _apply_sbatch_args(self, job: dict[str, object], args: list[str]) -> None: if option in self._DIRECTIVE_FIELDS: self._set_directive(job, self._DIRECTIVE_FIELDS[option], value, option) elif option == "--nodes": - self._set_directive(job, "nodes", str(value), option) + self._set_directive(job, "nodes", [int(item) for item in str(value).split("-", 1)], option) elif option == "--nodelist": - self._set_directive(job, "required_nodes", str(value).split(","), option) + self._set_directive(job, "nodelist", str(value), option) elif option == "--exclude": - self._set_directive(job, "excluded_nodes", str(value).split(","), option) + self._set_directive(job, "exclude_nodes", str(value), option) elif option == "--ntasks": self._set_directive(job, "tasks", int(value), option) elif option == "--ntasks-per-node": @@ -335,10 +333,10 @@ def _apply_sbatch_args(self, job: dict[str, object], args: list[str]) -> None: minutes = ( int(value) if str(value).isdigit() else math.ceil(parse_time_limit(str(value)).total_seconds() / 60) ) - self._set_directive(job, "time_limit", {"set": True, "number": minutes}, option) + self._set_directive(job, "time_limit", minutes, option) elif option in {"--gres", "--gpus-per-node"}: - tres = self._gpu_tres(str(value), from_gres=option == "--gres") - self._set_directive(job, "tres_per_node", tres, option) + gres = self._gpu_gres(str(value), from_gres=option == "--gres") + self._set_directive(job, "gres", gres, option) elif option in {"--chdir", "-D"}: self._set_directive(job, "current_working_directory", value, option) else: @@ -356,7 +354,7 @@ def _rest_job_description(self, script: str, script_path: Path) -> dict[str, obj self._apply_sbatch_args(job, args) job.setdefault("current_working_directory", str(script_path.parent.absolute())) - job["environment"] = [f"PATH={os.environ.get('PATH', '/usr/local/bin:/usr/bin:/bin')}"] + job["environment"] = {"PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin")} return job def _submit_sbatch_rest(self, script_path: Path, operation_name: str, *, wait: bool = False) -> int: @@ -451,7 +449,7 @@ def _rest_exit_code(cls, record: dict[str, Any]) -> str: return_code = cls._rest_number(exit_code.get("return_code")) signal_value = exit_code.get("signal") if isinstance(signal_value, dict): - signal_value = signal_value.get("id", signal_value) + signal_value = signal_value.get("id", signal_value.get("signal_id", signal_value)) signal = cls._rest_number(signal_value) return f"{return_code}:{signal}" @@ -729,7 +727,7 @@ def validate_install_environment(self) -> None: raise EnvironmentError("Required binary 'git' is not installed.") try: self._rest_request("GET", "slurm", "ping/") - self._rest_request("GET", "slurmdb", "ping/") + self._rest_request("GET", "slurmdb", "jobs/?start_time=now&skip_steps=true") except RuntimeError as exc: raise EnvironmentError(f"Failed to access the Slurm REST API: {exc}") from exc return diff --git a/tests/systems/slurm/test_system.py b/tests/systems/slurm/test_system.py index 1e79f59db..6907e0f2b 100644 --- a/tests/systems/slurm/test_system.py +++ b/tests/systems/slurm/test_system.py @@ -62,7 +62,7 @@ def test_slurm_api_request_expands_headers(rest_slurm_system: SlurmSystem, monke request.assert_called_once_with( "GET", - "https://slurm.example.com/slurm/v0.0.43/ping/", + "https://slurm.example.com/slurm/v0.0.38/ping/", headers={"X-SLURM-USER-NAME": "cloudai", "X-SLURM-USER-TOKEN": "secret"}, json=None, timeout=30, @@ -81,6 +81,7 @@ def test_submit_job_through_slurm_api(rest_slurm_system: SlurmSystem, tmp_path: #SBATCH --reservation=nightly #SBATCH --distribution=block #SBATCH -N 2 +#SBATCH --nodelist=node[01-02] #SBATCH --exclude=node03,node04 #SBATCH --gpus-per-node=8 #SBATCH --gres=gpu:8 @@ -106,13 +107,14 @@ def test_submit_job_through_slurm_api(rest_slurm_system: SlurmSystem, tmp_path: "account": "cloudai", "reservation": "nightly", "distribution": "block", - "nodes": "2", - "excluded_nodes": ["node03", "node04"], - "tres_per_node": "gres/gpu:8", + "nodes": [2], + "nodelist": "node[01-02]", + "exclude_nodes": "node03,node04", + "gres": "gpu:8", "tasks_per_node": 8, - "time_limit": {"set": True, "number": 21}, + "time_limit": 21, "current_working_directory": str(tmp_path), - "environment": [f"PATH={os.environ['PATH']}"], + "environment": {"PATH": os.environ["PATH"]}, } @@ -140,21 +142,16 @@ def test_slurm_api_job_lifecycle(rest_slurm_system: SlurmSystem): { "job_id": 42, "name": "rest-test", - "state": {"current": ["COMPLETED"]}, - "exit_code": {"return_code": {"number": 0}, "signal": {"id": {"number": 0}}}, - "time": {"start": {"number": 100}, "end": {"number": 120}, "elapsed": {"number": 20}}, + "state": {"current": "COMPLETED"}, + "exit_code": {"return_code": 0, "signal": {"signal_id": 0}}, + "time": {"start": 100, "end": 120, "elapsed": 20}, "nodes": "node[01-02]", - "submit_line": "sbatch job.sbatch", "steps": [ { "step": {"id": "batch", "name": "batch"}, - "state": ["COMPLETED"], - "exit_code": {"return_code": {"number": 0}, "signal": {"id": {"number": 0}}}, - "time": { - "start": {"number": 100}, - "end": {"number": 120}, - "elapsed": {"number": 20}, - }, + "state": "COMPLETED", + "exit_code": {"return_code": 0, "signal": {"signal_id": 0}}, + "time": {"start": 100, "end": 120, "elapsed": 20}, } ], } @@ -179,8 +176,8 @@ def test_slurm_api_job_lifecycle(rest_slurm_system: SlurmSystem): def test_slurm_api_nodes_cancel_and_validation(rest_slurm_system: SlurmSystem): nodes_response = { "nodes": [ - {"name": "node01", "partitions": ["main"], "state": ["IDLE", "DRAIN"], "gres": "gpu:8"}, - {"name": "node02", "partitions": ["main", "backup"], "state": ["ALLOCATED"]}, + {"name": "node01", "partitions": ["main"], "state": "IDLE+DRAIN", "gres": "gpu:8"}, + {"name": "node02", "partitions": ["main", "backup"], "state": "ALLOCATED"}, ] } jobs_response = { @@ -188,7 +185,7 @@ def test_slurm_api_nodes_cancel_and_validation(rest_slurm_system: SlurmSystem): { "job_id": 42, "partition": "main", - "job_state": ["RUNNING"], + "job_state": "RUNNING", "nodes": "node02", "user_name": "cloudai", } @@ -222,7 +219,7 @@ def request(_method: str, service: str, path: str, **_kwargs): assert rest_request.call_args_list[-3:] == [ call("DELETE", "slurm", "job/42"), call("GET", "slurm", "ping/"), - call("GET", "slurmdb", "ping/"), + call("GET", "slurmdb", "jobs/?start_time=now&skip_steps=true"), ] From a5500f92621154749bfa37e1be123cb68a237561 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Wed, 9 Sep 2026 16:31:37 +0200 Subject: [PATCH 3/5] Fix CI copyright years --- conf/common/system/example_slurm_cluster.toml | 2 +- src/cloudai/systems/slurm/__init__.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/conf/common/system/example_slurm_cluster.toml b/conf/common/system/example_slurm_cluster.toml index 17ca534ac..4bd2ea131 100644 --- a/conf/common/system/example_slurm_cluster.toml +++ b/conf/common/system/example_slurm_cluster.toml @@ -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"); diff --git a/src/cloudai/systems/slurm/__init__.py b/src/cloudai/systems/slurm/__init__.py index c7391b82c..d9621ff60 100644 --- a/src/cloudai/systems/slurm/__init__.py +++ b/src/cloudai/systems/slurm/__init__.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"); From 1421da16d76382918d9612010f027d90dc78e8a1 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Wed, 9 Sep 2026 18:16:31 +0200 Subject: [PATCH 4/5] Fix Slurm REST integration issues --- src/cloudai/systems/slurm/slurm_system.py | 21 +++++++++++++++++---- tests/systems/slurm/test_system.py | 4 ++-- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index 295bdb088..b32d014c4 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -259,13 +259,17 @@ def _rest_request( timeout=self._REST_TIMEOUT_SECONDS, verify=self.slurm_api.verify_certs, ) - response.raise_for_status() - data = response.json() + try: + data = response.json() + except ValueError: + response.raise_for_status() + raise if not isinstance(data, dict): raise RuntimeError(f"Slurm API returned a non-object response from {url}.") if errors := data.get("errors"): details = "; ".join(self._rest_message(error) for error in errors) raise RuntimeError(f"Slurm API request failed: {details}") + response.raise_for_status() for warning in data.get("warnings", []): logging.warning("Slurm API warning: %s", self._rest_message(warning)) return data @@ -320,7 +324,10 @@ def _apply_sbatch_args(self, job: dict[str, object], args: list[str]) -> None: if option in self._DIRECTIVE_FIELDS: self._set_directive(job, self._DIRECTIVE_FIELDS[option], value, option) elif option == "--nodes": - self._set_directive(job, "nodes", [int(item) for item in str(value).split("-", 1)], option) + node_counts = [int(item) for item in str(value).split("-", 1)] + if len(node_counts) == 1: + node_counts.append(node_counts[0]) + self._set_directive(job, "nodes", node_counts, option) elif option == "--nodelist": self._set_directive(job, "nodelist", str(value), option) elif option == "--exclude": @@ -727,7 +734,7 @@ def validate_install_environment(self) -> None: raise EnvironmentError("Required binary 'git' is not installed.") try: self._rest_request("GET", "slurm", "ping/") - self._rest_request("GET", "slurmdb", "jobs/?start_time=now&skip_steps=true") + self._rest_request("GET", "slurmdb", "clusters/") except RuntimeError as exc: raise EnvironmentError(f"Failed to access the Slurm REST API: {exc}") from exc return @@ -1143,6 +1150,9 @@ def scancel(self, job_id: int) -> None: Args: job_id (int): The ID of the job to cancel. """ + if job_id == 0: + return + if self.uses_slurm_api: self._rest_request("DELETE", "slurm", f"job/{job_id}") return @@ -1310,6 +1320,9 @@ def system_installables(self) -> list[Installable]: return [File(Path(__file__).parent.absolute() / "slurm-metadata.sh")] def complete_job(self, job: SlurmJob) -> list[str]: + if job.id == 0: + return [] + if self.uses_slurm_api: assert isinstance(job.id, int) rest_job = self._rest_accounting_job(job.id) diff --git a/tests/systems/slurm/test_system.py b/tests/systems/slurm/test_system.py index 6907e0f2b..52b060e6c 100644 --- a/tests/systems/slurm/test_system.py +++ b/tests/systems/slurm/test_system.py @@ -107,7 +107,7 @@ def test_submit_job_through_slurm_api(rest_slurm_system: SlurmSystem, tmp_path: "account": "cloudai", "reservation": "nightly", "distribution": "block", - "nodes": [2], + "nodes": [2, 2], "nodelist": "node[01-02]", "exclude_nodes": "node03,node04", "gres": "gpu:8", @@ -219,7 +219,7 @@ def request(_method: str, service: str, path: str, **_kwargs): assert rest_request.call_args_list[-3:] == [ call("DELETE", "slurm", "job/42"), call("GET", "slurm", "ping/"), - call("GET", "slurmdb", "jobs/?start_time=now&skip_steps=true"), + call("GET", "slurmdb", "clusters/"), ] From 443f6dd936d0a903ff085bb2ac9d4137b5059b9c Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Fri, 11 Sep 2026 20:55:42 +0200 Subject: [PATCH 5/5] Refactor Slurm REST client --- pyproject.toml | 1 + src/cloudai/systems/slurm/__init__.py | 3 +- .../systems/slurm/slurm_rest_client.py | 415 ++++++++++++++++++ src/cloudai/systems/slurm/slurm_system.py | 392 ++--------------- tests/systems/slurm/test_system.py | 13 +- uv.lock | 11 + 6 files changed, 470 insertions(+), 365 deletions(-) create mode 100644 src/cloudai/systems/slurm/slurm_rest_client.py diff --git a/pyproject.toml b/pyproject.toml index 3404704e6..67e7d19cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,7 @@ dependencies = [ "websockets~=16.0", "rich~=14.3", "requests~=2.33", + "tenacity~=9.1", "click~=8.3", "huggingface-hub~=1.4", "numpy>=2.4.6; python_version >= '3.14'", diff --git a/src/cloudai/systems/slurm/__init__.py b/src/cloudai/systems/slurm/__init__.py index d9621ff60..755feb33d 100644 --- a/src/cloudai/systems/slurm/__init__.py +++ b/src/cloudai/systems/slurm/__init__.py @@ -20,8 +20,9 @@ from .slurm_job import SlurmJob from .slurm_metadata import SlurmJobMetadata, SlurmStepMetadata, SlurmSystemMetadata from .slurm_node import SlurmNode, SlurmNodeState +from .slurm_rest_client import SlurmAPIConfig from .slurm_runner import SlurmRunner -from .slurm_system import SlurmAPIConfig, SlurmGroup, SlurmPartition, SlurmSystem, parse_node_list +from .slurm_system import SlurmGroup, SlurmPartition, SlurmSystem, parse_node_list __all__ = [ "SingleSbatchRunner", diff --git a/src/cloudai/systems/slurm/slurm_rest_client.py b/src/cloudai/systems/slurm/slurm_rest_client.py new file mode 100644 index 000000000..280e12c32 --- /dev/null +++ b/src/cloudai/systems/slurm/slurm_rest_client.py @@ -0,0 +1,415 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import logging +import math +import os +import re +import shlex +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, ClassVar + +import requests +from pydantic import BaseModel, ConfigDict, Field, field_validator +from tenacity import Retrying, before_sleep_log, retry_if_exception_type, stop_after_attempt, wait_fixed + +from cloudai.core import JobIdRetrievalError +from cloudai.util import parse_time_limit + +from .slurm_metadata import SlurmStepMetadata + +logger = logging.getLogger(__name__) + + +class SlurmAPIConfig(BaseModel): + """Connection details for a Slurm REST API endpoint.""" + + model_config = ConfigDict(extra="forbid") + + url: str + headers: dict[str, str] = Field(default_factory=dict) + verify_certs: bool = True + + @field_validator("url") + @classmethod + def _validate_url(cls, value: str) -> str: + value = value.strip().rstrip("/") + if not value: + raise ValueError("slurm_api.url must be non-blank") + return value + + +class SlurmRestClient: + """Translate CloudAI Slurm operations to slurmrestd v0.0.38 requests.""" + + API_VERSION: ClassVar[str] = "v0.0.38" + REQUEST_TIMEOUT_SECONDS: ClassVar[int] = 30 + TERMINAL_JOB_STATES: ClassVar[frozenset[str]] = frozenset( + { + "BOOT_FAIL", + "CANCELLED", + "COMPLETED", + "DEADLINE", + "FAILED", + "NODE_FAIL", + "OUT_OF_MEMORY", + "PREEMPTED", + "REVOKED", + "SPECIAL_EXIT", + "TIMEOUT", + } + ) + DIRECTIVE_FIELDS: ClassVar[dict[str, str]] = { + "--job-name": "name", + "-J": "name", + "--output": "standard_output", + "-o": "standard_output", + "--error": "standard_error", + "-e": "standard_error", + "--partition": "partition", + "-p": "partition", + "--account": "account", + "-A": "account", + "--reservation": "reservation", + "--distribution": "distribution", + } + + def __init__(self, config: SlurmAPIConfig, retry_pause_seconds: int) -> None: + self.config = config + self.retry_pause_seconds = retry_pause_seconds + + def _headers(self) -> dict[str, str]: + """Expand environment variables in configured headers.""" + headers: dict[str, str] = {} + for name, value in self.config.headers.items(): + expanded = os.path.expandvars(value) + if re.search(r"\$(?:[A-Za-z_][A-Za-z0-9_]*|\{[^}]+\})", expanded): + raise EnvironmentError(f"Environment variable referenced by Slurm API header '{name}' is not set.") + headers[name] = expanded + return headers + + @staticmethod + def _message(item: object) -> str: + """Extract useful text from Slurm error/warning objects; e.g. `{"error": "bad"}` becomes `"bad"`.""" + if not isinstance(item, dict): + return str(item) + return str(item.get("error") or item.get("description") or item) + + def _request_once(self, method: str, service: str, path: str, payload: dict[str, object] | None) -> dict[str, Any]: + """Send and validate one request; retry policy is applied by `request`.""" + url = f"{self.config.url}/{service}/{self.API_VERSION}/{path.lstrip('/')}" + response = requests.request( + method, + url, + headers=self._headers(), + json=payload, + timeout=self.REQUEST_TIMEOUT_SECONDS, + verify=self.config.verify_certs, + ) + try: + data = response.json() + except ValueError: + response.raise_for_status() + raise + if not isinstance(data, dict): + raise RuntimeError(f"Slurm API returned a non-object response from {url}.") + if errors := data.get("errors"): + details = "; ".join(self._message(error) for error in errors) + raise RuntimeError(f"Slurm API request failed: {details}") + response.raise_for_status() + for warning in data.get("warnings", []): + logger.warning("Slurm API warning: %s", self._message(warning)) + return data + + def request( + self, + method: str, + service: str, + path: str, + *, + payload: dict[str, object] | None = None, + retry_threshold: int = 1, + ) -> dict[str, Any]: + """Call slurmrestd, retrying failures up to `retry_threshold` attempts.""" + if retry_threshold < 1: + raise ValueError("retry_threshold must be at least 1") + + retrying = Retrying( + stop=stop_after_attempt(retry_threshold), + wait=wait_fixed(self.retry_pause_seconds), + retry=retry_if_exception_type((requests.RequestException, ValueError, RuntimeError)), + before_sleep=before_sleep_log(logger, logging.WARNING), + reraise=True, + ) + try: + return retrying(self._request_once, method, service, path, payload) + except (requests.RequestException, ValueError, RuntimeError) as exc: + raise RuntimeError(f"Slurm API request failed after {retry_threshold} attempt(s): {exc}") from exc + + @staticmethod + def _directive_value(args: list[str], index: int, option: str) -> tuple[str, int]: + """Read one SBATCH value and next index; e.g. `(["--time", "10"], 0, "--time")` returns `("10", 2)`.""" + token = args[index] + if "=" in token: + return token.split("=", 1)[1], index + 1 + if option == "--nodes" and token.startswith("-N") and token != "-N": + return token[2:], index + 1 + if index + 1 >= len(args): + raise ValueError(f"SBATCH directive '{option}' requires a value.") + return args[index + 1], index + 2 + + @staticmethod + def _gpu_gres(value: str, *, from_gres: bool) -> str: + """Normalize GPU requests; e.g. `--gpus-per-node=8` becomes REST GRES `gpu:8`.""" + return value if from_gres else f"gpu:{value}" + + @staticmethod + def _set_directive(job: dict[str, object], field: str, value: object, option: str) -> None: + """Set one REST field, rejecting conflicting aliases; e.g. `--gres` and `--gpus-per-node` must agree.""" + existing = job.get(field) + if existing is not None and existing != value: + raise ValueError(f"Conflicting SBATCH directives for '{option}'.") + job[field] = value + + def _apply_sbatch_args(self, job: dict[str, object], args: list[str]) -> None: # noqa: C901 + """Map tokenized SBATCH directives into v0.0.38 fields; e.g. `["-N", "2"]` sets `nodes=[2, 2]`.""" + index = 0 + while index < len(args): + token = args[index] + option = token.split("=", 1)[0] + if token.startswith("-N"): + option = "--nodes" + elif option == "-n": + option = "--ntasks" + elif option == "-D": + option = "--chdir" + value, index = self._directive_value(args, index, option) + + if option in self.DIRECTIVE_FIELDS: + self._set_directive(job, self.DIRECTIVE_FIELDS[option], value, option) + elif option == "--nodes": + node_counts = [int(item) for item in str(value).split("-", 1)] + if len(node_counts) == 1: + node_counts.append(node_counts[0]) + self._set_directive(job, "nodes", node_counts, option) + elif option == "--nodelist": + self._set_directive(job, "nodelist", str(value), option) + elif option == "--exclude": + self._set_directive(job, "exclude_nodes", str(value), option) + elif option == "--ntasks": + self._set_directive(job, "tasks", int(value), option) + elif option == "--ntasks-per-node": + self._set_directive(job, "tasks_per_node", int(value), option) + elif option == "--time": + minutes = ( + int(value) if str(value).isdigit() else math.ceil(parse_time_limit(str(value)).total_seconds() / 60) + ) + self._set_directive(job, "time_limit", minutes, option) + elif option in {"--gres", "--gpus-per-node"}: + gres = self._gpu_gres(str(value), from_gres=option == "--gres") + self._set_directive(job, "gres", gres, option) + elif option == "--chdir": + self._set_directive(job, "current_working_directory", value, option) + else: + raise ValueError(f"SBATCH directive '{option}' is not supported by CloudAI's Slurm REST transport.") + + def _job_description(self, script: str, script_path: Path) -> dict[str, object]: + """Build REST job properties from leading `#SBATCH` lines; script body remains unchanged.""" + job: dict[str, object] = {} + for line in script.splitlines(): + stripped = line.strip() + if not stripped or (stripped.startswith("#") and not stripped.startswith("#SBATCH")): + continue + if not stripped.startswith("#SBATCH"): + break + args = shlex.split(stripped.removeprefix("#SBATCH").strip()) + self._apply_sbatch_args(job, args) + + job.setdefault("current_working_directory", str(script_path.parent.absolute())) + job["environment"] = {"PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin")} + return job + + def submit_sbatch( + self, script_path: Path, operation_name: str, *, wait: bool = False, monitor_interval: int = 1 + ) -> int: + """Submit an SBATCH file and optionally wait for a terminal accounting state.""" + try: + script = script_path.read_text(encoding="utf-8") + data = self.request( + "POST", + "slurm", + "job/submit", + payload={"script": script, "job": self._job_description(script, script_path)}, + ) + except (OSError, RuntimeError, ValueError) as exc: + raise JobIdRetrievalError( + test_name=operation_name, + command=f"POST /slurm/{self.API_VERSION}/job/submit", + stdout="", + stderr=str(exc), + message="Failed to submit job through Slurm REST API.", + ) from exc + + job_id = data.get("job_id") + if not isinstance(job_id, int): + raise JobIdRetrievalError( + test_name=operation_name, + command=f"POST /slurm/{self.API_VERSION}/job/submit", + stdout=str(data), + stderr="", + message="Failed to retrieve job ID.", + ) + + if wait: + while not self.is_job_completed(job_id): + time.sleep(monitor_interval) + return job_id + + @staticmethod + def values(value: object) -> list[str]: + """Normalize Slurm scalar/list wrappers; e.g. `{"current": "IDLE+DRAIN"}` becomes `["IDLE", "DRAIN"]`.""" + if isinstance(value, dict): + value = value.get("current", []) + if isinstance(value, list): + return [str(item) for item in value] + if value is None: + return [] + return [item for item in re.split(r"[,+]", str(value)) if item] + + @classmethod + def states(cls, record: dict[str, Any]) -> list[str]: + """Extract normalized states; e.g. `{"job_state": "running+"}` becomes `["RUNNING"]`.""" + value = record.get("state", record.get("job_state")) + return [state.upper().rstrip("+") for state in cls.values(value)] + + @staticmethod + def _number(value: object) -> int: + """Decode Slurm number wrappers; e.g. `{"set": true, "number": 12}` becomes `12`.""" + if isinstance(value, dict): + if value.get("set") is False: + return 0 + value = value.get("number", 0) + if not isinstance(value, (str, int, float)): + return 0 + try: + return int(value or 0) + except (TypeError, ValueError): + return 0 + + @classmethod + def _time(cls, value: object) -> str: + """Normalize Slurm time values; e.g. epoch `100` becomes a UTC ISO-8601 timestamp.""" + if isinstance(value, str) and not value.isdigit(): + return value + timestamp = cls._number(value) + if not timestamp: + return "" + return datetime.fromtimestamp(timestamp, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + @classmethod + def _exit_code(cls, record: dict[str, Any]) -> str: + """Normalize composite exit status; e.g. return code `1` plus signal `9` becomes `"1:9"`.""" + exit_code = record.get("exit_code") + if isinstance(exit_code, str): + return exit_code + if not isinstance(exit_code, dict): + return "0:0" + return_code = cls._number(exit_code.get("return_code")) + signal_value = exit_code.get("signal") + if isinstance(signal_value, dict): + signal_value = signal_value.get("id", signal_value.get("signal_id", signal_value)) + signal = cls._number(signal_value) + return f"{return_code}:{signal}" + + @staticmethod + def _records(data: dict[str, Any], field: str) -> list[dict[str, Any]]: + """Read object records; e.g. `jobs` returns dictionary entries from `data["jobs"]`.""" + records = data.get(field, []) + if not isinstance(records, list): + raise RuntimeError(f"Slurm API returned an invalid {field} response.") + return [record for record in records if isinstance(record, dict)] + + def cluster_nodes(self) -> list[dict[str, Any]]: + """Return node records from slurmctld.""" + return self._records(self.request("GET", "slurm", "nodes/"), "nodes") + + def queue_jobs(self) -> list[dict[str, Any]]: + """Return current job records from slurmctld.""" + return self._records(self.request("GET", "slurm", "jobs/"), "jobs") + + def accounting_job(self, job_id: int, retry_threshold: int = 3) -> dict[str, Any] | None: + """Return one job from slurmdbd, retrying while accounting catches up.""" + data = self.request("GET", "slurmdb", f"job/{job_id}", retry_threshold=retry_threshold) + return next((job for job in self._records(data, "jobs") if self._number(job.get("job_id")) == job_id), None) + + def job_states(self, job_id: int, retry_threshold: int = 3) -> list[str]: + """Return job and step states from slurmdbd.""" + job = self.accounting_job(job_id, retry_threshold) + if job is None: + return [] + states = self.states(job) + steps = job.get("steps", []) + if isinstance(steps, list): + for step in steps: + if isinstance(step, dict): + states.extend(self.states(step)) + return states + + def is_job_completed(self, job_id: int, retry_threshold: int = 3) -> bool: + """Return whether accounting reports any terminal state and no running state.""" + states = self.job_states(job_id, retry_threshold) + if "RUNNING" in states: + return False + return any(state in self.TERMINAL_JOB_STATES for state in states) + + @classmethod + def step_metadata(cls, job: dict[str, Any]) -> list[SlurmStepMetadata]: + """Convert one accounting job and its steps to CloudAI metadata records.""" + job_id = cls._number(job.get("job_id")) + steps = job.get("steps", []) + records = [job, *(step for step in steps if isinstance(step, dict))] if isinstance(steps, list) else [job] + metadata: list[SlurmStepMetadata] = [] + for index, record in enumerate(records): + step = record.get("step", {}) if isinstance(record.get("step"), dict) else {} + times = record.get("time", {}) if isinstance(record.get("time"), dict) else {} + states = cls.states(record) + metadata.append( + SlurmStepMetadata( + job_id=job_id, + step_id="" if index == 0 else str(step.get("id", "")), + name=str(record.get("name", step.get("name", ""))), + state=states[0] if states else "", + exit_code=cls._exit_code(record), + start_time=cls._time(times.get("start")), + end_time=cls._time(times.get("end")), + elapsed_time_sec=cls._number(times.get("elapsed")), + submit_line=str(record.get("submit_line", "")), + ) + ) + return metadata + + def cancel(self, job_id: int) -> None: + """Cancel a Slurm job through slurmctld.""" + self.request("DELETE", "slurm", f"job/{job_id}") + + def validate(self) -> None: + """Verify access to slurmctld and slurmdbd endpoints used by CloudAI.""" + self.request("GET", "slurm", "ping/") + self.request("GET", "slurmdb", "clusters/") diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index b32d014c4..c6e799e6e 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -17,28 +17,25 @@ from __future__ import annotations import logging -import math -import os import re import shlex import shutil import subprocess import time from copy import copy -from datetime import datetime, timezone from pathlib import Path from typing import Any, ClassVar, Dict, Iterable, List, Optional, Tuple, Union -import requests from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator from cloudai.core import BaseJob, File, Installable, JobIdRetrievalError, System from cloudai.models.scenario import ReportConfig, parse_reports_spec -from cloudai.util import CommandShell, parse_time_limit +from cloudai.util import CommandShell from .slurm_job import SlurmJob from .slurm_metadata import SlurmStepMetadata from .slurm_node import SlurmNode, SlurmNodeState +from .slurm_rest_client import SlurmAPIConfig, SlurmRestClient class DataRepositoryConfig(BaseModel): @@ -48,24 +45,6 @@ class DataRepositoryConfig(BaseModel): verify_certs: bool = True -class SlurmAPIConfig(BaseModel): - """Connection details for a Slurm REST API endpoint.""" - - model_config = ConfigDict(extra="forbid") - - url: str - headers: dict[str, str] = Field(default_factory=dict) - verify_certs: bool = True - - @field_validator("url") - @classmethod - def _validate_url(cls, value: str) -> str: - value = value.strip().rstrip("/") - if not value: - raise ValueError("slurm_api.url must be non-blank") - return value - - def parse_node_list(node_list: str) -> List[str]: """ Expand a list of node names (with ranges) into a flat list of individual node names, keeping leading zeroes. @@ -125,7 +104,9 @@ class SlurmSystem(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.""" if self.uses_slurm_api: - return self._submit_sbatch_rest(script_path, operation_name, wait=wait) + return self._rest_client.submit_sbatch( + script_path, operation_name, wait=wait, monitor_interval=self.monitor_interval + ) wait_arg = " --wait" if wait else "" command = f"sbatch{wait_arg} {shlex.quote(str(script_path))}" @@ -171,37 +152,6 @@ def submit_sbatch(self, script_path: Path, operation_name: str, *, wait: bool = "--container-image", "--container-mounts", ) - _REST_API_VERSION: ClassVar[str] = "v0.0.38" - _REST_TIMEOUT_SECONDS: ClassVar[int] = 30 - _TERMINAL_JOB_STATES: ClassVar[frozenset[str]] = frozenset( - { - "BOOT_FAIL", - "CANCELLED", - "COMPLETED", - "DEADLINE", - "FAILED", - "NODE_FAIL", - "OUT_OF_MEMORY", - "PREEMPTED", - "REVOKED", - "SPECIAL_EXIT", - "TIMEOUT", - } - ) - _DIRECTIVE_FIELDS: ClassVar[dict[str, str]] = { - "--job-name": "name", - "-J": "name", - "--output": "standard_output", - "-o": "standard_output", - "--error": "standard_error", - "-e": "standard_error", - "--partition": "partition", - "-p": "partition", - "--account": "account", - "-A": "account", - "--reservation": "reservation", - "--distribution": "distribution", - } @field_validator("reports", mode="before") @classmethod @@ -220,200 +170,23 @@ def uses_slurm_api(self) -> bool: """Whether Slurm communication uses slurmrestd instead of local CLI tools.""" return self.slurm_api is not None - def _rest_headers(self) -> dict[str, str]: - assert self.slurm_api is not None - headers: dict[str, str] = {} - for name, value in self.slurm_api.headers.items(): - expanded = os.path.expandvars(value) - if re.search(r"\$(?:[A-Za-z_][A-Za-z0-9_]*|\{[^}]+\})", expanded): - raise EnvironmentError(f"Environment variable referenced by Slurm API header '{name}' is not set.") - headers[name] = expanded - return headers - - @staticmethod - def _rest_message(item: object) -> str: - if not isinstance(item, dict): - return str(item) - return str(item.get("error") or item.get("description") or item) - - def _rest_request( - self, - method: str, - service: str, - path: str, - *, - payload: dict[str, object] | None = None, - retry_threshold: int = 1, - ) -> dict[str, Any]: - assert self.slurm_api is not None - url = f"{self.slurm_api.url}/{service}/{self._REST_API_VERSION}/{path.lstrip('/')}" - last_error = "" - - for attempt in range(retry_threshold): - try: - response = requests.request( - method, - url, - headers=self._rest_headers(), - json=payload, - timeout=self._REST_TIMEOUT_SECONDS, - verify=self.slurm_api.verify_certs, - ) - try: - data = response.json() - except ValueError: - response.raise_for_status() - raise - if not isinstance(data, dict): - raise RuntimeError(f"Slurm API returned a non-object response from {url}.") - if errors := data.get("errors"): - details = "; ".join(self._rest_message(error) for error in errors) - raise RuntimeError(f"Slurm API request failed: {details}") - response.raise_for_status() - for warning in data.get("warnings", []): - logging.warning("Slurm API warning: %s", self._rest_message(warning)) - return data - except (requests.RequestException, ValueError, RuntimeError) as exc: - last_error = str(exc) - if attempt + 1 < retry_threshold: - logging.warning( - "Slurm API request failed; retrying (%d/%d): %s", - attempt + 1, - retry_threshold, - exc, - ) - time.sleep(self.status_retry_pause_seconds) - - raise RuntimeError(f"Slurm API request failed after {retry_threshold} attempt(s): {last_error}") - - @staticmethod - def _directive_value(args: list[str], index: int, option: str) -> tuple[str, int]: - token = args[index] - if "=" in token: - return token.split("=", 1)[1], index + 1 - if option == "--nodes" and token.startswith("-N") and token != "-N": - return token[2:], index + 1 - if index + 1 >= len(args): - raise ValueError(f"SBATCH directive '{option}' requires a value.") - return args[index + 1], index + 2 - - @staticmethod - def _gpu_gres(value: str, *, from_gres: bool) -> str: - return value if from_gres else f"gpu:{value}" - - @staticmethod - def _set_directive(job: dict[str, object], field: str, value: object, option: str) -> None: - existing = job.get(field) - if existing is not None and existing != value: - raise ValueError(f"Conflicting SBATCH directives for '{option}'.") - job[field] = value - - def _apply_sbatch_args(self, job: dict[str, object], args: list[str]) -> None: # noqa: C901 - index = 0 - while index < len(args): - token = args[index] - option = token.split("=", 1)[0] - if token.startswith("-N"): - option = "--nodes" - elif option == "-n": - option = "--ntasks" - elif option == "-D": - option = "--chdir" - value, index = self._directive_value(args, index, option) - - if option in self._DIRECTIVE_FIELDS: - self._set_directive(job, self._DIRECTIVE_FIELDS[option], value, option) - elif option == "--nodes": - node_counts = [int(item) for item in str(value).split("-", 1)] - if len(node_counts) == 1: - node_counts.append(node_counts[0]) - self._set_directive(job, "nodes", node_counts, option) - elif option == "--nodelist": - self._set_directive(job, "nodelist", str(value), option) - elif option == "--exclude": - self._set_directive(job, "exclude_nodes", str(value), option) - elif option == "--ntasks": - self._set_directive(job, "tasks", int(value), option) - elif option == "--ntasks-per-node": - self._set_directive(job, "tasks_per_node", int(value), option) - elif option == "--time": - minutes = ( - int(value) if str(value).isdigit() else math.ceil(parse_time_limit(str(value)).total_seconds() / 60) - ) - self._set_directive(job, "time_limit", minutes, option) - elif option in {"--gres", "--gpus-per-node"}: - gres = self._gpu_gres(str(value), from_gres=option == "--gres") - self._set_directive(job, "gres", gres, option) - elif option in {"--chdir", "-D"}: - self._set_directive(job, "current_working_directory", value, option) - else: - raise ValueError(f"SBATCH directive '{option}' is not supported by CloudAI's Slurm REST transport.") - - def _rest_job_description(self, script: str, script_path: Path) -> dict[str, object]: - job: dict[str, object] = {} - for line in script.splitlines(): - stripped = line.strip() - if not stripped or (stripped.startswith("#") and not stripped.startswith("#SBATCH")): - continue - if not stripped.startswith("#SBATCH"): - break - args = shlex.split(stripped.removeprefix("#SBATCH").strip()) - self._apply_sbatch_args(job, args) - - job.setdefault("current_working_directory", str(script_path.parent.absolute())) - job["environment"] = {"PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin")} - return job - - def _submit_sbatch_rest(self, script_path: Path, operation_name: str, *, wait: bool = False) -> int: - try: - script = script_path.read_text(encoding="utf-8") - data = self._rest_request( - "POST", - "slurm", - "job/submit", - payload={"script": script, "job": self._rest_job_description(script, script_path)}, - ) - except (OSError, RuntimeError, ValueError) as exc: - raise JobIdRetrievalError( - test_name=operation_name, - command=f"POST /slurm/{self._REST_API_VERSION}/job/submit", - stdout="", - stderr=str(exc), - message="Failed to submit job through Slurm REST API.", - ) from exc - - job_id = data.get("job_id") - if not isinstance(job_id, int): - raise JobIdRetrievalError( - test_name=operation_name, - command=f"POST /slurm/{self._REST_API_VERSION}/job/submit", - stdout=str(data), - stderr="", - message="Failed to retrieve job ID.", - ) - - if wait: - while not self._is_rest_job_completed(job_id): - time.sleep(self.monitor_interval) - return job_id + @property + def _rest_client(self) -> SlurmRestClient: + """Build a client from REST config, or fail if REST mode is disabled.""" + if self.slurm_api is None: + raise RuntimeError("Slurm REST API is not configured.") + return SlurmRestClient(self.slurm_api, self.status_retry_pause_seconds) @staticmethod - def _rest_values(value: object) -> list[str]: - if isinstance(value, dict): - value = value.get("current", []) - if isinstance(value, list): - return [str(item) for item in value] - if value is None: - return [] - return [item for item in re.split(r"[,+]", str(value)) if item] - - @classmethod - def _rest_states(cls, record: dict[str, Any]) -> list[str]: - value = record.get("state", record.get("job_state")) - return [state.upper().rstrip("+") for state in cls._rest_values(value)] + def _job_id(job: BaseJob) -> int: + """Return numeric Slurm ID; e.g. `SlurmJob(..., id=42)` returns `42`.""" + if not isinstance(job.id, int): + raise TypeError(f"Slurm job ID must be an integer, got {type(job.id).__name__}.") + return job.id def _rest_node_state(self, node: dict[str, Any]) -> SlurmNodeState: - states = [self.convert_state_to_enum(state) for state in self._rest_states(node)] + """Choose significant state from REST state flags; e.g. `IDLE+DRAIN` resolves to `DRAINED`.""" + states = [self.convert_state_to_enum(state) for state in self._rest_client.states(node)] ordinary_states = { SlurmNodeState.ALLOCATED, SlurmNodeState.ALLOCATED_COMPLETING, @@ -424,92 +197,6 @@ def _rest_node_state(self, node: dict[str, Any]) -> SlurmNodeState: fallback = states[0] if states else SlurmNodeState.UNKNOWN_STATE return next((state for state in states if state not in ordinary_states), fallback) - @staticmethod - def _rest_number(value: object) -> int: - if isinstance(value, dict): - if value.get("set") is False: - return 0 - value = value.get("number", 0) - if not isinstance(value, (str, int, float)): - return 0 - try: - return int(value or 0) - except (TypeError, ValueError): - return 0 - - @classmethod - def _rest_time(cls, value: object) -> str: - if isinstance(value, str) and not value.isdigit(): - return value - timestamp = cls._rest_number(value) - if not timestamp: - return "" - return datetime.fromtimestamp(timestamp, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") - - @classmethod - def _rest_exit_code(cls, record: dict[str, Any]) -> str: - exit_code = record.get("exit_code") - if isinstance(exit_code, str): - return exit_code - if not isinstance(exit_code, dict): - return "0:0" - return_code = cls._rest_number(exit_code.get("return_code")) - signal_value = exit_code.get("signal") - if isinstance(signal_value, dict): - signal_value = signal_value.get("id", signal_value.get("signal_id", signal_value)) - signal = cls._rest_number(signal_value) - return f"{return_code}:{signal}" - - def _rest_accounting_job(self, job_id: int, retry_threshold: int = 3) -> dict[str, Any] | None: - data = self._rest_request("GET", "slurmdb", f"job/{job_id}", retry_threshold=retry_threshold) - jobs = data.get("jobs", []) - if not isinstance(jobs, list): - raise RuntimeError("Slurm API returned an invalid jobs response.") - return next( - (item for item in jobs if isinstance(item, dict) and self._rest_number(item.get("job_id")) == job_id), - None, - ) - - def _rest_job_states(self, job_id: int, retry_threshold: int = 3) -> list[str]: - job = self._rest_accounting_job(job_id, retry_threshold) - if job is None: - return [] - states = self._rest_states(job) - for step in job.get("steps", []): - if isinstance(step, dict): - states.extend(self._rest_states(step)) - return states - - def _is_rest_job_completed(self, job_id: int, retry_threshold: int = 3) -> bool: - states = self._rest_job_states(job_id, retry_threshold) - if "RUNNING" in states: - return False - return any(state in self._TERMINAL_JOB_STATES for state in states) - - @classmethod - def _rest_step_metadata(cls, job: dict[str, Any]) -> list[SlurmStepMetadata]: - job_id = cls._rest_number(job.get("job_id")) - records = [job, *(step for step in job.get("steps", []) if isinstance(step, dict))] - metadata: list[SlurmStepMetadata] = [] - for index, record in enumerate(records): - step = record.get("step", {}) if isinstance(record.get("step"), dict) else {} - times = record.get("time", {}) if isinstance(record.get("time"), dict) else {} - states = cls._rest_states(record) - metadata.append( - SlurmStepMetadata( - job_id=job_id, - step_id="" if index == 0 else str(step.get("id", "")), - name=str(record.get("name", step.get("name", ""))), - state=states[0] if states else "", - exit_code=cls._rest_exit_code(record), - start_time=cls._rest_time(times.get("start")), - end_time=cls._rest_time(times.get("end")), - elapsed_time_sec=cls._rest_number(times.get("elapsed")), - submit_line=str(record.get("submit_line", "")), - ) - ) - return metadata - @property def groups(self) -> Dict[str, Dict[str, List[SlurmNode]]]: groups: Dict[str, Dict[str, List[SlurmNode]]] = {} @@ -540,17 +227,14 @@ def supports_gpu_directives(self) -> bool: if self.uses_slurm_api: try: - data = self._rest_request("GET", "slurm", "nodes/") + nodes = self._rest_client.cluster_nodes() except RuntimeError as exc: logging.warning("Error checking GPU support: %s", exc) self.supports_gpu_directives_cache = True return True self.supports_gpu_directives_cache = any( - "gpu" in str(node.get(field, "")).lower() - for node in data.get("nodes", []) - if isinstance(node, dict) - for field in ("gres", "tres") + "gpu" in str(node.get(field, "")).lower() for node in nodes for field in ("gres", "tres") ) return self.supports_gpu_directives_cache @@ -587,13 +271,12 @@ def update(self) -> None: def nodes_from_sinfo(self) -> list[SlurmNode]: if self.uses_slurm_api: - data = self._rest_request("GET", "slurm", "nodes/") nodes: list[SlurmNode] = [] - for node in data.get("nodes", []): - if not isinstance(node, dict) or not node.get("name"): + for node in self._rest_client.cluster_nodes(): + if not node.get("name"): continue state = self._rest_node_state(node) - for partition in self._rest_values(node.get("partitions")): + for partition in self._rest_client.values(node.get("partitions")): nodes.append(SlurmNode(name=str(node["name"]), partition=partition, state=state)) return nodes @@ -617,10 +300,9 @@ def nodes_from_sinfo(self) -> list[SlurmNode]: def nodes_from_squeue(self) -> list[SlurmNode]: if self.uses_slurm_api: - data = self._rest_request("GET", "slurm", "jobs/") nodes: list[SlurmNode] = [] - for job in data.get("jobs", []): - if not isinstance(job, dict) or not {"RUNNING", "PENDING"}.intersection(self._rest_states(job)): + for job in self._rest_client.queue_jobs(): + if not {"RUNNING", "PENDING"}.intersection(self._rest_client.states(job)): continue partition = str(job.get("partition", "")) user = str(job.get("user_name", job.get("user", "N/A"))) @@ -713,7 +395,7 @@ def submit_job(self, submission_command: str, test_name: str) -> int: stderr=f"Unsupported sbatch command arguments: {' '.join(unsupported)}", message="Failed to submit job through Slurm REST API.", ) - return self._submit_sbatch_rest(Path(args[-1]), test_name, wait=wait) + return self.submit_sbatch(Path(args[-1]), test_name, wait=wait) stdout, stderr = self.cmd_shell.execute(submission_command).communicate() job_id = self._parse_submitted_job_id(stdout) @@ -733,8 +415,7 @@ def validate_install_environment(self) -> None: if shutil.which("git") is None: raise EnvironmentError("Required binary 'git' is not installed.") try: - self._rest_request("GET", "slurm", "ping/") - self._rest_request("GET", "slurmdb", "clusters/") + self._rest_client.validate() except RuntimeError as exc: raise EnvironmentError(f"Failed to access the Slurm REST API: {exc}") from exc return @@ -770,8 +451,7 @@ def is_job_running(self, job: BaseJob, retry_threshold: int = 3) -> bool: cannot be determined after the specified number of retries. """ if self.uses_slurm_api: - assert isinstance(job.id, int) - return "RUNNING" in self._rest_job_states(job.id, retry_threshold) + return "RUNNING" in self._rest_client.job_states(self._job_id(job), retry_threshold) retry_count = 0 command = f"sacct -j {job.id} --format=State --noheader" @@ -824,8 +504,7 @@ def is_job_completed(self, job: BaseJob, retry_threshold: int = 3) -> bool: RuntimeError: If unable to determine job status after retries, or if a non-retryable error is encountered. """ if self.uses_slurm_api: - assert isinstance(job.id, int) - return self._is_rest_job_completed(job.id, retry_threshold) + return self._rest_client.is_job_completed(self._job_id(job), retry_threshold) retry_count = 0 command = f"sacct -j {job.id} --format=State --noheader" @@ -864,9 +543,8 @@ def is_job_completed(self, job: BaseJob, retry_threshold: int = 3) -> bool: def get_job_status(self, job: BaseJob, retry_threshold: int = 3) -> list[SlurmStepMetadata]: if self.uses_slurm_api: - assert isinstance(job.id, int) - rest_job = self._rest_accounting_job(job.id, retry_threshold) - return self._rest_step_metadata(rest_job) if rest_job else [] + rest_job = self._rest_client.accounting_job(self._job_id(job), retry_threshold) + return self._rest_client.step_metadata(rest_job) if rest_job else [] retry_count = 0 command = ( @@ -901,8 +579,7 @@ def kill(self, job: BaseJob) -> None: Args: job (BaseJob): The job to be terminated. """ - assert isinstance(job.id, int) - self.scancel(job.id) + self.scancel(self._job_id(job)) @classmethod def format_node_list(cls, node_names: List[str]) -> str: @@ -1154,7 +831,7 @@ def scancel(self, job_id: int) -> None: return if self.uses_slurm_api: - self._rest_request("DELETE", "slurm", f"job/{job_id}") + self._rest_client.cancel(job_id) return self.cmd_shell.execute(f"scancel {job_id}") @@ -1324,8 +1001,7 @@ def complete_job(self, job: SlurmJob) -> list[str]: return [] if self.uses_slurm_api: - assert isinstance(job.id, int) - rest_job = self._rest_accounting_job(job.id) + rest_job = self._rest_client.accounting_job(self._job_id(job)) spec = str(rest_job.get("nodes", "")) if rest_job else "" else: out, _ = self.fetch_command_output(f"sacct -j {job.id} -p --noheader -X --format=NodeList") diff --git a/tests/systems/slurm/test_system.py b/tests/systems/slurm/test_system.py index 52b060e6c..f5549e9c7 100644 --- a/tests/systems/slurm/test_system.py +++ b/tests/systems/slurm/test_system.py @@ -35,6 +35,7 @@ parse_node_list, ) from cloudai.systems.slurm.slurm_metadata import SlurmStepMetadata +from cloudai.systems.slurm.slurm_rest_client import SlurmRestClient from cloudai.workloads.nccl_test import NCCLCmdArgs, NCCLTestDefinition @@ -57,8 +58,8 @@ def test_slurm_api_request_expands_headers(rest_slurm_system: SlurmSystem, monke response = Mock() response.json.return_value = {"pings": [], "errors": [], "warnings": []} - with patch("cloudai.systems.slurm.slurm_system.requests.request", return_value=response) as request: - rest_slurm_system._rest_request("GET", "slurm", "ping/") + with patch("cloudai.systems.slurm.slurm_rest_client.requests.request", return_value=response) as request: + rest_slurm_system._rest_client.request("GET", "slurm", "ping/") request.assert_called_once_with( "GET", @@ -93,7 +94,7 @@ def test_submit_job_through_slurm_api(rest_slurm_system: SlurmSystem, tmp_path: script_path = tmp_path / "job.sbatch" script_path.write_text(script) - with patch.object(rest_slurm_system, "_rest_request", return_value={"job_id": 123}) as request: + with patch.object(SlurmRestClient, "request", return_value={"job_id": 123}) as request: job_id = rest_slurm_system.submit_job(f"sbatch {script_path}", "rest-test") assert job_id == 123 @@ -123,7 +124,7 @@ def test_slurm_api_rejects_unsupported_sbatch_directive(rest_slurm_system: Slurm script_path.write_text("#!/bin/bash\n#SBATCH --qos=high\nsrun true\n") with ( - patch.object(rest_slurm_system, "_rest_request") as request, + patch.object(SlurmRestClient, "request") as request, pytest.raises(JobIdRetrievalError, match="Failed to submit job through Slurm REST API"), ): rest_slurm_system.submit_sbatch(script_path, "rest-test") @@ -159,7 +160,7 @@ def test_slurm_api_job_lifecycle(rest_slurm_system: SlurmSystem): } job = SlurmJob(test_run=Mock(), id=42) - with patch.object(rest_slurm_system, "_rest_request", return_value=response): + with patch.object(SlurmRestClient, "request", return_value=response): assert rest_slurm_system.is_job_running(job) is False assert rest_slurm_system.is_job_completed(job) is True assert rest_slurm_system.complete_job(job) == ["node01", "node02"] @@ -201,7 +202,7 @@ def request(_method: str, service: str, path: str, **_kwargs): rest_slurm_system.supports_gpu_directives_cache = None with ( - patch.object(rest_slurm_system, "_rest_request", side_effect=request) as rest_request, + patch.object(SlurmRestClient, "request", side_effect=request) as rest_request, patch("cloudai.systems.slurm.slurm_system.shutil.which", return_value="/usr/bin/git"), ): assert rest_slurm_system.supports_gpu_directives is True diff --git a/uv.lock b/uv.lock index f8db48cf4..0e57f6efa 100644 --- a/uv.lock +++ b/uv.lock @@ -281,6 +281,7 @@ dependencies = [ { name = "requests" }, { name = "rich" }, { name = "tbparse" }, + { name = "tenacity" }, { name = "toml" }, { name = "websockets" }, ] @@ -361,6 +362,7 @@ requires-dist = [ { name = "sphinxext-opengraph", marker = "extra == 'docs'", specifier = "~=0.13" }, { name = "taplo", marker = "extra == 'dev'", specifier = "~=0.9.3" }, { name = "tbparse", specifier = "~=0.0.9" }, + { name = "tenacity", specifier = "~=9.1" }, { name = "toml", specifier = "~=0.10.2" }, { name = "vulture", marker = "extra == 'dev'", specifier = "==2.14" }, { name = "websockets", specifier = "~=16.0" }, @@ -2447,6 +2449,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/17/7a/56818acfbf03ef6eae4c3e0a7091ab6ffd8c3b44af93f2982297c07b50f2/tbparse-0.0.9-py3-none-any.whl", hash = "sha256:51a001728bc539a1efed9f03450b1e0151ad14011c8c52f156cf55dc1fbaa884", size = 19595, upload-time = "2024-08-16T04:37:48.546Z" }, ] +[[package]] +name = "tenacity" +version = "9.1.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/47/c6/ee486fd809e357697ee8a44d3d69222b344920433d3b6666ccd9b374630c/tenacity-9.1.4.tar.gz", hash = "sha256:adb31d4c263f2bd041081ab33b498309a57c77f9acf2db65aadf0898179cf93a", size = 49413, upload-time = "2026-02-07T10:45:33.841Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl", hash = "sha256:6095a360c919085f28c6527de529e76a06ad89b23659fa881ae0649b867a9d55", size = 28926, upload-time = "2026-02-07T10:45:32.24Z" }, +] + [[package]] name = "tensorboard" version = "2.20.0"