diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 33b904d9b..c291b0dea 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/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. + +
+ +```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] +``` + +
+ +`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`: + +
+ + ```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] + ``` + +
+ +??? 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:` 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 diff --git a/mkdocs/docs/reference/server/config.yml.md b/mkdocs/docs/reference/server/config.yml.md index 217e489b5..fdc3f4c7c 100644 --- a/mkdocs/docs/reference/server/config.yml.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index ca0843893..38c4054ae 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 8d57f73d6..477e9747e 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/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", diff --git a/src/dstack/_internal/core/backends/base/offers.py b/src/dstack/_internal/core/backends/base/offers.py index a7e0239c8..90d463234 100644 --- a/src/dstack/_internal/core/backends/base/offers.py +++ b/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) diff --git a/src/dstack/_internal/core/backends/configurators.py b/src/dstack/_internal/core/backends/configurators.py index cdeac7f60..bcc144cf5 100644 --- a/src/dstack/_internal/core/backends/configurators.py +++ b/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 diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index 6cc7c08a0..012299aa4 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/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, diff --git a/src/dstack/_internal/core/backends/kubernetes/utils.py b/src/dstack/_internal/core/backends/kubernetes/utils.py index a90f750f1..4a39bf564 100644 --- a/src/dstack/_internal/core/backends/kubernetes/utils.py +++ b/src/dstack/_internal/core/backends/kubernetes/utils.py @@ -13,10 +13,8 @@ Union, cast, ) -from uuid import UUID import yaml -from cachetools import TTLCache from kubernetes.client import V1Status, VersionApi from kubernetes.client.exceptions import ApiException from kubernetes.watch import Watch @@ -33,8 +31,6 @@ KubernetesProxyJumpConfig, ) from dstack._internal.core.models.common import CoreModel -from dstack._internal.core.models.instances import InstanceOffer -from dstack._internal.core.models.runs import Job, Run from dstack._internal.utils.logging import get_logger logger = get_logger(__name__) @@ -199,30 +195,6 @@ def kubeconfig_dict_to_kubeconfig(kubeconfig_dict: dict) -> Kubeconfig: return Kubeconfig.__response__.parse_obj(kubeconfig_dict) -class SkipOfferCache: - """ - `SkipOfferCache` is used to track (run/job, offer) pairs that failed to provision. - - The current implementation tracks _any_ job of the specific run (identified by `Run.id`) - on the specific cluster (identified by `InstanceOffer.region`, that is, a kubeconfig context). - """ - - def __init__(self, *, ttl: int, maxsize: int = 1000) -> None: - self._cache = TTLCache[tuple[UUID, str], 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 - - 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) - - def call_api_method( method: Callable[P, T], expected: Union[int, tuple[int, ...], list[int]], diff --git a/src/dstack/_internal/core/backends/models.py b/src/dstack/_internal/core/backends/models.py index c21141378..a7bb8c9ad 100644 --- a/src/dstack/_internal/core/backends/models.py +++ b/src/dstack/_internal/core/backends/models.py @@ -66,6 +66,11 @@ RunpodBackendConfig, RunpodBackendConfigWithCreds, ) +from dstack._internal.core.backends.slurm.models import ( + SlurmBackendConfig, + SlurmBackendConfigWithCreds, + SlurmBackendFileConfigWithCreds, +) from dstack._internal.core.backends.tensordock.models import ( TensorDockBackendConfig, TensorDockBackendConfigWithCreds, @@ -104,6 +109,7 @@ VastAIBackendConfig, VerdaBackendConfig, VultrBackendConfig, + SlurmBackendConfig, DstackBackendConfig, DstackBaseBackendConfig, ] @@ -130,6 +136,7 @@ TensorDockBackendConfigWithCreds, VastAIBackendConfigWithCreds, VultrBackendConfigWithCreds, + SlurmBackendConfigWithCreds, DstackBackendConfig, ] @@ -155,6 +162,7 @@ TensorDockBackendConfigWithCreds, VastAIBackendConfigWithCreds, VultrBackendConfigWithCreds, + SlurmBackendFileConfigWithCreds, ] diff --git a/src/dstack/_internal/core/backends/slurm/__init__.py b/src/dstack/_internal/core/backends/slurm/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/dstack/_internal/core/backends/slurm/backend.py b/src/dstack/_internal/core/backends/slurm/backend.py new file mode 100644 index 000000000..4361c9629 --- /dev/null +++ b/src/dstack/_internal/core/backends/slurm/backend.py @@ -0,0 +1,16 @@ +from dstack._internal.core.backends.base.backend import Backend +from dstack._internal.core.backends.slurm.compute import SlurmCompute +from dstack._internal.core.backends.slurm.models import SlurmConfig +from dstack._internal.core.models.backends.base import BackendType + + +class SlurmBackend(Backend): + TYPE = BackendType.SLURM + COMPUTE_CLASS = SlurmCompute + + def __init__(self, config: SlurmConfig): + self.config = config + self._compute = SlurmCompute(self.config) + + def compute(self) -> SlurmCompute: + return self._compute diff --git a/src/dstack/_internal/core/backends/slurm/client.py b/src/dstack/_internal/core/backends/slurm/client.py new file mode 100644 index 000000000..b3ea2740a --- /dev/null +++ b/src/dstack/_internal/core/backends/slurm/client.py @@ -0,0 +1,333 @@ +import dataclasses +import ipaddress +import re +import shlex +import uuid +from types import TracebackType +from typing import Optional + +import paramiko +from typing_extensions import Self + +from dstack._internal.core.backends.slurm.resources import Node, ResolvedNode +from dstack._internal.core.errors import ComputeError +from dstack._internal.utils.ssh import pkey_from_str + +DEFAULT_TIMEOUT = 10 + + +@dataclasses.dataclass +class ExecResult: + exit_status: int + stdout: bytes + stderr: bytes + + @property + def ok(self) -> bool: + return self.exit_status == 0 + + +class SlurmClientError(ComputeError): + pass + + +class SlurmClient: + _client: Optional[paramiko.SSHClient] = None + + def __init__( + self, + *, + hostname: str, + port: int, + user: str, + private_key: str, + timeout: Optional[float] = None, + ) -> None: + self._hostname = hostname + self._port = port + self._user = user + self._private_key = private_key + self._timeout = timeout or DEFAULT_TIMEOUT + + def __enter__(self) -> Self: + self.connect() + return self + + def __exit__( + self, + exc_type: Optional[type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> None: + self.close() + + def connect(self, *, timeout: Optional[float] = None) -> None: + """ + Connect to the SSH server. No-op if already connected. + """ + if self._client is not None: + return + try: + pkey = pkey_from_str(self._private_key) + except ValueError as e: + raise SlurmClientError(f"Failed to load private key: {e}") from e + self._client = paramiko.SSHClient() + self._client.set_missing_host_key_policy(paramiko.AutoAddPolicy) + try: + self._client.connect( + hostname=self._hostname, + port=self._port, + username=self._user, + pkey=pkey, + timeout=self._get_timeout(timeout), + ) + except (paramiko.SSHException, OSError) as e: + raise SlurmClientError(f"Failed to connect: {e}") from e + + def close(self) -> None: + """ + Disconnect from the SSH server. No-op if not connected. + """ + if self._client is None: + return + self._client.close() + self._client = None + + def ping(self) -> None: + res = self.exec("scontrol ping") + if not res.ok: + raise SlurmClientError(f"Failed to ping: {res}") + + def get_partitions(self) -> list[str]: + res = self.exec("sinfo -h -o '%P'") + if not res.ok: + raise SlurmClientError(f"Failed to get partitions: {res}") + return [ + part.rstrip("*") for line in res.stdout.decode().splitlines() if (part := line.strip()) + ] + + def get_nodes(self) -> list[Node]: + res = self.exec("scontrol show -o node") + if not res.ok: + raise SlurmClientError(f"Failed to get nodes: {res}") + nodes: list[Node] = [] + for node_line in filter(None, res.stdout.decode().splitlines()): + node_dict = _parse_scontrol_show_line(node_line) + try: + node = _build_node_from_dict(node_dict) + except ValueError as e: + raise SlurmClientError(f"Failed to parse node: {e}") from e + nodes.append(node) + return nodes + + def submit_batch_script(self, batch_script: str) -> str: + script = f"sbatch --parsable << 'EOF'\n{batch_script}\nEOF" + res = self.exec(script) + if not res.ok: + raise SlurmClientError(f"Failed to submit batch script: {res}") + output = res.stdout.decode().strip() + if not output: + raise SlurmClientError("Failed to submit batch script: sbatch output is empty") + # > --parsable + # > Outputs only the job ID number and the cluster name if present. + # > The values are separated by a semicolon. + job_id, _, _ = output.partition(";") + job_id = job_id.strip() + if not job_id: + raise SlurmClientError( + f"Failed to submit batch script: unexpected sbatch output: {output!r}" + ) + return job_id + + def get_job_state(self, job_id: str) -> Optional[str]: + res = self.exec(f"squeue -j {job_id} --only-job-state -h -o '%T'") + if not res.ok: + raise SlurmClientError(f"Failed to queue job state: {res}") + if state := res.stdout.decode().strip(): + return state + return None + + def get_job_partition(self, job_id: str) -> str: + res = self.exec(f"squeue -j {job_id} -h -o '%P'") + if not res.ok: + raise SlurmClientError(f"Failed to queue job partition: {res}") + return res.stdout.decode().strip() + + def get_job_nodes(self, job_id: str) -> list[ResolvedNode]: + res = self.exec(f""" + set -eu + nodelist=$(squeue -j {job_id} -h -o '%N') + scontrol show -o node="$nodelist" + """) + if not res.ok: + raise SlurmClientError(f"Failed to get job nodes: {res}") + + nodes: list[ResolvedNode] = [] + nodes_to_resolve: list[tuple[Node, str]] = [] + for node_line in filter(None, res.stdout.decode().splitlines()): + node_dict = _parse_scontrol_show_line(node_line) + try: + node = _build_node_from_dict(node_dict) + except ValueError as e: + raise SlurmClientError(f"Failed to parse node: {e}") from e + + addresses: list[str] = [] + if node_addr := node_dict.get("nodeaddr"): + addresses.append(node_addr) + if node_hostname := node_dict.get("nodehostname"): + addresses.append(node_hostname) + hostnames: list[str] = [] + ips: list[str] = [] + for address in addresses: + try: + ips.append(str(ipaddress.ip_address(address))) + except ValueError: + hostnames.append(address) + hostname = next(iter(hostnames), None) + ip = next(iter(ips), None) + if ip is None: + if hostname is None: + raise SlurmClientError(f"Failed to get hostname/IP: {node_line!r}") + nodes_to_resolve.append((node, hostname)) + else: + hostname = hostname or ip + nodes.append(ResolvedNode(**dataclasses.asdict(node), hostname=hostname, ip=ip)) + + if nodes_to_resolve: + # getent prints one line per address; keep only the first so that each hostname maps + # to exactly one output line, or '!' if resolution fails + res = self.exec(f""" + set -eu + for hostname in {shlex.join(hostname for _, hostname in nodes_to_resolve)}; do + if entry=$(getent hosts "$hostname"); then + echo "$entry" | head -n 1 + else + echo '!' + fi + done + """) + if not res.ok: + raise SlurmClientError(f"Failed to resolve IPs: {res}") + host_lines = res.stdout.decode().strip().splitlines() + if len(host_lines) != len(nodes_to_resolve): + raise SlurmClientError(f"Failed to resolve IPs: unexpected output: {res}") + for (node, hostname), host_line in zip(nodes_to_resolve, host_lines): + if host_line.startswith("!"): + raise SlurmClientError(f"Failed to resolve hostname {hostname} to IP: {res}") + ip = next(iter(host_line.split()), None) + if ip is None: + raise SlurmClientError(f"Failed to resolve hostname {hostname} to IP: {res}") + nodes.append(ResolvedNode(**dataclasses.asdict(node), hostname=hostname, ip=ip)) + + nodes.sort(key=lambda n: n.name) + return nodes + + def cancel_job(self, job_id: str) -> None: + res = self.exec(f"scancel {job_id}") + if not res.ok: + raise SlurmClientError(f"Failed to cancel job: {res}") + + def exec(self, script: str, *, timeout: Optional[float] = None) -> ExecResult: + """ + Execute a shell script using the user's login shell. + The client must be already connected. + """ + if self._client is None: + raise SlurmClientError("Not connected") + # boundary is used to strip pam's and/or shell's MOTD (or any other messages) + boundary = f"__dstack_boundary_{uuid.uuid4().hex}__" + _script = f"echo {boundary}; echo {boundary} >&2\n{script}\n" + try: + exit_status, stdout, stderr = self._exec(_script, self._get_timeout(timeout)) + except (paramiko.SSHException, OSError) as e: + raise SlurmClientError(f"Failed to exec {script!r}: {e}") from e + stdout = _strip_login_output(stdout, boundary) + stderr = _strip_login_output(stderr, boundary) + return ExecResult( + exit_status=exit_status, + stdout=stdout, + stderr=stderr, + ) + + def _exec(self, script: str, timeout: float) -> tuple[int, bytes, bytes]: + assert self._client is not None + transport = self._client.get_transport() + assert transport is not None + chan = transport.open_session(timeout=timeout) + # We use Channel.invoke_shell() ("shell" request) instead of Channel.exec_command() ("exec" + # request) to get a login shell as if the user is logged in interactively + try: + chan.settimeout(timeout) + chan.invoke_shell() + chan.sendall(script.encode()) + chan.shutdown_write() + stdout = chan.makefile("r", -1).read() + stderr = chan.makefile_stderr("r", -1).read() + exit_status = chan.recv_exit_status() + finally: + chan.close() + return exit_status, stdout, stderr + + def _get_timeout(self, timeout: Optional[float] = None) -> float: + return timeout or self._timeout + + +def _strip_login_output(output: bytes, boundary: str) -> bytes: + _, sep, rest = output.partition(boundary.encode() + b"\n") + return rest if sep else output + + +# A key starts at line-start or after whitespace: letters/digits/underscores then '=' +_SINFO_SHOW_KEY_REGEX = re.compile(r"(?:^|\s)(?P[A-Za-z_]\w*)=") + + +def _parse_scontrol_show_line(line: str, *, normalize_key: bool = True) -> dict[str, str]: + line = line.strip() + result = {} + matches = list(_SINFO_SHOW_KEY_REGEX.finditer(line)) + for next_index, match in enumerate(matches, 1): + key = match.group("key") + if normalize_key: + key = key.lower() + value_start = match.end() + value_end = matches[next_index].start() if next_index < len(matches) else len(line) + result[key] = line[value_start:value_end].strip() + return result + + +def _build_node_from_dict(node_dict: dict[str, str]) -> Node: + name = node_dict.get("nodename") + if not name: + raise ValueError("Missing node name") + + cpus_raw = node_dict.get("cpuefctv") + if not cpus_raw: + cpus_raw = node_dict.get("cputot") + if not cpus_raw: + raise ValueError("Failed to detect CPU count") + try: + cpus = int(cpus_raw) + except ValueError as e: + raise ValueError(f"Failed to parse CPU count: {e}") from e + + memory_raw = node_dict.get("realmemory") + if not memory_raw: + raise ValueError("Failed to detect memory") + try: + memory_mib = int(memory_raw) + except ValueError as e: + raise ValueError(f"Failed to parse memory: {e}") from e + + return Node( + name=name, + arch=node_dict.get("arch"), + cpus=cpus, + memory_mib=memory_mib, + gres=_split_by_comma(node_dict.get("gres", "")), + partitions=_split_by_comma(node_dict.get("partitions", "")), + ) + + +def _split_by_comma(value: str) -> list[str]: + # NB: empty items are unconditionally removed: "foo , , bar,,," -> ["foo", "bar"] + return [item for _item in value.split(",") if (item := _item.strip())] diff --git a/src/dstack/_internal/core/backends/slurm/cluster.py b/src/dstack/_internal/core/backends/slurm/cluster.py new file mode 100644 index 000000000..ecac2fe7f --- /dev/null +++ b/src/dstack/_internal/core/backends/slurm/cluster.py @@ -0,0 +1,153 @@ +import contextlib +import threading +import time +from collections import defaultdict +from typing import Optional + +from dstack._internal.core.backends.base.offers import gpu_matches_gpu_spec +from dstack._internal.core.backends.slurm.client import SlurmClient +from dstack._internal.core.backends.slurm.models import ( + SlurmBackendConfigWithCreds, + SlurmClusterConfigWithCreds, +) +from dstack._internal.core.backends.slurm.resources import GPUModel, Node +from dstack._internal.core.models.resources import GPUSpec + + +class SlurmCluster: + NODES_CACHE_TTL = 600 + PARTITIONS_CACHE_TTL = 600 + + def __init__(self, config: SlurmClusterConfigWithCreds) -> None: + self.region = self.name = config.name + self.hostname = config.hostname + self.port = config.port or 22 + self.user = config.user + + self._private_key = config.private_key.content + self._cpu_partitions = config.cpu_partitions + + gpu_model_to_partitions_map: defaultdict[GPUModel, set[str]] = defaultdict(set) + for gpu_partition_config in config.gpu_partitions or []: + gpu_model = GPUModel.from_string(gpu_partition_config.gpu) + for partition in gpu_partition_config.partitions: + gpu_model_to_partitions_map[gpu_model].add(partition) + self._gpu_model_to_partitions_map = dict(gpu_model_to_partitions_map) + + partition_to_gpu_model_map: dict[str, GPUModel] = {} + for gpu_model, partitions in self._gpu_model_to_partitions_map.items(): + for partition in partitions: + other_gpu_model = partition_to_gpu_model_map.get(partition) + if other_gpu_model is not None: + raise ValueError( + f"Multiple GPU models mapped in cluster {self.name!r}" + f" partition {partition!r}: {other_gpu_model}, {gpu_model}" + ) + partition_to_gpu_model_map[partition] = gpu_model + self._partition_to_gpu_model_map = partition_to_gpu_model_map + + now = time.monotonic() + self._nodes_cache: tuple[Node, ...] = () + self._nodes_cache_lock = threading.Lock() + self._nodes_cache_expiry = now + self._partitions_cache: list[str] = [] + self._partitions_cache_lock = threading.Lock() + self._partitions_cache_expiry = now + + def __str__(self) -> str: + return f"(name={self.name} hostname={self.hostname})" + + def __repr__(self) -> str: + return f"{self.__class__.__name__}{self}" + + def get_client(self, *, timeout: Optional[float] = None) -> SlurmClient: + return SlurmClient( + hostname=self.hostname, + port=self.port, + user=self.user, + private_key=self._private_key, + timeout=timeout, + ) + + def get_cpu_partitions(self) -> Optional[set[str]]: + """ + Returns all configured CPU partitions. The result should be filtered against partitions + actually present in the cluster, see `get_discovered_partitions()`. + + If `cpu_partitions` is not set in the config, `None` is returned. + if `cpu_partitions` is set to an empty array, an empty set is returned. + """ + if self._cpu_partitions is None: + return None + return set(self._cpu_partitions) + + def get_gpu_partitions(self) -> set[str]: + """ + Returns all configured GPU partitions. The result should be filtered against partitions + actually present in the cluster, see `get_discovered_partitions()`. + + If `gpu_partitions` is not set or set to an empty array in the config, an empty set + is returned. + """ + return set(self._partition_to_gpu_model_map.keys()) + + def filter_gpu_partitions(self, gpu_spec: GPUSpec) -> set[str]: + """ + Filter configured GPU partitions by GPUSpec. The result should be filtered against + partitions actually present in the cluster, see `get_discovered_partitions()`. + """ + filtered_partitions: set[str] = set() + for gpu_model, partitions in self._gpu_model_to_partitions_map.items(): + if gpu_matches_gpu_spec(gpu_model.to_gpu(), gpu_spec): + filtered_partitions.update(partitions) + return filtered_partitions + + def get_partition_gpu_model(self, partition: str) -> Optional[GPUModel]: + """ + Returns a GPU configured for the given partition, if any. + """ + return self._partition_to_gpu_model_map.get(partition) + + def get_discovered_nodes(self, client: Optional[SlurmClient] = None) -> tuple[Node, ...]: + """ + Returns all nodes discovered in the cluster. The result is cached. + """ + with self._nodes_cache_lock: + now = time.monotonic() + if now >= self._nodes_cache_expiry: + self._nodes_cache = tuple(self._discover_nodes(client)) + self._nodes_cache_expiry = now + self.NODES_CACHE_TTL + return self._nodes_cache + + def get_discovered_partitions(self, client: Optional[SlurmClient] = None) -> set[str]: + """ + Returns all partitions discovered in the cluster. The result is cached. + """ + with self._partitions_cache_lock: + now = time.monotonic() + if now >= self._partitions_cache_expiry: + self._partitions_cache = self._discover_partitions(client) + self._partitions_cache_expiry = now + self.PARTITIONS_CACHE_TTL + return set(self._partitions_cache) + + def _discover_nodes(self, client: Optional[SlurmClient]) -> list[Node]: + with self._get_client_context(client) as client: + client.connect() + return client.get_nodes() + + def _discover_partitions(self, client: Optional[SlurmClient]) -> list[str]: + with self._get_client_context(client) as client: + client.connect() + return client.get_partitions() + + def _get_client_context( + self, client: Optional[SlurmClient] + ) -> contextlib.AbstractContextManager[SlurmClient]: + if client is None: + client = self.get_client() + return contextlib.closing(client) + return contextlib.nullcontext(client) + + +def get_clusters_from_backend_config(config: SlurmBackendConfigWithCreds) -> list[SlurmCluster]: + return [SlurmCluster(c) for c in config.clusters] diff --git a/src/dstack/_internal/core/backends/slurm/compute.py b/src/dstack/_internal/core/backends/slurm/compute.py new file mode 100644 index 000000000..e3389272b --- /dev/null +++ b/src/dstack/_internal/core/backends/slurm/compute.py @@ -0,0 +1,468 @@ +import concurrent.futures +import contextlib +import shlex +import time +from functools import partial +from typing import Optional + +from dstack._internal.core.backends.base.compute import ( + Compute, + ComputeWithAllOffersCached, + ComputeWithGroupProvisioningSupport, + ComputeWithMultinodeSupport, + generate_unique_backend_name, + get_docker_commands, + normalize_arch, +) +from dstack._internal.core.backends.base.models import JobConfiguration +from dstack._internal.core.backends.base.offers import OfferModifier, RegionalSkipOfferCache +from dstack._internal.core.backends.slurm.cluster import ( + SlurmCluster, + get_clusters_from_backend_config, +) +from dstack._internal.core.backends.slurm.models import SlurmConfig +from dstack._internal.core.backends.slurm.resources import ( + GPUModel, + Node, + RequestedResources, + get_requested_resources_from_resources_spec, + parse_gres_gpu_count, +) +from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT +from dstack._internal.core.errors import ComputeError, SkipOffer +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.compute_groups import ComputeGroup, ComputeGroupProvisioningData +from dstack._internal.core.models.instances import ( + Disk, + Gpu, + InstanceAvailability, + InstanceOfferWithAvailability, + InstanceRuntime, + InstanceType, + Resources, + SSHConnectionParams, +) +from dstack._internal.core.models.placement import PlacementGroup +from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run +from dstack._internal.core.models.volumes import Volume +from dstack._internal.utils.docker import is_default_registry, parse_image_name +from dstack._internal.utils.logging import get_logger + +logger = get_logger(__name__) + +# An arbitrarily chosen sane limit; not enforced by Slurm +SLURM_JOB_NAME_MAX_LENGTH = 64 + +SLURM_JOB_OUTPUT = "/dev/null" + +PROVISIONING_TIMEOUT = 60 + +PRE_RUNNER_COMMANDS = [ + # enroot creates the rootfs directory group-writable but OpenSSH in strict mode requires + # all directories in the authorized_keys path not be group or world writable. + "chmod 0755 /", + # enroot is an unprivileged single-user runtime, where the user from the parent user ns + # (the one who starts the container) is mapped to either UID 0 (our case; we use + # Pyxis's --container-remap-root option, which is translated to enroot's --root option) + # or the same UID. + # OpenSSH cannot operate in such an environment (without tweaking build-time options) + # as it requires a separate unprivileged account (SSH_PRIVSEP_USER build-time option, + # sshd by default) for the privsep feature. + # We work around this limitation by replacing the privsep account with root. + "sed -i '/^sshd:/d' /etc/passwd", + "echo 'sshd:x:0:0:privsep:/dev/null:/sbin/nologin' >> /etc/passwd", +] + + +class SlurmCompute( + ComputeWithAllOffersCached, + ComputeWithMultinodeSupport, + ComputeWithGroupProvisioningSupport, + Compute, +): + def __init__(self, config: SlurmConfig): + super().__init__() + self._region_to_cluster_map = { + c.region: c for c in get_clusters_from_backend_config(config) + } + # NB: The current implementation of RegionalSkipOfferCache is suitable despite of + # lack of zones (partitions) support since it tracks unique runs (identified by Run.id), + # not their requirements, thus we only skip offers within the same run (-> the same + # configuration -> the same zones) + self._skip_offer_cache = RegionalSkipOfferCache(ttl=60) + + def get_all_offers_with_availability(self) -> list[InstanceOfferWithAvailability]: + offers: list[InstanceOfferWithAvailability] = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: + future_to_cluster_map: dict[ + concurrent.futures.Future[list[InstanceOfferWithAvailability]], SlurmCluster + ] = {} + for cluster in self._region_to_cluster_map.values(): + future = executor.submit(_get_cluster_offers, cluster) + future_to_cluster_map[future] = cluster + for future in concurrent.futures.as_completed(future_to_cluster_map): + try: + cluster_offers = future.result() + except ComputeError as e: + logger.warning( + "Failed to get offers from cluster %s: %s: %s", + future_to_cluster_map[future], + e.__class__.__name__, + e, + ) + continue + offers.extend(cluster_offers) + return offers + + def get_offers_modifiers(self, requirements: Requirements) -> list[OfferModifier]: + requested_resources = get_requested_resources_from_resources_spec(requirements.resources) + return [partial(self._offer_modifier, requested_resources)] + + def run_job( + self, + run: Run, + job: Job, + instance_offer: InstanceOfferWithAvailability, + project_ssh_public_key: str, + project_ssh_private_key: str, + volumes: list[Volume], + placement_group: Optional[PlacementGroup], + ) -> JobProvisioningData: + compute_provisioning_data = self._run_slurm_job( + run=run, + job=job, + instance_offer=instance_offer, + project_ssh_public_key=project_ssh_public_key, + ) + return compute_provisioning_data.job_provisioning_datas[0] + + def run_jobs( + self, + run: Run, + job_configurations: list[JobConfiguration], + instance_offer: InstanceOfferWithAvailability, + project_ssh_public_key: str, + project_ssh_private_key: str, + placement_group: Optional[PlacementGroup], + ) -> ComputeGroupProvisioningData: + master_job = job_configurations[0].job + return self._run_slurm_job( + run=run, + job=master_job, + instance_offer=instance_offer, + project_ssh_public_key=project_ssh_public_key, + ) + + def terminate_instance( + self, instance_id: str, region: str, backend_data: Optional[str] = None + ): + _, slurm_job_id, _ = _parse_instance_id(instance_id) + with self._get_cluster(region).get_client() as client: + client.cancel_job(slurm_job_id) + + def terminate_compute_group(self, compute_group: ComputeGroup): + region = compute_group.provisioning_data.region + slurm_job_id = compute_group.provisioning_data.compute_group_id + with self._get_cluster(region).get_client() as client: + client.cancel_job(slurm_job_id) + + def _run_slurm_job( + self, + run: Run, + job: Job, + instance_offer: InstanceOfferWithAvailability, + project_ssh_public_key: str, + ) -> ComputeGroupProvisioningData: + if job.job_spec.registry_auth is not None: + self._skip_offer_cache.add(run, job, instance_offer) + raise ComputeError("Private registries are not supported yet") + + region = instance_offer.region + cluster = self._get_cluster(region) + if self._skip_offer_cache.check(run, job, instance_offer): + raise SkipOffer(f"Cluster {cluster} has recently failed to schedule a similar job") + with ( + cluster.get_client() as client, + contextlib.ExitStack() as exit_stack, + ): + assert run.run_spec.run_name is not None + slurm_job_name = generate_unique_backend_name( + resource_name=run.run_spec.run_name, + project_name=run.project_name, + max_length=SLURM_JOB_NAME_MAX_LENGTH, + ) + + assert run.run_spec.ssh_key_pub is not None + authorized_keys = [project_ssh_public_key.strip(), run.run_spec.ssh_key_pub.strip()] + + node_count = job.job_spec.jobs_per_replica + resources_spec = job.job_spec.requirements.resources + requested_resources = get_requested_resources_from_resources_spec(resources_spec) + + partitions = _get_cluster_partitions(cluster, requested_resources) + if requested_resources.gpu_count > 0: + assert resources_spec.gpu is not None + partitions = partitions & cluster.filter_gpu_partitions(resources_spec.gpu) + logger.debug("Matching GPU partitions: %s", partitions) + else: + logger.debug("CPU partitions: %s", partitions) + requested_partitions = run.run_spec.configuration.availability_zones + if requested_partitions is not None: + partitions = partitions & set(requested_partitions) + logger.debug("Filtered partitions: %s", partitions) + if not partitions: + self._skip_offer_cache.add(run, job, instance_offer) + raise ComputeError("No matching partitions found") + + script_commands: list[str] = ["set -eu"] + sbatch_directives = [ + f"#SBATCH --job-name={slurm_job_name}", + f"#SBATCH --output={SLURM_JOB_OUTPUT}", + "#SBATCH --no-requeue", + f"#SBATCH --partition={','.join(partitions)}", + f"#SBATCH --nodes={node_count}", + "#SBATCH --ntasks-per-node=1", + f"#SBATCH --cpus-per-task={requested_resources.cpu_count}", + f"#SBATCH --mem={requested_resources.memory_mib}M", + ] + srun_options = [ + f"--cpus-per-task={requested_resources.cpu_count}", + f"--container-image={_build_image_uri(job.job_spec.image_name)}", + f"--container-name={slurm_job_name}", + "--container-remap-root", + "--container-writable", + "--no-container-mount-home", + ] + if requested_resources.gpu_count > 0: + sbatch_directives.append(f"#SBATCH --gres=gpu:{requested_resources.gpu_count}") + else: + # Force skip enroot's nvidia hook on CPU nodes even if NVIDIA_VISIBLE_DEVICES + # is set in the image to avoid failure: + # > nvidia-container-cli: initialization error: nvml error: driver not loaded + # > [ERROR] /etc/enroot/hooks.d/98-nvidia.sh exited with return code 1 + script_commands.append("export NVIDIA_VISIBLE_DEVICES=void") + srun_options.append("--container-env=NVIDIA_VISIBLE_DEVICES") + dstack_commands = get_docker_commands( + authorized_keys=authorized_keys, pre_runner_commands=PRE_RUNNER_COMMANDS + ) + srun_command = ["srun"] + srun_options + ["sh", "-c", " && ".join(dstack_commands)] + script_commands.append(shlex.join(srun_command)) + script_lines: list[str] = ["#!/bin/sh"] + sbatch_directives + script_commands + + slurm_job_id = client.submit_batch_script("\n".join(script_lines)) + exit_stack.callback(client.cancel_job, slurm_job_id) + + submitted_at = time.monotonic() + job_state: Optional[str] = None + while time.monotonic() - submitted_at < PROVISIONING_TIMEOUT: + job_state = client.get_job_state(slurm_job_id) + logger.debug("slurm_job_id: %s, state: %s", slurm_job_id, job_state) + if job_state is None: + self._skip_offer_cache.add(run, job, instance_offer) + raise ComputeError( + f"Slurm job {slurm_job_id} not found. It either failed or was canceled" + ) + if job_state.upper() == "RUNNING": + break + time.sleep(1) + else: + self._skip_offer_cache.add(run, job, instance_offer) + raise ComputeError( + f"Slurm job {slurm_job_id} didn't start in {PROVISIONING_TIMEOUT} seconds." + f" Last state: {job_state}" + ) + + job_nodes = client.get_job_nodes(slurm_job_id) + if len(job_nodes) != node_count: + raise ComputeError(f"Node count mismatch: len({job_nodes}) != {node_count}") + job_partition = client.get_job_partition(slurm_job_id) + + jpds = [ + JobProvisioningData( + backend=BackendType.SLURM, + instance_type=_build_instance_type(cluster, job_node, requested_resources), + instance_id=_build_instance_id(slurm_job_name, slurm_job_id, job_node.name), + hostname=job_node.hostname, + internal_ip=job_node.ip, + region=region, + availability_zone=job_partition, + price=0.0, + username="root", + ssh_port=DSTACK_RUNNER_SSH_PORT, + dockerized=False, + ssh_proxy=SSHConnectionParams( + hostname=cluster.hostname, + port=cluster.port, + username=cluster.user, + ), + backend_data=None, + ) + for job_node in job_nodes + ] + + res = client.exec(f""" + set -eu + if [ ! -e ~/.ssh/authorized_keys ]; then + mkdir -p ~/.ssh + chmod 700 ~/.ssh + touch ~/.ssh/authorized_keys + chmod 600 ~/.ssh/authorized_keys + fi + for key in {shlex.join(authorized_keys)}; do + if ! grep -qF "$key" ~/.ssh/authorized_keys; then + echo 'command="/bin/false"' "$key" >> ~/.ssh/authorized_keys + fi + done + """) + if not res.ok: + raise ComputeError(f"Failed to add authorized keys: {res}") + + exit_stack.pop_all() + + return ComputeGroupProvisioningData( + compute_group_id=slurm_job_id, + compute_group_name=slurm_job_name, + backend=BackendType.SLURM, + region=region, + job_provisioning_datas=jpds, + ) + + def _offer_modifier( + self, + requested_resources: RequestedResources, + offer: InstanceOfferWithAvailability, + ) -> Optional[InstanceOfferWithAvailability]: + resources = offer.instance.resources + if ( + resources.cpus < requested_resources.cpu_count + or resources.memory_mib < requested_resources.memory_mib + or len(resources.gpus) < requested_resources.gpu_count + ): + return None + + cluster = self._get_cluster(offer.region) + partitions = _get_cluster_partitions(cluster, requested_resources) + assert offer.availability_zones is not None + filtered_partitions = set(offer.availability_zones) & partitions + if not filtered_partitions: + return None + + offer_copy = offer.copy(deep=True) + _adjust_resources(offer_copy.instance.resources, requested_resources) + offer_copy.availability_zones = list(filtered_partitions) + return offer_copy + + def _get_cluster(self, region: str) -> SlurmCluster: + try: + return self._region_to_cluster_map[region] + except KeyError: + raise ComputeError(f"Unknown region: {region!r}") + + +def _get_cluster_offers(cluster: SlurmCluster) -> list[InstanceOfferWithAvailability]: + nodes = cluster.get_discovered_nodes() + return [ + InstanceOfferWithAvailability( + backend=BackendType.SLURM, + instance=_build_instance_type(cluster, node), + region=cluster.region, + price=0.0, + availability=InstanceAvailability.UNKNOWN, + availability_zones=node.partitions, + instance_runtime=InstanceRuntime.RUNNER, + ) + for node in nodes + ] + + +def _get_cluster_partitions( + cluster: SlurmCluster, requested_resources: RequestedResources +) -> set[str]: + discovered_partitions = cluster.get_discovered_partitions() + if requested_resources.gpu_count > 0: + gpu_partitions = cluster.get_gpu_partitions() + return discovered_partitions & gpu_partitions + cpu_partitions = cluster.get_cpu_partitions() + if cpu_partitions is not None: + return discovered_partitions & cpu_partitions + cpu_partitions = discovered_partitions - cluster.get_gpu_partitions() + return cpu_partitions + + +def _build_instance_type( + cluster: SlurmCluster, node: Node, requested_resources: Optional[RequestedResources] = None +) -> InstanceType: + gpus: list[Gpu] = [] + gpu_count: int = 0 + for gres in node.gres: + try: + gpu_count += parse_gres_gpu_count(gres) + except ValueError as e: + logger.warning("Failed to parse GPU GRES: %s: %s", node, e) + if gpu_count > 0: + node_gpu_models: set[GPUModel] = set() + for partition in node.partitions: + gpu_model = cluster.get_partition_gpu_model(partition) + if gpu_model is not None: + node_gpu_models.add(gpu_model) + if not node_gpu_models: + logger.warning("GPU GRES found but not mapped: %s", node) + elif len(node_gpu_models) > 1: + logger.warning("Multiple GPU models mapped: %s: %s", node, node_gpu_models) + else: + gpu = next(iter(node_gpu_models)).to_gpu() + gpus = [gpu] * gpu_count + + try: + cpu_arch = normalize_arch(node.arch).to_cpu_architecture() + except ValueError as e: + logger.warning("Failed to normalize CPU arch: %s: %s", node, e) + cpu_arch = None + + instance_type = InstanceType( + name=node.name, + resources=Resources( + cpu_arch=cpu_arch, + cpus=node.cpus, + memory_mib=node.memory_mib, + gpus=gpus, + spot=False, + disk=Disk(size_mib=0), + ), + ) + if requested_resources is not None: + # NB: unconditionally updates resources, may set values higher than the current ones + _adjust_resources(instance_type.resources, requested_resources) + return instance_type + + +def _adjust_resources(resources: Resources, requested_resources: RequestedResources) -> None: + resources.cpus = requested_resources.cpu_count + resources.memory_mib = requested_resources.memory_mib + resources.gpus = resources.gpus[: requested_resources.gpu_count] + resources.disk = Disk(size_mib=requested_resources.disk_mib) + + +def _build_image_uri(image_name: str) -> str: + # https://github.com/NVIDIA/pyxis/wiki/Usage#image-uri-formats + image = parse_image_name(image_name) + image_uri: str + if image.digest is not None: + image_uri = f"{image.repo}@{image.digest}" + else: + image_uri = f"{image.repo}:{image.tag}" + registry = image.registry + if registry is not None and is_default_registry(registry): + registry = None + if registry is not None: + image_uri = f"{registry}#{image_uri}" + return image_uri + + +def _build_instance_id(slurm_job_name: str, slurm_job_id: str, node_name: str) -> str: + return f"{slurm_job_name}:{slurm_job_id}:{node_name}" + + +def _parse_instance_id(compute_group_id: str) -> tuple[str, str, str]: + slurm_job_name, slurm_job_id, node_name = compute_group_id.split(":", maxsplit=2) + return slurm_job_name, slurm_job_id, node_name diff --git a/src/dstack/_internal/core/backends/slurm/configurator.py b/src/dstack/_internal/core/backends/slurm/configurator.py new file mode 100644 index 000000000..084dcf555 --- /dev/null +++ b/src/dstack/_internal/core/backends/slurm/configurator.py @@ -0,0 +1,73 @@ +import concurrent.futures + +from dstack._internal.core.backends.base.configurator import BackendRecord, Configurator +from dstack._internal.core.backends.slurm.backend import SlurmBackend +from dstack._internal.core.backends.slurm.cluster import ( + SlurmCluster, + get_clusters_from_backend_config, +) +from dstack._internal.core.backends.slurm.models import ( + SlurmBackendConfig, + SlurmBackendConfigWithCreds, + SlurmConfig, + SlurmStoredConfig, +) +from dstack._internal.core.errors import ServerClientError +from dstack._internal.core.models.backends.base import BackendType + + +class SlurmConfigurator( + Configurator[ + SlurmBackendConfig, + SlurmBackendConfigWithCreds, + ] +): + TYPE = BackendType.SLURM + BACKEND_CLASS = SlurmBackend + + def validate_config(self, config: SlurmBackendConfigWithCreds, default_creds_enabled: bool): + try: + clusters = get_clusters_from_backend_config(config) + except Exception as e: + raise ServerClientError(str(e)) + self._check_clusters(clusters) + + def create_backend( + self, project_name: str, config: SlurmBackendConfigWithCreds + ) -> BackendRecord: + return BackendRecord( + config=SlurmStoredConfig.__response__.parse_obj(config).json(), + auth="", + ) + + def get_backend_config_with_creds(self, record: BackendRecord) -> SlurmBackendConfigWithCreds: + config = self._get_config(record) + return SlurmBackendConfigWithCreds.__response__.parse_obj(config) + + def get_backend_config_without_creds(self, record: BackendRecord) -> SlurmBackendConfig: + config = self._get_config(record) + return SlurmBackendConfig.__response__.parse_obj(config) + + def get_backend(self, record: BackendRecord) -> SlurmBackend: + return SlurmBackend(self._get_config(record)) + + def _get_config(self, record: BackendRecord) -> SlurmConfig: + return SlurmConfig.__response__.parse_raw(record.config) + + def _check_clusters(self, clusters: list[SlurmCluster]) -> None: + error_messages: list[str] = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor: + future_to_cluster_map: dict[concurrent.futures.Future[None], SlurmCluster] = {} + for cluster in clusters: + future = executor.submit(self._check_cluster, cluster) + future_to_cluster_map[future] = cluster + for future in concurrent.futures.as_completed(future_to_cluster_map): + exc = future.exception() + if exc is not None: + error_messages.append(f"{future_to_cluster_map[future]}: {exc}") + if error_messages: + raise ServerClientError(f"Failed to check clusters: {', '.join(error_messages)}") + + def _check_cluster(self, cluster: SlurmCluster) -> None: + with cluster.get_client(timeout=10) as client: + client.ping() diff --git a/src/dstack/_internal/core/backends/slurm/models.py b/src/dstack/_internal/core/backends/slurm/models.py new file mode 100644 index 000000000..485fd96bc --- /dev/null +++ b/src/dstack/_internal/core/backends/slurm/models.py @@ -0,0 +1,134 @@ +from typing import Annotated, Literal, Optional, Union + +from pydantic import Field, root_validator + +from dstack._internal.core.backends.base.models import fill_data +from dstack._internal.core.models.common import CoreModel + + +class SlurmGPUPartitionConfig(CoreModel): + gpu: Annotated[ + str, + Field( + description=( + "The GPU model, in the `[vendor:]name[:memory]` format, e.g., `H200`, `A100:40GB`, `MI300X`" + ) + ), + ] + partitions: Annotated[ + list[str], Field(description="The list of partitions with the specified GPU model") + ] + + +class SlurmPrivateKeyConfig(CoreModel): + path: Annotated[str, Field(description="The path to the private key file")] = "" + content: Annotated[str, Field(description="The contents of the private key file")] + + +# `BaseSlurmClusterConfig` holds only non-sensitive fields safe to return in without-creds API +# responses. Connection details (hostname/port/user) and the private key are sensitive and live +# in `BaseSlurmClusterConfigWithCreds`/the creds-bearing configs instead. +# +# The credentialless `SlurmClusterConfig` and the creds-bearing configs are siblings, not +# parent/child. If the latter subclassed the former, a `list[SlurmClusterConfig]` field would +# accept a creds-bearing instance as-is (isinstance passthrough) and leak the sensitive fields +# into the without-creds API response instead of dropping them on re-validation. +class BaseSlurmClusterConfig(CoreModel): + name: Annotated[str, Field(description="The name of the cluster. Used as a region name")] + gpu_partitions: Annotated[ + Optional[list[SlurmGPUPartitionConfig]], + Field( + description=( + "The mapping of GPU models to partitions." + " Only partitions listed here are considered for GPU jobs." + " If not set, GPU jobs are not allowed" + ), + ), + ] = None + cpu_partitions: Annotated[ + Optional[list[str]], + Field( + description=( + "Partitions considered for CPU jobs." + " Defaults to all cluster partitions except those listed in `gpu_partitions`" + ), + ), + ] = None + + +class SlurmClusterConfig(BaseSlurmClusterConfig): + pass + + +class BaseSlurmClusterConfigWithCreds(BaseSlurmClusterConfig): + hostname: Annotated[str, Field(description="The hostname or IP address of the login node")] + port: Annotated[Optional[int], Field(description="The SSH port of the login node")] = None + user: Annotated[str, Field(description="The user to log in to the login node")] + + +class SlurmClusterConfigWithCreds(BaseSlurmClusterConfigWithCreds): + private_key: Annotated[SlurmPrivateKeyConfig, Field(description="The private key of the user")] + + +# Unlike other backends, `SlurmBackendConfigWithCreds` does not subclass `SlurmBackendConfig`: +# `clusters` differs in item type and `list` is invariant, so overriding it in a subclass is a +# type error. The two configs share only `type`, so they are kept as independent classes. +class SlurmBackendConfig(CoreModel): + type: Annotated[ + Literal["slurm"], + Field(description="The type of backend"), + ] = "slurm" + clusters: Annotated[list[SlurmClusterConfig], Field(description="Cluster configurations")] + + +class SlurmBackendConfigWithCreds(CoreModel): + type: Annotated[ + Literal["slurm"], + Field(description="The type of backend"), + ] = "slurm" + clusters: Annotated[ + list[SlurmClusterConfigWithCreds], Field(description="Cluster configurations") + ] + + +class SlurmPrivateKeyFileConfig(CoreModel): + path: Annotated[str, Field(description="The path to the private key file")] = "" + content: Annotated[ + Optional[str], + Field( + description=( + "The contents of the private key file." + " When configuring via `server/config.yml`, it's automatically filled from `path`." + " When configuring via UI, it has to be specified explicitly" + ) + ), + ] = None + + @root_validator + def fill_data(cls, values: dict) -> dict: + return fill_data(values, filename_field="path", data_field="content") + + +class SlurmClusterFileConfig(BaseSlurmClusterConfigWithCreds): + private_key: Annotated[ + SlurmPrivateKeyFileConfig, Field(description="The private key of the user") + ] + + +class SlurmBackendFileConfigWithCreds(CoreModel): + type: Annotated[ + Literal["slurm"], + Field(description="The type of backend"), + ] = "slurm" + clusters: Annotated[list[SlurmClusterFileConfig], Field(description="Cluster configurations")] + + +AnySlurmBackendConfig = Union[SlurmBackendConfig, SlurmBackendConfigWithCreds] + + +class SlurmStoredConfig(SlurmBackendConfigWithCreds): + pass + + +class SlurmConfig(SlurmStoredConfig): + pass diff --git a/src/dstack/_internal/core/backends/slurm/resources.py b/src/dstack/_internal/core/backends/slurm/resources.py new file mode 100644 index 000000000..9a481f6c6 --- /dev/null +++ b/src/dstack/_internal/core/backends/slurm/resources.py @@ -0,0 +1,146 @@ +import dataclasses +from typing import Optional + +import gpuhunt +from typing_extensions import Self + +from dstack._internal.core.models.instances import Gpu +from dstack._internal.core.models.resources import ( + DEFAULT_MEMORY_SIZE, + CPUSpec, + Memory, + ResourcesSpec, +) + + +@dataclasses.dataclass +class Node: + name: str + arch: Optional[str] + cpus: int + memory_mib: int + gres: list[str] + partitions: list[str] + + +@dataclasses.dataclass +class ResolvedNode(Node): + hostname: str + ip: str + + +def parse_gres_gpu_count(gres: str) -> int: + """ + Parse the number of GPUs from a single GRES entry, e.g. `gpu:8`, `gpu:tesla:2`, + `gpu:tesla:4(S:0-1)` (the `gpu[:type]:count[(S:sockets)]` format). + + Returns 0 for non-`gpu` GRES entries. Raises `ValueError` if it is a `gpu` entry + but the count cannot be parsed. + """ + if not gres.startswith("gpu:"): + return 0 + # Strip the optional socket affinity suffix, e.g. gpu:tesla:4(S:0-1) -> gpu:tesla:4 + spec = gres.split("(", maxsplit=1)[0] + return int(spec.rsplit(":", maxsplit=1)[-1]) + + +@dataclasses.dataclass(frozen=True) +class GPUModel: + vendor: gpuhunt.AcceleratorVendor + name: str + memory_mib: int + + def __str__(self) -> str: + return f"{self.vendor.value}:{self.name}:{round(self.memory_mib / 1024)}GB" + + @classmethod + def from_string(cls, s: str) -> Self: + fields = [f for _f in s.split(":") if (f := _f.strip())] + if not fields or len(fields) > 3: + raise ValueError("Invalid format") + + vendor_raw: Optional[str] = None + name_raw: Optional[str] = None + memory_raw: Optional[str] = None + + if len(fields) == 3: + [vendor_raw, name_raw, memory_raw] = fields + elif len(fields) == 1: + [name_raw] = fields + else: + try: + gpuhunt.AcceleratorVendor.cast(fields[0]) + except ValueError: + [name_raw, memory_raw] = fields + else: + [vendor_raw, name_raw] = fields + + vendor: Optional[gpuhunt.AcceleratorVendor] = None + if vendor_raw is not None: + vendor = gpuhunt.AcceleratorVendor(vendor_raw) + memory: Optional[Memory] = None + if memory_raw is not None: + memory = Memory.parse(memory_raw) + + if vendor is not None and memory is not None: + return cls( + vendor=vendor, + name=name_raw, + memory_mib=round(memory * 1024), + ) + accelerators = gpuhunt.find_accelerators( + names=[name_raw.replace(" ", "")], + vendors=[vendor] if vendor is not None else None, + ) + if memory is not None: + memory_gib = round(memory) + accelerators = [a for a in accelerators if a.memory == memory_gib] + if not accelerators: + raise ValueError(f"No matching GPU model found: {s}") + if len(accelerators) > 1: + raise ValueError(f"Multiple matching GPU models found: {s}: {accelerators}") + accelerator = accelerators[0] + return cls( + vendor=accelerator.vendor, + name=accelerator.name, + memory_mib=accelerator.memory * 1024, + ) + + def to_gpu(self) -> Gpu: + return Gpu(vendor=self.vendor, name=self.name, memory_mib=self.memory_mib) + + +@dataclasses.dataclass +class RequestedResources: + cpu_count: int + memory_mib: int + gpu_count: int + disk_mib: int + + +def get_requested_resources_from_resources_spec(spec: ResourcesSpec) -> RequestedResources: + assert isinstance(spec.cpu, CPUSpec) + # 1 is the default value of --cpus-per-task + cpu_count = spec.cpu.count.min or 1 + + # We cannot use 0 as a fallback since --mem=0 is a special case + # We cannot infer the default value which Slurm will use since it's partition-specific + # The easiest/safest option is to use some sane default + memory = spec.memory.min or DEFAULT_MEMORY_SIZE.min + assert memory + memory_mib = round(memory * 1024) + + gpu_count: int = 0 + if spec.gpu is not None: + gpu_count = spec.gpu.count.min or 0 + + disk_mib: int = 0 + if spec.disk is not None: + disk_mib = round((spec.disk.size.min or 0) * 1024) + + return RequestedResources( + cpu_count=cpu_count, + memory_mib=memory_mib, + gpu_count=gpu_count, + disk_mib=disk_mib, + ) diff --git a/src/dstack/_internal/core/models/backends/base.py b/src/dstack/_internal/core/models/backends/base.py index 80e241b31..419b2cc70 100644 --- a/src/dstack/_internal/core/models/backends/base.py +++ b/src/dstack/_internal/core/models/backends/base.py @@ -25,6 +25,7 @@ class BackendType(str, enum.Enum): VASTAI (BackendType): Vast.ai Marketplace VERDA (BackendType): Verda Cloud VULTR (BackendType): Vultr + SLURM (BackendType): Slurm """ AMDDEVCLOUD = "amddevcloud" @@ -50,3 +51,4 @@ class BackendType(str, enum.Enum): VASTAI = "vastai" VERDA = "verda" VULTR = "vultr" + SLURM = "slurm" diff --git a/src/dstack/_internal/server/background/pipeline_tasks/common.py b/src/dstack/_internal/server/background/pipeline_tasks/common.py index 0c204cbc3..56e0dc5f9 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/common.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/common.py @@ -14,6 +14,8 @@ def get_provisioning_timeout(backend_type: BackendType, instance_type_name: str) return timedelta(minutes=20) if backend_type == BackendType.KUBERNETES: return timedelta(minutes=20) + if backend_type == BackendType.SLURM: + return timedelta(minutes=20) if backend_type == BackendType.OCI and instance_type_name.startswith("BM."): return timedelta(minutes=20) if backend_type == BackendType.VULTR and instance_type_name.startswith("vbm"): diff --git a/src/tests/_internal/core/backends/base/test_offers.py b/src/tests/_internal/core/backends/base/test_offers.py new file mode 100644 index 000000000..8f3383c64 --- /dev/null +++ b/src/tests/_internal/core/backends/base/test_offers.py @@ -0,0 +1,94 @@ +import gpuhunt +import pytest + +from dstack._internal.core.backends.base.offers import gpu_matches_gpu_spec +from dstack._internal.core.models.instances import Gpu +from dstack._internal.core.models.resources import GPUSpec + +NVIDIA = gpuhunt.AcceleratorVendor.NVIDIA +AMD = gpuhunt.AcceleratorVendor.AMD + + +def make_gpu(name: str = "A100", memory_gib: int = 80, vendor=NVIDIA) -> Gpu: + return Gpu(name=name, memory_mib=memory_gib * 1024, vendor=vendor) + + +class TestGpuMatchesGpuSpec: + def test_empty_spec_matches_any_gpu(self): + assert gpu_matches_gpu_spec(make_gpu(vendor=AMD), GPUSpec()) + + def test_vendor_matches(self): + assert gpu_matches_gpu_spec(make_gpu(vendor=NVIDIA), GPUSpec(vendor="nvidia")) + + def test_vendor_mismatch(self): + assert not gpu_matches_gpu_spec(make_gpu(vendor=AMD), GPUSpec(vendor="nvidia")) + + def test_name_matches(self): + assert gpu_matches_gpu_spec(make_gpu(name="A100"), GPUSpec(name="A100")) + + @pytest.mark.parametrize( + ("gpu_name", "spec_name"), + [("A100", "a100"), ("a100", "A100"), ("H100", "h100")], + ) + def test_name_match_is_case_insensitive(self, gpu_name: str, spec_name: str): + assert gpu_matches_gpu_spec(make_gpu(name=gpu_name), GPUSpec(name=spec_name)) + + def test_name_matches_any_in_list(self): + assert gpu_matches_gpu_spec(make_gpu(name="A100"), GPUSpec(name=["V100", "A100"])) + + def test_name_mismatch(self): + assert not gpu_matches_gpu_spec(make_gpu(name="A100"), GPUSpec(name="V100")) + + def test_name_not_in_list(self): + assert not gpu_matches_gpu_spec(make_gpu(name="A100"), GPUSpec(name=["V100", "H100"])) + + def test_memory_within_range(self): + assert gpu_matches_gpu_spec(make_gpu(memory_gib=40), GPUSpec(memory="16GB..80GB")) + + def test_memory_equal_to_min(self): + assert gpu_matches_gpu_spec(make_gpu(memory_gib=16), GPUSpec(memory="16GB..")) + + def test_memory_equal_to_max(self): + assert gpu_matches_gpu_spec(make_gpu(memory_gib=80), GPUSpec(memory="..80GB")) + + def test_memory_below_min(self): + assert not gpu_matches_gpu_spec(make_gpu(memory_gib=15), GPUSpec(memory="16GB..")) + + def test_memory_above_max(self): + assert not gpu_matches_gpu_spec(make_gpu(memory_gib=100), GPUSpec(memory="..80GB")) + + def test_compute_capability_equal(self): + # A100 has compute capability 8.0 + assert gpu_matches_gpu_spec(make_gpu(name="A100"), GPUSpec(compute_capability="8.0")) + + def test_compute_capability_gpu_higher_than_required(self): + # A100 (8.0) satisfies the 7.0 minimum + assert gpu_matches_gpu_spec(make_gpu(name="A100"), GPUSpec(compute_capability="7.0")) + + def test_compute_capability_gpu_lower_than_required(self): + # A100 (8.0) does not satisfy the 9.0 minimum + assert not gpu_matches_gpu_spec(make_gpu(name="A100"), GPUSpec(compute_capability="9.0")) + + def test_compute_capability_non_nvidia_never_matches(self): + assert not gpu_matches_gpu_spec( + make_gpu(name="MI300X", vendor=AMD), GPUSpec(compute_capability="8.0") + ) + + def test_compute_capability_unknown_gpu_name_never_matches(self): + assert not gpu_matches_gpu_spec( + make_gpu(name="UnknownGPU"), GPUSpec(compute_capability="8.0") + ) + + def test_all_constraints_match(self): + spec = GPUSpec( + vendor="nvidia", + name="A100", + memory="40GB..80GB", + compute_capability="8.0", + ) + assert gpu_matches_gpu_spec(make_gpu(name="A100", memory_gib=80), spec) + + def test_single_failing_constraint_rejects_gpu(self): + # Everything matches except memory, which is below the minimum. + spec = GPUSpec(vendor="nvidia", name="A100", memory="40GB..80GB") + assert not gpu_matches_gpu_spec(make_gpu(name="A100", memory_gib=24), spec) diff --git a/src/tests/_internal/core/backends/slurm/__init__.py b/src/tests/_internal/core/backends/slurm/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/tests/_internal/core/backends/slurm/test_client.py b/src/tests/_internal/core/backends/slurm/test_client.py new file mode 100644 index 000000000..7c3a5e7d7 --- /dev/null +++ b/src/tests/_internal/core/backends/slurm/test_client.py @@ -0,0 +1,116 @@ +from dstack._internal.core.backends.slurm.client import _parse_scontrol_show_line + + +class TestParseScontrolShowLine: + # A single `scontrol show -o node` line for an allocated node. It exercises the tricky + # cases the parser must handle: + # - multi-word value (OS) containing spaces and `#` + # - values with embedded `=` that must not be split into new keys (CfgTRES/AllocTRES) + # - "(null)" and "N/A" sentinels kept verbatim + # - underscore in a key (MCS_label) + # - double spaces between some fields + SCONTROL_SHOW_NODE_LINE = ( + "NodeName=worker-1 Arch=x86_64 CoresPerSocket=1 CPUAlloc=4 CPUEfctv=4 CPUTot=4 " + "CPULoad=0.10 AvailableFeatures=(null) ActiveFeatures=(null) Gres=(null) GresDrain=N/A " + "NodeAddr=worker-1 NodeHostName=worker-1 Version=25.11.2 " + "OS=Linux 7.0.0-22-generic #22-Ubuntu SMP PREEMPT_DYNAMIC Mon May 25 15:54:34 UTC 2026 " + "RealMemory=15000 AllocMem=0 FreeMem=14769 Sockets=4 Boards=1 State=ALLOCATED " + "ThreadsPerCore=1 TmpDisk=0 Weight=1 Owner=N/A MCS_label=N/A Partitions=main,wrk-1 " + "BootTime=2026-06-16T10:02:01 SlurmdStartTime=2026-06-16T10:02:13 " + "LastBusyTime=2026-06-16T11:09:07 ResumeAfterTime=None " + "CfgTRES=cpu=4,mem=15000M,billing=4 AllocTRES=cpu=4,mem=15000M,billing=4 " + "CurrentWatts=0 AveWatts=0" + ) + + def test_parses_full_node_line(self): + assert _parse_scontrol_show_line(self.SCONTROL_SHOW_NODE_LINE) == { + "nodename": "worker-1", + "arch": "x86_64", + "corespersocket": "1", + "cpualloc": "4", + "cpuefctv": "4", + "cputot": "4", + "cpuload": "0.10", + "availablefeatures": "(null)", + "activefeatures": "(null)", + "gres": "(null)", + "gresdrain": "N/A", + "nodeaddr": "worker-1", + "nodehostname": "worker-1", + "version": "25.11.2", + "os": "Linux 7.0.0-22-generic #22-Ubuntu SMP PREEMPT_DYNAMIC Mon May 25 15:54:34 UTC 2026", + "realmemory": "15000", + "allocmem": "0", + "freemem": "14769", + "sockets": "4", + "boards": "1", + "state": "ALLOCATED", + "threadspercore": "1", + "tmpdisk": "0", + "weight": "1", + "owner": "N/A", + "mcs_label": "N/A", + "partitions": "main,wrk-1", + "boottime": "2026-06-16T10:02:01", + "slurmdstarttime": "2026-06-16T10:02:13", + "lastbusytime": "2026-06-16T11:09:07", + "resumeaftertime": "None", + "cfgtres": "cpu=4,mem=15000M,billing=4", + "alloctres": "cpu=4,mem=15000M,billing=4", + "currentwatts": "0", + "avewatts": "0", + } + + def test_keeps_multi_word_value_intact(self): + result = _parse_scontrol_show_line(self.SCONTROL_SHOW_NODE_LINE) + + # The whole `OS=...` value is captured up to the next key, with surrounding + # whitespace (including the trailing double space) stripped. + assert ( + result["os"] + == "Linux 7.0.0-22-generic #22-Ubuntu SMP PREEMPT_DYNAMIC Mon May 25 15:54:34 UTC 2026" + ) + + def test_does_not_split_value_on_embedded_equals(self): + result = _parse_scontrol_show_line(self.SCONTROL_SHOW_NODE_LINE) + + # `cpu=...`, `mem=...`, `billing=...` are part of the value, not new keys, + # because they are not preceded by whitespace. + assert result["cfgtres"] == "cpu=4,mem=15000M,billing=4" + assert result["alloctres"] == "cpu=4,mem=15000M,billing=4" + assert "cpu" not in result + assert "mem" not in result + assert "billing" not in result + + def test_keeps_null_and_na_sentinels_verbatim(self): + result = _parse_scontrol_show_line(self.SCONTROL_SHOW_NODE_LINE) + + assert result["gres"] == "(null)" + assert result["availablefeatures"] == "(null)" + assert result["gresdrain"] == "N/A" + assert result["owner"] == "N/A" + + def test_normalizes_keys_to_lowercase_by_default(self): + result = _parse_scontrol_show_line(self.SCONTROL_SHOW_NODE_LINE) + + assert "nodename" in result + assert "mcs_label" in result + assert "NodeName" not in result + + def test_preserves_key_case_when_normalize_key_is_false(self): + result = _parse_scontrol_show_line(self.SCONTROL_SHOW_NODE_LINE, normalize_key=False) + + assert result["NodeName"] == "worker-1" + assert result["MCS_label"] == "N/A" + assert result["CfgTRES"] == "cpu=4,mem=15000M,billing=4" + assert "nodename" not in result + + def test_strips_surrounding_whitespace_on_the_line(self): + assert _parse_scontrol_show_line(" NodeName=worker-1 CPUTot=4 \n") == { + "nodename": "worker-1", + "cputot": "4", + } + + def test_returns_empty_dict_for_blank_line(self): + assert _parse_scontrol_show_line("") == {} + assert _parse_scontrol_show_line(" \n") == {} diff --git a/src/tests/_internal/core/backends/slurm/test_configurator.py b/src/tests/_internal/core/backends/slurm/test_configurator.py new file mode 100644 index 000000000..bf576a80e --- /dev/null +++ b/src/tests/_internal/core/backends/slurm/test_configurator.py @@ -0,0 +1,56 @@ +from dstack._internal.core.backends.base.configurator import BackendRecord +from dstack._internal.core.backends.slurm.configurator import SlurmConfigurator +from dstack._internal.core.backends.slurm.models import ( + SlurmBackendConfigWithCreds, + SlurmClusterConfig, + SlurmClusterConfigWithCreds, +) + +PRIVATE_KEY = "-----BEGIN OPENSSH PRIVATE KEY-----\nSECRET\n-----END OPENSSH PRIVATE KEY-----" + + +HOSTNAME = "login.secret.example" +PORT = 2222 +USER = "alice" + + +def _make_backend_record() -> BackendRecord: + config = SlurmBackendConfigWithCreds( + clusters=[ + { + "name": "c1", + "hostname": HOSTNAME, + "port": PORT, + "user": USER, + "private_key": {"content": PRIVATE_KEY}, + } + ] + ) + return SlurmConfigurator().create_backend("proj", config) + + +class TestSlurmConfiguratorCreds: + def test_without_creds_strips_sensitive_fields(self): + record = _make_backend_record() + config = SlurmConfigurator().get_backend_config_without_creds(record) + cluster = config.clusters[0] + # The credentialless cluster config exposes only non-sensitive fields; connection + # details and the private key must not be present at all. + assert set(cluster.__fields__) == {"name", "gpu_partitions", "cpu_partitions"} + rendered = config.json() + for secret in (PRIVATE_KEY, HOSTNAME, str(PORT), USER): + assert secret not in rendered + + def test_with_creds_keeps_sensitive_fields(self): + record = _make_backend_record() + cluster = SlurmConfigurator().get_backend_config_with_creds(record).clusters[0] + assert cluster.private_key.content == PRIVATE_KEY + assert cluster.hostname == HOSTNAME + assert cluster.port == PORT + assert cluster.user == USER + + def test_with_and_without_cluster_configs_are_not_subclasses(self): + # A subclass relationship would let a `list[SlurmClusterConfig]` field accept a + # creds-bearing instance as-is and leak the private key into without-creds responses. + assert not issubclass(SlurmClusterConfigWithCreds, SlurmClusterConfig) + assert not issubclass(SlurmClusterConfig, SlurmClusterConfigWithCreds) diff --git a/src/tests/_internal/core/backends/slurm/test_resources.py b/src/tests/_internal/core/backends/slurm/test_resources.py new file mode 100644 index 000000000..36137f71b --- /dev/null +++ b/src/tests/_internal/core/backends/slurm/test_resources.py @@ -0,0 +1,107 @@ +import gpuhunt +import pytest + +from dstack._internal.core.backends.slurm.resources import GPUModel, parse_gres_gpu_count + +NVIDIA = gpuhunt.AcceleratorVendor.NVIDIA +AMD = gpuhunt.AcceleratorVendor.AMD + + +class TestGPUModelStr: + def test_nvidia(self): + assert ( + str(GPUModel(vendor=NVIDIA, name="A100", memory_mib=80 * 1024)) == "nvidia:A100:80GB" + ) + + def test_amd(self): + assert ( + str(GPUModel(vendor=AMD, name="MI300X", memory_mib=192 * 1024)) == "amd:MI300X:192GB" + ) + + def test_roundtrips_through_from_string(self): + model = GPUModel(vendor=NVIDIA, name="A100", memory_mib=80 * 1024) + assert GPUModel.from_string(str(model)) == model + + +class TestGPUModelFromString: + def test_vendor_name_memory_uses_values_verbatim(self): + # With vendor and memory both given, gpuhunt is bypassed, so an unknown name is accepted + # and the name is kept verbatim (no space normalization). + assert GPUModel.from_string("nvidia:Custom Accel:16GB") == GPUModel( + vendor=NVIDIA, name="Custom Accel", memory_mib=16 * 1024 + ) + + def test_name_only_looks_up_nvidia(self): + assert GPUModel.from_string("H100") == GPUModel( + vendor=NVIDIA, name="H100", memory_mib=80 * 1024 + ) + + def test_name_only_looks_up_amd(self): + assert GPUModel.from_string("MI300X") == GPUModel( + vendor=AMD, name="MI300X", memory_mib=192 * 1024 + ) + + def test_name_only_normalizes_spaces_for_lookup(self): + assert GPUModel.from_string("RTX 4090") == GPUModel( + vendor=NVIDIA, name="RTX4090", memory_mib=24 * 1024 + ) + + def test_vendor_and_name(self): + assert GPUModel.from_string("nvidia:H100") == GPUModel( + vendor=NVIDIA, name="H100", memory_mib=80 * 1024 + ) + + @pytest.mark.parametrize("memory_gb", [40, 80]) + def test_name_and_memory_disambiguates_variants(self, memory_gb: int): + assert GPUModel.from_string(f"A100:{memory_gb}GB") == GPUModel( + vendor=NVIDIA, name="A100", memory_mib=memory_gb * 1024 + ) + + def test_memory_is_rounded_to_gib_for_matching(self): + assert GPUModel.from_string("A100:80.4GB") == GPUModel( + vendor=NVIDIA, name="A100", memory_mib=80 * 1024 + ) + + def test_raises_when_no_gpu_matches_name(self): + with pytest.raises(ValueError, match="No matching GPU model found"): + GPUModel.from_string("ThisGpuDoesNotExist") + + def test_raises_when_multiple_gpus_match_name(self): + with pytest.raises(ValueError, match="Multiple matching GPU models found"): + GPUModel.from_string("A100") + + def test_raises_when_vendor_does_not_match(self): + with pytest.raises(ValueError, match="No matching GPU model found"): + GPUModel.from_string("amd:A100") + + def test_raises_when_memory_matches_no_variant(self): + with pytest.raises(ValueError, match="No matching GPU model found"): + GPUModel.from_string("A100:24GB") + + @pytest.mark.parametrize("s", ["", " ", ":::", "a:b:c:d", "nvidia:A100:80GB:extra"]) + def test_raises_on_invalid_format(self, s: str): + with pytest.raises(ValueError, match="Invalid format"): + GPUModel.from_string(s) + + +class TestParseGresGpuCount: + def test_count_only(self): + assert parse_gres_gpu_count("gpu:8") == 8 + + def test_type_and_count(self): + assert parse_gres_gpu_count("gpu:tesla:2") == 2 + + def test_count_with_socket_affinity(self): + assert parse_gres_gpu_count("gpu:8(S:0)") == 8 + + def test_type_and_count_with_socket_affinity(self): + assert parse_gres_gpu_count("gpu:tesla:4(S:0-1)") == 4 + + @pytest.mark.parametrize("gres", ["mps:200", "mem:1024", "gpu", ""]) + def test_returns_zero_for_non_gpu_gres(self, gres: str): + assert parse_gres_gpu_count(gres) == 0 + + @pytest.mark.parametrize("gres", ["gpu:tesla", "gpu:", "gpu:x(S:0)"]) + def test_raises_when_count_is_not_an_integer(self, gres: str): + with pytest.raises(ValueError): + parse_gres_gpu_count(gres) diff --git a/src/tests/_internal/server/routers/test_backends.py b/src/tests/_internal/server/routers/test_backends.py index 62eff12a3..23b51c5fc 100644 --- a/src/tests/_internal/server/routers/test_backends.py +++ b/src/tests/_internal/server/routers/test_backends.py @@ -98,6 +98,7 @@ async def test_returns_backend_types(self, client: AsyncClient): "nebius", "oci", "runpod", + "slurm", "vastai", "verda", "vultr",