Normalize measured GPU VRAM in SSH and Kubernetes

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

@@ -47,6 +47,7 @@
     AMD_GPU_NAME_TO_DEVICE_IDS,
     AMD_GPU_NODE_TAINT,
     LABEL_VALUE_MAX_LENGTH,
+    NVIDIA_GPU_MEMORY_LABEL,
     NVIDIA_GPU_NODE_TAINT,
     NVIDIA_GPU_PRODUCT_LABEL,
     OBJECT_NAME_MAX_LENGTH,
@@ -790,19 +791,19 @@ def _get_pod_spec_parameters_for_gpu(
 def _get_nvidia_gpu_node_affinity(
     gpu_spec: GPUSpec, nodes: list[client.V1Node]
 ) -> client.V1NodeAffinity:
-    matching_gpu_label_values: set[str] = set()
+    matching_gpu_labels: set[tuple[str, str]] = set()
     for node in nodes:
         labels = get_node_labels(node)
-        gpu = get_nvidia_gpu_from_node_labels(labels)
+        gpu = get_nvidia_gpu_from_node_labels(labels, node_name=get_node_name(node))
         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:
+            matching_gpu_labels.add(
+                (labels[NVIDIA_GPU_PRODUCT_LABEL], labels[NVIDIA_GPU_MEMORY_LABEL])
+            )
+    if not matching_gpu_labels:
         raise ComputeError(
             f"NVIDIA GPU is requested but no matching GPU labels found: {gpu_spec=}"
         )
-    logger.debug(
-        "Selecting nodes by labels %s for NVIDIA %s", matching_gpu_label_values, gpu_spec.name
-    )
+    logger.debug("Selecting nodes by labels %s for NVIDIA %s", matching_gpu_labels, gpu_spec.name)
     return client.V1NodeAffinity(
         required_during_scheduling_ignored_during_execution=client.V1NodeSelector(
             node_selector_terms=[
@@ -811,10 +812,16 @@ def _get_nvidia_gpu_node_affinity(
                         client.V1NodeSelectorRequirement(
                             key=NVIDIA_GPU_PRODUCT_LABEL,
                             operator=Operator.IN,
-                            values=list(matching_gpu_label_values),
+                            values=[product],
+                        ),
+                        client.V1NodeSelectorRequirement(
+                            key=NVIDIA_GPU_MEMORY_LABEL,
+                            operator=Operator.IN,
+                            values=[memory],
                         ),
                     ],
-                ),
+                )
+                for product, memory in sorted(matching_gpu_labels)
             ],
         ),
     )

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

@@ -8,7 +8,7 @@
 from typing import Callable, Literal, Optional, Union, cast, get_args
 
 import gpuhunt
-from gpuhunt import KNOWN_AMD_GPUS, KNOWN_NVIDIA_GPUS, AcceleratorVendor
+from gpuhunt import KNOWN_AMD_GPUS, AcceleratorVendor
 from kubernetes.client import CoreV1Api, V1Node, V1Taint
 from kubernetes.utils.quantity import parse_quantity as _parse_quantity
 from typing_extensions import Self
@@ -27,6 +27,7 @@
 from dstack._internal.core.models.resources import Memory, ResourcesSpec
 from dstack._internal.utils import docker as docker_utils
 from dstack._internal.utils.common import get_or_error
+from dstack._internal.utils.gpu import convert_nvidia_gpu_name
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -47,6 +48,7 @@
 NVIDIA_GPU_RESOURCE = "nvidia.com/gpu"
 NVIDIA_GPU_NODE_TAINT = NVIDIA_GPU_RESOURCE
 NVIDIA_GPU_PRODUCT_LABEL = f"{NVIDIA_GPU_RESOURCE}.product"
+NVIDIA_GPU_MEMORY_LABEL = f"{NVIDIA_GPU_RESOURCE}.memory"
 
 AMD_GPU_RESOURCE = "amd.com/gpu"
 AMD_GPU_NODE_TAINT = AMD_GPU_RESOURCE
