diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index c291b0dea..92d0f72e9 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1247,9 +1247,9 @@ projects: If you use ranges with [`resources`](../concepts/tasks.md#resources) (e.g. `gpu: 1..8` or `memory: 64GB..`) in fleet or run configurations, other backends collect and try all offers that satisfy the range. The `kubernetes` backend handles it differently. - + * For `gpu`, if you specify a range (e.g. `gpu: 4..8`), the `kubernetes` backend only provisions pods with the GPU count equal to the lower limit (`4`). The upper limit of the GPU range is always ignored. - * For other resources such as `cpu`, `memory`, and `disk`, the `kubernetes` backend passes the lower and upper limits of the range as Kubernetes [requests and limits](https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/#requests-and-limits) respectively. If the upper limit is not set, the Kubernetes limit is also not set. + * For other resources such as `cpu`, `memory`, and `disk`, the `kubernetes` backend passes the lower and upper limits of the range as Kubernetes [requests and limits](https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/#requests-and-limits) respectively. If the lower limit is not set, the Kubernetes request is set to `0`, unlike Kubernetes, where the request defaults to the limit if not set. If the upper limit is not set, the Kubernetes limit is also not set. Example: @@ -1260,7 +1260,7 @@ projects: ide: vscode resources: - cpu: 32..64 + cpu: ..64 memory: 1024GB disk: 100GB.. gpu: nvidia:4..8 @@ -1272,7 +1272,7 @@ projects: | Resource | Request | Limit | |---------------------|----------|-----------| - | `cpu` | `32` | `64` | + | `cpu` | `0` | `64` | | `memory` | `1024Gi` | `1024Gi` | | `ephemeral-storage` | `100Gi` | _not set_ | | `nvidia.com/gpu` | `4` | `4` | diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index a9572fa33..0deb74d80 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -7,6 +7,7 @@ from contextlib import ExitStack from decimal import Decimal from enum import Enum +from functools import partial from typing import List, Optional from gpuhunt import AcceleratorVendor @@ -15,7 +16,7 @@ from dstack._internal.core.backends.base.compute import ( Compute, - ComputeWithFilteredOffersCached, + ComputeWithAllOffersCached, ComputeWithGatewaySupport, ComputeWithInstanceVolumesSupport, ComputeWithMultinodeSupport, @@ -29,27 +30,33 @@ get_dstack_gateway_commands, merge_tags, ) -from dstack._internal.core.backends.base.offers import RegionalSkipOfferCache, gpu_matches_gpu_spec +from dstack._internal.core.backends.base.offers import ( + OfferModifier, + RegionalSkipOfferCache, + gpu_matches_gpu_spec, +) from dstack._internal.core.backends.kubernetes.api_client import API_CLIENT_EXCEPTIONS from dstack._internal.core.backends.kubernetes.models import KubernetesConfig from dstack._internal.core.backends.kubernetes.resources import ( 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, PodPhase, + ResourceLimits, + ResourceRequests, TaintEffect, + adjust_resources_by_resource_requests, build_base_labels, build_dockerconfigjson, filter_invalid_labels, format_memory, get_amd_gpu_from_node_labels, - get_gpu_request_from_gpu_spec, get_instance_offer_from_node, get_instance_offers, get_node_labels, @@ -80,7 +87,7 @@ SSHConnectionParams, ) from dstack._internal.core.models.placement import PlacementGroup -from dstack._internal.core.models.resources import CPUSpec, GPUSpec +from dstack._internal.core.models.resources import GPUSpec from dstack._internal.core.models.routers import AnyGatewayRouterConfig from dstack._internal.core.models.runs import ( Job, @@ -125,7 +132,7 @@ def load(cls, raw: str) -> Self: class KubernetesCompute( - ComputeWithFilteredOffersCached, + ComputeWithAllOffersCached, ComputeWithPrivilegedSupport, ComputeWithInstanceVolumesSupport, ComputeWithVolumeSupport, @@ -138,9 +145,7 @@ def __init__(self, config: KubernetesConfig): self.region_cluster_map = {c.region: c for c in get_clusters_from_backend_config(config)} self.skip_offer_cache = RegionalSkipOfferCache(ttl=60) - def get_offers_by_requirements( - self, requirements: Requirements - ) -> list[InstanceOfferWithAvailability]: + def get_all_offers_with_availability(self) -> list[InstanceOfferWithAvailability]: offers: list[InstanceOfferWithAvailability] = [] with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: future_cluster_map: dict[ @@ -148,7 +153,7 @@ def get_offers_by_requirements( ] = {} for region, cluster in self.region_cluster_map.items(): api = client.CoreV1Api(cluster.api_client) - future = executor.submit(get_instance_offers, api, region, requirements) + future = executor.submit(get_instance_offers, api, region) future_cluster_map[future] = cluster for future in concurrent.futures.as_completed(future_cluster_map): try: @@ -164,6 +169,10 @@ def get_offers_by_requirements( offers.extend(cluster_offers) return offers + def get_offers_modifiers(self, requirements: Requirements) -> list[OfferModifier]: + resource_requests = ResourceRequests.from_resources_spec(requirements.resources) + return [partial(_offer_modifier, resource_requests)] + def run_job( self, run: Run, @@ -383,18 +392,15 @@ 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)) - # In the original offer, the resources have already been adjusted according to - # the run configuration resource requirements, see get_offers_by_requirements() - original_resources = provisioning_data.instance_type.resources - instance_offer = get_instance_offer_from_node( - node=node, - region=cluster.region, - cpu_request=original_resources.cpus, - memory_mib_request=original_resources.memory_mib, - gpu_request=len(original_resources.gpus), - disk_mib_request=original_resources.disk.size_mib, - ) + instance_offer = get_instance_offer_from_node(node=node, region=cluster.region) 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 {} + ) + adjust_resources_by_resource_requests( + instance_offer.instance.resources, resource_requests, force=True + ) provisioning_data.instance_type = instance_offer.instance provisioning_data.price = instance_offer.price @@ -728,7 +734,7 @@ def delete_volume(self, volume: Volume): def _get_pod_spec_parameters_for_gpu( api: client.CoreV1Api, gpu_spec: GPUSpec -) -> tuple[str, client.V1NodeAffinity, str]: +) -> tuple[AnyKubernetesGPUResource, client.V1NodeAffinity, str]: 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 @@ -736,10 +742,10 @@ def _get_pod_spec_parameters_for_gpu( # 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 NVIDIA_GPU_RESOURCE, node_affinity, NVIDIA_GPU_NODE_TAINT + return KubernetesResource.NVIDIA_GPU, node_affinity, NVIDIA_GPU_NODE_TAINT if gpu_vendor == AcceleratorVendor.AMD: node_affinity = _get_amd_gpu_node_affinity(gpu_spec, nodes) - return AMD_GPU_RESOURCE, node_affinity, AMD_GPU_NODE_TAINT + return KubernetesResource.AMD_GPU, node_affinity, AMD_GPU_NODE_TAINT raise ComputeError(f"Unsupported GPU vendor: {gpu_vendor}") @@ -802,6 +808,14 @@ def _get_amd_gpu_node_affinity( ) +def _offer_modifier( + resource_requests: ResourceRequests, offer: InstanceOfferWithAvailability +) -> InstanceOfferWithAvailability: + offer_copy = offer.copy(deep=True) + adjust_resources_by_resource_requests(offer_copy.instance.resources, resource_requests) + return offer_copy + + def _create_jump_pod_service_if_not_exists( api: client.CoreV1Api, namespace: str, @@ -1100,8 +1114,6 @@ def _create_job_pod( volumes: list[Volume], authorized_keys: list[str], ) -> None: - resources_requests: dict[str, str] = {} - resources_limits: dict[str, str] = {} node_affinity: Optional[client.V1NodeAffinity] = None tolerations: list[client.V1Toleration] = [] volumes_: list[client.V1Volume] = [] @@ -1109,18 +1121,14 @@ def _create_job_pod( env_vars: list[client.V1EnvVar] = [] resources_spec = job_spec.requirements.resources - assert isinstance(resources_spec.cpu, CPUSpec) - if (cpu_min := resources_spec.cpu.count.min) is not None: - resources_requests["cpu"] = str(cpu_min) - if (cpu_max := resources_spec.cpu.count.max) is not None: - resources_limits["cpu"] = str(cpu_max) - gpu_spec = resources_spec.gpu - if gpu_spec is not None and (gpu_request := get_gpu_request_from_gpu_spec(gpu_spec)) > 0: + resource_requests = ResourceRequests.from_resources_spec(resources_spec) + resource_limits = ResourceLimits.from_resources_spec(resources_spec) + gpu_resource: Optional[AnyKubernetesGPUResource] = 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) - logger.debug("Requesting GPU resource: %s=%d", gpu_resource, gpu_request) - resources_requests[gpu_resource] = str(gpu_request) - # Limit must be set (GPU resources cannot be overcommitted) and must be equal to request. - resources_limits[gpu_resource] = str(gpu_request) + 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]: tolerations.append( @@ -1131,15 +1139,6 @@ def _create_job_pod( # into the image (NVIDIA images and images based on them including dstackai/base) # See https://github.com/NVIDIA/k8s-device-plugin/issues/61 env_vars.append(client.V1EnvVar(name="NVIDIA_VISIBLE_DEVICES", value="void")) - if (memory_min := resources_spec.memory.min) is not None: - resources_requests["memory"] = format_memory(memory_min) - if (memory_max := resources_spec.memory.max) is not None: - resources_limits["memory"] = format_memory(memory_max) - if (disk_spec := resources_spec.disk) is not None: - if (disk_min := disk_spec.size.min) is not None: - resources_requests["ephemeral-storage"] = format_memory(disk_min) - if (disk_max := disk_spec.size.max) is not None: - resources_limits["ephemeral-storage"] = format_memory(disk_max) if (shm_size := resources_spec.shm_size) is not None: shm_volume_name = "dev-shm" volumes_.append( @@ -1246,8 +1245,8 @@ def _create_job_pod( ), ), resources=client.V1ResourceRequirements( - requests=resources_requests, - limits=resources_limits, + requests=resource_requests.to_kubernetes_map(gpu_resource), + limits=resource_limits.to_kubernetes_map(gpu_resource), ), volume_mounts=volume_mounts, env=env_vars, diff --git a/src/dstack/_internal/core/backends/kubernetes/resources.py b/src/dstack/_internal/core/backends/kubernetes/resources.py index 2d3ff02de..4dcfedc99 100644 --- a/src/dstack/_internal/core/backends/kubernetes/resources.py +++ b/src/dstack/_internal/core/backends/kubernetes/resources.py @@ -5,8 +5,9 @@ from collections.abc import Mapping from decimal import Decimal from enum import Enum -from typing import Callable, Literal, Optional, Union, cast +from typing import Callable, Literal, Optional, Union, cast, get_args +import gpuhunt from gpuhunt import KNOWN_AMD_GPUS, KNOWN_NVIDIA_GPUS, AcceleratorVendor # XXX: kubernetes.utils is missing in the stubs package @@ -15,7 +16,6 @@ from typing_extensions import Self from dstack._internal.core.backends.base.compute import normalize_arch -from dstack._internal.core.backends.base.offers import filter_offers_by_requirements from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.instances import ( Disk, @@ -26,8 +26,7 @@ InstanceType, Resources, ) -from dstack._internal.core.models.resources import CPUSpec, GPUSpec, Memory -from dstack._internal.core.models.runs import Requirements +from dstack._internal.core.models.resources import CPUSpec, Memory, ResourcesSpec from dstack._internal.utils import docker as docker_utils from dstack._internal.utils.common import get_or_error from dstack._internal.utils.logging import get_logger @@ -96,6 +95,19 @@ class KubernetesResource(str, Enum): NVIDIA_GPU = NVIDIA_GPU_RESOURCE AMD_GPU = AMD_GPU_RESOURCE + @classmethod + def from_gpu_vendor(cls, vendor: gpuhunt.AcceleratorVendor) -> "AnyKubernetesGPUResource": + match vendor: + case gpuhunt.AcceleratorVendor.NVIDIA: + return KubernetesResource.NVIDIA_GPU + case gpuhunt.AcceleratorVendor.AMD: + return KubernetesResource.AMD_GPU + raise ValueError(f"Unsupported accelerator vendor: {vendor}") + + +AnyKubernetesGPUResource = Literal[KubernetesResource.NVIDIA_GPU, KubernetesResource.AMD_GPU] +GPU_RESOURCES: tuple[AnyKubernetesGPUResource, ...] = get_args(AnyKubernetesGPUResource) + @dataclasses.dataclass class KubernetesResources: @@ -135,6 +147,125 @@ def __sub__(self, other: Self) -> Self: return type(self)(**dct) +@dataclasses.dataclass(frozen=True) +class ResourceRequestsLimits: + cpu: Optional[int] + memory_mib: Optional[int] + disk_mib: Optional[int] + gpu: int + + def to_kubernetes_map( + self, gpu_resource: Optional[AnyKubernetesGPUResource] = None + ) -> dict[str, str]: + dct: dict[str, str] = {} + if self.cpu is not None: + dct[KubernetesResource.CPU.value] = str(self.cpu) + if self.memory_mib is not None: + dct[KubernetesResource.MEMORY.value] = f"{self.memory_mib}Mi" + if self.disk_mib is not None: + dct[KubernetesResource.EPHEMERAL_STORAGE.value] = f"{self.disk_mib}Mi" + if self.gpu > 0: + if gpu_resource is None: + raise ValueError("gpu_resource is not specified") + dct[gpu_resource.value] = str(self.gpu) + return dct + + +@dataclasses.dataclass(frozen=True) +class ResourceRequests(ResourceRequestsLimits): + cpu: int + memory_mib: int + disk_mib: int + + @classmethod + def from_resources_spec(cls, spec: ResourcesSpec) -> Self: + assert isinstance(spec.cpu, CPUSpec) + cpu = spec.cpu.count.min or 0 + memory_mib: int = 0 + if spec.memory.min is not None: + memory_mib = round(spec.memory.min * 1024) + disk_mib: int = 0 + if spec.disk is not None and spec.disk.size.min is not None: + disk_mib = round(spec.disk.size.min * 1024) + gpu: int = 0 + if spec.gpu is not None: + gpu = spec.gpu.count.min or 0 + return cls( + cpu=cpu, + memory_mib=memory_mib, + disk_mib=disk_mib, + gpu=gpu, + ) + + @classmethod + def from_kubernetes_map(cls, map_: Mapping[str, str]) -> Self: + cpu_qty = map_.get(KubernetesResource.CPU.value, "0") + cpu = round(parse_quantity(cpu_qty)) + memory_qty = map_.get(KubernetesResource.MEMORY.value, "0") + memory_mib = round(parse_quantity(memory_qty) / 2**20) + 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 + return cls( + cpu=cpu, + memory_mib=memory_mib, + disk_mib=disk_mib, + gpu=gpu, + ) + + +@dataclasses.dataclass(frozen=True) +class ResourceLimits(ResourceRequestsLimits): + @classmethod + def from_resources_spec(cls, spec: ResourcesSpec) -> Self: + assert isinstance(spec.cpu, CPUSpec) + cpu = spec.cpu.count.max + memory_mib: Optional[int] = None + if spec.memory.max is not None: + memory_mib = round(spec.memory.max * 1024) + disk_mib: Optional[int] = None + if spec.disk is not None: + if spec.disk.size.max is not None: + disk_mib = round(spec.disk.size.max * 1024) + gpu: int = 0 + if spec.gpu is not None: + # GPU resources cannot be overcommitted, limit must be equal to request + gpu = spec.gpu.count.min or 0 + assert isinstance(spec.cpu, CPUSpec) + return cls( + cpu=cpu, + memory_mib=memory_mib, + disk_mib=disk_mib, + gpu=gpu, + ) + + +def adjust_resources_by_resource_requests( + resources: Resources, + resource_requests: ResourceRequests, + *, + force: bool = False, +) -> None: + cpu = resource_requests.cpu + if not force: + cpu = min(resources.cpus, cpu) + resources.cpus = cpu + memory_mib = resource_requests.memory_mib + if not force: + memory_mib = min(resources.memory_mib, memory_mib) + resources.memory_mib = memory_mib + resources.gpus = resources.gpus[: resource_requests.gpu] + disk_mib = resource_requests.disk_mib + if not force: + disk_mib = min(resources.disk.size_mib, disk_mib) + resources.disk = Disk(size_mib=disk_mib) + + def build_base_labels( *, component: Literal["ssh-proxy", "job", "gateway", "volume"], @@ -225,10 +356,6 @@ def format_memory(memory: Memory) -> str: return f"{float(memory)}Gi" -def get_gpu_request_from_gpu_spec(gpu_spec: GPUSpec) -> int: - return gpu_spec.count.min or 0 - - def get_node_name(node: V1Node) -> Optional[str]: if (metadata := node.metadata) is None: return None @@ -258,20 +385,7 @@ def is_taint_tolerated(taint: V1Taint) -> bool: return taint.key in (NVIDIA_GPU_NODE_TAINT, AMD_GPU_NODE_TAINT) -def get_instance_offers( - api: CoreV1Api, region: str, requirements: Requirements -) -> list[InstanceOfferWithAvailability]: - resources_spec = requirements.resources - assert isinstance(resources_spec.cpu, CPUSpec) - cpu_request = resources_spec.cpu.count.min or 0 - memory_mib_request = round((resources_spec.memory.min or 0) * 1024) - gpu_request = 0 - if (gpu_spec := resources_spec.gpu) is not None: - gpu_request = get_gpu_request_from_gpu_spec(gpu_spec) - disk_mib_request = 0 - if (disk_spec := resources_spec.disk) is not None: - disk_mib_request = round((disk_spec.size.min or 0) * 1024) - +def get_instance_offers(api: CoreV1Api, region: str) -> list[InstanceOfferWithAvailability]: nodes_allocated_resources = _get_nodes_allocated_resources(api) offers: list[InstanceOfferWithAvailability] = [] for node in api.list_node().items: @@ -282,24 +396,14 @@ def get_instance_offers( node_name=node_name, node_allocated_resources=nodes_allocated_resources.get(node_name), region=region, - cpu_request=cpu_request, - memory_mib_request=memory_mib_request, - gpu_request=gpu_request, - disk_mib_request=disk_mib_request, ) if offer is not None: - offers.extend(filter_offers_by_requirements([offer], requirements)) + offers.append(offer) return offers def get_instance_offer_from_node( - node: V1Node, - *, - region: str, - cpu_request: int, - memory_mib_request: int, - gpu_request: int, - disk_mib_request: int, + node: V1Node, region: str ) -> Optional[InstanceOfferWithAvailability]: node_name = get_node_name(node) if node_name is None: @@ -309,10 +413,6 @@ def get_instance_offer_from_node( node_name=node_name, node_allocated_resources=None, region=region, - cpu_request=cpu_request, - memory_mib_request=memory_mib_request, - gpu_request=gpu_request, - disk_mib_request=disk_mib_request, ) @@ -365,10 +465,6 @@ def _get_instance_offer_from_node( node_name: str, node_allocated_resources: Optional[KubernetesResources], region: str, - cpu_request: int, - memory_mib_request: int, - gpu_request: int, - disk_mib_request: int, ) -> Optional[InstanceOfferWithAvailability]: try: node_spec = get_or_error(node.spec) @@ -398,11 +494,11 @@ def _get_instance_offer_from_node( instance=InstanceType( name=node_name, resources=Resources( - cpus=min(cpu_request, cpu), + cpus=cpu, cpu_arch=cpu_arch, - memory_mib=min(memory_mib_request, memory_mib), - gpus=gpus[:gpu_request], - disk=Disk(size_mib=min(disk_mib_request, disk_mib)), + memory_mib=memory_mib, + gpus=gpus, + disk=Disk(size_mib=disk_mib), spot=False, ), ), diff --git a/src/tests/_internal/core/backends/kubernetes/test_resources.py b/src/tests/_internal/core/backends/kubernetes/test_resources.py index 1839a4191..e1c55dc96 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_resources.py +++ b/src/tests/_internal/core/backends/kubernetes/test_resources.py @@ -4,12 +4,24 @@ from gpuhunt import AcceleratorVendor from dstack._internal.core.backends.kubernetes.resources import ( + KubernetesResource, + ResourceLimits, + ResourceRequests, + adjust_resources_by_resource_requests, get_amd_gpu_from_node_labels, get_nvidia_gpu_from_node_labels, validate_label_key, validate_label_value, ) -from dstack._internal.core.models.instances import Gpu +from dstack._internal.core.models.instances import Disk, Gpu, Resources +from dstack._internal.core.models.resources import ( + CPUSpec, + DiskSpec, + GPUSpec, + Memory, + Range, + ResourcesSpec, +) class TestGetNvidiaGPUFromNodeLabels: @@ -55,6 +67,101 @@ def test_returns_none_if_multiple_gpu_models(self, caplog: pytest.LogCaptureFixt assert "Multiple AMD GPU models detected" in caplog.text +class TestResourceRequests: + def test_from_resources_spec_uses_lower_bounds_and_defaults_unset_to_zero(self): + spec = ResourcesSpec( + cpu=CPUSpec(count=Range[int](min=None, max=64)), + memory=Range[Memory](min=Memory.parse("1024GB"), max=Memory.parse("1024GB")), + disk=DiskSpec(size=Range[Memory](min=Memory.parse("100GB"), max=None)), + gpu=GPUSpec(count=Range[int](min=4, max=8)), + ) + assert ResourceRequests.from_resources_spec(spec) == ResourceRequests( + cpu=0, memory_mib=1024 * 1024, disk_mib=100 * 1024, gpu=4 + ) + + def test_to_kubernetes_map(self): + requests = ResourceRequests(cpu=0, memory_mib=1024 * 1024, disk_mib=100 * 1024, gpu=4) + assert requests.to_kubernetes_map(KubernetesResource.NVIDIA_GPU) == { + "cpu": "0", + "memory": "1048576Mi", + "ephemeral-storage": "102400Mi", + "nvidia.com/gpu": "4", + } + + def test_kubernetes_map_round_trip(self): + requests = ResourceRequests(cpu=8, memory_mib=16384, disk_mib=51200, gpu=2) + map_ = requests.to_kubernetes_map(KubernetesResource.AMD_GPU) + assert ResourceRequests.from_kubernetes_map(map_) == requests + + +class TestResourceLimits: + def test_from_resources_spec_uses_upper_bounds_and_leaves_unset_as_none(self): + spec = ResourcesSpec( + cpu=CPUSpec(count=Range[int](min=None, max=64)), + memory=Range[Memory](min=Memory.parse("1024GB"), max=Memory.parse("1024GB")), + disk=DiskSpec(size=Range[Memory](min=Memory.parse("100GB"), max=None)), + gpu=GPUSpec(count=Range[int](min=4, max=8)), + ) + # GPU limit must equal the request (min), since GPUs cannot be overcommitted. + assert ResourceLimits.from_resources_spec(spec) == ResourceLimits( + cpu=64, memory_mib=1024 * 1024, disk_mib=None, gpu=4 + ) + + def test_to_kubernetes_map_omits_unset_resources(self): + limits = ResourceLimits(cpu=64, memory_mib=1024 * 1024, disk_mib=None, gpu=4) + assert limits.to_kubernetes_map(KubernetesResource.NVIDIA_GPU) == { + "cpu": "64", + "memory": "1048576Mi", + "nvidia.com/gpu": "4", + } + + +def _node_resources() -> Resources: + return Resources( + cpus=64, + memory_mib=256 * 1024, + gpus=[Gpu(name="H100", memory_mib=80 * 1024)] * 8, + disk=Disk(size_mib=1000 * 1024), + spot=False, + ) + + +class TestAdjustResourcesByResourceRequests: + def test_clamps_each_resource_to_the_smaller_of_node_and_request(self): + resources = _node_resources() + adjust_resources_by_resource_requests( + resources, + ResourceRequests(cpu=8, memory_mib=16 * 1024, disk_mib=100 * 1024, gpu=2), + ) + assert resources.cpus == 8 + assert resources.memory_mib == 16 * 1024 + assert resources.disk == Disk(size_mib=100 * 1024) + assert len(resources.gpus) == 2 + + def test_does_not_exceed_node_resources(self): + resources = _node_resources() + adjust_resources_by_resource_requests( + resources, + ResourceRequests(cpu=128, memory_mib=512 * 1024, disk_mib=4000 * 1024, gpu=16), + ) + assert resources.cpus == 64 + assert resources.memory_mib == 256 * 1024 + assert resources.disk == Disk(size_mib=1000 * 1024) + assert len(resources.gpus) == 8 + + def test_force_sets_resources_to_the_request(self): + resources = _node_resources() + adjust_resources_by_resource_requests( + resources, + ResourceRequests(cpu=8, memory_mib=16 * 1024, disk_mib=100 * 1024, gpu=2), + force=True, + ) + assert resources.cpus == 8 + assert resources.memory_mib == 16 * 1024 + assert resources.disk == Disk(size_mib=100 * 1024) + assert len(resources.gpus) == 2 + + class TestLabelValidation: @pytest.mark.parametrize( "key",