From 8ea541b76f85e890c15e8d71301c31ab16089abd Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 8 Sep 2026 19:07:40 +0200 Subject: [PATCH 01/16] 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 02/16] 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 03/16] 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 04/16] 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 05/16] 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" From c0a95ec1f38d85047aa2ad70020f12f531561952 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Sat, 12 Sep 2026 19:24:08 +0200 Subject: [PATCH 06/16] Clean up Slurm REST client API --- .../systems/slurm/slurm_rest_client.py | 98 ++++++++++--------- tests/systems/slurm/test_system.py | 14 +-- 2 files changed, 58 insertions(+), 54 deletions(-) diff --git a/src/cloudai/systems/slurm/slurm_rest_client.py b/src/cloudai/systems/slurm/slurm_rest_client.py index 280e12c32..79d2022b4 100644 --- a/src/cloudai/systems/slurm/slurm_rest_client.py +++ b/src/cloudai/systems/slurm/slurm_rest_client.py @@ -16,38 +16,38 @@ from __future__ import annotations +import datetime import logging import math import os +import pathlib import re import shlex import time -from datetime import datetime, timezone -from pathlib import Path from typing import Any, ClassVar +import pydantic 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 +import tenacity -from cloudai.core import JobIdRetrievalError -from cloudai.util import parse_time_limit +import cloudai.core +import cloudai.util from .slurm_metadata import SlurmStepMetadata logger = logging.getLogger(__name__) -class SlurmAPIConfig(BaseModel): +class SlurmAPIConfig(pydantic.BaseModel): """Connection details for a Slurm REST API endpoint.""" - model_config = ConfigDict(extra="forbid") + model_config = pydantic.ConfigDict(extra="forbid") url: str - headers: dict[str, str] = Field(default_factory=dict) + headers: dict[str, str] = pydantic.Field(default_factory=dict) verify_certs: bool = True - @field_validator("url") + @pydantic.field_validator("url") @classmethod def _validate_url(cls, value: str) -> str: value = value.strip().rstrip("/") @@ -59,9 +59,9 @@ def _validate_url(cls, value: str) -> str: 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( + _API_VERSION: ClassVar[str] = "v0.0.38" + _REQUEST_TIMEOUT_SECONDS: ClassVar[int] = 30 + _TERMINAL_JOB_STATES: ClassVar[frozenset[str]] = frozenset( { "BOOT_FAIL", "CANCELLED", @@ -76,7 +76,7 @@ class SlurmRestClient: "TIMEOUT", } ) - DIRECTIVE_FIELDS: ClassVar[dict[str, str]] = { + _DIRECTIVE_FIELDS: ClassVar[dict[str, str]] = { "--job-name": "name", "-J": "name", "--output": "standard_output", @@ -92,13 +92,13 @@ class SlurmRestClient: } def __init__(self, config: SlurmAPIConfig, retry_pause_seconds: int) -> None: - self.config = config - self.retry_pause_seconds = retry_pause_seconds + 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(): + 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.") @@ -113,15 +113,14 @@ def _message(item: object) -> str: 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('/')}" + 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, + timeout=self._REQUEST_TIMEOUT_SECONDS, + verify=self._config.verify_certs, ) try: data = response.json() @@ -138,7 +137,7 @@ def _request_once(self, method: str, service: str, path: str, payload: dict[str, logger.warning("Slurm API warning: %s", self._message(warning)) return data - def request( + def _request( self, method: str, service: str, @@ -151,11 +150,11 @@ def request( 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), + retrying = tenacity.Retrying( + stop=tenacity.stop_after_attempt(retry_threshold), + wait=tenacity.wait_fixed(self._retry_pause_seconds), + retry=tenacity.retry_if_exception_type((requests.RequestException, ValueError, RuntimeError)), + before_sleep=tenacity.before_sleep_log(logger, logging.WARNING), reraise=True, ) try: @@ -182,10 +181,13 @@ def _gpu_gres(value: str, *, from_gres: bool) -> str: @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.""" + """Merge directives mapped to one REST field; e.g. `--gres=gpu:8` and `--gpus-per-node=8` must agree.""" existing = job.get(field) if existing is not None and existing != value: - raise ValueError(f"Conflicting SBATCH directives for '{option}'.") + raise ValueError( + f"Conflicting SBATCH directives for '{option}': REST field '{field}' " + f"is already {existing!r}, got {value!r}." + ) job[field] = value def _apply_sbatch_args(self, job: dict[str, object], args: list[str]) -> None: # noqa: C901 @@ -202,8 +204,8 @@ def _apply_sbatch_args(self, job: dict[str, object], args: list[str]) -> None: 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) + 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: @@ -219,7 +221,9 @@ def _apply_sbatch_args(self, job: dict[str, object], args: list[str]) -> None: 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) + int(value) + if str(value).isdigit() + else math.ceil(cloudai.util.parse_time_limit(str(value)).total_seconds() / 60) ) self._set_directive(job, "time_limit", minutes, option) elif option in {"--gres", "--gpus-per-node"}: @@ -230,7 +234,7 @@ def _apply_sbatch_args(self, job: dict[str, object], args: list[str]) -> None: 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]: + def _job_description(self, script: str, script_path: pathlib.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(): @@ -247,21 +251,21 @@ def _job_description(self, script: str, script_path: Path) -> dict[str, object]: return job def submit_sbatch( - self, script_path: Path, operation_name: str, *, wait: bool = False, monitor_interval: int = 1 + self, script_path: pathlib.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( + 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( + raise cloudai.core.JobIdRetrievalError( test_name=operation_name, - command=f"POST /slurm/{self.API_VERSION}/job/submit", + command=f"POST /slurm/{self._API_VERSION}/job/submit", stdout="", stderr=str(exc), message="Failed to submit job through Slurm REST API.", @@ -269,9 +273,9 @@ def submit_sbatch( job_id = data.get("job_id") if not isinstance(job_id, int): - raise JobIdRetrievalError( + raise cloudai.core.JobIdRetrievalError( test_name=operation_name, - command=f"POST /slurm/{self.API_VERSION}/job/submit", + command=f"POST /slurm/{self._API_VERSION}/job/submit", stdout=str(data), stderr="", message="Failed to retrieve job ID.", @@ -321,7 +325,7 @@ def _time(cls, value: object) -> str: timestamp = cls._number(value) if not timestamp: return "" - return datetime.fromtimestamp(timestamp, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + return datetime.datetime.fromtimestamp(timestamp, tz=datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") @classmethod def _exit_code(cls, record: dict[str, Any]) -> str: @@ -348,15 +352,15 @@ def _records(data: dict[str, Any], field: str) -> list[dict[str, Any]]: def cluster_nodes(self) -> list[dict[str, Any]]: """Return node records from slurmctld.""" - return self._records(self.request("GET", "slurm", "nodes/"), "nodes") + 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") + 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) + 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]: @@ -377,7 +381,7 @@ def is_job_completed(self, job_id: int, retry_threshold: int = 3) -> bool: 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) + return any(state in self._TERMINAL_JOB_STATES for state in states) @classmethod def step_metadata(cls, job: dict[str, Any]) -> list[SlurmStepMetadata]: @@ -407,9 +411,9 @@ def step_metadata(cls, job: dict[str, Any]) -> list[SlurmStepMetadata]: def cancel(self, job_id: int) -> None: """Cancel a Slurm job through slurmctld.""" - self.request("DELETE", "slurm", f"job/{job_id}") + 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/") + self._request("GET", "slurm", "ping/") + self._request("GET", "slurmdb", "clusters/") diff --git a/tests/systems/slurm/test_system.py b/tests/systems/slurm/test_system.py index f5549e9c7..5bf6cfd12 100644 --- a/tests/systems/slurm/test_system.py +++ b/tests/systems/slurm/test_system.py @@ -56,14 +56,14 @@ def test_slurm_api_request_expands_headers(rest_slurm_system: SlurmSystem, monke monkeypatch.setenv("SLURM_USER", "cloudai") monkeypatch.setenv("SLURM_JWT", "secret") response = Mock() - response.json.return_value = {"pings": [], "errors": [], "warnings": []} + response.json.return_value = {"nodes": [], "errors": [], "warnings": []} with patch("cloudai.systems.slurm.slurm_rest_client.requests.request", return_value=response) as request: - rest_slurm_system._rest_client.request("GET", "slurm", "ping/") + rest_slurm_system._rest_client.cluster_nodes() request.assert_called_once_with( "GET", - "https://slurm.example.com/slurm/v0.0.38/ping/", + "https://slurm.example.com/slurm/v0.0.38/nodes/", headers={"X-SLURM-USER-NAME": "cloudai", "X-SLURM-USER-TOKEN": "secret"}, json=None, timeout=30, @@ -94,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(SlurmRestClient, "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 @@ -124,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(SlurmRestClient, "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") @@ -160,7 +160,7 @@ def test_slurm_api_job_lifecycle(rest_slurm_system: SlurmSystem): } job = SlurmJob(test_run=Mock(), id=42) - with patch.object(SlurmRestClient, "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"] @@ -202,7 +202,7 @@ def request(_method: str, service: str, path: str, **_kwargs): rest_slurm_system.supports_gpu_directives_cache = None with ( - patch.object(SlurmRestClient, "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 From 3e214788db2a07252e1cf545c1f5f7f449bfaeb6 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Sat, 12 Sep 2026 19:36:50 +0200 Subject: [PATCH 07/16] Simplify Slurm REST job construction --- src/cloudai/systems/slurm/slurm_rest_client.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/src/cloudai/systems/slurm/slurm_rest_client.py b/src/cloudai/systems/slurm/slurm_rest_client.py index 79d2022b4..ece718a3d 100644 --- a/src/cloudai/systems/slurm/slurm_rest_client.py +++ b/src/cloudai/systems/slurm/slurm_rest_client.py @@ -47,14 +47,6 @@ class SlurmAPIConfig(pydantic.BaseModel): headers: dict[str, str] = pydantic.Field(default_factory=dict) verify_certs: bool = True - @pydantic.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.""" @@ -113,7 +105,7 @@ def _message(item: object) -> str: 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]: - url = f"{self._config.url}/{service}/{self._API_VERSION}/{path.lstrip('/')}" + url = f"{self._config.url.rstrip('/')}/{service}/{self._API_VERSION}/{path.lstrip('/')}" response = requests.request( method, url, @@ -234,7 +226,7 @@ def _apply_sbatch_args(self, job: dict[str, object], args: list[str]) -> None: else: raise ValueError(f"SBATCH directive '{option}' is not supported by CloudAI's Slurm REST transport.") - def _job_description(self, script: str, script_path: pathlib.Path) -> dict[str, object]: + def _make_job(self, script: str, script_path: pathlib.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(): @@ -260,7 +252,7 @@ def submit_sbatch( "POST", "slurm", "job/submit", - payload={"script": script, "job": self._job_description(script, script_path)}, + payload={"script": script, "job": self._make_job(script, script_path)}, ) except (OSError, RuntimeError, ValueError) as exc: raise cloudai.core.JobIdRetrievalError( From 2037a606121324f84259e66c7da2184c8d462184 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Sat, 12 Sep 2026 19:49:13 +0200 Subject: [PATCH 08/16] Simplify SBATCH directive parsing --- .../systems/slurm/slurm_rest_client.py | 119 +++++++++--------- tests/systems/slurm/test_system.py | 3 +- 2 files changed, 64 insertions(+), 58 deletions(-) diff --git a/src/cloudai/systems/slurm/slurm_rest_client.py b/src/cloudai/systems/slurm/slurm_rest_client.py index ece718a3d..6e01e4210 100644 --- a/src/cloudai/systems/slurm/slurm_rest_client.py +++ b/src/cloudai/systems/slurm/slurm_rest_client.py @@ -17,6 +17,7 @@ from __future__ import annotations import datetime +import getopt import logging import math import os @@ -82,6 +83,25 @@ class SlurmRestClient: "--reservation": "reservation", "--distribution": "distribution", } + _SHORT_DIRECTIVES: ClassVar[str] = "J:o:e:p:A:N:n:D:" + _LONG_DIRECTIVES: ClassVar[list[str]] = [ + "job-name=", + "output=", + "error=", + "partition=", + "account=", + "reservation=", + "distribution=", + "nodes=", + "nodelist=", + "exclude=", + "ntasks=", + "ntasks-per-node=", + "time=", + "gres=", + "gpus-per-node=", + "chdir=", + ] def __init__(self, config: SlurmAPIConfig, retry_pause_seconds: int) -> None: self._config = config @@ -154,17 +174,16 @@ def _request( 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 + @classmethod + def _parse_sbatch_line(cls, line: str) -> list[tuple[str, str]]: + """Parse one SBATCH line; e.g. `--nodes 2 --time=10` becomes two directives.""" + try: + directives, values = getopt.gnu_getopt(shlex.split(line), cls._SHORT_DIRECTIVES, cls._LONG_DIRECTIVES) + except getopt.GetoptError as exc: + raise ValueError(f"Invalid SBATCH directive: {exc}.") from exc + if values: + raise ValueError(f"Unexpected SBATCH value(s): {' '.join(values)}.") + return directives @staticmethod def _gpu_gres(value: str, *, from_gres: bool) -> str: @@ -182,49 +201,37 @@ def _set_directive(job: dict[str, object], field: str, value: object, option: st ) 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(cloudai.util.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 _apply_sbatch_directive(self, job: dict[str, object], option: str, value: str) -> None: # noqa: C901 + """Map one SBATCH directive to its v0.0.38 job field; e.g. `--nodes=2` sets `nodes=[2, 2]`.""" + option = {"-N": "--nodes", "-n": "--ntasks", "-D": "--chdir"}.get(option, 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 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", value, option) + elif option == "--exclude": + self._set_directive(job, "exclude_nodes", 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 value.isdigit() else math.ceil(cloudai.util.parse_time_limit(value).total_seconds() / 60) + ) + self._set_directive(job, "time_limit", minutes, option) + elif option in {"--gres", "--gpus-per-node"}: + gres = self._gpu_gres(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 _make_job(self, script: str, script_path: pathlib.Path) -> dict[str, object]: """Build REST job properties from leading `#SBATCH` lines; script body remains unchanged.""" @@ -235,8 +242,8 @@ def _make_job(self, script: str, script_path: pathlib.Path) -> dict[str, object] continue if not stripped.startswith("#SBATCH"): break - args = shlex.split(stripped.removeprefix("#SBATCH").strip()) - self._apply_sbatch_args(job, args) + for option, value in self._parse_sbatch_line(stripped.removeprefix("#SBATCH").strip()): + self._apply_sbatch_directive(job, option, value) job.setdefault("current_working_directory", str(script_path.parent.absolute())) job["environment"] = {"PATH": os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin")} diff --git a/tests/systems/slurm/test_system.py b/tests/systems/slurm/test_system.py index 5bf6cfd12..bdcd77165 100644 --- a/tests/systems/slurm/test_system.py +++ b/tests/systems/slurm/test_system.py @@ -74,8 +74,7 @@ def test_slurm_api_request_expands_headers(rest_slurm_system: SlurmSystem, monke 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 --job-name=rest-test --output=/shared/stdout.txt #SBATCH --error=/shared/stderr.txt #SBATCH --partition=gpu #SBATCH --account=cloudai From 33b496e308d455ef9adaf8ed966a4142e56750a1 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Sat, 12 Sep 2026 19:55:30 +0200 Subject: [PATCH 09/16] Simplify SBATCH directive mapping --- src/cloudai/systems/slurm/slurm_rest_client.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/src/cloudai/systems/slurm/slurm_rest_client.py b/src/cloudai/systems/slurm/slurm_rest_client.py index 6e01e4210..b7578c4c3 100644 --- a/src/cloudai/systems/slurm/slurm_rest_client.py +++ b/src/cloudai/systems/slurm/slurm_rest_client.py @@ -82,6 +82,9 @@ class SlurmRestClient: "-A": "account", "--reservation": "reservation", "--distribution": "distribution", + "--nodelist": "nodelist", + "--exclude": "exclude_nodes", + "--chdir": "current_working_directory", } _SHORT_DIRECTIVES: ClassVar[str] = "J:o:e:p:A:N:n:D:" _LONG_DIRECTIVES: ClassVar[list[str]] = [ @@ -185,11 +188,6 @@ def _parse_sbatch_line(cls, line: str) -> list[tuple[str, str]]: raise ValueError(f"Unexpected SBATCH value(s): {' '.join(values)}.") return directives - @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: """Merge directives mapped to one REST field; e.g. `--gres=gpu:8` and `--gpus-per-node=8` must agree.""" @@ -201,7 +199,7 @@ def _set_directive(job: dict[str, object], field: str, value: object, option: st ) job[field] = value - def _apply_sbatch_directive(self, job: dict[str, object], option: str, value: str) -> None: # noqa: C901 + def _apply_sbatch_directive(self, job: dict[str, object], option: str, value: str) -> None: """Map one SBATCH directive to its v0.0.38 job field; e.g. `--nodes=2` sets `nodes=[2, 2]`.""" option = {"-N": "--nodes", "-n": "--ntasks", "-D": "--chdir"}.get(option, option) @@ -212,10 +210,6 @@ def _apply_sbatch_directive(self, job: dict[str, object], option: str, value: st 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", value, option) - elif option == "--exclude": - self._set_directive(job, "exclude_nodes", value, option) elif option == "--ntasks": self._set_directive(job, "tasks", int(value), option) elif option == "--ntasks-per-node": @@ -226,10 +220,8 @@ def _apply_sbatch_directive(self, job: dict[str, object], option: str, value: st ) self._set_directive(job, "time_limit", minutes, option) elif option in {"--gres", "--gpus-per-node"}: - gres = self._gpu_gres(value, from_gres=option == "--gres") + gres = value if option == "--gres" else f"gpu:{value}" 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.") From 82e71144e8282d74e6c6e199dc99085fbd95a8de Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Sat, 12 Sep 2026 20:05:18 +0200 Subject: [PATCH 10/16] Simplify Slurm REST client methods --- src/cloudai/systems/slurm/slurm_rest_client.py | 6 ++---- src/cloudai/systems/slurm/slurm_system.py | 6 ++---- tests/systems/slurm/test_system.py | 1 - 3 files changed, 4 insertions(+), 9 deletions(-) diff --git a/src/cloudai/systems/slurm/slurm_rest_client.py b/src/cloudai/systems/slurm/slurm_rest_client.py index b7578c4c3..4d7afba64 100644 --- a/src/cloudai/systems/slurm/slurm_rest_client.py +++ b/src/cloudai/systems/slurm/slurm_rest_client.py @@ -342,21 +342,19 @@ def _records(data: dict[str, Any], field: str) -> list[dict[str, Any]]: 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: + def get_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) + job = self.get_job(job_id, retry_threshold) if job is None: return [] states = self.states(job) diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index c6e799e6e..943e5063a 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -412,8 +412,6 @@ 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_client.validate() except RuntimeError as exc: @@ -543,7 +541,7 @@ 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: - rest_job = self._rest_client.accounting_job(self._job_id(job), retry_threshold) + rest_job = self._rest_client.get_job(self._job_id(job), retry_threshold) return self._rest_client.step_metadata(rest_job) if rest_job else [] retry_count = 0 @@ -1001,7 +999,7 @@ def complete_job(self, job: SlurmJob) -> list[str]: return [] if self.uses_slurm_api: - rest_job = self._rest_client.accounting_job(self._job_id(job)) + rest_job = self._rest_client.get_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 bdcd77165..a3c194947 100644 --- a/tests/systems/slurm/test_system.py +++ b/tests/systems/slurm/test_system.py @@ -202,7 +202,6 @@ def request(_method: str, service: str, path: str, **_kwargs): rest_slurm_system.supports_gpu_directives_cache = None with ( 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 assert [(node.name, node.partition, node.state) for node in rest_slurm_system.nodes_from_sinfo()] == [ From d807d996d9bbee0d22c22892cb076da34dce9bce Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Sat, 12 Sep 2026 20:24:35 +0200 Subject: [PATCH 11/16] Refine Slurm REST read operations --- .../systems/slurm/slurm_rest_client.py | 242 ++++++++++-------- src/cloudai/systems/slurm/slurm_system.py | 95 ++++--- 2 files changed, 203 insertions(+), 134 deletions(-) diff --git a/src/cloudai/systems/slurm/slurm_rest_client.py b/src/cloudai/systems/slurm/slurm_rest_client.py index 4d7afba64..b65380f2f 100644 --- a/src/cloudai/systems/slurm/slurm_rest_client.py +++ b/src/cloudai/systems/slurm/slurm_rest_client.py @@ -277,92 +277,65 @@ def submit_sbatch( 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.datetime.fromtimestamp(timestamp, tz=datetime.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 self._records(self._request("GET", "slurm", "nodes/"), "nodes") + nodes = self._request("GET", "slurm", "nodes/").get("nodes", []) + if not isinstance(nodes, list): + raise RuntimeError("Slurm API returned an invalid nodes response.") + return [node for node in nodes if isinstance(node, dict)] + + def has_gpus(self) -> bool: + return any( + "gpu" in str(node.get(field, "")).lower() for node in self.cluster_nodes() for field in ("gres", "tres") + ) def queue_jobs(self) -> list[dict[str, Any]]: - return self._records(self._request("GET", "slurm", "jobs/"), "jobs") - - def get_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) + jobs = self._request("GET", "slurm", "jobs/").get("jobs", []) + if not isinstance(jobs, list): + raise RuntimeError("Slurm API returned an invalid jobs response.") + return [job for job in jobs if isinstance(job, dict)] + + def _get_job(self, job_id: int, retry_threshold: int = 3) -> dict[str, Any] | None: + jobs = self._request("GET", "slurmdb", f"job/{job_id}", retry_threshold=retry_threshold).get("jobs", []) + if not isinstance(jobs, list): + raise RuntimeError("Slurm API returned an invalid jobs response.") + + for job in jobs: + if not isinstance(job, dict): + continue + response_job_id = job.get("job_id") + if isinstance(response_job_id, dict): + if response_job_id.get("set") is False: + continue + response_job_id = response_job_id.get("number", 0) + if not isinstance(response_job_id, (str, int, float)): + continue + try: + if int(response_job_id) == job_id: + return job + except (TypeError, ValueError): + continue + return None def job_states(self, job_id: int, retry_threshold: int = 3) -> list[str]: """Return job and step states from slurmdbd.""" - job = self.get_job(job_id, retry_threshold) + job = self._get_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)) + records = [job, *(step for step in steps if isinstance(step, dict))] if isinstance(steps, list) else [job] + states: list[str] = [] + for record in records: + raw_states = record.get("state", record.get("job_state")) + if isinstance(raw_states, dict): + raw_states = raw_states.get("current", []) + if isinstance(raw_states, list): + record_states = raw_states + elif raw_states is None: + record_states = [] + else: + record_states = re.split(r"[,+]", str(raw_states)) + states.extend(str(state).upper().rstrip("+") for state in record_states if state) return states def is_job_completed(self, job_id: int, retry_threshold: int = 3) -> bool: @@ -372,31 +345,100 @@ def is_job_completed(self, job_id: int, retry_threshold: int = 3) -> bool: 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")) + @staticmethod + def _make_step_metadata(job_id: int, record: dict[str, Any], *, is_job: bool) -> SlurmStepMetadata: # noqa: C901 + step = record.get("step", {}) if isinstance(record.get("step"), dict) else {} + times = record.get("time", {}) if isinstance(record.get("time"), dict) else {} + + raw_states = record.get("state", record.get("job_state")) + if isinstance(raw_states, dict): + raw_states = raw_states.get("current", []) + if isinstance(raw_states, list): + states = [str(state).upper().rstrip("+") for state in raw_states if state] + elif raw_states is None: + states = [] + else: + states = [state.upper().rstrip("+") for state in re.split(r"[,+]", str(raw_states)) if state] + + exit_code = record.get("exit_code") + if isinstance(exit_code, str): + formatted_exit_code = exit_code + elif isinstance(exit_code, dict): + return_code = exit_code.get("return_code", 0) + if isinstance(return_code, dict): + return_code = return_code.get("number", 0) if return_code.get("set") is not False else 0 + signal = exit_code.get("signal", 0) + if isinstance(signal, dict): + signal = signal.get("id", signal.get("signal_id", signal.get("number", 0))) + try: + parsed_return_code = int(return_code) if isinstance(return_code, (str, int, float)) else 0 + parsed_signal = int(signal) if isinstance(signal, (str, int, float)) else 0 + formatted_exit_code = f"{parsed_return_code}:{parsed_signal}" + except (TypeError, ValueError): + formatted_exit_code = "0:0" + else: + formatted_exit_code = "0:0" + + formatted_times: list[str] = [] + for field in ("start", "end"): + raw_time = times.get(field) + if isinstance(raw_time, str) and not raw_time.isdigit(): + formatted_times.append(raw_time) + continue + if isinstance(raw_time, dict): + raw_time = raw_time.get("number", 0) if raw_time.get("set") is not False else 0 + try: + timestamp = int(raw_time) if isinstance(raw_time, (str, int, float)) else 0 + except (TypeError, ValueError): + timestamp = 0 + formatted_times.append( + datetime.datetime.fromtimestamp(timestamp, tz=datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + if timestamp + else "" + ) + + elapsed = times.get("elapsed", 0) + if isinstance(elapsed, dict): + elapsed = elapsed.get("number", 0) if elapsed.get("set") is not False else 0 + try: + elapsed_seconds = int(elapsed) if isinstance(elapsed, (str, int, float)) else 0 + except (TypeError, ValueError): + elapsed_seconds = 0 + + return SlurmStepMetadata( + job_id=job_id, + step_id="" if is_job else str(step.get("id", "")), + name=str(record.get("name", step.get("name", ""))), + state=states[0] if states else "", + exit_code=formatted_exit_code, + start_time=formatted_times[0], + end_time=formatted_times[1], + elapsed_time_sec=elapsed_seconds, + submit_line=str(record.get("submit_line", "")), + ) + + def get_job_status(self, job_id: int, retry_threshold: int = 3) -> list[SlurmStepMetadata]: + job = self._get_job(job_id, retry_threshold) + if job is None: + return [] + + response_job_id = job.get("job_id") + if isinstance(response_job_id, dict): + response_job_id = response_job_id.get("number", 0) if response_job_id.get("set") is not False else 0 + try: + metadata_job_id = int(response_job_id) if isinstance(response_job_id, (str, int, float)) else 0 + except (TypeError, ValueError): + metadata_job_id = 0 + 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 + return [ + self._make_step_metadata(metadata_job_id, record, is_job=index == 0) for index, record in enumerate(records) + ] + + def get_job_nodes(self, job_id: int) -> str: + job = self._get_job(job_id) + return str(job.get("nodes", "")) if job else "" def cancel(self, job_id: int) -> None: """Cancel a Slurm job through slurmctld.""" diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index 943e5063a..021a39d1c 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -186,7 +186,17 @@ def _job_id(job: BaseJob) -> int: def _rest_node_state(self, node: dict[str, Any]) -> SlurmNodeState: """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)] + raw_states = node.get("state", node.get("job_state")) + if isinstance(raw_states, dict): + raw_states = raw_states.get("current", []) + if isinstance(raw_states, list): + state_names = [str(state) for state in raw_states] + elif raw_states is None: + state_names = [] + else: + state_names = [state for state in re.split(r"[,+]", str(raw_states)) if state] + + states = [self.convert_state_to_enum(state.upper().rstrip("+")) for state in state_names] ordinary_states = { SlurmNodeState.ALLOCATED, SlurmNodeState.ALLOCATED_COMPLETING, @@ -197,6 +207,51 @@ 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) + def _nodes_from_rest(self) -> list[SlurmNode]: + nodes: list[SlurmNode] = [] + for node in self._rest_client.cluster_nodes(): + if not node.get("name"): + continue + + partitions = node.get("partitions") + if isinstance(partitions, dict): + partitions = partitions.get("current", []) + if isinstance(partitions, list): + partition_names = [str(partition) for partition in partitions] + elif partitions is None: + partition_names = [] + else: + partition_names = [partition for partition in re.split(r"[,+]", str(partitions)) if partition] + + state = self._rest_node_state(node) + nodes.extend( + SlurmNode(name=str(node["name"]), partition=partition, state=state) for partition in partition_names + ) + return nodes + + def _allocated_nodes_from_rest(self) -> list[SlurmNode]: + nodes: list[SlurmNode] = [] + for job in self._rest_client.queue_jobs(): + raw_states = job.get("state", job.get("job_state")) + if isinstance(raw_states, dict): + raw_states = raw_states.get("current", []) + if isinstance(raw_states, list): + states = {str(state).upper().rstrip("+") for state in raw_states} + elif raw_states is None: + states = set() + else: + states = {state.upper().rstrip("+") for state in re.split(r"[,+]", str(raw_states)) if state} + if not {"RUNNING", "PENDING"}.intersection(states): + continue + + partition = str(job.get("partition", "")) + user = str(job.get("user_name", job.get("user", "N/A"))) + nodes.extend( + SlurmNode(name=name, partition=partition, state=SlurmNodeState.ALLOCATED, user=user) + for name in parse_node_list(str(job.get("nodes", ""))) + ) + return nodes + @property def groups(self) -> Dict[str, Dict[str, List[SlurmNode]]]: groups: Dict[str, Dict[str, List[SlurmNode]]] = {} @@ -227,15 +282,11 @@ def supports_gpu_directives(self) -> bool: if self.uses_slurm_api: try: - nodes = self._rest_client.cluster_nodes() + self.supports_gpu_directives_cache = self._rest_client.has_gpus() 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 nodes for field in ("gres", "tres") - ) return self.supports_gpu_directives_cache stdout, stderr = self.fetch_command_output("scontrol show config") @@ -271,14 +322,7 @@ def update(self) -> None: def nodes_from_sinfo(self) -> list[SlurmNode]: if self.uses_slurm_api: - nodes: list[SlurmNode] = [] - 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_client.values(node.get("partitions")): - nodes.append(SlurmNode(name=str(node["name"]), partition=partition, state=state)) - return nodes + return self._nodes_from_rest() sinfo_output, _ = self.fetch_command_output("sinfo --noheader -o '%P|%t|%u|%N'") nodes: list[SlurmNode] = [] @@ -300,22 +344,7 @@ def nodes_from_sinfo(self) -> list[SlurmNode]: def nodes_from_squeue(self) -> list[SlurmNode]: if self.uses_slurm_api: - nodes: list[SlurmNode] = [] - 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"))) - 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 + return self._allocated_nodes_from_rest() squeue_output, _ = self.fetch_command_output("squeue --states=running,pending --noheader -o '%P|%T|%N|%u'") nodes: list[SlurmNode] = [] @@ -541,8 +570,7 @@ 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: - rest_job = self._rest_client.get_job(self._job_id(job), retry_threshold) - return self._rest_client.step_metadata(rest_job) if rest_job else [] + return self._rest_client.get_job_status(self._job_id(job), retry_threshold) retry_count = 0 command = ( @@ -999,8 +1027,7 @@ def complete_job(self, job: SlurmJob) -> list[str]: return [] if self.uses_slurm_api: - rest_job = self._rest_client.get_job(self._job_id(job)) - spec = str(rest_job.get("nodes", "")) if rest_job else "" + spec = self._rest_client.get_job_nodes(self._job_id(job)) 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 From c94766e408576ec250d89cbd37e45015445068f9 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 14 Sep 2026 15:45:21 +0200 Subject: [PATCH 12/16] Use slurmctld for REST job status --- .../systems/slurm/slurm_rest_client.py | 154 ++++++------------ src/cloudai/systems/slurm/slurm_system.py | 2 +- tests/systems/slurm/test_system.py | 28 +--- 3 files changed, 56 insertions(+), 128 deletions(-) diff --git a/src/cloudai/systems/slurm/slurm_rest_client.py b/src/cloudai/systems/slurm/slurm_rest_client.py index b65380f2f..90c428247 100644 --- a/src/cloudai/systems/slurm/slurm_rest_client.py +++ b/src/cloudai/systems/slurm/slurm_rest_client.py @@ -244,7 +244,7 @@ def _make_job(self, script: str, script_path: pathlib.Path) -> dict[str, object] def submit_sbatch( self, script_path: pathlib.Path, operation_name: str, *, wait: bool = False, monitor_interval: int = 1 ) -> int: - """Submit an SBATCH file and optionally wait for a terminal accounting state.""" + """Submit an SBATCH file and optionally wait for a terminal job state.""" try: script = script_path.read_text(encoding="utf-8") data = self._request( @@ -295,7 +295,7 @@ def queue_jobs(self) -> list[dict[str, Any]]: return [job for job in jobs if isinstance(job, dict)] def _get_job(self, job_id: int, retry_threshold: int = 3) -> dict[str, Any] | None: - jobs = self._request("GET", "slurmdb", f"job/{job_id}", retry_threshold=retry_threshold).get("jobs", []) + jobs = self._request("GET", "slurm", f"job/{job_id}", retry_threshold=retry_threshold).get("jobs", []) if not isinstance(jobs, list): raise RuntimeError("Slurm API returned an invalid jobs response.") @@ -303,10 +303,6 @@ def _get_job(self, job_id: int, retry_threshold: int = 3) -> dict[str, Any] | No if not isinstance(job, dict): continue response_job_id = job.get("job_id") - if isinstance(response_job_id, dict): - if response_job_id.get("set") is False: - continue - response_job_id = response_job_id.get("number", 0) if not isinstance(response_job_id, (str, int, float)): continue try: @@ -316,106 +312,18 @@ def _get_job(self, job_id: int, retry_threshold: int = 3) -> dict[str, Any] | No continue return None - def job_states(self, job_id: int, retry_threshold: int = 3) -> list[str]: - """Return job and step states from slurmdbd.""" + def get_job_state(self, job_id: int, retry_threshold: int = 3) -> str: + """Return current job state from slurmctld.""" job = self._get_job(job_id, retry_threshold) if job is None: - return [] + return "" - steps = job.get("steps", []) - records = [job, *(step for step in steps if isinstance(step, dict))] if isinstance(steps, list) else [job] - states: list[str] = [] - for record in records: - raw_states = record.get("state", record.get("job_state")) - if isinstance(raw_states, dict): - raw_states = raw_states.get("current", []) - if isinstance(raw_states, list): - record_states = raw_states - elif raw_states is None: - record_states = [] - else: - record_states = re.split(r"[,+]", str(raw_states)) - states.extend(str(state).upper().rstrip("+") for state in record_states if state) - return states + state = job.get("job_state") + return str(state).upper().rstrip("+") if state else "" 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) - - @staticmethod - def _make_step_metadata(job_id: int, record: dict[str, Any], *, is_job: bool) -> SlurmStepMetadata: # noqa: C901 - step = record.get("step", {}) if isinstance(record.get("step"), dict) else {} - times = record.get("time", {}) if isinstance(record.get("time"), dict) else {} - - raw_states = record.get("state", record.get("job_state")) - if isinstance(raw_states, dict): - raw_states = raw_states.get("current", []) - if isinstance(raw_states, list): - states = [str(state).upper().rstrip("+") for state in raw_states if state] - elif raw_states is None: - states = [] - else: - states = [state.upper().rstrip("+") for state in re.split(r"[,+]", str(raw_states)) if state] - - exit_code = record.get("exit_code") - if isinstance(exit_code, str): - formatted_exit_code = exit_code - elif isinstance(exit_code, dict): - return_code = exit_code.get("return_code", 0) - if isinstance(return_code, dict): - return_code = return_code.get("number", 0) if return_code.get("set") is not False else 0 - signal = exit_code.get("signal", 0) - if isinstance(signal, dict): - signal = signal.get("id", signal.get("signal_id", signal.get("number", 0))) - try: - parsed_return_code = int(return_code) if isinstance(return_code, (str, int, float)) else 0 - parsed_signal = int(signal) if isinstance(signal, (str, int, float)) else 0 - formatted_exit_code = f"{parsed_return_code}:{parsed_signal}" - except (TypeError, ValueError): - formatted_exit_code = "0:0" - else: - formatted_exit_code = "0:0" - - formatted_times: list[str] = [] - for field in ("start", "end"): - raw_time = times.get(field) - if isinstance(raw_time, str) and not raw_time.isdigit(): - formatted_times.append(raw_time) - continue - if isinstance(raw_time, dict): - raw_time = raw_time.get("number", 0) if raw_time.get("set") is not False else 0 - try: - timestamp = int(raw_time) if isinstance(raw_time, (str, int, float)) else 0 - except (TypeError, ValueError): - timestamp = 0 - formatted_times.append( - datetime.datetime.fromtimestamp(timestamp, tz=datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") - if timestamp - else "" - ) - - elapsed = times.get("elapsed", 0) - if isinstance(elapsed, dict): - elapsed = elapsed.get("number", 0) if elapsed.get("set") is not False else 0 - try: - elapsed_seconds = int(elapsed) if isinstance(elapsed, (str, int, float)) else 0 - except (TypeError, ValueError): - elapsed_seconds = 0 - - return SlurmStepMetadata( - job_id=job_id, - step_id="" if is_job else str(step.get("id", "")), - name=str(record.get("name", step.get("name", ""))), - state=states[0] if states else "", - exit_code=formatted_exit_code, - start_time=formatted_times[0], - end_time=formatted_times[1], - elapsed_time_sec=elapsed_seconds, - submit_line=str(record.get("submit_line", "")), - ) + """Return whether slurmctld reports a terminal job state.""" + return self.get_job_state(job_id, retry_threshold) in self._TERMINAL_JOB_STATES def get_job_status(self, job_id: int, retry_threshold: int = 3) -> list[SlurmStepMetadata]: job = self._get_job(job_id, retry_threshold) @@ -423,17 +331,48 @@ def get_job_status(self, job_id: int, retry_threshold: int = 3) -> list[SlurmSte return [] response_job_id = job.get("job_id") - if isinstance(response_job_id, dict): - response_job_id = response_job_id.get("number", 0) if response_job_id.get("set") is not False else 0 try: metadata_job_id = int(response_job_id) if isinstance(response_job_id, (str, int, float)) else 0 except (TypeError, ValueError): metadata_job_id = 0 - steps = job.get("steps", []) - records = [job, *(step for step in steps if isinstance(step, dict))] if isinstance(steps, list) else [job] + raw_exit_code = job.get("exit_code") + return_code = 0 + signal = 0 + if isinstance(raw_exit_code, int) and 0 <= raw_exit_code <= 0xFFFF: + if os.WIFEXITED(raw_exit_code): + return_code = os.WEXITSTATUS(raw_exit_code) + elif os.WIFSIGNALED(raw_exit_code): + signal = os.WTERMSIG(raw_exit_code) + + raw_start_time = job.get("start_time") + raw_end_time = job.get("end_time") + start_timestamp = int(raw_start_time) if isinstance(raw_start_time, (str, int, float)) else 0 + end_timestamp = int(raw_end_time) if isinstance(raw_end_time, (str, int, float)) else 0 + start_time = ( + datetime.datetime.fromtimestamp(start_timestamp, tz=datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + if start_timestamp + else "" + ) + end_time = ( + datetime.datetime.fromtimestamp(end_timestamp, tz=datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + if end_timestamp + else "" + ) + elapsed_seconds = max(end_timestamp - start_timestamp, 0) if start_timestamp and end_timestamp else 0 + return [ - self._make_step_metadata(metadata_job_id, record, is_job=index == 0) for index, record in enumerate(records) + SlurmStepMetadata( + job_id=metadata_job_id, + step_id="", + name=str(job.get("name", "")), + state=str(job.get("job_state", "")).upper().rstrip("+"), + exit_code=f"{return_code}:{signal}", + start_time=start_time, + end_time=end_time, + elapsed_time_sec=elapsed_seconds, + submit_line=str(job.get("command", "")), + ) ] def get_job_nodes(self, job_id: int) -> str: @@ -445,6 +384,5 @@ def cancel(self, job_id: int) -> None: self._request("DELETE", "slurm", f"job/{job_id}") def validate(self) -> None: - """Verify access to slurmctld and slurmdbd endpoints used by CloudAI.""" + """Verify access to slurmctld 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 021a39d1c..968165a9e 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -478,7 +478,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: - return "RUNNING" in self._rest_client.job_states(self._job_id(job), retry_threshold) + return self._rest_client.get_job_state(self._job_id(job), retry_threshold) == "RUNNING" retry_count = 0 command = f"sacct -j {job.id} --format=State --noheader" diff --git a/tests/systems/slurm/test_system.py b/tests/systems/slurm/test_system.py index a3c194947..411212163 100644 --- a/tests/systems/slurm/test_system.py +++ b/tests/systems/slurm/test_system.py @@ -142,34 +142,25 @@ def test_slurm_api_job_lifecycle(rest_slurm_system: SlurmSystem): { "job_id": 42, "name": "rest-test", - "state": {"current": "COMPLETED"}, - "exit_code": {"return_code": 0, "signal": {"signal_id": 0}}, - "time": {"start": 100, "end": 120, "elapsed": 20}, + "job_state": "FAILED", + "exit_code": 256, + "start_time": 100, + "end_time": 120, "nodes": "node[01-02]", - "steps": [ - { - "step": {"id": "batch", "name": "batch"}, - "state": "COMPLETED", - "exit_code": {"return_code": 0, "signal": {"signal_id": 0}}, - "time": {"start": 100, "end": 120, "elapsed": 20}, - } - ], } ] } job = SlurmJob(test_run=Mock(), id=42) - with patch.object(SlurmRestClient, "_request", return_value=response): + with patch.object(SlurmRestClient, "_request", return_value=response) as request: 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 request.call_args_list == [call("GET", "slurm", "job/42", retry_threshold=3)] * 4 + assert [(item.step_id, item.name, item.state) for item in metadata] == [("", "rest-test", "FAILED")] + assert metadata[0].exit_code == "1:0" assert metadata[0].elapsed_time_sec == 20 @@ -215,10 +206,9 @@ def request(_method: str, service: str, path: str, **_kwargs): rest_slurm_system.scancel(42) rest_slurm_system.validate_install_environment() - assert rest_request.call_args_list[-3:] == [ + assert rest_request.call_args_list[-2:] == [ call("DELETE", "slurm", "job/42"), call("GET", "slurm", "ping/"), - call("GET", "slurmdb", "clusters/"), ] From be3aee2d1962f4408d6023f4388e04e206de723c Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 14 Sep 2026 16:15:02 +0200 Subject: [PATCH 13/16] Simplify Slurm REST response handling --- src/cloudai/systems/slurm/__init__.py | 4 +- src/cloudai/systems/slurm/slurm_node.py | 31 ++++- .../systems/slurm/slurm_rest_client.py | 131 ++++++++++++------ src/cloudai/systems/slurm/slurm_system.py | 110 +-------------- tests/systems/slurm/test_system.py | 17 ++- 5 files changed, 139 insertions(+), 154 deletions(-) diff --git a/src/cloudai/systems/slurm/__init__.py b/src/cloudai/systems/slurm/__init__.py index 755feb33d..da876d60c 100644 --- a/src/cloudai/systems/slurm/__init__.py +++ b/src/cloudai/systems/slurm/__init__.py @@ -19,10 +19,10 @@ from .slurm_installer import SlurmInstaller from .slurm_job import SlurmJob from .slurm_metadata import SlurmJobMetadata, SlurmStepMetadata, SlurmSystemMetadata -from .slurm_node import SlurmNode, SlurmNodeState +from .slurm_node import SlurmNode, SlurmNodeState, parse_node_list from .slurm_rest_client import SlurmAPIConfig from .slurm_runner import SlurmRunner -from .slurm_system import SlurmGroup, SlurmPartition, SlurmSystem, parse_node_list +from .slurm_system import SlurmGroup, SlurmPartition, SlurmSystem __all__ = [ "SingleSbatchRunner", diff --git a/src/cloudai/systems/slurm/slurm_node.py b/src/cloudai/systems/slurm/slurm_node.py index cca703231..917937c10 100644 --- a/src/cloudai/systems/slurm/slurm_node.py +++ b/src/cloudai/systems/slurm/slurm_node.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,11 +14,40 @@ # See the License for the specific language governing permissions and # limitations under the License. +import re from enum import Enum from pydantic import BaseModel, ConfigDict +def parse_node_list(node_list: str) -> list[str]: + """Expand a Slurm node list such as `node[01-03]` into individual names.""" + node_list = node_list.strip() + nodes = [] + if not node_list: + return [] + + components = re.split(r",\s*(?![^[]*\])", node_list) + for component in components: + if "[" not in component: + nodes.append(component) + continue + + header, node_number = component.split("[") + for node_range in node_number.rstrip("]").split(","): + if "-" not in node_range: + nodes.append(f"{header}{node_range}") + continue + + start_node, end_node = node_range.split("-") + width = len(end_node) + nodes.extend( + f"{header}{node_number:0{width}d}" for node_number in range(int(start_node), int(end_node) + 1) + ) + + return nodes + + class SlurmNodeState(Enum): """ Enumeration of possible states for a Slurm compute node, as defined by the Slurm workload manager. diff --git a/src/cloudai/systems/slurm/slurm_rest_client.py b/src/cloudai/systems/slurm/slurm_rest_client.py index 90c428247..6b9c6e4d2 100644 --- a/src/cloudai/systems/slurm/slurm_rest_client.py +++ b/src/cloudai/systems/slurm/slurm_rest_client.py @@ -25,7 +25,7 @@ import re import shlex import time -from typing import Any, ClassVar +from typing import Any, ClassVar, cast import pydantic import requests @@ -35,6 +35,7 @@ import cloudai.util from .slurm_metadata import SlurmStepMetadata +from .slurm_node import SlurmNode, SlurmNodeState, parse_node_list logger = logging.getLogger(__name__) @@ -277,40 +278,95 @@ def submit_sbatch( time.sleep(monitor_interval) return job_id - def cluster_nodes(self) -> list[dict[str, Any]]: - nodes = self._request("GET", "slurm", "nodes/").get("nodes", []) + def _cluster_nodes(self) -> list[dict[str, Any]]: + nodes = self._request("GET", "slurm", "nodes/").get("nodes") if not isinstance(nodes, list): raise RuntimeError("Slurm API returned an invalid nodes response.") - return [node for node in nodes if isinstance(node, dict)] + return cast(list[dict[str, Any]], nodes) + + @staticmethod + def _node_state(state: str, state_flags: list[str]) -> SlurmNodeState: + """Combine the v0.0.38 node state and flags; e.g. `idle` plus `DRAIN` means `DRAINED`.""" + state = state.upper() + flags = {flag.upper() for flag in state_flags} + + if "NOT_RESPONDING" in flags: + return SlurmNodeState.NOT_RESPONDING + if "DRAIN" in flags: + if state in {"ALLOCATED", "MIXED"} or "COMPLETING" in flags: + return SlurmNodeState.DRAINING + return SlurmNodeState.DRAINED + if "FAIL" in flags: + return SlurmNodeState.FAILING if state == "ALLOCATED" or "COMPLETING" in flags else SlurmNodeState.FAIL + + flag_states = { + "INVALID_REG": SlurmNodeState.INVALID_REGISTRATION, + "MAINTENANCE": SlurmNodeState.MAINTENANCE, + "POWER_DOWN": SlurmNodeState.PENDING_POWER_DOWN_STATE, + "POWER_UP": SlurmNodeState.BEING_POWERED_UP_OR_CONFIGURED, + "POWERED_DOWN": SlurmNodeState.POWERED_DOWN_STATE, + "POWERING_DOWN": SlurmNodeState.POWERING_DOWN_STATE, + "POWERING_UP": SlurmNodeState.POWERING_UP_STATE, + "REBOOT_REQUESTED": SlurmNodeState.REBOOT_REQUESTED, + "REBOOT_ISSUED": SlurmNodeState.REBOOT_ISSUED_STATE, + "PERFCTRS": SlurmNodeState.USING_NETWORK_PERFORMANCE_COUNTERS, + "PLANNED": SlurmNodeState.PLANNED_STATE, + "RESERVED": SlurmNodeState.RESERVED, + } + for flag in state_flags: + if node_state := flag_states.get(flag.upper()): + return node_state + + if "COMPLETING" in flags: + return SlurmNodeState.ALLOCATED_COMPLETING if state == "ALLOCATED" else SlurmNodeState.COMPLETING + + try: + return SlurmNodeState(state) + except ValueError: + return SlurmNodeState.UNKNOWN_STATE + + def get_nodes(self) -> list[SlurmNode]: + nodes: list[SlurmNode] = [] + for node in self._cluster_nodes(): + state = self._node_state(node["state"], node["state_flags"]) + nodes.extend( + SlurmNode(name=node["name"], partition=partition, state=state) for partition in node["partitions"] + ) + return nodes def has_gpus(self) -> bool: return any( - "gpu" in str(node.get(field, "")).lower() for node in self.cluster_nodes() for field in ("gres", "tres") + "gpu" in str(node.get(field, "")).lower() for node in self._cluster_nodes() for field in ("gres", "tres") ) - def queue_jobs(self) -> list[dict[str, Any]]: - jobs = self._request("GET", "slurm", "jobs/").get("jobs", []) + def _queue_jobs(self) -> list[dict[str, Any]]: + jobs = self._request("GET", "slurm", "jobs/").get("jobs") if not isinstance(jobs, list): raise RuntimeError("Slurm API returned an invalid jobs response.") - return [job for job in jobs if isinstance(job, dict)] + return cast(list[dict[str, Any]], jobs) + + def get_allocated_nodes(self) -> list[SlurmNode]: + nodes: list[SlurmNode] = [] + for job in self._queue_jobs(): + if job["job_state"].upper().rstrip("+") not in {"RUNNING", "PENDING"}: + continue + + nodes.extend( + SlurmNode( + name=name, + partition=job["partition"] or "", + state=SlurmNodeState.ALLOCATED, + user=job["user_name"] or "N/A", + ) + for name in parse_node_list(job["nodes"] or "") + ) + return nodes def _get_job(self, job_id: int, retry_threshold: int = 3) -> dict[str, Any] | None: - jobs = self._request("GET", "slurm", f"job/{job_id}", retry_threshold=retry_threshold).get("jobs", []) + jobs = self._request("GET", "slurm", f"job/{job_id}", retry_threshold=retry_threshold).get("jobs") if not isinstance(jobs, list): raise RuntimeError("Slurm API returned an invalid jobs response.") - - for job in jobs: - if not isinstance(job, dict): - continue - response_job_id = job.get("job_id") - if not isinstance(response_job_id, (str, int, float)): - continue - try: - if int(response_job_id) == job_id: - return job - except (TypeError, ValueError): - continue - return None + return cast(dict[str, Any], jobs[0]) if jobs else None def get_job_state(self, job_id: int, retry_threshold: int = 3) -> str: """Return current job state from slurmctld.""" @@ -318,8 +374,7 @@ def get_job_state(self, job_id: int, retry_threshold: int = 3) -> str: if job is None: return "" - state = job.get("job_state") - return str(state).upper().rstrip("+") if state else "" + return job["job_state"].upper().rstrip("+") def is_job_completed(self, job_id: int, retry_threshold: int = 3) -> bool: """Return whether slurmctld reports a terminal job state.""" @@ -330,25 +385,17 @@ def get_job_status(self, job_id: int, retry_threshold: int = 3) -> list[SlurmSte if job is None: return [] - response_job_id = job.get("job_id") - try: - metadata_job_id = int(response_job_id) if isinstance(response_job_id, (str, int, float)) else 0 - except (TypeError, ValueError): - metadata_job_id = 0 - - raw_exit_code = job.get("exit_code") + raw_exit_code = job["exit_code"] return_code = 0 signal = 0 - if isinstance(raw_exit_code, int) and 0 <= raw_exit_code <= 0xFFFF: + if 0 <= raw_exit_code <= 0xFFFF: if os.WIFEXITED(raw_exit_code): return_code = os.WEXITSTATUS(raw_exit_code) elif os.WIFSIGNALED(raw_exit_code): signal = os.WTERMSIG(raw_exit_code) - raw_start_time = job.get("start_time") - raw_end_time = job.get("end_time") - start_timestamp = int(raw_start_time) if isinstance(raw_start_time, (str, int, float)) else 0 - end_timestamp = int(raw_end_time) if isinstance(raw_end_time, (str, int, float)) else 0 + start_timestamp = job["start_time"] + end_timestamp = job["end_time"] start_time = ( datetime.datetime.fromtimestamp(start_timestamp, tz=datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") if start_timestamp @@ -363,21 +410,23 @@ def get_job_status(self, job_id: int, retry_threshold: int = 3) -> list[SlurmSte return [ SlurmStepMetadata( - job_id=metadata_job_id, + job_id=job["job_id"], step_id="", - name=str(job.get("name", "")), - state=str(job.get("job_state", "")).upper().rstrip("+"), + name=job["name"], + state=job["job_state"].upper().rstrip("+"), exit_code=f"{return_code}:{signal}", start_time=start_time, end_time=end_time, elapsed_time_sec=elapsed_seconds, - submit_line=str(job.get("command", "")), + submit_line=job.get("command") or "", ) ] def get_job_nodes(self, job_id: int) -> str: job = self._get_job(job_id) - return str(job.get("nodes", "")) if job else "" + if job is None: + return "" + return job.get("nodes") or "" def cancel(self, job_id: int) -> None: """Cancel a Slurm job through slurmctld.""" diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index 968165a9e..06151f2e4 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -34,7 +34,7 @@ from .slurm_job import SlurmJob from .slurm_metadata import SlurmStepMetadata -from .slurm_node import SlurmNode, SlurmNodeState +from .slurm_node import SlurmNode, SlurmNodeState, parse_node_list from .slurm_rest_client import SlurmAPIConfig, SlurmRestClient @@ -45,42 +45,6 @@ class DataRepositoryConfig(BaseModel): verify_certs: bool = True -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. - - Args: - node_list (str): A list of node names, possibly including ranges. - - Returns: - List[str]: A flat list of expanded node names with preserved zeroes. - """ - node_list = node_list.strip() - nodes = [] - if not node_list: - return [] - - components = re.split(r",\s*(?![^[]*\])", node_list) - for component in components: - if "[" not in component: - nodes.append(component) - else: - header, node_number = component.split("[") - node_number = node_number.replace("]", "") - ranges = node_number.split(",") - for r in ranges: - if "-" in r: - start_node, end_node = r.split("-") - number_of_digits = len(end_node) - nodes.extend( - [f"{header}{str(i).zfill(number_of_digits)}" for i in range(int(start_node), int(end_node) + 1)] - ) - else: - nodes.append(f"{header}{r}") - - return nodes - - class SlurmGroup(BaseModel): """Represents a group of nodes within a partition.""" @@ -184,74 +148,6 @@ def _job_id(job: BaseJob) -> 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: - """Choose significant state from REST state flags; e.g. `IDLE+DRAIN` resolves to `DRAINED`.""" - raw_states = node.get("state", node.get("job_state")) - if isinstance(raw_states, dict): - raw_states = raw_states.get("current", []) - if isinstance(raw_states, list): - state_names = [str(state) for state in raw_states] - elif raw_states is None: - state_names = [] - else: - state_names = [state for state in re.split(r"[,+]", str(raw_states)) if state] - - states = [self.convert_state_to_enum(state.upper().rstrip("+")) for state in state_names] - 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) - - def _nodes_from_rest(self) -> list[SlurmNode]: - nodes: list[SlurmNode] = [] - for node in self._rest_client.cluster_nodes(): - if not node.get("name"): - continue - - partitions = node.get("partitions") - if isinstance(partitions, dict): - partitions = partitions.get("current", []) - if isinstance(partitions, list): - partition_names = [str(partition) for partition in partitions] - elif partitions is None: - partition_names = [] - else: - partition_names = [partition for partition in re.split(r"[,+]", str(partitions)) if partition] - - state = self._rest_node_state(node) - nodes.extend( - SlurmNode(name=str(node["name"]), partition=partition, state=state) for partition in partition_names - ) - return nodes - - def _allocated_nodes_from_rest(self) -> list[SlurmNode]: - nodes: list[SlurmNode] = [] - for job in self._rest_client.queue_jobs(): - raw_states = job.get("state", job.get("job_state")) - if isinstance(raw_states, dict): - raw_states = raw_states.get("current", []) - if isinstance(raw_states, list): - states = {str(state).upper().rstrip("+") for state in raw_states} - elif raw_states is None: - states = set() - else: - states = {state.upper().rstrip("+") for state in re.split(r"[,+]", str(raw_states)) if state} - if not {"RUNNING", "PENDING"}.intersection(states): - continue - - partition = str(job.get("partition", "")) - user = str(job.get("user_name", job.get("user", "N/A"))) - nodes.extend( - SlurmNode(name=name, partition=partition, state=SlurmNodeState.ALLOCATED, user=user) - for name in parse_node_list(str(job.get("nodes", ""))) - ) - return nodes - @property def groups(self) -> Dict[str, Dict[str, List[SlurmNode]]]: groups: Dict[str, Dict[str, List[SlurmNode]]] = {} @@ -322,7 +218,7 @@ def update(self) -> None: def nodes_from_sinfo(self) -> list[SlurmNode]: if self.uses_slurm_api: - return self._nodes_from_rest() + return self._rest_client.get_nodes() sinfo_output, _ = self.fetch_command_output("sinfo --noheader -o '%P|%t|%u|%N'") nodes: list[SlurmNode] = [] @@ -344,7 +240,7 @@ def nodes_from_sinfo(self) -> list[SlurmNode]: def nodes_from_squeue(self) -> list[SlurmNode]: if self.uses_slurm_api: - return self._allocated_nodes_from_rest() + return self._rest_client.get_allocated_nodes() squeue_output, _ = self.fetch_command_output("squeue --states=running,pending --noheader -o '%P|%T|%N|%u'") nodes: list[SlurmNode] = [] diff --git a/tests/systems/slurm/test_system.py b/tests/systems/slurm/test_system.py index 411212163..789217aa9 100644 --- a/tests/systems/slurm/test_system.py +++ b/tests/systems/slurm/test_system.py @@ -59,7 +59,7 @@ def test_slurm_api_request_expands_headers(rest_slurm_system: SlurmSystem, monke response.json.return_value = {"nodes": [], "errors": [], "warnings": []} with patch("cloudai.systems.slurm.slurm_rest_client.requests.request", return_value=response) as request: - rest_slurm_system._rest_client.cluster_nodes() + rest_slurm_system._rest_client.get_nodes() request.assert_called_once_with( "GET", @@ -167,8 +167,19 @@ 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", + "state_flags": ["DRAIN"], + "gres": "gpu:8", + }, + { + "name": "node02", + "partitions": ["main", "backup"], + "state": "allocated", + "state_flags": [], + }, ] } jobs_response = { From 7610af2437e5ee5159c2c96440e0d114721f1d45 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 14 Sep 2026 16:18:58 +0200 Subject: [PATCH 14/16] Preserve Slurm node list parser --- src/cloudai/systems/slurm/slurm_node.py | 39 +++++++++++++++---------- 1 file changed, 24 insertions(+), 15 deletions(-) diff --git a/src/cloudai/systems/slurm/slurm_node.py b/src/cloudai/systems/slurm/slurm_node.py index 917937c10..d2dade15c 100644 --- a/src/cloudai/systems/slurm/slurm_node.py +++ b/src/cloudai/systems/slurm/slurm_node.py @@ -16,12 +16,21 @@ import re from enum import Enum +from typing import List from pydantic import BaseModel, ConfigDict -def parse_node_list(node_list: str) -> list[str]: - """Expand a Slurm node list such as `node[01-03]` into individual names.""" +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. + + Args: + node_list (str): A list of node names, possibly including ranges. + + Returns: + List[str]: A flat list of expanded node names with preserved zeroes. + """ node_list = node_list.strip() nodes = [] if not node_list: @@ -31,19 +40,19 @@ def parse_node_list(node_list: str) -> list[str]: for component in components: if "[" not in component: nodes.append(component) - continue - - header, node_number = component.split("[") - for node_range in node_number.rstrip("]").split(","): - if "-" not in node_range: - nodes.append(f"{header}{node_range}") - continue - - start_node, end_node = node_range.split("-") - width = len(end_node) - nodes.extend( - f"{header}{node_number:0{width}d}" for node_number in range(int(start_node), int(end_node) + 1) - ) + else: + header, node_number = component.split("[") + node_number = node_number.replace("]", "") + ranges = node_number.split(",") + for r in ranges: + if "-" in r: + start_node, end_node = r.split("-") + number_of_digits = len(end_node) + nodes.extend( + [f"{header}{str(i).zfill(number_of_digits)}" for i in range(int(start_node), int(end_node) + 1)] + ) + else: + nodes.append(f"{header}{r}") return nodes From e8779353f21bd1d957d393813e17af10972687de Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 14 Sep 2026 16:53:43 +0200 Subject: [PATCH 15/16] Correct Slurm REST service requirements --- doc/USER_GUIDE.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/doc/USER_GUIDE.rst b/doc/USER_GUIDE.rst index 627b6eeee..64d916250 100644 --- a/doc/USER_GUIDE.rst +++ b/doc/USER_GUIDE.rst @@ -86,8 +86,8 @@ Field Descriptions Slurm REST API ~~~~~~~~~~~~~~ -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. +CloudAI uses the Slurm 22.05 REST API v0.0.38 when ``slurm_api`` is configured. The ``slurm`` endpoint backed by +``slurmctld`` must be enabled by the service; ``slurmdbd`` is not required. .. code-block:: toml From ef1b1affaa5ce56c9a172175f0cba8831808208a Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 14 Sep 2026 20:42:14 +0200 Subject: [PATCH 16/16] Use canonical Slurm node fields --- src/cloudai/systems/slurm/slurm_rest_client.py | 4 ++-- tests/systems/slurm/test_system.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/cloudai/systems/slurm/slurm_rest_client.py b/src/cloudai/systems/slurm/slurm_rest_client.py index 6b9c6e4d2..b3328456d 100644 --- a/src/cloudai/systems/slurm/slurm_rest_client.py +++ b/src/cloudai/systems/slurm/slurm_rest_client.py @@ -83,8 +83,8 @@ class SlurmRestClient: "-A": "account", "--reservation": "reservation", "--distribution": "distribution", - "--nodelist": "nodelist", - "--exclude": "exclude_nodes", + "--nodelist": "required_nodes", + "--exclude": "excluded_nodes", "--chdir": "current_working_directory", } _SHORT_DIRECTIVES: ClassVar[str] = "J:o:e:p:A:N:n:D:" diff --git a/tests/systems/slurm/test_system.py b/tests/systems/slurm/test_system.py index 789217aa9..175f8c5c0 100644 --- a/tests/systems/slurm/test_system.py +++ b/tests/systems/slurm/test_system.py @@ -108,8 +108,8 @@ def test_submit_job_through_slurm_api(rest_slurm_system: SlurmSystem, tmp_path: "reservation": "nightly", "distribution": "block", "nodes": [2, 2], - "nodelist": "node[01-02]", - "exclude_nodes": "node03,node04", + "required_nodes": "node[01-02]", + "excluded_nodes": "node03,node04", "gres": "gpu:8", "tasks_per_node": 8, "time_limit": 21,