diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index e488520fb2..7651258219 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1397,6 +1397,23 @@ 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" + 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, e.g., `gpu: A100:10GB` for one slice or `gpu: A100:10GB:2` for two. + 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 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. ### 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)