diff --git a/conf/common/system/example_slurm_cluster.toml b/conf/common/system/example_slurm_cluster.toml index a66815c81..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"); @@ -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..64d916250 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 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 + + [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..67e7d19cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,8 @@ dependencies = [ "jinja2~=3.1.6", "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 1a92a318a..da876d60c 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"); @@ -19,12 +19,14 @@ 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", + "SlurmAPIConfig", "SlurmCommandGenStrategy", "SlurmGroup", "SlurmInstaller", diff --git a/src/cloudai/systems/slurm/slurm_node.py b/src/cloudai/systems/slurm/slurm_node.py index cca703231..d2dade15c 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,49 @@ # See the License for the specific language governing permissions and # limitations under the License. +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 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 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 new file mode 100644 index 000000000..b3328456d --- /dev/null +++ b/src/cloudai/systems/slurm/slurm_rest_client.py @@ -0,0 +1,437 @@ +# 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 datetime +import getopt +import logging +import math +import os +import pathlib +import re +import shlex +import time +from typing import Any, ClassVar, cast + +import pydantic +import requests +import tenacity + +import cloudai.core +import cloudai.util + +from .slurm_metadata import SlurmStepMetadata +from .slurm_node import SlurmNode, SlurmNodeState, parse_node_list + +logger = logging.getLogger(__name__) + + +class SlurmAPIConfig(pydantic.BaseModel): + """Connection details for a Slurm REST API endpoint.""" + + model_config = pydantic.ConfigDict(extra="forbid") + + url: str + headers: dict[str, str] = pydantic.Field(default_factory=dict) + verify_certs: bool = True + + +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", + "--nodelist": "required_nodes", + "--exclude": "excluded_nodes", + "--chdir": "current_working_directory", + } + _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 + 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]: + url = f"{self._config.url.rstrip('/')}/{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 = 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: + 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 + + @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 _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.""" + existing = job.get(field) + if existing is not None and existing != value: + 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_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) + + 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 == "--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 = value if option == "--gres" else f"gpu:{value}" + self._set_directive(job, "gres", gres, 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.""" + 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 + 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")} + return job + + 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 job state.""" + try: + script = script_path.read_text(encoding="utf-8") + data = self._request( + "POST", + "slurm", + "job/submit", + payload={"script": script, "job": self._make_job(script, script_path)}, + ) + except (OSError, RuntimeError, ValueError) as exc: + raise cloudai.core.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 cloudai.core.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 + + 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 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") + ) + + 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 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") + if not isinstance(jobs, list): + raise RuntimeError("Slurm API returned an invalid jobs response.") + 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.""" + job = self._get_job(job_id, retry_threshold) + if job is None: + return "" + + 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.""" + 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) + if job is None: + return [] + + raw_exit_code = job["exit_code"] + return_code = 0 + signal = 0 + 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) + + 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 + 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 [ + SlurmStepMetadata( + job_id=job["job_id"], + step_id="", + 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=job.get("command") or "", + ) + ] + + def get_job_nodes(self, job_id: int) -> str: + job = self._get_job(job_id) + if job is None: + return "" + return job.get("nodes") or "" + + 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 endpoints used by CloudAI.""" + self._request("GET", "slurm", "ping/") diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index f27367e12..06151f2e4 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -34,7 +34,8 @@ 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 class DataRepositoryConfig(BaseModel): @@ -44,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.""" @@ -102,6 +67,11 @@ 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._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))}" return self.submit_job(command, operation_name) @@ -123,6 +93,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 @@ -158,6 +129,25 @@ 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 + + @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 _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 + @property def groups(self) -> Dict[str, Dict[str, List[SlurmNode]]]: groups: Dict[str, Dict[str, List[SlurmNode]]] = {} @@ -186,6 +176,15 @@ 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: + 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 + 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 +217,9 @@ 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: + return self._rest_client.get_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 +239,9 @@ def nodes_from_sinfo(self) -> list[SlurmNode]: return nodes def nodes_from_squeue(self) -> list[SlurmNode]: + if self.uses_slurm_api: + 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] = [] for line in squeue_output.split("\n"): @@ -294,6 +299,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(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 +336,13 @@ 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: + try: + self._rest_client.validate() + 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 +373,9 @@ 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: + 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" @@ -388,6 +426,9 @@ 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: + 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" @@ -424,6 +465,9 @@ 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: + return self._rest_client.get_job_status(self._job_id(job), retry_threshold) + retry_count = 0 command = ( f"sacct -j {job.id} --format=JobID,JobName,State,ExitCode,Start,End,ElapsedRAW,SubmitLine " @@ -457,8 +501,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: @@ -706,6 +749,12 @@ 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_client.cancel(job_id) + return self.cmd_shell.execute(f"scancel {job_id}") def fetch_command_output(self, command: str) -> Tuple[str, str]: @@ -870,8 +919,14 @@ 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 job.id == 0: + return [] + + if self.uses_slurm_api: + 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 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..175f8c5c0 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,16 +26,203 @@ from cloudai.core import BaseJob, JobIdRetrievalError, TestRun from cloudai.models.scenario import ReportConfig from cloudai.systems.slurm import ( + SlurmAPIConfig, SlurmCommandGenStrategy, + SlurmJob, SlurmNode, SlurmNodeState, SlurmSystem, 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 +@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 = {"nodes": [], "errors": [], "warnings": []} + + with patch("cloudai.systems.slurm.slurm_rest_client.requests.request", return_value=response) as request: + rest_slurm_system._rest_client.get_nodes() + + request.assert_called_once_with( + "GET", + "https://slurm.example.com/slurm/v0.0.38/nodes/", + 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 --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 --nodelist=node[01-02] +#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(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 + 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, 2], + "required_nodes": "node[01-02]", + "excluded_nodes": "node03,node04", + "gres": "gpu:8", + "tasks_per_node": 8, + "time_limit": 21, + "current_working_directory": str(tmp_path), + "environment": {"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(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") + + 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", + "job_state": "FAILED", + "exit_code": 256, + "start_time": 100, + "end_time": 120, + "nodes": "node[01-02]", + } + ] + } + job = SlurmJob(test_run=Mock(), id=42) + + 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 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 + + +def test_slurm_api_nodes_cancel_and_validation(rest_slurm_system: SlurmSystem): + nodes_response = { + "nodes": [ + { + "name": "node01", + "partitions": ["main"], + "state": "idle", + "state_flags": ["DRAIN"], + "gres": "gpu:8", + }, + { + "name": "node02", + "partitions": ["main", "backup"], + "state": "allocated", + "state_flags": [], + }, + ] + } + 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(SlurmRestClient, "_request", side_effect=request) as rest_request, + ): + 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[-2:] == [ + call("DELETE", "slurm", "job/42"), + call("GET", "slurm", "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..0e57f6efa 100644 --- a/uv.lock +++ b/uv.lock @@ -278,8 +278,10 @@ dependencies = [ { name = "pandas" }, { name = "pydantic" }, { name = "pyyaml" }, + { name = "requests" }, { name = "rich" }, { name = "tbparse" }, + { name = "tenacity" }, { name = "toml" }, { name = "websockets" }, ] @@ -346,6 +348,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" }, @@ -359,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" }, @@ -2445,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"