diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 20894f9de..96f4157b2 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/mkdocs/docs/concepts/backends.md @@ -1257,6 +1257,29 @@ projects: The jump host can be a GPU node or a CPU-only node — it makes no difference. The only requirement is that both the `dstack` server and CLI can reach `hostname:port`. + The jump pod is managed by a ReplicaSet, so Kubernetes recreates it if it's evicted or its node is removed. To make eviction less likely, the jump pod's CPU and memory requests equal its limits, which gives it the [`Guaranteed`](https://kubernetes.io/docs/concepts/workloads/pods/pod-qos/#guaranteed) QoS class. Set `proxy_jump.cpu` and `proxy_jump.memory` to change the defaults, `500m` and `1Gi`: + + ```yaml + contexts: + - name: gpu-cluster-a + proxy_jump: + cpu: 1 + memory: 2Gi + ``` + + Each SSH session through the jump pod takes about 2 MiB of memory. Each `dstack` server replica keeps at least one session open per running job, and each `dstack attach` adds one more. If the jump pod exceeds its memory limit, Kubernetes restarts it. + + When the jump pod restarts or is recreated, its SSH sessions drop. The `dstack` server reconnects automatically, but `dstack attach` may fail for runs submitted before that. If reliable SSH access to runs is important, deploy [`dstack-sshproxy`](../guides/server-deployment.md#ssh-proxy). + + !!! warning "Changing `cpu` and `memory`" + New values don't apply to an existing jump pod. To apply them, update the jump pod's ReplicaSet in the context's namespace and delete the pod so that Kubernetes recreates it: + + ```shell + kubectl -n set resources replicaset dstack--ssh-jump-pod \ + --requests=cpu=1,memory=2Gi --limits=cpu=1,memory=2Gi + kubectl -n delete pod -l app.kubernetes.io/instance=dstack--ssh-jump-pod + ``` + !!! info "Region and namespace" Each enabled context becomes its own `dstack` region, named after the context. When creating a `dstack` [volume](volumes.md) or [gateway](gateways.md), the `region` field selects which cluster the resource is provisioned in. diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index c9cd0e7f9..f9c392927 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -41,7 +41,11 @@ 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.models import ( + JUMP_POD_DEFAULT_CPU, + JUMP_POD_DEFAULT_MEMORY, + KubernetesConfig, +) from dstack._internal.core.backends.kubernetes.resources import ( AMD_GPU_DEVICE_ID_LABEL_PREFIX, AMD_GPU_NAME_TO_DEVICE_IDS, @@ -116,6 +120,10 @@ JUMP_POD_IMAGE = "testcontainers/sshd:1.3.0@sha256:c50c0f59554dcdb2d9e5e705112144428ae9d04ac0af6322b365a18e24213a6a" JUMP_POD_SSH_PORT = 22 JUMP_POD_USER = "root" +# Under disk pressure, the kubelet first evicts pods whose disk usage exceeds their +# ephemeral-storage request. Covers the writable layer and the logs with the default log +# rotation (5 x 10Mi). +JUMP_POD_EPHEMERAL_STORAGE_REQUEST = "64Mi" JOB_POD_SCHEDULING_TIMEOUT = 10 @@ -245,6 +253,8 @@ def run_job( jump_pod_service_name=jump_pod_service_name, jump_pod_port=cluster.proxy_jump.port, authorized_keys=project_authorized_keys, + cpu_request_limit=cluster.proxy_jump.cpu or JUMP_POD_DEFAULT_CPU, + memory_request_limit=cluster.proxy_jump.memory or JUMP_POD_DEFAULT_MEMORY, ) pod_name = generate_unique_instance_name_for_job( @@ -876,6 +886,8 @@ def _create_jump_pod_service_if_not_exists( jump_pod_service_name: str, jump_pod_port: Optional[int], authorized_keys: list[str], + cpu_request_limit: str, + memory_request_limit: str, ) -> None: base_labels = build_base_labels( component="ssh-proxy", @@ -977,6 +989,16 @@ def _create_jump_pod_service_if_not_exists( if not tolerations: logger.warning("No appropriate node found, the jump pod may never be scheduled") commands = _get_jump_pod_commands(authorized_keys) + # Equal CPU and memory requests and limits for the Guaranteed QoS class + limits = { + KubernetesResource.CPU.value: cpu_request_limit, + KubernetesResource.MEMORY.value: memory_request_limit, + } + requests = { + **limits, + # No limit, as exceeding it gets the pod evicted + KubernetesResource.EPHEMERAL_STORAGE.value: JUMP_POD_EPHEMERAL_STORAGE_REQUEST, + } replica_set = client.V1ReplicaSet( metadata=client.V1ObjectMeta( name=jump_pod_name, @@ -1001,6 +1023,10 @@ def _create_jump_pod_service_if_not_exists( container_port=JUMP_POD_SSH_PORT, ) ], + resources=client.V1ResourceRequirements( + requests=requests, + limits=limits, + ), ) ], restart_policy=RestartPolicy.ALWAYS, @@ -1093,6 +1119,18 @@ def _check_and_configure_jump_pod_service( if _is_pod_ready(jump_pod): break else: + for jump_pod in jump_pods: + # E.g., no node has enough free resources or all nodes have untolerated taints. + # Provisioning waits for the jump pod, so without this warning, the only symptom is + # runs hitting the provisioning timeout. + condition = _get_pod_unscheduled_condition(jump_pod) + if condition is not None: + logger.warning( + "Jump pod %s cannot be scheduled: %s: %s", + get_object_name(jump_pod), + condition.reason, + condition.message, + ) logger.debug("Jump pod %s is not ready yet", jump_pod_name) return None @@ -1448,13 +1486,20 @@ def _get_unscheduled_pod_reason_message( name=pod_name, namespace=namespace, ) - if pod is not None and pod.status is not None and pod.status.conditions: - for cond in pod.status.conditions: - if cond.type == "PodScheduled" and cond.status == "False": - return cond.reason, cond.message + if pod is not None and (condition := _get_pod_unscheduled_condition(pod)) is not None: + return condition.reason, condition.message return None, None +def _get_pod_unscheduled_condition(pod: client.V1Pod) -> Optional[client.V1PodCondition]: + if pod.status is None: + return None + for condition in pod.status.conditions or []: + if condition.type == "PodScheduled" and condition.status == "False": + return condition + return None + + def _is_pod_ready(pod: client.V1Pod) -> bool: if pod.metadata is None: return False diff --git a/src/dstack/_internal/core/backends/kubernetes/configurator.py b/src/dstack/_internal/core/backends/kubernetes/configurator.py index 8a4e2e17f..89cab248c 100644 --- a/src/dstack/_internal/core/backends/kubernetes/configurator.py +++ b/src/dstack/_internal/core/backends/kubernetes/configurator.py @@ -1,3 +1,5 @@ +from typing import Optional + from dstack._internal.core.backends.base.configurator import ( BackendRecord, Configurator, @@ -8,8 +10,10 @@ KubernetesBackendConfig, KubernetesBackendConfigWithCreds, KubernetesConfig, + KubernetesProxyJumpConfig, KubernetesStoredConfig, ) +from dstack._internal.core.backends.kubernetes.resources import parse_quantity, validate_quantity from dstack._internal.core.backends.kubernetes.utils import ( check_cluster, get_clusters_from_backend_config, @@ -21,6 +25,12 @@ logger = get_logger(__name__) +# Sanity floors that catch unit mistakes, e.g., `memory: 512` (bytes). Measured with the jump pod +# image: below 100m, SSH handshakes through the jump pod slow down sharply (~0.5s at 50m, ~2s at +# 10m); 32Mi fits the container (~9Mi) and a few sessions (~1.6Mi each). +_JUMP_POD_MIN_CPU = "100m" +_JUMP_POD_MIN_MEMORY = "32Mi" + class KubernetesConfigurator( Configurator[ @@ -35,6 +45,7 @@ def validate_config( self, config: KubernetesBackendConfigWithCreds, default_creds_enabled: bool ): self._check_config_contexts(config) + self._check_config_proxy_jumps(config) try: clusters = get_clusters_from_backend_config(config, request_timeout=10, retries=0) except Exception as e: @@ -90,3 +101,32 @@ def _check_config_contexts(self, config: KubernetesBackendConfig): seen.add(name) if duplicates: raise ServerClientError(f"duplicate contexts: {', '.join(sorted(duplicates))}") + + def _check_config_proxy_jumps(self, config: KubernetesBackendConfig): + if config.proxy_jump is not None: + self._check_proxy_jump(config.proxy_jump) + for context in config.contexts or []: + if not isinstance(context, str) and context.proxy_jump is not None: + self._check_proxy_jump(context.proxy_jump, context_name=context.name) + + def _check_proxy_jump( + self, proxy_jump: KubernetesProxyJumpConfig, *, context_name: Optional[str] = None + ): + for field, quantity, min_quantity in [ + ("cpu", proxy_jump.cpu, _JUMP_POD_MIN_CPU), + ("memory", proxy_jump.memory, _JUMP_POD_MIN_MEMORY), + ]: + if quantity is None: + continue + try: + validate_quantity(quantity) + except ValueError: + reason = "must be a Kubernetes quantity" + else: + if parse_quantity(quantity) >= parse_quantity(min_quantity): + continue + reason = f"must be at least {min_quantity}" + msg = f"invalid proxy_jump.{field} value {quantity!r}" + if context_name is not None: + msg = f"{msg} in context {context_name!r}" + raise ServerClientError(f"{msg}: {reason}") diff --git a/src/dstack/_internal/core/backends/kubernetes/models.py b/src/dstack/_internal/core/backends/kubernetes/models.py index c9df6ac88..ac126f9fa 100644 --- a/src/dstack/_internal/core/backends/kubernetes/models.py +++ b/src/dstack/_internal/core/backends/kubernetes/models.py @@ -6,6 +6,12 @@ from dstack._internal.core.backends.base.models import fill_data from dstack._internal.core.models.common import CoreModel +# Measured with the jump pod image (OpenSSH 9.9). The container takes ~9 MiB plus ~1.6 MiB per +# tunneled SSH session (two sshd-session processes), so 1Gi fits ~600 sessions. Each running job +# holds one or more sessions per server replica. At 500m, it tunnels ~130-180 MiB/s. +JUMP_POD_DEFAULT_CPU = "500m" +JUMP_POD_DEFAULT_MEMORY = "1Gi" + class KubernetesProxyJumpConfig(CoreModel): hostname: Annotated[ @@ -14,6 +20,24 @@ class KubernetesProxyJumpConfig(CoreModel): port: Annotated[ Optional[int], Field(description="Any port accessible outside of the cluster") ] = None + cpu: Annotated[ + Optional[str], + Field( + description=( + "The jump pod's CPU request and limit, as a Kubernetes quantity," + f" e.g., `1`, `500m`. Defaults to `{JUMP_POD_DEFAULT_CPU}`" + ) + ), + ] = None + memory: Annotated[ + Optional[str], + Field( + description=( + "The jump pod's memory request and limit, as a Kubernetes quantity," + f" e.g., `2Gi`, `512Mi`. Defaults to `{JUMP_POD_DEFAULT_MEMORY}`" + ) + ), + ] = None class KubernetesContextConfig(CoreModel): diff --git a/src/dstack/_internal/core/backends/kubernetes/resources.py b/src/dstack/_internal/core/backends/kubernetes/resources.py index 37aa36ef6..95eb5b46e 100644 --- a/src/dstack/_internal/core/backends/kubernetes/resources.py +++ b/src/dstack/_internal/core/backends/kubernetes/resources.py @@ -44,6 +44,11 @@ LABEL_VALUE_MAX_LENGTH = 63 LABEL_VALUE_REGEX = re.compile(r"^(?:[A-Za-z0-9](?:[A-Za-z0-9_.-]*[A-Za-z0-9])?)?$") +# https://kubernetes.io/docs/reference/kubernetes-api/common-definitions/quantity/ +QUANTITY_REGEX = re.compile( + r"^[+-]?(?:[0-9]+(?:\.[0-9]*)?|\.[0-9]+)(?:[KMGTPE]i|[numkMGTPE]|[eE][+-]?[0-9]+)?$" +) + NVIDIA_GPU_RESOURCE = "nvidia.com/gpu" NVIDIA_GPU_NODE_TAINT = NVIDIA_GPU_RESOURCE NVIDIA_GPU_PRODUCT_LABEL = f"{NVIDIA_GPU_RESOURCE}.product" @@ -345,6 +350,13 @@ def build_dockerconfigjson(image_name: str, username: str, password: str) -> str parse_quantity = cast(Callable[[Union[str, int, float, Decimal]], Decimal], _parse_quantity) +def validate_quantity(quantity: str) -> None: + # Stricter than parse_quantity(), which also accepts some strings that the API rejects, + # e.g., `1 Gi`, `1K`, `500mi` + if QUANTITY_REGEX.fullmatch(quantity) is None: + raise ValueError("Invalid quantity") + + def format_memory(memory: Memory) -> str: return f"{float(memory)}Gi" diff --git a/src/dstack/_internal/core/compatibility/backends.py b/src/dstack/_internal/core/compatibility/backends.py index 97ef46c38..725df7127 100644 --- a/src/dstack/_internal/core/compatibility/backends.py +++ b/src/dstack/_internal/core/compatibility/backends.py @@ -1,5 +1,9 @@ from dstack._internal.core.backends.gcp.models import GCPBackendConfigWithCreds from dstack._internal.core.backends.hotaisle.models import HotAisleBackendConfigWithCreds +from dstack._internal.core.backends.kubernetes.models import ( + KubernetesBackendConfigWithCreds, + KubernetesProxyJumpConfig, +) from dstack._internal.core.backends.models import AnyBackendConfigWithCreds from dstack._internal.core.models.common import IncludeExcludeStrDictType @@ -15,4 +19,28 @@ def get_backend_config_excludes(config: AnyBackendConfigWithCreds) -> IncludeExc excludes["tpu"] = True if isinstance(config, HotAisleBackendConfigWithCreds) and config.bare_metal is None: excludes["bare_metal"] = True + if isinstance(config, KubernetesBackendConfigWithCreds): + if config.proxy_jump is not None: + excludes["proxy_jump"] = _get_kubernetes_proxy_jump_excludes([config.proxy_jump]) + if config.contexts is not None: + proxy_jumps = [ + context.proxy_jump + for context in config.contexts + if not isinstance(context, str) and context.proxy_jump is not None + ] + excludes["contexts"] = { + "__all__": {"proxy_jump": _get_kubernetes_proxy_jump_excludes(proxy_jumps)} + } + return excludes + + +def _get_kubernetes_proxy_jump_excludes( + proxy_jumps: list[KubernetesProxyJumpConfig], +) -> IncludeExcludeStrDictType: + # A field is excluded only if it's unset in all `proxy_jumps` + excludes: IncludeExcludeStrDictType = {} + if all(proxy_jump.cpu is None for proxy_jump in proxy_jumps): + excludes["cpu"] = True + if all(proxy_jump.memory is None for proxy_jump in proxy_jumps): + excludes["memory"] = True return excludes diff --git a/src/tests/_internal/core/backends/kubernetes/test_compute.py b/src/tests/_internal/core/backends/kubernetes/test_compute.py index 3d1a7a906..c8ffe10d0 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_compute.py +++ b/src/tests/_internal/core/backends/kubernetes/test_compute.py @@ -1,3 +1,4 @@ +import logging from collections.abc import Generator from datetime import datetime, timezone from typing import Optional @@ -8,6 +9,7 @@ from kubernetes.client.exceptions import ApiException from dstack._internal.core.backends.kubernetes.compute import ( + JUMP_POD_EPHEMERAL_STORAGE_REQUEST, KubernetesCompute, _check_and_configure_jump_pod_service, _create_jump_pod_service_if_not_exists, @@ -142,7 +144,18 @@ def _make_pod( ready: bool = True, node_name: Optional[str] = "node-a", terminating: bool = False, + unschedulable_message: Optional[str] = None, ) -> client.V1Pod: + conditions = [client.V1PodCondition(type="Ready", status="True" if ready else "False")] + if unschedulable_message is not None: + conditions.append( + client.V1PodCondition( + type="PodScheduled", + status="False", + reason="Unschedulable", + message=unschedulable_message, + ) + ) return client.V1Pod( metadata=client.V1ObjectMeta( name=name, @@ -152,12 +165,7 @@ def _make_pod( ), ), spec=client.V1PodSpec(containers=[], node_name=node_name), - status=client.V1PodStatus( - phase=phase, - conditions=[ - client.V1PodCondition(type="Ready", status="True" if ready else "False"), - ], - ), + status=client.V1PodStatus(phase=phase, conditions=conditions), ) @@ -202,7 +210,14 @@ def make_apis( _mock_read(apps_api.read_namespaced_replica_set, replica_set) return api, apps_api - def call(self, api: Mock, apps_api: Mock, jump_pod_port: Optional[int] = None) -> None: + def call( + self, + api: Mock, + apps_api: Mock, + jump_pod_port: Optional[int] = None, + cpu: str = "500m", + memory: str = "1Gi", + ) -> None: _create_jump_pod_service_if_not_exists( api=api, apps_api=apps_api, @@ -212,15 +227,28 @@ def call(self, api: Mock, apps_api: Mock, jump_pod_port: Optional[int] = None) - jump_pod_service_name=JUMP_POD_SERVICE_NAME, jump_pod_port=jump_pod_port, authorized_keys=["ssh-ed25519 AAAA project"], + cpu_request_limit=cpu, + memory_request_limit=memory, ) - def assert_replica_set_created(self, apps_api: Mock) -> None: + def assert_replica_set_created( + self, apps_api: Mock, *, cpu: str = "500m", memory: str = "1Gi" + ) -> None: apps_api.create_namespaced_replica_set.assert_called_once() replica_set = apps_api.create_namespaced_replica_set.call_args.kwargs["body"] assert replica_set.metadata.name == JUMP_POD_NAME assert replica_set.spec.replicas == 1 assert replica_set.spec.selector.match_labels == JUMP_POD_SELECTOR assert replica_set.spec.template.metadata.labels == JUMP_POD_LABELS + resources = replica_set.spec.template.spec.containers[0].resources + # Equal CPU and memory requests and limits for the Guaranteed QoS class, and + # an ephemeral-storage request without a limit + assert resources.limits == {"cpu": cpu, "memory": memory} + assert resources.requests == { + "cpu": cpu, + "memory": memory, + "ephemeral-storage": JUMP_POD_EPHEMERAL_STORAGE_REQUEST, + } def assert_service_created( self, api: Mock, *, selector: dict[str, str], node_port: Optional[int] @@ -250,6 +278,13 @@ def test_creates_replica_set_and_service_if_not_found(self): self.assert_replica_set_created(apps_api) self.assert_service_created(api, selector=JUMP_POD_SELECTOR, node_port=None) + def test_sets_jump_pod_resources(self): + api, apps_api = self.make_apis() + + self.call(api, apps_api, cpu="1", memory="2Gi") + + self.assert_replica_set_created(apps_api, cpu="1", memory="2Gi") + def test_does_nothing_if_replica_set_and_service_exist(self): api, apps_api = self.make_apis( replica_set=client.V1ReplicaSet( @@ -408,12 +443,39 @@ def test_returns_ssh_proxy_on_node_of_ready_pod(self): pytest.param(_make_pod(terminating=True), id="terminating"), ], ) - def test_waits_if_no_pod_is_ready(self, run_ssh_command_mock: Mock, pod: client.V1Pod): + def test_waits_if_no_pod_is_ready( + self, caplog: pytest.LogCaptureFixture, run_ssh_command_mock: Mock, pod: client.V1Pod + ): + caplog.set_level(logging.WARNING) api = self.make_api(pods=[pod]) assert self.call(api) is None api.read_namespaced_pod.assert_not_called() run_ssh_command_mock.assert_not_called() + assert "cannot be scheduled" not in caplog.text + + def test_warns_if_pod_cannot_be_scheduled( + self, caplog: pytest.LogCaptureFixture, run_ssh_command_mock: Mock + ): + caplog.set_level(logging.WARNING) + api = self.make_api( + pods=[ + _make_pod( + name="pending", + phase="Pending", + ready=False, + node_name=None, + unschedulable_message="0/3 nodes are available: 3 Insufficient memory.", + ), + ] + ) + + assert self.call(api) is None + run_ssh_command_mock.assert_not_called() + assert ( + "Jump pod pending cannot be scheduled: Unschedulable:" + " 0/3 nodes are available: 3 Insufficient memory." + ) in caplog.text def test_waits_if_no_pods(self, run_ssh_command_mock: Mock): # E.g., the ReplicaSet has not created a pod yet. diff --git a/src/tests/_internal/core/backends/kubernetes/test_configurator.py b/src/tests/_internal/core/backends/kubernetes/test_configurator.py index 36f32a08b..a8b983355 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_configurator.py +++ b/src/tests/_internal/core/backends/kubernetes/test_configurator.py @@ -1,3 +1,4 @@ +import re from unittest.mock import Mock import pytest @@ -76,6 +77,85 @@ def test_validate_config_duplicate_contexts(self): with pytest.raises(ServerClientError, match="duplicate contexts: ctx-1, ctx-3"): KubernetesConfigurator().validate_config(config, default_creds_enabled=True) + @pytest.mark.usefixtures("get_clusters_mock") + @pytest.mark.parametrize( + ["cpu", "memory"], + [ + pytest.param("100m", "32Mi", id="minimum"), + pytest.param("1", "2Gi", id="integer-cpu"), + pytest.param(0.5, "1.5Gi", id="number-cpu"), + ], + ) + def test_validate_config_valid_proxy_jump_resources(self, cpu, memory): + config = KubernetesBackendConfigWithCreds( + kubeconfig=KubeconfigConfig(data="mocked", filename="-"), + contexts=[ + KubernetesContextConfig( + name="ctx", + proxy_jump=KubernetesProxyJumpConfig(cpu=cpu, memory=memory), + ), + ], + ) + KubernetesConfigurator().validate_config(config, default_creds_enabled=True) + + @pytest.mark.usefixtures("get_clusters_mock") + @pytest.mark.parametrize( + ["proxy_jump", "expected_error"], + [ + pytest.param( + KubernetesProxyJumpConfig(cpu="50m"), + "invalid proxy_jump.cpu value '50m': must be at least 100m", + id="cpu-below-minimum", + ), + pytest.param( + KubernetesProxyJumpConfig(cpu="-1"), + "invalid proxy_jump.cpu value '-1': must be at least 100m", + id="negative-cpu", + ), + pytest.param( + KubernetesProxyJumpConfig(memory="512"), + "invalid proxy_jump.memory value '512': must be at least 32Mi", + id="memory-in-bytes", + ), + pytest.param( + KubernetesProxyJumpConfig(memory="1 Gi"), + "invalid proxy_jump.memory value '1 Gi': must be a Kubernetes quantity", + id="invalid-memory-quantity", + ), + ], + ) + def test_validate_config_invalid_proxy_jump_resources( + self, proxy_jump: KubernetesProxyJumpConfig, expected_error: str + ): + config = KubernetesBackendConfigWithCreds( + kubeconfig=KubeconfigConfig(data="mocked", filename="-"), + proxy_jump=proxy_jump, + namespace="ns", + ) + with pytest.raises(ServerClientError, match=re.escape(expected_error)): + KubernetesConfigurator().validate_config(config, default_creds_enabled=True) + + @pytest.mark.usefixtures("get_clusters_mock") + def test_validate_config_invalid_context_proxy_jump_resources(self): + config = KubernetesBackendConfigWithCreds( + kubeconfig=KubeconfigConfig(data="mocked", filename="-"), + contexts=[ + "ctx-1", + KubernetesContextConfig( + name="ctx-2", + proxy_jump=KubernetesProxyJumpConfig(cpu="1", memory="500mi"), + ), + ], + ) + with pytest.raises( + ServerClientError, + match=re.escape( + "invalid proxy_jump.memory value '500mi' in context 'ctx-2':" + " must be a Kubernetes quantity" + ), + ): + KubernetesConfigurator().validate_config(config, default_creds_enabled=True) + def test_validate_config_cluster_check_failed( self, monkeypatch: pytest.MonkeyPatch, get_clusters_mock: Mock ): diff --git a/src/tests/_internal/core/backends/kubernetes/test_resources.py b/src/tests/_internal/core/backends/kubernetes/test_resources.py index e1c55dc96..bd2324d40 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_resources.py +++ b/src/tests/_internal/core/backends/kubernetes/test_resources.py @@ -12,6 +12,7 @@ get_nvidia_gpu_from_node_labels, validate_label_key, validate_label_value, + validate_quantity, ) from dstack._internal.core.models.instances import Disk, Gpu, Resources from dstack._internal.core.models.resources import ( @@ -218,3 +219,45 @@ 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) + + +class TestValidateQuantity: + @pytest.mark.parametrize( + "quantity", + [ + pytest.param("1", id="integer"), + pytest.param("0.5", id="decimal"), + pytest.param(".5", id="no-integer-part"), + pytest.param("1.", id="no-fractional-part"), + pytest.param("500m", id="milli"), + pytest.param("100u", id="micro"), + pytest.param("1k", id="decimal-si"), + pytest.param("1E", id="exa"), + pytest.param("512Mi", id="binary-si"), + pytest.param("1Ei", id="binary-si-exbi"), + pytest.param("1e3", id="exponent"), + pytest.param("1.5E-3", id="signed-exponent"), + # A valid format, the value is checked separately + pytest.param("-1", id="negative"), + ], + ) + def test_valid(self, quantity: str): + validate_quantity(quantity) + + # parse_quantity() accepts some of these, but the API rejects them + @pytest.mark.parametrize( + "quantity", + [ + pytest.param("1 Gi", id="space"), + pytest.param("1K", id="uppercase-kilo"), + pytest.param("500mi", id="lowercase-binary-si"), + pytest.param("1gi", id="lowercase-gibi"), + pytest.param("1GB", id="bytes-unit"), + pytest.param("1e", id="exponent-without-number"), + pytest.param("Gi", id="no-number"), + pytest.param("", id="empty"), + ], + ) + def test_invalid(self, quantity: str): + with pytest.raises(ValueError, match="Invalid quantity"): + validate_quantity(quantity) diff --git a/src/tests/_internal/core/compatibility/test_backends.py b/src/tests/_internal/core/compatibility/test_backends.py index dfc3273ee..f360cadbe 100644 --- a/src/tests/_internal/core/compatibility/test_backends.py +++ b/src/tests/_internal/core/compatibility/test_backends.py @@ -7,6 +7,12 @@ HotAisleAPIKeyCreds, HotAisleBackendConfigWithCreds, ) +from dstack._internal.core.backends.kubernetes.models import ( + KubeconfigConfig, + KubernetesBackendConfigWithCreds, + KubernetesContextConfig, + KubernetesProxyJumpConfig, +) from dstack._internal.core.backends.models import AnyBackendConfigWithCreds from dstack._internal.core.compatibility.backends import get_backend_config_excludes @@ -42,3 +48,62 @@ def test_keeps_set_hotaisle_bare_metal(self): ) assert _request_body(config)["bare_metal"] is True + + def test_excludes_unset_kubernetes_proxy_jump_resources(self): + config = KubernetesBackendConfigWithCreds( + kubeconfig=KubeconfigConfig(data="mocked"), + proxy_jump=KubernetesProxyJumpConfig(hostname="203.0.113.1"), + ) + + assert _request_body(config)["proxy_jump"] == {"hostname": "203.0.113.1", "port": None} + + def test_keeps_set_kubernetes_proxy_jump_resources(self): + config = KubernetesBackendConfigWithCreds( + kubeconfig=KubeconfigConfig(data="mocked"), + proxy_jump=KubernetesProxyJumpConfig(cpu="1"), + ) + + assert _request_body(config)["proxy_jump"] == {"hostname": None, "port": None, "cpu": "1"} + + def test_excludes_unset_kubernetes_context_proxy_jump_resources(self): + config = KubernetesBackendConfigWithCreds( + kubeconfig=KubeconfigConfig(data="mocked"), + contexts=[ + "ctx-1", + KubernetesContextConfig(name="ctx-2"), + KubernetesContextConfig( + name="ctx-3", proxy_jump=KubernetesProxyJumpConfig(hostname="203.0.113.1") + ), + ], + ) + + assert _request_body(config)["contexts"] == [ + "ctx-1", + {"name": "ctx-2", "proxy_jump": None}, + {"name": "ctx-3", "proxy_jump": {"hostname": "203.0.113.1", "port": None}}, + ] + + def test_keeps_kubernetes_context_proxy_jump_resources_set_in_any_context(self): + # A field set in one context is sent for all contexts, as only newer servers accept it + config = KubernetesBackendConfigWithCreds( + kubeconfig=KubeconfigConfig(data="mocked"), + contexts=[ + KubernetesContextConfig( + name="ctx-1", proxy_jump=KubernetesProxyJumpConfig(cpu="1") + ), + KubernetesContextConfig( + name="ctx-2", proxy_jump=KubernetesProxyJumpConfig(memory="2Gi") + ), + ], + ) + + assert _request_body(config)["contexts"] == [ + { + "name": "ctx-1", + "proxy_jump": {"hostname": None, "port": None, "cpu": "1", "memory": None}, + }, + { + "name": "ctx-2", + "proxy_jump": {"hostname": None, "port": None, "cpu": None, "memory": "2Gi"}, + }, + ]