From 23a0df3a4dc87696d1cd8dc3e072ebcec9f6f0de Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 07:28:46 +0000 Subject: [PATCH 1/9] Support NVIDIA MIG slices in the Kubernetes backend Offer MIG slices created by the NVIDIA GPU Operator (single and mixed strategies) as GPUs of the parent model with the profile's memory, e.g. a 1g.10gb slice of an A100 is A100:10GB. - Single strategy: a node whose nvidia.com/gpu.product is -MIG- is offered as that many slices. - Mixed strategy: each nvidia.com/mig- resource on a node is offered separately next to the node's whole GPUs; the pod requests that resource and its node affinity uses the parent product label. - Runs get slices only when the GPU memory spec has an upper bound, so gpu: A100 never lands on a slice; --full-offers shows everything. - After a failed scheduling attempt, skip only offers for the same GPU resource in the cluster, so a 3g.20gb failure doesn't rule out the cluster's 4g.20gb slices or whole GPUs. Tested on GKE 1.36 with GPU Operator v26.7.2 (driver 595) and A100 40GB nodes across single, mixed and MIG-only layouts. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EDgQBSHAm1rxrhSUXWziy3 --- mkdocs/docs/concepts/backends.md | 36 +++ .../core/backends/kubernetes/compute.py | 114 +++++-- .../core/backends/kubernetes/resources.py | 260 +++++++++++++--- .../core/backends/kubernetes/test_compute.py | 200 ++++++++++++- .../backends/kubernetes/test_resources.py | 282 ++++++++++++++++++ 5 files changed, 832 insertions(+), 60 deletions(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index e488520fb2..a30e6327e3 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1397,6 +1397,42 @@ projects: This applies to offers shown in `dstack apply` (run plans), during provisioning, and in `dstack offer`. Unlike other backends, offers for the `kubernetes` backend always reflect the lower limit of the range. +??? info "Multi-Instance GPU (MIG)" + The `kubernetes` backend supports NVIDIA [MIG](https://docs.nvidia.com/datacenter/tesla/mig-user-guide/latest/index.html) slices + configured with the NVIDIA GPU Operator's [MIG Manager](https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/gpu-operator-mig.html), + with either the `single` or the `mixed` [MIG strategy](https://docs.nvidia.com/datacenter/cloud-native/kubernetes/latest/index.html). + + A MIG slice is offered as a GPU of the same model with the slice's memory, taken from the MIG profile name. + For example, a `1g.10gb` slice of an A100 is `A100:10GB`. + + To run on a slice, set an upper limit on GPU memory: + +
+ + ```yaml + type: dev-environment + ide: vscode + + resources: + gpu: A100:10GB + ``` + +
+ + Without an upper limit on GPU memory (e.g., `gpu: A100` or `gpu: 24GB..`), runs only get whole GPUs. + To see slices in `dstack offer`, set the limit (e.g., `--gpu A100:10GB`) or pass `--full-offers`. + + With the `mixed` strategy, a node that has whole GPUs and slices of several sizes appears as several offers, one per GPU type. + With `--full-offers`, each of them shows the node's full CPU, memory, and disk, which they share. + + Limitations: + + * NVIDIA doesn't support NCCL with MIG, so a run with several slices (e.g., `gpu: A100:10GB:2`) can't use NCCL-based multi-GPU training. + Several slices suit independent processes, one per slice (a single process can use several slices only with NVIDIA driver R570 or newer). + To run several copies of a model, use a service with `replicas`, each on one slice. + * Slices with the same memory (e.g., `3g.40gb` and `4g.40gb`) and media-engine variants (e.g., `1g.10gb+me`) can't be requested specifically. + * `dstack metrics` doesn't show GPU metrics for slices. + > To learn more, see the [Lambda](../examples/clusters/lambda/#kubernetes) and [Crusoe](../examples/clusters/crusoe/#kubernetes) examples. ### Slurm diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index 980cd67615..37e98e862c 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -9,6 +9,7 @@ from enum import Enum from functools import partial from typing import List, Optional +from uuid import UUID from gpuhunt import AcceleratorVendor from kubernetes import client @@ -36,8 +37,8 @@ merge_tags, ) from dstack._internal.core.backends.base.offers import ( + BaseSkipOfferCache, OfferModifier, - RegionalSkipOfferCache, gpu_matches_gpu_spec, ) from dstack._internal.core.backends.kubernetes.api_client import API_CLIENT_EXCEPTIONS @@ -46,12 +47,14 @@ AMD_GPU_DEVICE_ID_LABEL_PREFIX, AMD_GPU_NAME_TO_DEVICE_IDS, AMD_GPU_NODE_TAINT, + AMD_GPU_RESOURCE, LABEL_VALUE_MAX_LENGTH, NVIDIA_GPU_NODE_TAINT, NVIDIA_GPU_PRODUCT_LABEL, + NVIDIA_GPU_RESOURCE, OBJECT_NAME_MAX_LENGTH, - AnyKubernetesGPUResource, - KubernetesResource, + KubernetesOfferBackendData, + KubernetesResources, PodPhase, ResourceLimits, ResourceRequests, @@ -62,12 +65,14 @@ filter_invalid_labels, format_memory, get_amd_gpu_from_node_labels, + get_gpu_resource_from_kubernetes_map, get_instance_offer_from_node, get_instance_offers, get_node_labels, get_node_name, get_nvidia_gpu_from_node_labels, is_hard_taint, + is_nvidia_mig_product, is_taint_tolerated, parse_quantity, ) @@ -88,6 +93,7 @@ GatewayReplicaProvisioningData, ) from dstack._internal.core.models.instances import ( + InstanceOffer, InstanceOfferWithAvailability, SSHConnectionParams, ) @@ -157,7 +163,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 = RegionalSkipOfferCache(ttl=60) + self.skip_offer_cache = _KubernetesSkipOfferCache(ttl=60) def get_all_offers_with_availability( self, unallocated_resources: bool @@ -191,7 +197,8 @@ def get_offers_modifiers( if full_offers: return [] resource_requests = ResourceRequests.from_resources_spec(requirements.resources) - return [partial(_offer_modifier, resource_requests)] + mig_allowed = _is_mig_allowed(requirements.resources.gpu) + return [partial(_offer_modifier, resource_requests, mig_allowed)] def run_job( self, @@ -284,6 +291,9 @@ def run_job( job_spec=job.job_spec, volumes=volumes, requirements=requirements, + mig_resource=KubernetesOfferBackendData.load( + instance_offer.backend_data + ).mig_resource, authorized_keys=authorized_keys, ) exit_stack.callback( @@ -427,16 +437,19 @@ def update_provisioning_data( provisioning_data.hostname = get_or_error(service_spec.cluster_ip) pod_spec = get_or_error(pod.spec) node = api.read_node(name=get_or_error(pod_spec.node_name)) + resource_requirements = get_or_error(pod_spec.containers[0].resources) + requests_map = resource_requirements.requests or {} # TODO: Set provisioning_data.gpu_driver from the node labels: # nvidia.com/cuda.driver-version.full set by GPU Feature Discovery # (or nvidia.com/cuda.driver.{major,minor,rev} set by older GFD versions), # amd.com/gpu.driver-version set by the AMD GPU Operator. - instance_offer = get_instance_offer_from_node(node=node, region=cluster.region) + instance_offer = get_instance_offer_from_node( + node=node, + region=cluster.region, + gpu_resource=get_gpu_resource_from_kubernetes_map(requests_map), + ) if instance_offer is not None: - resource_requirements = get_or_error(pod_spec.containers[0].resources) - resource_requests = ResourceRequests.from_kubernetes_map( - resource_requirements.requests or {} - ) + resource_requests = ResourceRequests.from_kubernetes_map(requests_map) adjust_resources_by_resource_requests( instance_offer.instance.resources, resource_requests, force=True ) @@ -771,29 +784,45 @@ def delete_volume(self, volume: Volume): def _get_pod_spec_parameters_for_gpu( - api: client.CoreV1Api, gpu_spec: GPUSpec -) -> tuple[AnyKubernetesGPUResource, client.V1NodeAffinity, str]: + api: client.CoreV1Api, gpu_spec: GPUSpec, mig_resource: Optional[str] +) -> tuple[str, client.V1NodeAffinity, str]: + """ + Returns the GPU resource to request, the node affinity, and the node taint to tolerate. + + Args: + mig_resource: The `mixed` strategy MIG resource of the offer, if any. If set, it is + requested instead of `nvidia.com/gpu`. + """ nodes = api.list_node().items gpu_vendor = gpu_spec.vendor # If no vendor specified, we assume it's NVIDIA. Technically, it's possible to request either # NVIDIA or AMD in the run configuration using only GPU names (e.g.,`gpu: H100,MI300X:8`), # but we ignore such configurations as it's hard to translate them to K8s request. if gpu_vendor is None or gpu_vendor == AcceleratorVendor.NVIDIA: - node_affinity = _get_nvidia_gpu_node_affinity(gpu_spec, nodes) - return KubernetesResource.NVIDIA_GPU, node_affinity, NVIDIA_GPU_NODE_TAINT + node_affinity = _get_nvidia_gpu_node_affinity(gpu_spec, nodes, mig_resource) + gpu_resource = mig_resource if mig_resource is not None else NVIDIA_GPU_RESOURCE + return gpu_resource, node_affinity, NVIDIA_GPU_NODE_TAINT if gpu_vendor == AcceleratorVendor.AMD: node_affinity = _get_amd_gpu_node_affinity(gpu_spec, nodes) - return KubernetesResource.AMD_GPU, node_affinity, AMD_GPU_NODE_TAINT + return AMD_GPU_RESOURCE, node_affinity, AMD_GPU_NODE_TAINT raise ComputeError(f"Unsupported GPU vendor: {gpu_vendor}") def _get_nvidia_gpu_node_affinity( - gpu_spec: GPUSpec, nodes: list[client.V1Node] + gpu_spec: GPUSpec, nodes: list[client.V1Node], mig_resource: Optional[str] ) -> client.V1NodeAffinity: + mig_allowed = _is_mig_allowed(gpu_spec) matching_gpu_label_values: set[str] = set() for node in nodes: labels = get_node_labels(node) - gpu = get_nvidia_gpu_from_node_labels(labels) + if mig_resource is None: + # `single` strategy MIG slices are also `nvidia.com/gpu` + if is_nvidia_mig_product(labels) and not mig_allowed: + continue + elif not _node_has_allocatable_resource(node, mig_resource): + continue + # With the `mixed` strategy, the product label is the parent GPU's product name + gpu = get_nvidia_gpu_from_node_labels(labels, mig_resource) 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: @@ -847,13 +876,53 @@ def _get_amd_gpu_node_affinity( def _offer_modifier( - resource_requests: ResourceRequests, offer: InstanceOfferWithAvailability -) -> InstanceOfferWithAvailability: + resource_requests: ResourceRequests, mig_allowed: bool, offer: InstanceOfferWithAvailability +) -> Optional[InstanceOfferWithAvailability]: + backend_data = KubernetesOfferBackendData.load(offer.backend_data) + if backend_data.mig: + if resource_requests.gpu == 0: + # Without GPUs, a `mixed` strategy MIG offer is a duplicate of the node's main offer + if backend_data.mig_resource is not None: + return None + elif not mig_allowed: + return None offer_copy = offer.model_copy(deep=True) adjust_resources_by_resource_requests(offer_copy.instance.resources, resource_requests) return offer_copy +def _is_mig_allowed(gpu_spec: Optional[GPUSpec]) -> bool: + """ + Returns whether a run may get MIG slices: only if it caps GPU memory, e.g., `A100:10GB` or + `..20GB`, so that, e.g., `A100` or `24GB..` only get whole GPUs. + """ + return gpu_spec is not None and gpu_spec.memory is not None and gpu_spec.memory.max is not None + + +def _node_has_allocatable_resource(node: client.V1Node, resource: str) -> bool: + if node.status is None or node.status.allocatable is None: + return False + return KubernetesResources.from_kubernetes_map(node.status.allocatable)[resource] > 0 + + +class _KubernetesSkipOfferCache(BaseSkipOfferCache[tuple[UUID, str, Optional[str]]]): + """ + Like `RegionalSkipOfferCache`, but tracks `mixed` strategy MIG offers separately by their + MIG resource. + + A pod for any other offer requests `nvidia.com/gpu` or `amd.com/gpu` on any node with a + matching GPU, so if it is not scheduled, no such offer in the region can be. A pod for a + `mixed` strategy MIG offer requests only that offer's MIG resource, so other MIG resources + and whole GPUs in the region may still be available. + """ + + def _build_key( + self, run: Run, job: Job, offer: InstanceOffer + ) -> tuple[UUID, str, Optional[str]]: + mig_resource = KubernetesOfferBackendData.load(offer.backend_data).mig_resource + return (run.id, offer.region, mig_resource) + + def _create_jump_pod_service_if_not_exists( api: client.CoreV1Api, namespace: str, @@ -1159,6 +1228,7 @@ def _create_job_pod( job_spec: JobSpec, volumes: list[Volume], requirements: Requirements, + mig_resource: Optional[str], authorized_keys: list[str], ) -> None: node_affinity: Optional[client.V1NodeAffinity] = None @@ -1170,11 +1240,13 @@ def _create_job_pod( resources_spec = requirements.resources resource_requests = ResourceRequests.from_resources_spec(resources_spec) resource_limits = ResourceLimits.from_resources_spec(resources_spec) - gpu_resource: Optional[AnyKubernetesGPUResource] = None + gpu_resource: Optional[str] = None if resource_requests.gpu > 0: gpu_spec = resources_spec.gpu assert gpu_spec is not None - gpu_resource, node_affinity, node_taint = _get_pod_spec_parameters_for_gpu(api, gpu_spec) + gpu_resource, node_affinity, node_taint = _get_pod_spec_parameters_for_gpu( + api, gpu_spec, mig_resource + ) logger.debug("Requesting GPU resource: %s=%d", gpu_resource, resource_requests.gpu) # It should be NoSchedule, but we also add NoExecute toleration just in case. for effect in [TaintEffect.NO_SCHEDULE, TaintEffect.NO_EXECUTE]: diff --git a/src/dstack/_internal/core/backends/kubernetes/resources.py b/src/dstack/_internal/core/backends/kubernetes/resources.py index b049eccc8c..40d7dc73f4 100644 --- a/src/dstack/_internal/core/backends/kubernetes/resources.py +++ b/src/dstack/_internal/core/backends/kubernetes/resources.py @@ -1,11 +1,13 @@ import base64 import dataclasses +import functools import json +import operator import re from collections.abc import Mapping from decimal import Decimal from enum import Enum -from typing import Callable, Literal, Optional, Union, cast, get_args +from typing import Any, Callable, Literal, Optional, Union, cast, get_args import gpuhunt from gpuhunt import KNOWN_AMD_GPUS, KNOWN_NVIDIA_GPUS, AcceleratorVendor @@ -15,6 +17,7 @@ from dstack._internal.core.backends.base.compute import normalize_arch from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.common import CoreModel, validate_extra_ignore from dstack._internal.core.models.instances import ( Disk, Gpu, @@ -47,6 +50,19 @@ NVIDIA_GPU_RESOURCE = "nvidia.com/gpu" NVIDIA_GPU_NODE_TAINT = NVIDIA_GPU_RESOURCE NVIDIA_GPU_PRODUCT_LABEL = f"{NVIDIA_GPU_RESOURCE}.product" +# With the `mixed` MIG strategy, the NVIDIA device plugin advertises MIG slices as a separate +# resource per MIG profile, e.g., nvidia.com/mig-1g.10gb (`+` in profile names is replaced +# with `.`, e.g., nvidia.com/mig-1g.10gb.me for 1g.10gb+me). GPUs with MIG disabled stay +# nvidia.com/gpu. See https://docs.nvidia.com/datacenter/cloud-native/kubernetes/latest/index.html +NVIDIA_MIG_RESOURCE_PREFIX = "nvidia.com/mig-" +# With the `single` MIG strategy, MIG slices are advertised as nvidia.com/gpu, and GPU Feature +# Discovery appends the MIG profile to the product label, e.g., A100-SXM4-80GB-MIG-1g.10gb. +# An invalid MIG layout is labeled as A100-SXM4-80GB-MIG-INVALID. +NVIDIA_MIG_PRODUCT_SEPARATOR = "-MIG-" +# The g.gb part of a MIG profile name, e.g., 1g.10gb. Suffixes such as +# +me, +me.all, -me, or +gfx don't change the memory. +# See https://docs.nvidia.com/datacenter/tesla/mig-user-guide/latest/supported-mig-profiles.html +MIG_PROFILE_MEMORY_REGEX = re.compile(r"^\d+g\.(\d+)gb") AMD_GPU_RESOURCE = "amd.com/gpu" AMD_GPU_NODE_TAINT = AMD_GPU_RESOURCE @@ -114,6 +130,8 @@ class KubernetesResources: ephemeral_storage: Decimal = Decimal("0") nvidia_gpu: Decimal = Decimal("0") amd_gpu: Decimal = Decimal("0") + nvidia_mig: dict[str, Decimal] = dataclasses.field(default_factory=dict) + """`mixed` strategy MIG slices by resource name, e.g., `nvidia.com/mig-1g.10gb`.""" @classmethod def from_kubernetes_map(cls, map_: Mapping[str, str]) -> Self: @@ -121,9 +139,16 @@ def from_kubernetes_map(cls, map_: Mapping[str, str]) -> Self: for resource in KubernetesResource: if (qty := map_.get(resource.value)) is not None: dct[resource.name.lower()] = parse_quantity(qty) - return cls(**dct) + nvidia_mig = { + resource: parse_quantity(qty) + for resource, qty in map_.items() + if resource.startswith(NVIDIA_MIG_RESOURCE_PREFIX) + } + return cls(**dct, nvidia_mig=nvidia_mig) def __getitem__(self, key: str) -> Decimal: + if key.startswith(NVIDIA_MIG_RESOURCE_PREFIX): + return self.nvidia_mig.get(key, Decimal("0")) try: resource = KubernetesResource(key) except ValueError: @@ -131,18 +156,21 @@ def __getitem__(self, key: str) -> Decimal: return getattr(self, resource.name.lower()) def __add__(self, other: Self) -> Self: - dct: dict[str, Decimal] = dataclasses.asdict(self) - qty: Decimal - for field, qty in dataclasses.asdict(other).items(): - dct[field] += qty - return type(self)(**dct) + return self._combine(other, operator.add) def __sub__(self, other: Self) -> Self: - dct: dict[str, Decimal] = dataclasses.asdict(self) - qty: Decimal - for field, qty in dataclasses.asdict(other).items(): - dct[field] -= qty - return type(self)(**dct) + return self._combine(other, operator.sub) + + def _combine(self, other: Self, op: Callable[[Decimal, Decimal], Decimal]) -> Self: + dct: dict[str, Decimal] = {} + for resource in KubernetesResource: + field = resource.name.lower() + dct[field] = op(getattr(self, field), getattr(other, field)) + nvidia_mig = { + resource: op(self[resource], other[resource]) + for resource in self.nvidia_mig.keys() | other.nvidia_mig.keys() + } + return type(self)(**dct, nvidia_mig=nvidia_mig) @dataclasses.dataclass(frozen=True) @@ -153,8 +181,13 @@ class ResourceRequestsLimits: gpu: int def to_kubernetes_map( - self, gpu_resource: Optional[AnyKubernetesGPUResource] = None + self, gpu_resource: Optional[Union[AnyKubernetesGPUResource, str]] = None ) -> dict[str, str]: + """ + Args: + gpu_resource: `nvidia.com/gpu`, `amd.com/gpu`, or a `mixed` strategy MIG resource, + e.g., `nvidia.com/mig-1g.10gb`. Required if GPUs are requested. + """ dct: dict[str, str] = {} if self.cpu is not None: dct[KubernetesResource.CPU.value] = str(self.cpu) @@ -165,7 +198,9 @@ def to_kubernetes_map( if self.gpu > 0: if gpu_resource is None: raise ValueError("gpu_resource is not specified") - dct[gpu_resource.value] = str(self.gpu) + if isinstance(gpu_resource, KubernetesResource): + gpu_resource = gpu_resource.value + dct[gpu_resource] = str(self.gpu) return dct @@ -203,11 +238,9 @@ def from_kubernetes_map(cls, map_: Mapping[str, str]) -> Self: disk_qty = map_.get(KubernetesResource.EPHEMERAL_STORAGE.value, "0") disk_mib = round(parse_quantity(disk_qty) / 2**20) gpu: int = 0 - for gpu_resource in GPU_RESOURCES: - gpu_qty = map_.get(gpu_resource) - if gpu_qty is not None: - gpu = round(parse_quantity(gpu_qty)) - break + gpu_resource = get_gpu_resource_from_kubernetes_map(map_) + if gpu_resource is not None: + gpu = round(parse_quantity(map_[gpu_resource])) return cls( cpu=cpu, memory_mib=memory_mib, @@ -240,6 +273,20 @@ def from_resources_spec(cls, spec: ResourcesSpec) -> Self: ) +def get_gpu_resource_from_kubernetes_map(map_: Mapping[str, str]) -> Optional[str]: + """ + Returns the GPU resource in a container's resource requests or limits map: `nvidia.com/gpu`, + `amd.com/gpu`, or a `mixed` strategy MIG resource, e.g., `nvidia.com/mig-1g.10gb`. + """ + for gpu_resource in GPU_RESOURCES: + if gpu_resource.value in map_: + return gpu_resource.value + for resource in map_: + if resource.startswith(NVIDIA_MIG_RESOURCE_PREFIX): + return resource + return None + + def adjust_resources_by_resource_requests( resources: Resources, resource_requests: ResourceRequests, @@ -378,6 +425,23 @@ def is_taint_tolerated(taint: V1Taint) -> bool: return taint.key in (NVIDIA_GPU_NODE_TAINT, AMD_GPU_NODE_TAINT) +class KubernetesOfferBackendData(CoreModel): + mig: bool = False + """Whether the offer's GPUs are NVIDIA MIG slices rather than whole GPUs.""" + mig_resource: Optional[str] = None + """The `mixed` strategy MIG resource that pods request for the offer's GPUs, e.g., + `nvidia.com/mig-1g.10gb`. Not set for other offers, including `single` strategy MIG slices, + whose pods request `nvidia.com/gpu` or `amd.com/gpu`.""" + + @classmethod + def load(cls, data: Mapping[str, Any]) -> Self: + return validate_extra_ignore(cls, data) + + def dump(self) -> dict[str, Any]: + # Offers without MIG slices keep empty backend data. + return self.model_dump(exclude_defaults=True) + + def get_instance_offers( api: CoreV1Api, region: str, unallocated_resources: bool ) -> list[InstanceOfferWithAvailability]: @@ -388,36 +452,56 @@ def get_instance_offers( for node in api.list_node().items: if (node_name := get_node_name(node)) is None: continue - offer = _get_instance_offer_from_node( - node=node, - node_name=node_name, - node_allocated_resources=( - None - if nodes_allocated_resources is None - else nodes_allocated_resources.get(node_name) - ), - region=region, + offers.extend( + _get_instance_offers_from_node( + node=node, + node_name=node_name, + node_allocated_resources=( + None + if nodes_allocated_resources is None + else nodes_allocated_resources.get(node_name) + ), + region=region, + ) ) - if offer is not None: - offers.append(offer) return offers def get_instance_offer_from_node( - node: V1Node, region: str + node: V1Node, region: str, gpu_resource: Optional[str] = None ) -> Optional[InstanceOfferWithAvailability]: + """ + Returns the node's offer for `gpu_resource`: the offer for the `mixed` strategy MIG resource + if `gpu_resource` is one, otherwise the node's main offer. + """ node_name = get_node_name(node) if node_name is None: return None - return _get_instance_offer_from_node( + mig_resource: Optional[str] = None + if gpu_resource is not None and gpu_resource.startswith(NVIDIA_MIG_RESOURCE_PREFIX): + mig_resource = gpu_resource + offers = _get_instance_offers_from_node( node=node, node_name=node_name, node_allocated_resources=None, region=region, ) + for offer in offers: + if KubernetesOfferBackendData.load(offer.backend_data).mig_resource == mig_resource: + return offer + return None + +def get_nvidia_gpu_from_node_labels( + labels: dict[str, str], mig_resource: Optional[str] = None +) -> Optional[Gpu]: + """ + Returns the GPU that the node provides via `nvidia.com/gpu`, or via the `mixed` strategy MIG + resource `mig_resource` if it is set. -def get_nvidia_gpu_from_node_labels(labels: dict[str, str]) -> Optional[Gpu]: + A MIG slice is returned with the parent GPU's name and the memory from the MIG profile name, + e.g., A100 with 10GB for 1g.10gb, the same way whole GPUs have their nominal memory. + """ # We rely on https://github.com/NVIDIA/k8s-device-plugin/tree/main/docs/gpu-feature-discovery # to detect gpus. Note that "nvidia.com/gpu.product" is not a short gpu name like "T4" or # "A100" but a product name like "Tesla-T4" or "A100-SXM4-40GB". @@ -425,6 +509,11 @@ def get_nvidia_gpu_from_node_labels(labels: dict[str, str]) -> Optional[Gpu]: gpu_product = labels.get(NVIDIA_GPU_PRODUCT_LABEL) if gpu_product is None: return None + # With the `single` MIG strategy, the product label ends with -MIG-. With the + # `mixed` strategy, it is the parent GPU's product name. + gpu_product, _, mig_profile = gpu_product.partition(NVIDIA_MIG_PRODUCT_SEPARATOR) + if mig_resource is not None: + mig_profile = mig_resource.removeprefix(NVIDIA_MIG_RESOURCE_PREFIX) gpu_product = gpu_product.replace("RTX-", "RTX") for gpu_name in NVIDIA_GPU_NAMES: if gpu_name.lower() in gpu_product.lower().split("-"): @@ -436,9 +525,22 @@ def get_nvidia_gpu_from_node_labels(labels: dict[str, str]) -> Optional[Gpu]: # A100 may come in two variants if "40GB" in gpu_product: gpu_memory = 40 * 1024 + if mig_profile: + mig_memory = _get_mig_memory_mib(mig_profile) + if mig_memory is None: + logger.warning("Failed to detect memory of MIG profile %s", mig_profile) + return None + gpu_memory = mig_memory return Gpu(vendor=AcceleratorVendor.NVIDIA, name=gpu_name, memory_mib=gpu_memory) +def is_nvidia_mig_product(labels: dict[str, str]) -> bool: + """ + Returns whether the node's `nvidia.com/gpu` resource provides `single` strategy MIG slices. + """ + return NVIDIA_MIG_PRODUCT_SEPARATOR in labels.get(NVIDIA_GPU_PRODUCT_LABEL, "") + + def get_amd_gpu_from_node_labels(labels: dict[str, str]) -> Optional[Gpu]: # (AMDGPUInfo.name, AMDGPUInfo.memory) pairs gpus: set[tuple[str, int]] = set() @@ -461,17 +563,21 @@ def get_amd_gpu_from_node_labels(labels: dict[str, str]) -> Optional[Gpu]: return None -def _get_instance_offer_from_node( +def _get_instance_offers_from_node( node: V1Node, node_name: str, node_allocated_resources: Optional[KubernetesResources], region: str, -) -> Optional[InstanceOfferWithAvailability]: +) -> list[InstanceOfferWithAvailability]: + """ + Returns the node's main offer, with the node's `nvidia.com/gpu` or `amd.com/gpu` GPUs, if any, + followed by an offer for each `mixed` strategy MIG resource with available slices. + """ try: node_spec = get_or_error(node.spec) if any(is_hard_taint(t) and not is_taint_tolerated(t) for t in node_spec.taints or []): logger.debug("Node %s: untolerated taint(s) found, skipping", node_name) - return None + return [] node_status = get_or_error(node.status) allocatable = get_or_error(node_status.allocatable) _cpu_arch: Optional[str] = None @@ -480,7 +586,7 @@ def _get_instance_offer_from_node( cpu_arch = normalize_arch(_cpu_arch).to_cpu_architecture() except ValueError as e: logger.exception("Failed to process node %s: %s: %s", node_name, type(e).__name__, e) - return None + return [] node_resources = KubernetesResources.from_kubernetes_map(allocatable) if node_allocated_resources is not None: @@ -488,8 +594,55 @@ def _get_instance_offer_from_node( cpu = max(0, int(node_resources.cpu)) memory_mib = max(0, int(node_resources.memory / 2**20)) disk_mib = max(0, int(node_resources.ephemeral_storage / 2**20)) + labels = get_node_labels(node) gpus = _get_gpus_from_node(node, node_name, node_resources) + make_offer = functools.partial( + _make_instance_offer, + node_name=node_name, + region=region, + cpu=cpu, + cpu_arch=cpu_arch, + memory_mib=memory_mib, + disk_mib=disk_mib, + ) + offers = [ + make_offer( + gpus=gpus, + backend_data=KubernetesOfferBackendData( + mig=( + len(gpus) > 0 + and gpus[0].vendor == AcceleratorVendor.NVIDIA + and is_nvidia_mig_product(labels) + ), + ), + ) + ] + # With the `mixed` MIG strategy, each MIG profile is a separate resource, and a pod can + # request only one of them, so each gets its own offer. The offers share the node's CPU, + # memory, and disk. + for mig_resource, mig_gpus in _get_nvidia_mig_gpus_from_node( + labels, node_name, node_resources + ): + offers.append( + make_offer( + gpus=mig_gpus, + backend_data=KubernetesOfferBackendData(mig=True, mig_resource=mig_resource), + ) + ) + return offers + +def _make_instance_offer( + *, + node_name: str, + region: str, + cpu: int, + cpu_arch: gpuhunt.CPUArchitecture, + memory_mib: int, + disk_mib: int, + gpus: list[Gpu], + backend_data: KubernetesOfferBackendData, +) -> InstanceOfferWithAvailability: return InstanceOfferWithAvailability( backend=BackendType.KUBERNETES, instance=InstanceType( @@ -507,6 +660,7 @@ def _get_instance_offer_from_node( region=region, availability=InstanceAvailability.AVAILABLE, instance_runtime=InstanceRuntime.RUNNER, + backend_data=backend_data.dump(), ) @@ -535,6 +689,38 @@ def _get_gpus_from_node( return [] +def _get_nvidia_mig_gpus_from_node( + labels: dict[str, str], node_name: str, node_resources: KubernetesResources +) -> list[tuple[str, list[Gpu]]]: + """ + Returns `(mig_resource, gpus)` pairs for the node's `mixed` strategy MIG resources with + available slices. + """ + mig_gpus: list[tuple[str, list[Gpu]]] = [] + for mig_resource, mig_qty in sorted(node_resources.nvidia_mig.items()): + mig_count = int(mig_qty) + if mig_count < 1: + continue + gpu = get_nvidia_gpu_from_node_labels(labels, mig_resource) + if gpu is None: + logger.warning( + "Node %s: MIG resource found, but failed to detect its GPU: %s=%d", + node_name, + mig_resource, + mig_count, + ) + continue + mig_gpus.append((mig_resource, [gpu] * mig_count)) + return mig_gpus + + +def _get_mig_memory_mib(mig_profile: str) -> Optional[int]: + match = MIG_PROFILE_MEMORY_REGEX.match(mig_profile) + if match is None: + return None + return int(match.group(1)) * 1024 + + def _get_nodes_allocated_resources(api: CoreV1Api) -> dict[str, KubernetesResources]: nodes_allocated_resources: dict[str, KubernetesResources] = {} for pod in api.list_pod_for_all_namespaces().items: diff --git a/src/tests/_internal/core/backends/kubernetes/test_compute.py b/src/tests/_internal/core/backends/kubernetes/test_compute.py index 39ef170850..79c51b6ac2 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_compute.py +++ b/src/tests/_internal/core/backends/kubernetes/test_compute.py @@ -1,7 +1,20 @@ +import uuid +from typing import Optional from unittest.mock import MagicMock, patch -from dstack._internal.core.backends.kubernetes.compute import KubernetesCompute +import pytest +from kubernetes.client import V1Node, V1NodeSpec, V1NodeStatus, V1ObjectMeta + +from dstack._internal.core.backends.kubernetes.compute import ( + KubernetesCompute, + _get_pod_spec_parameters_for_gpu, + _KubernetesSkipOfferCache, +) from dstack._internal.core.backends.kubernetes.models import KubeconfigConfig, KubernetesConfig +from dstack._internal.core.backends.kubernetes.resources import ( + KubernetesOfferBackendData, + get_instance_offers, +) from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.instances import ( Disk, @@ -11,7 +24,7 @@ InstanceType, Resources, ) -from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.resources import GPUSpec, ResourcesSpec from dstack._internal.core.models.runs import Requirements @@ -92,3 +105,186 @@ def test_get_offers_without_full_offers_adjusts_to_requested_slice(): assert adjusted_resources.memory_mib == 8 * 1024 assert len(adjusted_resources.gpus) == 1 assert adjusted_resources.disk.size_mib == 100 * 1024 + + +def _node(name: str, gpu_product: str, gpu_allocatable: dict[str, str]) -> V1Node: + # Labels and resources as published by NVIDIA GPU Feature Discovery and the device plugin + return V1Node( + metadata=V1ObjectMeta(name=name, labels={"nvidia.com/gpu.product": gpu_product}), + spec=V1NodeSpec(), + status=V1NodeStatus( + allocatable={ + "cpu": "200", + "memory": "1800Gi", + "ephemeral-storage": "1000Gi", + **gpu_allocatable, + } + ), + ) + + +def _api(nodes: list[V1Node]) -> MagicMock: + api = MagicMock() + api.list_node.return_value.items = nodes + api.list_pod_for_all_namespaces.return_value.items = [] + return api + + +# A `mixed` strategy MIG node: 4 GPUs with MIG disabled, 2 GPUs as 7 x 1g.10gb, +# 2 GPUs as 2 x 3g.40gb +_MIXED_NODE = _node( + "mixed", + "NVIDIA-A100-SXM4-80GB", + {"nvidia.com/gpu": "4", "nvidia.com/mig-1g.10gb": "14", "nvidia.com/mig-3g.40gb": "4"}, +) +_SINGLE_NODE = _node("single", "NVIDIA-A100-SXM4-80GB-MIG-1g.10gb", {"nvidia.com/gpu": "56"}) + + +class TestGetOffersWithMig: + @pytest.mark.parametrize( + ["gpu", "full_offers", "expected"], + [ + # Without a GPU memory cap, runs only get whole GPUs + pytest.param("A100", False, [("A100", 80, 1)], id="name"), + pytest.param("24GB..", False, [("A100", 80, 1)], id="memory-without-cap"), + pytest.param("A100:10GB", False, [("A100", 10, 1)], id="slice"), + pytest.param("A100:10GB:2", False, [("A100", 10, 2)], id="two-slices"), + pytest.param( + "10GB..40GB", False, [("A100", 10, 1), ("A100", 40, 1)], id="memory-range" + ), + # The mixed node's MIG offers would duplicate its main offer + pytest.param("0", False, [(None, 0, 0)], id="no-gpu"), + # Full offers show what the node has, including slices + pytest.param( + "A100", + True, + [("A100", 80, 4), ("A100", 10, 14), ("A100", 40, 4)], + id="full-offers", + ), + pytest.param( + "0..", + True, + [("A100", 80, 4), ("A100", 10, 14), ("A100", 40, 4)], + id="full-offers-no-gpu", + ), + ], + ) + def test_mixed_strategy_node( + self, gpu: str, full_offers: bool, expected: list[tuple[Optional[str], int, int]] + ): + assert _get_offer_gpus(_MIXED_NODE, gpu, full_offers) == expected + + @pytest.mark.parametrize( + ["gpu", "expected"], + [ + pytest.param("A100", [], id="name"), + pytest.param("A100:10GB", [("A100", 10, 1)], id="slice"), + pytest.param("0", [(None, 0, 0)], id="no-gpu"), + ], + ) + def test_single_strategy_node(self, gpu: str, expected: list[tuple[Optional[str], int, int]]): + assert _get_offer_gpus(_SINGLE_NODE, gpu, full_offers=False) == expected + + +def _get_offer_gpus( + node: V1Node, gpu: str, full_offers: bool +) -> list[tuple[Optional[str], int, int]]: + compute = _compute() + compute.get_all_offers_with_availability = MagicMock( + return_value=get_instance_offers(_api([node]), "ctx", unallocated_resources=False) + ) + requirements = Requirements( + resources=ResourcesSpec(cpu="2..", memory="8GB..", disk="100GB..", gpu=gpu) + ) + offers = compute.get_offers(requirements, full_offers=full_offers, unallocated_resources=False) + result: list[tuple[Optional[str], int, int]] = [] + for offer in offers: + gpus = offer.instance.resources.gpus + if gpus: + result.append((gpus[0].name, gpus[0].memory_mib // 1024, len(gpus))) + else: + result.append((None, 0, 0)) + return result + + +class TestGetPodSpecParametersForGpu: + _NODES = [ + _node("whole", "NVIDIA-A100-SXM4-80GB", {"nvidia.com/gpu": "8"}), + _SINGLE_NODE, + _MIXED_NODE, + _node("mixed-h100", "NVIDIA-H100-80GB-HBM3", {"nvidia.com/mig-1g.10gb": "7"}), + ] + + @pytest.mark.parametrize( + ["gpu", "mig_resource", "expected_resource", "expected_products"], + [ + pytest.param( + "A100", + None, + "nvidia.com/gpu", + # The mixed node's product label is the parent GPU's, same as the whole node's + ["NVIDIA-A100-SXM4-80GB"], + id="whole-gpu-excludes-single-strategy-slices", + ), + pytest.param( + "A100:10GB", + None, + "nvidia.com/gpu", + ["NVIDIA-A100-SXM4-80GB-MIG-1g.10gb"], + id="single-strategy-slice", + ), + pytest.param( + "A100:10GB", + "nvidia.com/mig-1g.10gb", + "nvidia.com/mig-1g.10gb", + # Only nodes with the MIG resource and the requested GPU model + ["NVIDIA-A100-SXM4-80GB"], + id="mixed-strategy-slice", + ), + ], + ) + def test_nvidia( + self, + gpu: str, + mig_resource: Optional[str], + expected_resource: str, + expected_products: list[str], + ): + gpu_resource, node_affinity, node_taint = _get_pod_spec_parameters_for_gpu( + _api(self._NODES), GPUSpec.model_validate(gpu), mig_resource + ) + assert gpu_resource == expected_resource + assert node_taint == "nvidia.com/gpu" + node_selector = node_affinity.required_during_scheduling_ignored_during_execution + expression = node_selector.node_selector_terms[0].match_expressions[0] + assert expression.key == "nvidia.com/gpu.product" + assert sorted(expression.values) == expected_products + + +class TestKubernetesSkipOfferCache: + def test_mig_offers_are_skipped_separately_by_mig_resource(self): + cache = _KubernetesSkipOfferCache(ttl=60) + run = MagicMock(id=uuid.uuid4()) + job = MagicMock() + whole_offer = _node_offer() + other_whole_offer = _node_offer() + mig_offer = _mig_offer("nvidia.com/mig-1g.10gb") + other_mig_offer = _mig_offer("nvidia.com/mig-3g.40gb") + + cache.add(run, job, mig_offer) + assert cache.check(run, job, mig_offer) + assert not cache.check(run, job, other_mig_offer) + assert not cache.check(run, job, whole_offer) + + # As before, a failed offer that isn't a mixed strategy MIG offer skips all such offers + # in the region + cache.add(run, job, whole_offer) + assert cache.check(run, job, other_whole_offer) + assert not cache.check(run, job, other_mig_offer) + assert not cache.check(MagicMock(id=uuid.uuid4()), job, other_whole_offer) + + +def _mig_offer(mig_resource: str) -> InstanceOfferWithAvailability: + offer = _node_offer() + offer.backend_data = KubernetesOfferBackendData(mig=True, mig_resource=mig_resource).dump() + return offer diff --git a/src/tests/_internal/core/backends/kubernetes/test_resources.py b/src/tests/_internal/core/backends/kubernetes/test_resources.py index e1c55dc96f..f0b171b730 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_resources.py +++ b/src/tests/_internal/core/backends/kubernetes/test_resources.py @@ -1,15 +1,35 @@ import logging +from decimal import Decimal +from typing import Optional +from unittest.mock import MagicMock import pytest from gpuhunt import AcceleratorVendor +from kubernetes.client import ( + V1Container, + V1Node, + V1NodeSpec, + V1NodeStatus, + V1ObjectMeta, + V1Pod, + V1PodSpec, + V1PodStatus, + V1ResourceRequirements, +) from dstack._internal.core.backends.kubernetes.resources import ( + KubernetesOfferBackendData, KubernetesResource, + KubernetesResources, ResourceLimits, ResourceRequests, adjust_resources_by_resource_requests, get_amd_gpu_from_node_labels, + get_gpu_resource_from_kubernetes_map, + get_instance_offer_from_node, + get_instance_offers, get_nvidia_gpu_from_node_labels, + is_nvidia_mig_product, validate_label_key, validate_label_value, ) @@ -37,6 +57,71 @@ def test_returns_correct_memory_for_different_A100(self): {"nvidia.com/gpu.product": "A100-SXM4-80GB"} ) == Gpu(vendor=AcceleratorVendor.NVIDIA, name="A100", memory_mib=80 * 1024) + @pytest.mark.parametrize( + ["product", "expected_name", "expected_memory_gb"], + [ + # The label from NVIDIA's "MIG Support in Kubernetes" docs + pytest.param("A100-SXM4-40GB-MIG-1g.5gb", "A100", 5, id="a100-40gb"), + # The label set by GPU Operator v26.7.2 on GKE + pytest.param("NVIDIA-A100-SXM4-40GB-MIG-1g.5gb", "A100", 5, id="a100-40gb-gke"), + pytest.param("NVIDIA-A100-SXM4-80GB-MIG-1g.10gb", "A100", 10, id="a100-80gb"), + pytest.param("NVIDIA-H100-80GB-HBM3-MIG-3g.40gb", "H100", 40, id="h100"), + # GPU Feature Discovery replaces `+` with `.` in labels: 1g.5gb+me -> 1g.5gb.me + pytest.param("NVIDIA-A100-SXM4-40GB-MIG-1g.5gb.me", "A100", 5, id="me-suffix"), + pytest.param("NVIDIA-A100-SXM4-80GB-MIG-1g.10gb-SHARED", "A100", 10, id="shared"), + ], + ) + def test_returns_mig_profile_memory_for_single_strategy_slice( + self, product: str, expected_name: str, expected_memory_gb: int + ): + assert get_nvidia_gpu_from_node_labels({"nvidia.com/gpu.product": product}) == Gpu( + vendor=AcceleratorVendor.NVIDIA, + name=expected_name, + memory_mib=expected_memory_gb * 1024, + ) + + def test_returns_none_for_invalid_single_strategy_layout( + self, caplog: pytest.LogCaptureFixture + ): + caplog.set_level(logging.WARNING) + labels = {"nvidia.com/gpu.product": "NVIDIA-A100-SXM4-80GB-MIG-INVALID"} + assert get_nvidia_gpu_from_node_labels(labels) is None + assert "Failed to detect memory of MIG profile INVALID" in caplog.text + + @pytest.mark.parametrize( + ["mig_resource", "expected_memory_gb"], + [ + pytest.param("nvidia.com/mig-1g.10gb", 10, id="1g"), + pytest.param("nvidia.com/mig-3g.40gb", 40, id="3g"), + # The device plugin replaces `+` with `.`: 1g.10gb+me -> mig-1g.10gb.me + pytest.param("nvidia.com/mig-1g.10gb.me", 10, id="me-suffix"), + ], + ) + def test_returns_mig_resource_memory_for_mixed_strategy_slice( + self, mig_resource: str, expected_memory_gb: int + ): + # With the mixed strategy, the product label is the parent GPU's + labels = {"nvidia.com/gpu.product": "NVIDIA-A100-SXM4-80GB"} + assert get_nvidia_gpu_from_node_labels(labels, mig_resource) == Gpu( + vendor=AcceleratorVendor.NVIDIA, name="A100", memory_mib=expected_memory_gb * 1024 + ) + + +class TestIsNvidiaMigProduct: + def test_true_for_single_strategy_slice(self): + labels = {"nvidia.com/gpu.product": "NVIDIA-A100-SXM4-80GB-MIG-1g.10gb"} + assert is_nvidia_mig_product(labels) + + @pytest.mark.parametrize( + "labels", + [ + pytest.param({}, id="no-labels"), + pytest.param({"nvidia.com/gpu.product": "NVIDIA-A100-SXM4-80GB"}, id="whole-gpu"), + ], + ) + def test_false_for_whole_gpu(self, labels: dict[str, str]): + assert not is_nvidia_mig_product(labels) + class TestGetAMDGPUFromNodeLabels: def test_returns_no_gpus_if_no_labels(self): @@ -93,6 +178,53 @@ def test_kubernetes_map_round_trip(self): map_ = requests.to_kubernetes_map(KubernetesResource.AMD_GPU) assert ResourceRequests.from_kubernetes_map(map_) == requests + def test_kubernetes_map_round_trip_with_mig_resource(self): + requests = ResourceRequests(cpu=8, memory_mib=16384, disk_mib=51200, gpu=2) + map_ = requests.to_kubernetes_map("nvidia.com/mig-1g.10gb") + assert map_["nvidia.com/mig-1g.10gb"] == "2" + assert ResourceRequests.from_kubernetes_map(map_) == requests + + +class TestGetGpuResourceFromKubernetesMap: + @pytest.mark.parametrize( + ["map_", "expected"], + [ + pytest.param({"cpu": "1"}, None, id="no-gpu"), + pytest.param({"cpu": "1", "nvidia.com/gpu": "1"}, "nvidia.com/gpu", id="nvidia"), + pytest.param({"cpu": "1", "amd.com/gpu": "1"}, "amd.com/gpu", id="amd"), + pytest.param( + {"cpu": "1", "nvidia.com/mig-1g.10gb": "1"}, "nvidia.com/mig-1g.10gb", id="mig" + ), + ], + ) + def test_returns_gpu_resource(self, map_: dict[str, str], expected: Optional[str]): + assert get_gpu_resource_from_kubernetes_map(map_) == expected + + +class TestKubernetesResources: + def test_from_kubernetes_map_includes_mig_resources(self): + resources = KubernetesResources.from_kubernetes_map( + {"cpu": "8", "nvidia.com/gpu": "4", "nvidia.com/mig-1g.10gb": "14"} + ) + assert resources.nvidia_gpu == 4 + assert resources["nvidia.com/mig-1g.10gb"] == 14 + assert resources["nvidia.com/mig-3g.40gb"] == 0 + + def test_subtracts_mig_resources(self): + allocatable = KubernetesResources.from_kubernetes_map( + {"cpu": "8", "nvidia.com/mig-1g.10gb": "14", "nvidia.com/mig-3g.40gb": "4"} + ) + allocated = KubernetesResources.from_kubernetes_map( + {"cpu": "2", "nvidia.com/mig-1g.10gb": "5"} + ) + assert allocatable - allocated == KubernetesResources( + cpu=Decimal(6), + nvidia_mig={ + "nvidia.com/mig-1g.10gb": Decimal(9), + "nvidia.com/mig-3g.40gb": Decimal(4), + }, + ) + class TestResourceLimits: def test_from_resources_spec_uses_upper_bounds_and_leaves_unset_as_none(self): @@ -218,3 +350,153 @@ def test_valid_value(self, value: str): def test_invalid_value(self, value: str, expected_error: str): with pytest.raises(ValueError, match=expected_error): validate_label_value(value) + + +def _node(name: str, gpu_product: str, gpu_allocatable: dict[str, str]) -> V1Node: + # Labels and resources as published by NVIDIA GPU Feature Discovery and the device plugin + return V1Node( + metadata=V1ObjectMeta(name=name, labels={"nvidia.com/gpu.product": gpu_product}), + spec=V1NodeSpec(), + status=V1NodeStatus( + allocatable={ + "cpu": "200", + "memory": "1800Gi", + "ephemeral-storage": "1000Gi", + **gpu_allocatable, + } + ), + ) + + +def _pod(node_name: str, requests: dict[str, str]) -> V1Pod: + return V1Pod( + spec=V1PodSpec( + node_name=node_name, + containers=[ + V1Container(name="c", resources=V1ResourceRequirements(requests=requests)) + ], + ), + status=V1PodStatus(phase="Running"), + ) + + +def _api(nodes: list[V1Node], pods: Optional[list[V1Pod]] = None) -> MagicMock: + api = MagicMock() + api.list_node.return_value.items = nodes + api.list_pod_for_all_namespaces.return_value.items = pods or [] + return api + + +# 4 GPUs with MIG disabled, 2 GPUs as 7 x 1g.10gb, 2 GPUs as 2 x 3g.40gb +_MIXED_NODE_GPU_ALLOCATABLE = { + "nvidia.com/gpu": "4", + "nvidia.com/mig-1g.10gb": "14", + "nvidia.com/mig-3g.40gb": "4", +} + + +def _offer_gpus(offer) -> tuple[Optional[str], int, int, KubernetesOfferBackendData]: + gpus = offer.instance.resources.gpus + return ( + gpus[0].name if gpus else None, + gpus[0].memory_mib // 1024 if gpus else 0, + len(gpus), + KubernetesOfferBackendData.load(offer.backend_data), + ) + + +class TestGetInstanceOffers: + def test_node_without_mig_has_one_offer_with_empty_backend_data(self): + api = _api([_node("node", "NVIDIA-A100-SXM4-80GB", {"nvidia.com/gpu": "8"})]) + offers = get_instance_offers(api, "ctx", unallocated_resources=False) + assert len(offers) == 1 + assert offers[0].backend_data == {} + assert _offer_gpus(offers[0])[:3] == ("A100", 80, 8) + + def test_single_strategy_node_offers_slices(self): + api = _api([_node("node", "NVIDIA-A100-SXM4-80GB-MIG-1g.10gb", {"nvidia.com/gpu": "56"})]) + offers = get_instance_offers(api, "ctx", unallocated_resources=False) + assert [_offer_gpus(o) for o in offers] == [ + ("A100", 10, 56, KubernetesOfferBackendData(mig=True)), + ] + + def test_mixed_strategy_node_has_offer_per_mig_resource(self): + api = _api([_node("node", "NVIDIA-A100-SXM4-80GB", _MIXED_NODE_GPU_ALLOCATABLE)]) + offers = get_instance_offers(api, "ctx", unallocated_resources=False) + assert [_offer_gpus(o) for o in offers] == [ + ("A100", 80, 4, KubernetesOfferBackendData()), + ( + "A100", + 10, + 14, + KubernetesOfferBackendData(mig=True, mig_resource="nvidia.com/mig-1g.10gb"), + ), + ( + "A100", + 40, + 4, + KubernetesOfferBackendData(mig=True, mig_resource="nvidia.com/mig-3g.40gb"), + ), + ] + # The offers share the node's CPU, memory, and disk + assert {o.instance.name for o in offers} == {"node"} + assert {o.instance.resources.cpus for o in offers} == {200} + + def test_mixed_strategy_node_skips_mig_resources_with_no_slices(self): + # After a MIG layout change, the device plugin keeps the old resources at 0 + # (as seen on GKE with GPU Operator v26.7.2) + gpu_allocatable = { + "nvidia.com/gpu": "0", + "nvidia.com/mig-1g.5gb": "0", + "nvidia.com/mig-3g.20gb": "1", + "nvidia.com/mig-4g.20gb": "1", + } + api = _api([_node("node", "NVIDIA-A100-SXM4-40GB", gpu_allocatable)]) + offers = get_instance_offers(api, "ctx", unallocated_resources=False) + assert [_offer_gpus(o) for o in offers] == [ + (None, 0, 0, KubernetesOfferBackendData()), + ( + "A100", + 20, + 1, + KubernetesOfferBackendData(mig=True, mig_resource="nvidia.com/mig-3g.20gb"), + ), + ( + "A100", + 20, + 1, + KubernetesOfferBackendData(mig=True, mig_resource="nvidia.com/mig-4g.20gb"), + ), + ] + + def test_unallocated_resources_subtract_mig_slices(self): + api = _api( + [_node("node", "NVIDIA-A100-SXM4-80GB", _MIXED_NODE_GPU_ALLOCATABLE)], + pods=[ + _pod("node", {"cpu": "50", "nvidia.com/gpu": "2"}), + _pod("node", {"nvidia.com/mig-1g.10gb": "5"}), + _pod("node", {"nvidia.com/mig-3g.40gb": "4"}), + ], + ) + offers = get_instance_offers(api, "ctx", unallocated_resources=True) + # No 3g.40gb offer: all of its slices are in use + assert [_offer_gpus(o)[:3] for o in offers] == [("A100", 80, 2), ("A100", 10, 9)] + assert {o.instance.resources.cpus for o in offers} == {150} + + +class TestGetInstanceOfferFromNode: + @pytest.mark.parametrize( + ["gpu_resource", "expected_memory_gb", "expected_count"], + [ + pytest.param(None, 80, 4, id="no-gpu"), + pytest.param("nvidia.com/gpu", 80, 4, id="whole-gpu"), + pytest.param("nvidia.com/mig-3g.40gb", 40, 4, id="mig"), + ], + ) + def test_returns_offer_for_gpu_resource( + self, gpu_resource: Optional[str], expected_memory_gb: int, expected_count: int + ): + node = _node("node", "NVIDIA-A100-SXM4-80GB", _MIXED_NODE_GPU_ALLOCATABLE) + offer = get_instance_offer_from_node(node, "ctx", gpu_resource) + assert offer is not None + assert _offer_gpus(offer)[1:3] == (expected_memory_gb, expected_count) From 204664a14dec70a34421c8663af66592cbfa40b0 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 07:50:32 +0000 Subject: [PATCH 2/9] Shorten the Kubernetes MIG docs note Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EDgQBSHAm1rxrhSUXWziy3 --- mkdocs/docs/concepts/backends.md | 37 +++++++------------------------- 1 file changed, 8 insertions(+), 29 deletions(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index a30e6327e3..a78f8d3d98 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1397,40 +1397,19 @@ projects: This applies to offers shown in `dstack apply` (run plans), during provisioning, and in `dstack offer`. Unlike other backends, offers for the `kubernetes` backend always reflect the lower limit of the range. -??? info "Multi-Instance GPU (MIG)" +??? info "MIG" The `kubernetes` backend supports NVIDIA [MIG](https://docs.nvidia.com/datacenter/tesla/mig-user-guide/latest/index.html) slices - configured with the NVIDIA GPU Operator's [MIG Manager](https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/gpu-operator-mig.html), - with either the `single` or the `mixed` [MIG strategy](https://docs.nvidia.com/datacenter/cloud-native/kubernetes/latest/index.html). + set up by the GPU Operator's [MIG Manager](https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/gpu-operator-mig.html), + with the `single` or `mixed` [strategy](https://docs.nvidia.com/datacenter/cloud-native/kubernetes/latest/index.html). + A slice is offered as a GPU with the slice's memory, e.g., a `1g.10gb` slice of an A100 is `A100:10GB`. - A MIG slice is offered as a GPU of the same model with the slice's memory, taken from the MIG profile name. - For example, a `1g.10gb` slice of an A100 is `A100:10GB`. - - To run on a slice, set an upper limit on GPU memory: - -
- - ```yaml - type: dev-environment - ide: vscode - - resources: - gpu: A100:10GB - ``` - -
- - Without an upper limit on GPU memory (e.g., `gpu: A100` or `gpu: 24GB..`), runs only get whole GPUs. - To see slices in `dstack offer`, set the limit (e.g., `--gpu A100:10GB`) or pass `--full-offers`. - - With the `mixed` strategy, a node that has whole GPUs and slices of several sizes appears as several offers, one per GPU type. - With `--full-offers`, each of them shows the node's full CPU, memory, and disk, which they share. + A run gets slices only if it sets an upper limit on GPU memory (e.g., `gpu: A100:10GB`); with `gpu: A100`, it gets whole GPUs only. + The same applies to `dstack offer`, unless you pass `--full-offers`. Limitations: - * NVIDIA doesn't support NCCL with MIG, so a run with several slices (e.g., `gpu: A100:10GB:2`) can't use NCCL-based multi-GPU training. - Several slices suit independent processes, one per slice (a single process can use several slices only with NVIDIA driver R570 or newer). - To run several copies of a model, use a service with `replicas`, each on one slice. - * Slices with the same memory (e.g., `3g.40gb` and `4g.40gb`) and media-engine variants (e.g., `1g.10gb+me`) can't be requested specifically. + * NVIDIA doesn't support NCCL with MIG, so several slices per run (e.g., `gpu: A100:10GB:2`) suit independent processes, not multi-GPU training. + * Slices with the same memory (e.g., `3g.40gb` and `4g.40gb`) and `+me` variants can't be requested specifically. * `dstack metrics` doesn't show GPU metrics for slices. > To learn more, see the [Lambda](../examples/clusters/lambda/#kubernetes) and [Crusoe](../examples/clusters/crusoe/#kubernetes) examples. From f6a6743091212b7b60bf33573910499f3d983871 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 07:52:01 +0000 Subject: [PATCH 3/9] Drop the NCCL note from the MIG docs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EDgQBSHAm1rxrhSUXWziy3 --- mkdocs/docs/concepts/backends.md | 1 - 1 file changed, 1 deletion(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index a78f8d3d98..9b7ca24952 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1408,7 +1408,6 @@ projects: Limitations: - * NVIDIA doesn't support NCCL with MIG, so several slices per run (e.g., `gpu: A100:10GB:2`) suit independent processes, not multi-GPU training. * Slices with the same memory (e.g., `3g.40gb` and `4g.40gb`) and `+me` variants can't be requested specifically. * `dstack metrics` doesn't show GPU metrics for slices. From 5800d419b4d0acff48605d9136502f7d17fbc3b5 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 07:52:40 +0000 Subject: [PATCH 4/9] Clarify the MIG slice matching limitation Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EDgQBSHAm1rxrhSUXWziy3 --- mkdocs/docs/concepts/backends.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 9b7ca24952..384f790277 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1408,7 +1408,7 @@ projects: Limitations: - * Slices with the same memory (e.g., `3g.40gb` and `4g.40gb`) and `+me` variants can't be requested specifically. + * Slices are matched by memory only, so `gpu: A100:40GB` can get either a `3g.40gb` or a `4g.40gb` slice. * `dstack metrics` doesn't show GPU metrics for slices. > To learn more, see the [Lambda](../examples/clusters/lambda/#kubernetes) and [Crusoe](../examples/clusters/crusoe/#kubernetes) examples. From 416b9a188a10a0e3e994c3751526d96828a582c5 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 07:53:26 +0000 Subject: [PATCH 5/9] Describe MIG slice offers in the docs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EDgQBSHAm1rxrhSUXWziy3 --- mkdocs/docs/concepts/backends.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 384f790277..d69da94f0c 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1404,7 +1404,9 @@ projects: A slice is offered as a GPU with the slice's memory, e.g., a `1g.10gb` slice of an A100 is `A100:10GB`. A run gets slices only if it sets an upper limit on GPU memory (e.g., `gpu: A100:10GB`); with `gpu: A100`, it gets whole GPUs only. - The same applies to `dstack offer`, unless you pass `--full-offers`. + + Offers in `dstack offer` and `dstack apply` follow the same rule; `--full-offers` shows all of them. + A node's whole GPUs and each of its slice sizes are separate offers, e.g., `A100:80GB:4` and `A100:10GB:14`. Limitations: From bb73bd6ba2cce2d6dfcad42bab844d13021b13a5 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 07:55:38 +0000 Subject: [PATCH 6/9] Rewrite the MIG docs note in dstack terms Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EDgQBSHAm1rxrhSUXWziy3 --- mkdocs/docs/concepts/backends.md | 31 +++++++++++++++++++++++-------- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index d69da94f0c..8836af5083 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1398,19 +1398,34 @@ projects: This applies to offers shown in `dstack apply` (run plans), during provisioning, and in `dstack offer`. Unlike other backends, offers for the `kubernetes` backend always reflect the lower limit of the range. ??? info "MIG" - The `kubernetes` backend supports NVIDIA [MIG](https://docs.nvidia.com/datacenter/tesla/mig-user-guide/latest/index.html) slices - set up by the GPU Operator's [MIG Manager](https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/gpu-operator-mig.html), - with the `single` or `mixed` [strategy](https://docs.nvidia.com/datacenter/cloud-native/kubernetes/latest/index.html). - A slice is offered as a GPU with the slice's memory, e.g., a `1g.10gb` slice of an A100 is `A100:10GB`. + If the cluster's NVIDIA GPUs are partitioned with [MIG](https://docs.nvidia.com/datacenter/tesla/mig-user-guide/latest/index.html), + `dstack` treats each MIG slice as a GPU with the slice's memory. For example, a `1g.10gb` slice of an A100 is an `A100:10GB` GPU. - A run gets slices only if it sets an upper limit on GPU memory (e.g., `gpu: A100:10GB`); with `gpu: A100`, it gets whole GPUs only. + To run on slices, specify the GPU memory: - Offers in `dstack offer` and `dstack apply` follow the same rule; `--full-offers` shows all of them. - A node's whole GPUs and each of its slice sizes are separate offers, e.g., `A100:80GB:4` and `A100:10GB:14`. +
+ + ```yaml + type: dev-environment + ide: vscode + + resources: + # One 10GB slice; A100:10GB:2 for two slices + gpu: A100:10GB + ``` + +
+ + If the GPU memory isn't specified or has no maximum (e.g., `gpu: A100` or `gpu: 24GB..`), the run gets whole GPUs only. + `dstack offer` and `dstack apply` show offers the same way, unless you pass `--full-offers`. + Slices are listed separately from whole GPUs, e.g., `A100:10GB:14` and `A100:80GB:4` on the same node. + + MIG must be configured with the GPU Operator's [MIG Manager](https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/gpu-operator-mig.html) + (the `single` or `mixed` [strategy](https://docs.nvidia.com/datacenter/cloud-native/kubernetes/latest/index.html)). Limitations: - * Slices are matched by memory only, so `gpu: A100:40GB` can get either a `3g.40gb` or a `4g.40gb` slice. + * Slices with the same memory are the same GPU for `dstack`, so `gpu: A100:40GB` can get either a `3g.40gb` or a `4g.40gb` slice. * `dstack metrics` doesn't show GPU metrics for slices. > To learn more, see the [Lambda](../examples/clusters/lambda/#kubernetes) and [Crusoe](../examples/clusters/crusoe/#kubernetes) examples. From a53f095b7e2cc3e7c9c7cb9f2a4d3e5903ac5b17 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 07:56:19 +0000 Subject: [PATCH 7/9] State that a GPU name alone never gets MIG slices Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EDgQBSHAm1rxrhSUXWziy3 --- mkdocs/docs/concepts/backends.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 8836af5083..f4f15e45c6 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1416,7 +1416,7 @@ projects: - If the GPU memory isn't specified or has no maximum (e.g., `gpu: A100` or `gpu: 24GB..`), the run gets whole GPUs only. + With just a GPU name (e.g., `gpu: A100`) or memory without a maximum (e.g., `gpu: 24GB..`), the run gets whole GPUs only, never slices. `dstack offer` and `dstack apply` show offers the same way, unless you pass `--full-offers`. Slices are listed separately from whole GPUs, e.g., `A100:10GB:14` and `A100:80GB:4` on the same node. From 74afd8d69682cfcd061b2c78e3f1ca680b7f15c8 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 07:57:01 +0000 Subject: [PATCH 8/9] Inline the MIG docs example Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EDgQBSHAm1rxrhSUXWziy3 --- mkdocs/docs/concepts/backends.md | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index f4f15e45c6..1d3c150d88 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1401,21 +1401,7 @@ projects: If the cluster's NVIDIA GPUs are partitioned with [MIG](https://docs.nvidia.com/datacenter/tesla/mig-user-guide/latest/index.html), `dstack` treats each MIG slice as a GPU with the slice's memory. For example, a `1g.10gb` slice of an A100 is an `A100:10GB` GPU. - To run on slices, specify the GPU memory: - -
- - ```yaml - type: dev-environment - ide: vscode - - resources: - # One 10GB slice; A100:10GB:2 for two slices - gpu: A100:10GB - ``` - -
- + To run on slices, specify the GPU memory, e.g., `gpu: A100:10GB` for one slice or `gpu: A100:10GB:2` for two. With just a GPU name (e.g., `gpu: A100`) or memory without a maximum (e.g., `gpu: 24GB..`), the run gets whole GPUs only, never slices. `dstack offer` and `dstack apply` show offers the same way, unless you pass `--full-offers`. Slices are listed separately from whole GPUs, e.g., `A100:10GB:14` and `A100:80GB:4` on the same node. From 8f5fea67a20c63891470b48dff5860cc668d0f55 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 07:57:26 +0000 Subject: [PATCH 9/9] Restore the earlier wording of the whole-GPU sentence Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EDgQBSHAm1rxrhSUXWziy3 --- mkdocs/docs/concepts/backends.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 1d3c150d88..7651258219 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1402,7 +1402,7 @@ projects: `dstack` treats each MIG slice as a GPU with the slice's memory. For example, a `1g.10gb` slice of an A100 is an `A100:10GB` GPU. To run on slices, specify the GPU memory, e.g., `gpu: A100:10GB` for one slice or `gpu: A100:10GB:2` for two. - With just a GPU name (e.g., `gpu: A100`) or memory without a maximum (e.g., `gpu: 24GB..`), the run gets whole GPUs only, never slices. + If the GPU memory isn't specified or has no maximum (e.g., `gpu: A100` or `gpu: 24GB..`), the run gets whole GPUs only. `dstack offer` and `dstack apply` show offers the same way, unless you pass `--full-offers`. Slices are listed separately from whole GPUs, e.g., `A100:10GB:14` and `A100:80GB:4` on the same node.