@@ -57,9 +59,6 @@
 # https://github.com/ROCm/k8s-device-plugin/issues/112
 AMD_GPU_DEVICE_ID_LABEL_PREFIX = f"beta.{AMD_GPU_RESOURCE}.device-id."
 
-NVIDIA_GPU_NAME_TO_GPU_INFO = {gpu.name: gpu for gpu in KNOWN_NVIDIA_GPUS}
-NVIDIA_GPU_NAMES = NVIDIA_GPU_NAME_TO_GPU_INFO.keys()
-
 AMD_GPU_DEVICE_ID_TO_GPU_INFO = {
     device_id: gpu_info for gpu_info in KNOWN_AMD_GPUS for device_id in gpu_info.device_ids
 }
@@ -417,26 +416,40 @@
     )
 
 
-def get_nvidia_gpu_from_node_labels(labels: dict[str, str]) -> Optional[Gpu]:
-    # 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".
-    # Thus, we convert the product name to a known gpu name.
+def get_nvidia_gpu_from_node_labels(
+    labels: dict[str, str], *, node_name: Optional[str] = None
+) -> Optional[Gpu]:
+    """Read the NVIDIA GFD product and per-device MiB; normalize whole-GPU VRAM.
+
+    Missing or invalid memory warns and returns None. Unknown models and unmatched
+    capacities retain measured VRAM, with a warning from gpuhunt.
+    """
+    # https://github.com/NVIDIA/k8s-device-plugin/tree/main/docs/gpu-feature-discovery
     gpu_product = labels.get(NVIDIA_GPU_PRODUCT_LABEL)
-    if gpu_product is None:
+    if not gpu_product:
         return None
-    gpu_product = gpu_product.replace("RTX-", "RTX")
-    for gpu_name in NVIDIA_GPU_NAMES:
-        if gpu_name.lower() in gpu_product.lower().split("-"):
-            break
-    else:
+    context = f"provider=kubernetes, node={node_name}" if node_name else "provider=kubernetes"
+    try:
+        gpu_memory_mib = int(labels.get(NVIDIA_GPU_MEMORY_LABEL, "0"))
+    except ValueError:
+        gpu_memory_mib = 0
+    if gpu_memory_mib <= 0:
+        logger.warning(
+            "Missing or invalid %s for NVIDIA GPU %r (%s): %r",
+            NVIDIA_GPU_MEMORY_LABEL,
+            gpu_product,
+            context,
+            labels.get(NVIDIA_GPU_MEMORY_LABEL),
+        )
         return None
-    gpu_info = NVIDIA_GPU_NAME_TO_GPU_INFO[gpu_name]
-    gpu_memory = gpu_info.memory * 1024
-    # A100 may come in two variants
-    if "40GB" in gpu_product:
-        gpu_memory = 40 * 1024
-    return Gpu(vendor=AcceleratorVendor.NVIDIA, name=gpu_name, memory_mib=gpu_memory)
+    gpu_product = gpu_product.removesuffix("-SHARED")
+    gpu_name = convert_nvidia_gpu_name(gpu_product.replace("-", " "))
+    gpu_memory_gib = gpuhunt.get_gpu_memory_gib(gpu_name, gpu_memory_mib / 1024, context=context)
+    return Gpu(
+        vendor=AcceleratorVendor.NVIDIA,
+        name=gpu_name,
+        memory_mib=int(gpu_memory_gib * 1024),
+    )
 
 
 def get_amd_gpu_from_node_labels(labels: dict[str, str]) -> Optional[Gpu]:
@@ -514,17 +527,17 @@
     node: V1Node, node_name: str, node_resources: KubernetesResources
 ) -> list[Gpu]:
     labels = get_node_labels(node)
