Kubernetes: run the jump pod with Guaranteed QoS

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 <namespace> set resources replicaset dstack-<project>-ssh-jump-pod \
+          --requests=cpu=1,memory=2Gi --limits=cpu=1,memory=2Gi
+        kubectl -n <namespace> delete pod -l app.kubernetes.io/instance=dstack-<project>-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.
 

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

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}")

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):

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"
 

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