From 94fb749ddb1431969598131864481fce11bff0fe Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Mon, 2 Nov 2026 14:14:41 +0000 Subject: [PATCH] Kubernetes: run the jump pod with Guaranteed QoS The jump pod requested no resources, so it was BestEffort, the first pod the kubelet evicts under node pressure. Since #4447, its ReplicaSet recreates it, but every eviction still drops all SSH sessions through the jump pod. * The jump pod container now requests and limits CPU and memory equally, which gives it the Guaranteed QoS class. The defaults, 500m and 1Gi, come from measuring the jump pod image: ~9 MiB plus ~1.6 MiB per tunneled SSH session, and 130-180 MiB/s of tunneled traffic at 500m. * The new proxy_jump.cpu and proxy_jump.memory properties override the defaults. Values must be Kubernetes quantities of at least 100m and 32Mi. * The container also requests 64Mi of ephemeral storage, without a limit. QoS doesn't cover disk pressure, under which the kubelet first evicts pods whose disk usage exceeds their ephemeral-storage request. * A jump pod that can't be scheduled, e.g., due to insufficient resources, is now logged as a warning with the scheduler's reason. Before, the only symptom was runs hitting the provisioning timeout. Notes: * New values don't apply to an existing jump pod. The docs describe how to apply them with kubectl. * A legacy jump pod that the ReplicaSet adopts on upgrade stays BestEffort until it's recreated. Fixes: https://github.com/dstackai/dstack/issues/4407 Co-Authored-By: Claude Opus 5.5 --- mkdocs/docs/concepts/backends.md | 23 ++++++ .../core/backends/kubernetes/compute.py | 55 +++++++++++-- .../core/backends/kubernetes/configurator.py | 40 ++++++++++ .../core/backends/kubernetes/models.py | 24 ++++++ .../core/backends/kubernetes/resources.py | 12 +++ .../_internal/core/compatibility/backends.py | 28 +++++++ .../core/backends/kubernetes/test_compute.py | 80 ++++++++++++++++--- .../backends/kubernetes/test_configurator.py | 80 +++++++++++++++++++ .../backends/kubernetes/test_resources.py | 43 ++++++++++ .../core/compatibility/test_backends.py | 65 +++++++++++++++ 10 files changed, 436 insertions(+), 14 deletions(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 20894f9def..96f4157b2a 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 c9cd0e7f9f..f9c3929273 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 8a4e2e17f0..89cab248c8 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 c9df6ac88c..ac126f9fad 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 37aa36ef6b..95eb5b46eb 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 97ef46c386..725df7127c 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 3d1a7a9064..c8ffe10d09 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 36f32a08b9..a8b9833550 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 e1c55dc96f..bd2324d400 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 dfc3273ee2..f360cadbe6 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"}, + }, + ]