-    for gpu_resource, gpu_getter in (
-        (NVIDIA_GPU_RESOURCE, get_nvidia_gpu_from_node_labels),
-        (AMD_GPU_RESOURCE, get_amd_gpu_from_node_labels),
-    ):
+    for gpu_resource in (NVIDIA_GPU_RESOURCE, AMD_GPU_RESOURCE):
         gpu_count = int(node_resources[gpu_resource])
         if gpu_count < 1:
             continue
-        gpu = gpu_getter(labels)
+        if gpu_resource == NVIDIA_GPU_RESOURCE:
+            gpu = get_nvidia_gpu_from_node_labels(labels, node_name=node_name)
+        else:
+            gpu = get_amd_gpu_from_node_labels(labels)
         if gpu is None:
             logger.warning(
-                "Node %s: GPU resource found, but failed to detect its model: %s=%d",
+                "Node %s: GPU resource found, but failed to detect its model or memory: %s=%d",
                 node_name,
                 gpu_resource,
                 gpu_count,

src/dstack/_internal/server/background/pipeline_tasks/instances/ssh_deploy.py

@@ -135,7 +135,9 @@ async def add_ssh_instance(instance_model: InstanceModel) -> ProcessResult:
         )
         return result
 
-    instance_type = host_info_to_instance_type(host_info, arch)
+    instance_type = host_info_to_instance_type(
+        host_info, arch, context=f"provider=ssh, instance={fmt(instance_model)}"
+    )
     try:
         instance_network, internal_ip = _resolve_ssh_instance_network(instance_model, host_info)
     except _SSHInstanceNetworkResolutionError as exc:

src/dstack/_internal/server/services/ssh_fleets/provisioning.py

@@ -6,7 +6,7 @@
 from typing import Any, Dict, Generator, List, Optional
 
 import paramiko
-from gpuhunt import AcceleratorVendor, correct_gpu_memory_gib
+from gpuhunt import AcceleratorVendor, get_gpu_memory_gib
 
 from dstack._internal.core.backends.base.authorized_keys import (
     get_add_authorized_keys_script,
@@ -265,7 +265,9 @@
     return out
 
 
-def host_info_to_instance_type(host_info: Dict[str, Any], arch: GoArchType) -> InstanceType:
+def host_info_to_instance_type(
+    host_info: Dict[str, Any], arch: GoArchType, *, context: Optional[str] = None
+) -> InstanceType:
     gpu_count = host_info.get("gpu_count", 0)
     if gpu_count > 0:
         gpu_vendor = AcceleratorVendor.cast(host_info.get("gpu_vendor", "nvidia"))
@@ -283,7 +285,10 @@
         else:
             # newer shim versions report gpu_memory as an integer
             gpu_memory_mib = float(gpu_memory_mib)
-        gpu_memory_mib = correct_gpu_memory_gib(gpu_name, gpu_memory_mib) * 1024
+        gpu_memory_mib = int(
+            get_gpu_memory_gib(gpu_name, gpu_memory_mib / 1024, context=context or "provider=ssh")
+            * 1024
+        )
         gpus = [Gpu(vendor=gpu_vendor, name=gpu_name, memory_mib=gpu_memory_mib)] * gpu_count
     else:
         gpus = []

src/tests/_internal/core/backends/kubernetes/test_compute.py

@@ -1,7 +1,15 @@
 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_nvidia_gpu_node_affinity,
+)
 from dstack._internal.core.backends.kubernetes.models import KubeconfigConfig, KubernetesConfig
