Support NVIDIA MIG slices in the Kubernetes backend

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

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