Add Slurm backend

mkdocs/docs/concepts/backends.md

@@ -1281,6 +1281,109 @@ projects:
 
 > To learn more, see the [Lambda](../examples/clusters/lambda/#kubernetes) and [Crusoe](../examples/clusters/crusoe/#kubernetes) examples.
 
+### Slurm
+
+`dstack` can orchestrate container-based runs across your [Slurm](https://slurm.schedmd.com/) clusters. A single `slurm` backend can manage one or many clusters — `dstack` connects to each cluster's login node over SSH and submits runs as Slurm jobs. Each cluster becomes its own `dstack` region, named after the cluster.
+
+<div editor-title="~/.dstack/server/config.yml">
+
+```yaml
+projects:
+- name: main
+  backends:
+  - type: slurm
+
+    clusters:
+    - name: gpu-cluster-a
+      hostname: login.example.com
+      user: admin
+      private_key:
+        path: ~/.ssh/id_rsa
+
+      gpu_partitions:
+      - gpu: H100
+        partitions: [gpu]
+```
+
+</div>
+
+`dstack` logs in to `hostname` as `user` using `private_key`, and uses this login node both to submit jobs and as an SSH jump host to reach the containers — no additional setup is required.
+
+!!! info "Prerequisites"
+    `dstack` runs each job as a Slurm batch job that launches the run's container via the [Pyxis](https://github.com/NVIDIA/pyxis) SPANK plugin and the [enroot](https://github.com/NVIDIA/enroot) container runtime. Both must be installed and configured on the cluster's compute nodes.
+
+!!! info "Partitions"
+    `dstack` selects which Slurm partitions to submit to based on whether a run requests GPUs:
+
+    - `gpu_partitions` maps a GPU model, in the `[vendor:]name[:memory]` format (e.g. `H100`, `A100:40GB`, `MI300X`), to the partitions that provide it. Only partitions listed here are used for GPU runs. If `gpu_partitions` is not set, GPU runs are not allowed.
+    - `cpu_partitions` lists the partitions used for CPU-only runs. If not set, it defaults to all cluster partitions except those listed in `gpu_partitions`.
+
+    ```yaml
+    clusters:
+    - name: gpu-cluster-a
+      hostname: login.example.com
+      user: admin
+      private_key:
+        path: ~/.ssh/id_rsa
+      gpu_partitions:
+      - gpu: H100
+        partitions: [gpu-h100]
+      - gpu: A100:40GB
+        partitions: [gpu-a100]
+      cpu_partitions: [cpu]
+    ```
+
+    Partitions are exposed as availability zones, so a run can target specific partitions via the `availability_zones` property in its configuration:
+
+    ```yaml
+    type: task
+    availability_zones: [gpu-h100]
+    ```
+
+    The partitions selected by default (the `gpu_partitions` for a GPU run, or the `cpu_partitions` for a CPU-only run) act as a bounding set: `availability_zones` can only narrow this set, not extend it — a partition outside it is never used. When `availability_zones` is not set, all default-selected partitions are eligible.
+
+!!! info "GPU model syntax"
+    The `gpu` field accepts the `[vendor:]name[:memory]` format. In most cases the name alone is enough: `dstack` matches it against its list of known GPU models, inferring the vendor and memory size and normalizing the name so it stays consistent with other backends (e.g. `RTX 4090` becomes `RTX4090`).
+
+    Specify the memory size only when a model comes in several memory variants — for example `A100:40GB` versus `A100:80GB` — otherwise the name is ambiguous and configuration fails.
+
+    The full `vendor:name:memory` form (e.g. `nvidia:A100:80GB`) bypasses this matching and normalization, so it can describe models `dstack` doesn't know about. It's used verbatim and requires all three fields. Prefer the short form where possible — it's easier and keeps GPU names consistent across backends.
+
+!!! info "Region"
+    Each enabled cluster becomes its own `dstack` region, named after the cluster's `name`. Use the `region` field in fleet or run configurations to target a specific cluster.
+
+??? info "User interface"
+    If you are configuring the `slurm` backend on the [project settings page](projects.md#backends), specify the contents of the private key in `content` instead of referencing a `path`:
+
+    <div editor-title="~/.dstack/server/config.yml">
+
+    ```yaml
+    type: slurm
+
+    clusters:
+    - name: gpu-cluster-a
+      hostname: login.example.com
+      user: admin
+      private_key:
+        content: |
+          -----BEGIN OPENSSH PRIVATE KEY-----
+          ...
+          -----END OPENSSH PRIVATE KEY-----
+      gpu_partitions:
+      - gpu: H100
+        partitions: [gpu]
+    ```
+
+    </div>
+
+??? info "Resources and offers"
+    Like the `kubernetes` backend, if you use ranges with [`resources`](../concepts/tasks.md#resources) (e.g. `gpu: 1..8` or `memory: 64GB..`) in fleet or run configurations, the `slurm` backend only requests the lower limit of each range — the upper limit is ignored. The requested CPU, memory, and GPU counts are passed to Slurm as `--cpus-per-task`, `--mem`, and `--gres=gpu:<count>` respectively.
+
+!!! warning "Known limitations"
+    - **One `dstack` job per node.** A `dstack` job can share a node with non-`dstack` Slurm jobs, but two `dstack` jobs cannot run on the same node — the second one fails to start and its run eventually fails by provisioning timeout. `dstack` does not control node placement, so this can happen whenever the Slurm scheduler places a second `dstack` job on a node that already runs one. This limitation is expected to be lifted in a future release.
+    - **Runs always execute as `root`.** `enroot` is a single-user container runtime: it can only map the invoking host user inside the container to either the same UID/GID or to `root` (`0:0`). `dstack` uses the latter, so runs always execute as `root`, and both the image's default user and the run configuration's `user` property are ignored. The single-user model is a fundamental `enroot` limitation, not specific to the `dstack` integration.
+    - **Private registries via `registry_auth` are not supported.** The `registry_auth` run configuration property is rejected. You can still pull images from a private registry by preconfiguring `enroot` credentials on the cluster's compute nodes — see the [enroot import documentation](https://github.com/NVIDIA/enroot/blob/main/doc/cmd/import.md#description).
+
 ### Runpod
 
 Log into your [Runpod](https://www.runpod.io/console/) console, click Settings in the sidebar, expand the `API Keys` section, and click

mkdocs/docs/reference/server/config.yml.md

@@ -176,6 +176,34 @@ to configure [backends](../../concepts/backends.md) and other [server-level sett
         type:
             required: true
 
+##### `projects[n].backends[type=slurm]` { #slurm data-toc-label="slurm" }
+
+#SCHEMA# dstack._internal.core.backends.slurm.models.SlurmBackendFileConfigWithCreds
+    overrides:
+        show_root_heading: false
+        type:
+            required: true
+        item_id_prefix: slurm-
+
+###### `projects[n].backends[type=slurm].clusters[n]` { #slurm-clusters data-toc-label="clusters" }
+
+#SCHEMA# dstack._internal.core.backends.slurm.models.SlurmClusterFileConfig
+    overrides:
+        show_root_heading: false
+        item_id_prefix: slurm-clusters-
+
+###### `projects[n].backends[type=slurm].clusters[n].gpu_partitions[n]` { #slurm-clusters-gpu_partitions data-toc-label="gpu_partitions" }
+
+#SCHEMA# dstack._internal.core.backends.slurm.models.SlurmGPUPartitionConfig
+    overrides:
+        show_root_heading: false
+
+###### `projects[n].backends[type=slurm].clusters[n].private_key` { #slurm-clusters-private_key data-toc-label="private_key" }
+
+#SCHEMA# dstack._internal.core.backends.slurm.models.SlurmPrivateKeyFileConfig
+    overrides:
+        show_root_heading: false
+
 ##### `projects[n].backends[type=vastai]` { #vastai data-toc-label="vastai" }
 
 #SCHEMA# dstack._internal.core.backends.vastai.models.VastAIBackendConfigWithCreds

pyproject.toml

@@ -32,7 +32,7 @@ dependencies = [
     "python-multipart>=0.0.16",
     "filelock",
     "psutil",
-    "gpuhunt==0.1.25",
+    "gpuhunt==0.1.26",
     "argcomplete>=3.5.0",
     "ignore-python>=0.2.0",
     "orjson",
@@ -104,6 +104,7 @@ include = [
     "src/dstack/_internal/core/backends/aws",
     "src/dstack/_internal/core/backends/kubernetes",
     "src/dstack/_internal/core/backends/runpod",
+    "src/dstack/_internal/core/backends/slurm",
     "src/dstack/_internal/cli/services/configurators",
     "src/dstack/_internal/cli/commands",
     "src/tests/_internal/server/background/pipeline_tasks",

src/dstack/_internal/core/backends/base/compute.py

@@ -982,6 +982,7 @@ def get_gateway_user_data(
 def get_docker_commands(
     authorized_keys: list[str],
     bin_path: Optional[PathLike] = None,
+    pre_runner_commands: Optional[Iterable[str]] = None,
 ) -> list[str]:
     dstack_runner_binary_path = get_dstack_runner_binary_path(bin_path)
     commands = [
@@ -1002,6 +1003,9 @@ def get_docker_commands(
         ": )",
     ]
 
+    if pre_runner_commands is not None:
+        commands.extend(pre_runner_commands)
+
     runner_command = [
         dstack_runner_binary_path,
         "--log-level",

src/dstack/_internal/core/backends/base/offers.py

@@ -1,8 +1,11 @@
+from abc import ABC, abstractmethod
 from collections.abc import Iterable, Iterator
 from dataclasses import asdict
-from typing import Callable, List, Optional, TypeVar
+from typing import Callable, Generic, List, Literal, Optional, TypeVar
+from uuid import UUID
 
 import gpuhunt
+from cachetools import TTLCache
 from pydantic import parse_obj_as
 
 from dstack._internal.core.models.backends.base import BackendType
@@ -14,8 +17,8 @@
     InstanceType,
     Resources,
 )
-from dstack._internal.core.models.resources import DEFAULT_DISK, CPUSpec, Memory, Range
-from dstack._internal.core.models.runs import Requirements
+from dstack._internal.core.models.resources import DEFAULT_DISK, CPUSpec, GPUSpec, Memory, Range
+from dstack._internal.core.models.runs import Job, Requirements, Run
 from dstack._internal.utils.common import get_or_error
 
 # Offers not supported by all dstack versions are hidden behind one or more flags.
@@ -31,6 +34,14 @@
 ]
 
 
+# NvidiaGPUInfo.name in KNOWN_NVIDIA_GPUS is not unique -- there are multiple variants
+# with the same name but different amount of memory, but Compute Capability of the variants
+# is always the same, so it's safe to use 1:1 mapping
+NVIDIA_GPU_NAME_TO_COMPUTE_CAPABILITY_MAP = {
+    gpu_info.name: gpu_info.compute_capability for gpu_info in gpuhunt.KNOWN_NVIDIA_GPUS
+}
+
+
 def get_catalog_offers(
     backend: BackendType,
     locations: Optional[List[str]] = None,
@@ -241,3 +252,68 @@ def modifier(offer: InstanceOfferWithAvailability) -> Optional[InstanceOfferWith
         return offer_copy
 
     return modifier
+
+
+def gpu_matches_gpu_spec(gpu: Gpu, gpu_spec: GPUSpec) -> bool:
+    if gpu_spec.vendor is not None and gpu.vendor != gpu_spec.vendor:
+        return False
+    if gpu_spec.name is not None and gpu.name.lower() not in map(str.lower, gpu_spec.name):
+        return False
+    if gpu_spec.memory is not None:
+        min_memory_gib = gpu_spec.memory.min
+        if min_memory_gib is not None and gpu.memory_mib < min_memory_gib * 1024:
+            return False
+        max_memory_gib = gpu_spec.memory.max
+        if max_memory_gib is not None and gpu.memory_mib > max_memory_gib * 1024:
+            return False
+    if gpu_spec.compute_capability is not None:
+        if gpu.vendor != gpuhunt.AcceleratorVendor.NVIDIA:
+            return False
+        compute_capability = NVIDIA_GPU_NAME_TO_COMPUTE_CAPABILITY_MAP.get(gpu.name)
+        if compute_capability is None:
+            return False
+        if compute_capability < gpu_spec.compute_capability:
+            return False
+    return True
+
+
+OfferKeyT = TypeVar("OfferKeyT")
+
+
+class BaseSkipOfferCache(Generic[OfferKeyT], ABC):
+    """
+    A base class for a cache to track (run/job, offer) pairs that failed to provision.
+
+    Implementations can be used to skip offers based on, e.g., a region, an instance type,
+    a region/type pair, etc.
+
+    Subclasses must implement `_build_key()`.
+    """
+
+    def __init__(self, *, ttl: int, maxsize: int = 1000) -> None:
+        self._cache = TTLCache[OfferKeyT, Literal[True]](maxsize=maxsize, ttl=ttl)
+
+    def add(self, run: Run, job: Job, offer: InstanceOffer) -> None:
+        self._cache[self._build_key(run, job, offer)] = True
+
+    def check(self, run: Run, job: Job, offer: InstanceOffer) -> bool:
+        return self._build_key(run, job, offer) in self._cache
+
+    @abstractmethod
+    def _build_key(self, run: Run, job: Job, offer: InstanceOffer) -> OfferKeyT:
+        pass
+
+
+class RegionalSkipOfferCache(BaseSkipOfferCache[tuple[UUID, str]]):
+    """
+    `RegionalSkipOfferRegionCache` tracks failed provisioning attempts based on the offer's region.
+
+    The current implementation tracks _any_ job of the specific run (identified by `Run.id`)
+    in the specific region (identified by `InstanceOffer.region`).
+    """
+
+    def _build_key(self, run: Run, job: Job, offer: InstanceOffer) -> tuple[UUID, str]:
+        # The current implementation uses only Run.id ignoring the job/job spec.
+        # A more sophisticated implementation could use some parts of the job spec
+        # (e.g., requirements, volumes) instead.
+        return (run.id, offer.region)

src/dstack/_internal/core/backends/configurators.py

@@ -137,6 +137,12 @@
 except ImportError:
     pass
 
+try:
+    from dstack._internal.core.backends.slurm.configurator import SlurmConfigurator
+
+    _CONFIGURATOR_CLASSES.append(SlurmConfigurator)
+except ImportError:
+    pass
 
 try:
     from dstack._internal.core.backends.vastai.configurator import VastAIConfigurator

src/dstack/_internal/core/backends/kubernetes/compute.py

@@ -29,6 +29,7 @@
     get_dstack_gateway_commands,
     merge_tags,
 )
+from dstack._internal.core.backends.base.offers import RegionalSkipOfferCache, gpu_matches_gpu_spec
 from dstack._internal.core.backends.kubernetes.api_client import API_CLIENT_EXCEPTIONS
 from dstack._internal.core.backends.kubernetes.models import KubernetesConfig
 from dstack._internal.core.backends.kubernetes.resources import (
@@ -37,7 +38,6 @@
     AMD_GPU_NODE_TAINT,
     AMD_GPU_RESOURCE,
     LABEL_VALUE_MAX_LENGTH,
-    NVIDIA_GPU_NAME_TO_GPU_INFO,
     NVIDIA_GPU_NODE_TAINT,
     NVIDIA_GPU_PRODUCT_LABEL,
     NVIDIA_GPU_RESOURCE,
@@ -62,7 +62,6 @@
 from dstack._internal.core.backends.kubernetes.utils import (
     LEGACY_CURRENT_CONTEXT_REGION,
     Cluster,
-    SkipOfferCache,
     call_api_method,
     get_clusters_from_backend_config,
     try_delete_object_if_exists,
@@ -77,7 +76,6 @@
     GatewayProvisioningData,
 )
 from dstack._internal.core.models.instances import (
-    Gpu,
     InstanceOfferWithAvailability,
     SSHConnectionParams,
 )
@@ -138,7 +136,7 @@ class KubernetesCompute(
     def __init__(self, config: KubernetesConfig):
         super().__init__()
         self.region_cluster_map = {c.region: c for c in get_clusters_from_backend_config(config)}
-        self.skip_offer_cache = SkipOfferCache(ttl=60)
+        self.skip_offer_cache = RegionalSkipOfferCache(ttl=60)
 
     def get_offers_by_requirements(
         self, requirements: Requirements
@@ -752,7 +750,7 @@ def _get_nvidia_gpu_node_affinity(
     for node in nodes:
         labels = get_node_labels(node)
         gpu = get_nvidia_gpu_from_node_labels(labels)
-        if gpu is not None and _gpu_matches_gpu_spec(gpu, gpu_spec):
+        if gpu is not None and gpu_matches_gpu_spec(gpu, gpu_spec):
             matching_gpu_label_values.add(labels[NVIDIA_GPU_PRODUCT_LABEL])
     if not matching_gpu_label_values:
         raise ComputeError(
@@ -785,7 +783,7 @@ def _get_amd_gpu_node_affinity(
     for node in nodes:
         labels = get_node_labels(node)
         gpu = get_amd_gpu_from_node_labels(labels)
-        if gpu is not None and _gpu_matches_gpu_spec(gpu, gpu_spec):
+        if gpu is not None and gpu_matches_gpu_spec(gpu, gpu_spec):
             matching_device_ids.update(AMD_GPU_NAME_TO_DEVICE_IDS[gpu.name])
     return client.V1NodeAffinity(
         required_during_scheduling_ignored_during_execution=client.V1NodeSelector(
@@ -804,29 +802,6 @@ def _get_amd_gpu_node_affinity(
     )
 
 
-def _gpu_matches_gpu_spec(gpu: Gpu, gpu_spec: GPUSpec) -> bool:
-    if gpu_spec.vendor is not None and gpu.vendor != gpu_spec.vendor:
-        return False
-    if gpu_spec.name is not None and gpu.name.lower() not in map(str.lower, gpu_spec.name):
-        return False
-    if gpu_spec.memory is not None:
-        min_memory_gib = gpu_spec.memory.min
-        if min_memory_gib is not None and gpu.memory_mib < min_memory_gib * 1024:
-            return False
-        max_memory_gib = gpu_spec.memory.max
-        if max_memory_gib is not None and gpu.memory_mib > max_memory_gib * 1024:
-            return False
-    if gpu_spec.compute_capability is not None:
-        if gpu.vendor != AcceleratorVendor.NVIDIA:
-            return False
-        gpu_info = NVIDIA_GPU_NAME_TO_GPU_INFO.get(gpu.name)
-        if gpu_info is None:
-            return False
-        if gpu_info.compute_capability < gpu_spec.compute_capability:
-            return False
-    return True
-
-
 def _create_jump_pod_service_if_not_exists(
     api: client.CoreV1Api,
     namespace: str,