+from dstack._internal.core.backends.kubernetes.resources import get_instance_offer_from_node
+from dstack._internal.core.errors import ComputeError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.instances import (
     Disk,
@@ -15,6 +23,63 @@
 from dstack._internal.core.models.runs import Requirements
 
 
+class TestGetNvidiaGpuNodeAffinity:
+    def test_memory_requirement_is_preserved_from_offer_to_node_selection(self):
+        nodes = [
+            _gpu_node("NVIDIA-RTX-A2000", "6144"),
+            _gpu_node("NVIDIA-RTX-A2000", "12288"),
+        ]
+        offers = [get_instance_offer_from_node(node, "ctx") for node in nodes]
+        assert [offer.instance.resources.gpus[0].memory_mib for offer in offers] == [6144, 12288]
+        assert all(len(offer.instance.resources.gpus) == 2 for offer in offers)
+
+        affinity = _get_nvidia_gpu_node_affinity(ResourcesSpec(gpu="A2000:12GB").gpu, nodes)
+        terms = affinity.required_during_scheduling_ignored_during_execution.node_selector_terms
+        assert len(terms) == 1
+        assert {rule.key: rule.values for rule in terms[0].match_expressions} == {
+            "nvidia.com/gpu.product": ["NVIDIA-RTX-A2000"],
+            "nvidia.com/gpu.memory": ["12288"],
+        }
+
+    def test_product_and_memory_pairs_cannot_be_combined_across_nodes(self):
+        nodes = [_gpu_node("A100-SXM4-40GB", "39538"), _gpu_node("A100-SXM4-80GB", "80896")]
+        affinity = _get_nvidia_gpu_node_affinity(ResourcesSpec(gpu="A100").gpu, nodes)
+        terms = affinity.required_during_scheduling_ignored_during_execution.node_selector_terms
+        assert [{rule.key: rule.values for rule in term.match_expressions} for term in terms] == [
+            {
+                "nvidia.com/gpu.product": ["A100-SXM4-40GB"],
+                "nvidia.com/gpu.memory": ["39538"],
+            },
+            {
+                "nvidia.com/gpu.product": ["A100-SXM4-80GB"],
+                "nvidia.com/gpu.memory": ["80896"],
+            },
+        ]
+
+    def test_does_not_select_a_smaller_capacity(self):
+        with pytest.raises(ComputeError, match="no matching GPU labels"):
+            _get_nvidia_gpu_node_affinity(
+                ResourcesSpec(gpu="A2000:12GB").gpu, [_gpu_node("NVIDIA-RTX-A2000", "6144")]
+            )
+
+    def test_missing_product_warns_when_a_node_has_gpu_resources(self, caplog):
+        node = _gpu_node("", "6144")
+        offer = get_instance_offer_from_node(node, "ctx")
+        assert offer.instance.resources.gpus == []
+        assert "GPU resource found, but failed to detect its model or memory" in caplog.text
+
+
+def _gpu_node(product: str, memory_mib: str) -> V1Node:
+    return V1Node(
+        metadata=V1ObjectMeta(
+            name="gpu-node",
+            labels={"nvidia.com/gpu.product": product, "nvidia.com/gpu.memory": memory_mib},
+        ),
+        spec=V1NodeSpec(),
+        status=V1NodeStatus(allocatable={"cpu": "8", "memory": "64Gi", "nvidia.com/gpu": "2"}),
+    )
+
+
 def _compute() -> KubernetesCompute:
     with patch(
         "dstack._internal.core.backends.kubernetes.compute.get_clusters_from_backend_config",

src/tests/_internal/core/backends/kubernetes/test_resources.py

@@ -28,14 +28,54 @@ class TestGetNvidiaGPUFromNodeLabels:
     def test_returns_none_if_no_labels(self):
         assert get_nvidia_gpu_from_node_labels({}) is None
 
-    def test_returns_correct_memory_for_different_A100(self):
+    @pytest.mark.parametrize(
+        "product,memory_mib,name,expected_gib",
+        [
+            ("A100-SXM4-40GB", 39538, "A100", 40),
+            ("A100-SXM4-80GB", 79 * 1024, "A100", 80),
+            ("Tesla-P100-PCIE-12GB", 12 * 1024, "P100", 12),
+            ("Tesla-P100-PCIE-16GB", 16 * 1024, "P100", 16),
+            ("NVIDIA-RTX-A2000", 6 * 1024, "A2000", 6),
+            ("NVIDIA-RTX-A2000", 12 * 1024, "A2000", 12),
+            ("Quadro-RTX-4000", 8 * 1024, "RTX4000", 8),
+            ("NVIDIA-RTX-4000-Ada-Generation", 20 * 1024, "RTX4000Ada", 20),
+            ("Quadro-RTX-5000", 16 * 1024, "RTX5000", 16),
+            ("NVIDIA-RTX-5000-Ada-Generation", 32 * 1024, "RTX5000Ada", 32),
+            ("NVIDIA-L4-SHARED", 22888, "L4", 24),
+        ],
+    )
+    def test_uses_measured_memory_and_preserves_model_generation(
+        self, product, memory_mib, name, expected_gib
+    ):
         assert get_nvidia_gpu_from_node_labels(
-            {"nvidia.com/gpu.product": "A100-SXM4-40GB"}
-        ) == Gpu(vendor=AcceleratorVendor.NVIDIA, name="A100", memory_mib=40 * 1024)
+            {"nvidia.com/gpu.product": product, "nvidia.com/gpu.memory": str(memory_mib)}
+        ) == Gpu(vendor=AcceleratorVendor.NVIDIA, name=name, memory_mib=expected_gib * 1024)
 
-        assert get_nvidia_gpu_from_node_labels(
-            {"nvidia.com/gpu.product": "A100-SXM4-80GB"}
-        ) == Gpu(vendor=AcceleratorVendor.NVIDIA, name="A100", memory_mib=80 * 1024)
+    @pytest.mark.parametrize("memory", [None, "", "invalid", "0", "-1"])
+    def test_warns_and_returns_none_if_memory_is_missing_or_invalid(self, memory, caplog):
+        labels = {"nvidia.com/gpu.product": "NVIDIA-A100"}
+        if memory is not None:
+            labels["nvidia.com/gpu.memory"] = memory
+        assert get_nvidia_gpu_from_node_labels(labels) is None
+        assert "Missing or invalid nvidia.com/gpu.memory" in caplog.text
+
+    @pytest.mark.parametrize(
+        "product,name,memory_mib,warning",
+        [
+            ("NVIDIA-Future-GPU", "FutureGPU", 12345, "Unknown GPU FutureGPU"),
+            ("NVIDIA-GeForce-RTX-4090", "RTX4090", 49140, "no registered capacity matches"),
+        ],
+    )
+    def test_warns_and_preserves_unresolved_measured_memory(
+        self, product, name, memory_mib, warning, caplog
+    ):
+        gpu = get_nvidia_gpu_from_node_labels(
+            {"nvidia.com/gpu.product": product, "nvidia.com/gpu.memory": str(memory_mib)},
+            node_name="gpu-node",
+        )
+        assert gpu == Gpu(vendor=AcceleratorVendor.NVIDIA, name=name, memory_mib=memory_mib)
+        assert warning in caplog.text
+        assert "provider=kubernetes, node=gpu-node" in caplog.text
 
 
 class TestGetAMDGPUFromNodeLabels:

src/tests/_internal/core/backends/slurm/test_resources.py

@@ -36,6 +36,12 @@ def test_name_only_looks_up_nvidia(self):
             vendor=NVIDIA, name="H100", memory_mib=80 * 1024
         )
 
+    @pytest.mark.parametrize("spec", ["P100", "nvidia:P100"])
+    def test_p100_without_memory_keeps_existing_capacity(self, spec):
+        assert GPUModel.from_string(spec) == GPUModel(
+            vendor=NVIDIA, name="P100", memory_mib=16 * 1024
+        )
+
     def test_name_only_looks_up_amd(self):
         assert GPUModel.from_string("MI300X") == GPUModel(
             vendor=AMD, name="MI300X", memory_mib=192 * 1024