diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index 980cd67615..fa59005f68 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/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) ], ), ) diff --git a/src/dstack/_internal/core/backends/kubernetes/resources.py b/src/dstack/_internal/core/backends/kubernetes/resources.py index b049eccc8c..a38b064417 100644 --- a/src/dstack/_internal/core/backends/kubernetes/resources.py +++ b/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_instance_offer_from_node( ) -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 @@ def _get_gpus_from_node( 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, diff --git a/src/dstack/_internal/server/background/pipeline_tasks/instances/ssh_deploy.py b/src/dstack/_internal/server/background/pipeline_tasks/instances/ssh_deploy.py index 26fb082260..bbf65263bd 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/instances/ssh_deploy.py +++ b/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: diff --git a/src/dstack/_internal/server/services/ssh_fleets/provisioning.py b/src/dstack/_internal/server/services/ssh_fleets/provisioning.py index 6feae9ca32..76887a7b5a 100644 --- a/src/dstack/_internal/server/services/ssh_fleets/provisioning.py +++ b/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 @@ def _get_shim_healthcheck(client: paramiko.SSHClient) -> Optional[str]: 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 @@ def host_info_to_instance_type(host_info: Dict[str, Any], arch: GoArchType) -> I 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 = [] diff --git a/src/tests/_internal/core/backends/kubernetes/test_compute.py b/src/tests/_internal/core/backends/kubernetes/test_compute.py index 39ef170850..4c374db794 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_compute.py +++ b/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", diff --git a/src/tests/_internal/core/backends/kubernetes/test_resources.py b/src/tests/_internal/core/backends/kubernetes/test_resources.py index e1c55dc96f..9d0c504100 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_resources.py +++ b/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: diff --git a/src/tests/_internal/core/backends/slurm/test_resources.py b/src/tests/_internal/core/backends/slurm/test_resources.py index 36137f71bc..320415fa07 100644 --- a/src/tests/_internal/core/backends/slurm/test_resources.py +++ b/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 diff --git a/src/tests/_internal/server/services/ssh_fleets/test_provisioning.py b/src/tests/_internal/server/services/ssh_fleets/test_provisioning.py new file mode 100644 index 0000000000..1dc4304569 --- /dev/null +++ b/src/tests/_internal/server/services/ssh_fleets/test_provisioning.py @@ -0,0 +1,37 @@ +import pytest + +from dstack._internal.core.backends.base.compute import GoArchType +from dstack._internal.server.services.ssh_fleets.provisioning import host_info_to_instance_type + + +class TestHostInfoToInstanceType: + @pytest.mark.parametrize("legacy_memory", [False, True]) + @pytest.mark.parametrize( + "product,reported_mib,expected_mib,warns", + [ + ("NVIDIA A100-SXM4-40GB", 39936, 40 * 1024, False), + ("NVIDIA A100-SXM4-80GB", 80896, 80 * 1024, False), + ("NVIDIA GeForce RTX 4090", 49140, 49140, True), + ("NVIDIA Future GPU", 12345, 12345, True), + ], + ) + def test_normalizes_known_capacity_and_preserves_unresolved_measurements( + self, product, reported_mib, expected_mib, warns, legacy_memory, caplog + ): + host_info = { + "gpu_count": 2, + "gpu_name": product, + "gpu_memory": f"{reported_mib} MiB" if legacy_memory else reported_mib, + "cpus": 16, + "memory": 64 * 1024**3, + "disk_size": 256 * 1024**3, + } + instance = host_info_to_instance_type( + host_info, GoArchType.AMD64, context="provider=ssh, instance=test-host" + ) + assert len(instance.resources.gpus) == 2 + assert all(gpu.memory_mib == expected_mib for gpu in instance.resources.gpus) + assert bool(caplog.records) == warns + if warns: + assert "keeping reported VRAM" in caplog.text + assert "provider=ssh, instance=test-host" in caplog.text