[Slurm] Rest API

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}"

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
 ---------------

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'",

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",

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.