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..627b6eeee 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. 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..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..755feb33d 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"); @@ -20,11 +20,13 @@ 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 SlurmGroup, SlurmPartition, SlurmSystem, parse_node_list __all__ = [ "SingleSbatchRunner", + "SlurmAPIConfig", "SlurmCommandGenStrategy", "SlurmGroup", "SlurmInstaller", 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 f27367e12..c6e799e6e 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -35,6 +35,7 @@ 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): @@ -102,6 +103,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 +129,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 +165,38 @@ 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 + + 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)] + 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) + @property def groups(self) -> Dict[str, Dict[str, List[SlurmNode]]]: groups: Dict[str, Dict[str, List[SlurmNode]]] = {} @@ -186,6 +225,19 @@ 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: + 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 nodes 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 +270,16 @@ 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: + 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 + 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 +299,24 @@ def nodes_from_sinfo(self) -> list[SlurmNode]: return nodes 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 + 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 +374,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 +411,15 @@ 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: + 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 +450,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 "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" @@ -388,6 +503,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 +542,10 @@ 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: + 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 = ( f"sacct -j {job.id} --format=JobID,JobName,State,ExitCode,Start,End,ElapsedRAW,SubmitLine " @@ -457,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: @@ -706,6 +827,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 +997,15 @@ 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: + 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") + 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..f5549e9c7 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,204 @@ 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 = {"pings": [], "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/") + + request.assert_called_once_with( + "GET", + "https://slurm.example.com/slurm/v0.0.38/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 --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], + "nodelist": "node[01-02]", + "exclude_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", + "state": {"current": "COMPLETED"}, + "exit_code": {"return_code": 0, "signal": {"signal_id": 0}}, + "time": {"start": 100, "end": 120, "elapsed": 20}, + "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): + 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(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()] == [ + ("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", "clusters/"), + ] + + @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"