Kubernetes: add proxy_jump.node_selector and tolerations

mkdocs/docs/concepts/backends.md

@@ -1271,12 +1271,32 @@ projects:
 
     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`"
+    To run the jump pod only on specific nodes, e.g., nodes with an `ExternalIP` or CPU-only nodes, label them and set `proxy_jump.node_selector`. If these nodes are [tainted](https://kubernetes.io/docs/concepts/scheduling-eviction/taint-and-toleration/), also set `proxy_jump.tolerations`:
+
+    ```yaml
+    contexts:
+    - name: gpu-cluster-a
+      proxy_jump:
+        node_selector:
+        - example.com/jump-host
+        tolerations:
+        - example.com/dedicated=jump-host:NoSchedule
+    ```
+
+    Each `node_selector` item is `key=value`, `key=` (an empty value), or `key` (any value), and a node must match all of them. Since a `nodeSelector` can't match any value, `dstack` applies `node_selector` as a required [node affinity](https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#node-affinity) term. If none of the selected nodes can run the jump pod, it stays pending and runs can't be provisioned, so label several nodes.
+
+    Each `tolerations` item is `key[=value][:effect]`, the `kubectl taint` format, so a taint added with `kubectl taint` can be copied as is. Without `=value`, any value is tolerated; without `:effect`, any effect. If `tolerations` isn't set, `dstack` tolerates the GPU node taints when every node is tainted, but not if `node_selector` is set.
+
+    !!! warning "Changing `cpu`, `memory`, `node_selector`, and `tolerations`"
         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
+        # cpu and memory
         kubectl -n <namespace> set resources replicaset dstack-<project>-ssh-jump-pod \
           --requests=cpu=1,memory=2Gi --limits=cpu=1,memory=2Gi
+        # node_selector and tolerations: edit spec.template.spec.affinity and .tolerations
+        kubectl -n <namespace> edit replicaset dstack-<project>-ssh-jump-pod
+        # Recreate the pod
         kubectl -n <namespace> delete pod -l app.kubernetes.io/instance=dstack-<project>-ssh-jump-pod
         ```
 

src/dstack/_internal/core/backends/kubernetes/compute.py

@@ -60,6 +60,7 @@
     OBJECT_NAME_MAX_LENGTH,
     AnyKubernetesGPUResource,
     KubernetesResource,
+    Operator,
     PodPhase,
     ResourceLimits,
     ResourceRequests,
@@ -80,7 +81,9 @@
     get_registry_auth_secret_name,
     is_hard_taint,
     is_taint_tolerated,
+    parse_node_selector_requirement,
     parse_quantity,
+    parse_toleration,
 )
 from dstack._internal.core.backends.kubernetes.utils import (
     LEGACY_CURRENT_CONTEXT_REGION,
@@ -139,11 +142,6 @@
 JOB_POD_SCHEDULING_TIMEOUT = 10
 
 
-class Operator(str, Enum):
-    EXISTS = "Exists"
-    IN = "In"
-
-
 class RestartPolicy(str, Enum):
     """
     This policy determines how Kubernetes reacts to containers exiting due to errors
@@ -267,6 +265,8 @@ def run_job(
             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,
+            node_selector=cluster.proxy_jump.node_selector,
+            tolerations=cluster.proxy_jump.tolerations,
         )
 
         pod_name = generate_unique_instance_name_for_job(
@@ -888,6 +888,8 @@ def _create_jump_pod_service_if_not_exists(
     authorized_keys: list[str],
     cpu_request_limit: str,
     memory_request_limit: str,
+    node_selector: Optional[list[str]],
+    tolerations: Optional[list[str]],
 ) -> None:
     base_labels = build_base_labels(
         component="ssh-proxy",
@@ -959,35 +961,19 @@ def _create_jump_pod_service_if_not_exists(
         # Otherwise, the existing legacy pod will be adopted by a newly created ReplicaSet.
 
     if should_create_replica_set:
-        # False if we found at least one node without any "hard" taint, that is, if we don't need
-        # to specify the toleration.
-        toleration_required = True
-        # (key, effect) pairs.
-        tolerated_taints: set[tuple[str, str]] = set()
-        for node in api.list_node().items:
-            if (node_spec := node.spec) is None:
-                continue
-            # True if the node has at least one NoExecute or NoSchedule taint.
-            has_hard_taint = False
-            taints = node_spec.taints or []
-            for taint in taints:
-                # A "soft" taint, ignore.
-                if not is_hard_taint(taint):
-                    continue
-                has_hard_taint = True
-                if is_taint_tolerated(taint):
-                    tolerated_taints.add((taint.key, taint.effect))
-            if not has_hard_taint:
-                toleration_required = False
-                break
-        tolerations: list[client.V1Toleration] = []
-        if toleration_required:
-            for key, effect in tolerated_taints:
-                tolerations.append(
-                    client.V1Toleration(key=key, operator=Operator.EXISTS, effect=effect)
-                )
-            if not tolerations:
-                logger.warning("No appropriate node found, the jump pod may never be scheduled")
+        affinity: Optional[client.V1Affinity] = None
+        # Not for an empty `node_selector`, as an empty node selector term matches no nodes
+        if node_selector:
+            affinity = _get_jump_pod_affinity(node_selector)
+        pod_tolerations: list[client.V1Toleration]
+        if tolerations is not None:
+            pod_tolerations = [parse_toleration(toleration) for toleration in tolerations]
+        elif node_selector is not None:
+            # Autodetection considers all nodes, not only the selected ones, so the user
+            # specifies tolerations for the selected nodes, if any are needed.
+            pod_tolerations = []
+        else:
+            pod_tolerations = _get_jump_pod_tolerations_from_node_taints(api)
         commands = _get_jump_pod_commands(authorized_keys)
         # Equal CPU and memory requests and limits for the Guaranteed QoS class
         limits = {
@@ -1030,7 +1016,8 @@ def _create_jump_pod_service_if_not_exists(
                             )
                         ],
                         restart_policy=RestartPolicy.ALWAYS,
-                        tolerations=tolerations,
+                        affinity=affinity,
+                        tolerations=pod_tolerations,
                     ),
                 ),
             ),
@@ -1220,6 +1207,60 @@ def _check_and_configure_jump_pod_service(
     )
 
 
+def _get_jump_pod_affinity(node_selector: list[str]) -> client.V1Affinity:
+    # A nodeSelector cannot express `Exists`. A single node affinity term matches like
+    # a nodeSelector: a node must meet all requirements.
+    return client.V1Affinity(
+        node_affinity=client.V1NodeAffinity(
+            required_during_scheduling_ignored_during_execution=client.V1NodeSelector(
+                node_selector_terms=[
+                    client.V1NodeSelectorTerm(
+                        match_expressions=[
+                            parse_node_selector_requirement(requirement)
+                            for requirement in node_selector
+                        ],
+                    ),
+                ],
+            ),
+        ),
+    )
+
+
+def _get_jump_pod_tolerations_from_node_taints(
+    api: client.CoreV1Api,
+) -> list[client.V1Toleration]:
+    # False if we found at least one node without any "hard" taint, that is, if we don't need
+    # to specify the toleration.
+    toleration_required = True
+    # (key, effect) pairs.
+    tolerated_taints: set[tuple[str, str]] = set()
+    for node in api.list_node().items:
+        if (node_spec := node.spec) is None:
+            continue
+        # True if the node has at least one NoExecute or NoSchedule taint.
+        has_hard_taint = False
+        taints = node_spec.taints or []
+        for taint in taints:
+            # A "soft" taint, ignore.
+            if not is_hard_taint(taint):
+                continue
+            has_hard_taint = True
+            if is_taint_tolerated(taint):
+                tolerated_taints.add((taint.key, taint.effect))
+        if not has_hard_taint:
+            toleration_required = False
+            break
+    tolerations: list[client.V1Toleration] = []
+    if toleration_required:
+        for key, effect in tolerated_taints:
+            tolerations.append(
+                client.V1Toleration(key=key, operator=Operator.EXISTS, effect=effect)
+            )
+        if not tolerations:
+            logger.warning("No appropriate node found, the jump pod may never be scheduled")
+    return tolerations
+
+
 def _get_jump_pod_commands(authorized_keys: list[str]) -> list[str]:
     authorized_keys_content = "\n".join(authorized_keys)
     commands = [

src/dstack/_internal/core/backends/kubernetes/configurator.py

@@ -13,7 +13,12 @@
     KubernetesProxyJumpConfig,
     KubernetesStoredConfig,
 )
-from dstack._internal.core.backends.kubernetes.resources import parse_quantity, validate_quantity
+from dstack._internal.core.backends.kubernetes.resources import (
+    parse_node_selector_requirement,
+    parse_quantity,
+    parse_toleration,
+    validate_quantity,
+)
 from dstack._internal.core.backends.kubernetes.utils import (
     check_cluster,
     get_clusters_from_backend_config,
@@ -112,6 +117,7 @@ def _check_config_proxy_jumps(self, config: KubernetesBackendConfig):
     def _check_proxy_jump(
         self, proxy_jump: KubernetesProxyJumpConfig, *, context_name: Optional[str] = None
     ):
+        in_context = f" in context {context_name!r}" if context_name is not None else ""
         for field, quantity, min_quantity in [
             ("cpu", proxy_jump.cpu, _JUMP_POD_MIN_CPU),
             ("memory", proxy_jump.memory, _JUMP_POD_MIN_MEMORY),
@@ -126,7 +132,17 @@ def _check_proxy_jump(
                 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}")
+            raise ServerClientError(
+                f"invalid proxy_jump.{field} value {quantity!r}{in_context}: {reason}"
+            )
+        for field, items, parse in [
+            ("node_selector", proxy_jump.node_selector, parse_node_selector_requirement),
+            ("tolerations", proxy_jump.tolerations, parse_toleration),
+        ]:
+            for item in items or []:
+                try:
+                    parse(item)
+                except ValueError as e:
+                    raise ServerClientError(
+                        f"invalid proxy_jump.{field} item {item!r}{in_context}: {e}"
+                    )

src/dstack/_internal/core/backends/kubernetes/models.py

@@ -38,6 +38,27 @@ class KubernetesProxyJumpConfig(CoreModel):
             )
         ),
     ] = None
+    node_selector: Annotated[
+        Optional[list[str]],
+        Field(
+            description=(
+                "The labels that the jump pod's node must have."
+                " Each item is `key=value`, `key=` (an empty value), or `key` (any value)"
+            )
+        ),
+    ] = None
+    tolerations: Annotated[
+        Optional[list[str]],
+        Field(
+            description=(
+                "The node taints that the jump pod tolerates."
+                " Each item is `key[=value][:effect]`, as in `kubectl taint`."
+                " Without `=value`, any value is tolerated; without `:effect`, any effect."
+                " If not set, defaults to the GPU node taints when every node is tainted,"
+                " or to none if `node_selector` is set"
+            )
+        ),
+    ] = None
 
 
 class KubernetesContextConfig(CoreModel):

src/dstack/_internal/core/backends/kubernetes/resources.py

@@ -9,7 +9,14 @@
 
 import gpuhunt
 from gpuhunt import KNOWN_AMD_GPUS, KNOWN_NVIDIA_GPUS, AcceleratorVendor
-from kubernetes.client import CoreV1Api, V1Node, V1ObjectMeta, V1Taint
+from kubernetes.client import (
+    CoreV1Api,
+    V1Node,
+    V1NodeSelectorRequirement,
+    V1ObjectMeta,
+    V1Taint,
+    V1Toleration,
+)
 from kubernetes.utils.quantity import parse_quantity as _parse_quantity
 from typing_extensions import Self
 
@@ -28,7 +35,7 @@
 )
 from dstack._internal.core.models.resources import Memory, ResourcesSpec
 from dstack._internal.utils import docker as docker_utils
-from dstack._internal.utils.common import get_or_error
+from dstack._internal.utils.common import get_or_error, list_enum_values_for_annotation
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -93,6 +100,12 @@ class TaintEffect(str, Enum):
     PREFER_NO_SCHEDULE = "PreferNoSchedule"
 
 
+class Operator(str, Enum):
+    EQUAL = "Equal"
+    EXISTS = "Exists"
+    IN = "In"
+
+
 class KubernetesResource(str, Enum):
     CPU = "cpu"
     MEMORY = "memory"
@@ -377,6 +390,48 @@ def validate_label_value(value: str) -> None:
         raise ValueError("Invalid value")
 
 
+def parse_node_selector_requirement(requirement: str) -> V1NodeSelectorRequirement:
+    """
+    Parses `key=value`, `key=` (an empty value), or `key` (any value).
+    """
+    key, value_sep, value = requirement.partition("=")
+    _validate_key(key)
+    if not value_sep:
+        return V1NodeSelectorRequirement(key=key, operator=Operator.EXISTS)
+    validate_label_value(value)
+    return V1NodeSelectorRequirement(key=key, operator=Operator.IN, values=[value])
+
+
+def parse_toleration(toleration: str) -> V1Toleration:
+    """
+    Parses `key[=value][:effect]`, the `kubectl taint` format with an optional effect.
+    Without `=value`, tolerates any value; without `:effect`, tolerates any effect.
+    """
+    key_value, effect_sep, effect = toleration.partition(":")
+    key, value_sep, value = key_value.partition("=")
+    # An empty key would tolerate all taints
+    _validate_key(key)
+    taint_effect: Optional[TaintEffect] = None
+    if effect_sep:
+        try:
+            taint_effect = TaintEffect(effect)
+        except ValueError:
+            raise ValueError(
+                f"Invalid effect, must be one of: {list_enum_values_for_annotation(TaintEffect)}"
+            )
+    if not value_sep:
+        return V1Toleration(key=key, operator=Operator.EXISTS, effect=taint_effect)
+    validate_label_value(value)
+    return V1Toleration(key=key, operator=Operator.EQUAL, value=value, effect=taint_effect)
+
+
+def _validate_key(key: str) -> None:
+    try:
+        validate_label_key(key)
+    except ValueError as e:
+        raise ValueError(f"Invalid key: {e}") from e
+
+
 def build_dockerconfigjson(image_name: str, username: str, password: str) -> str:
     registry = docker_utils.parse_image_name(image_name).registry
     if registry is None or docker_utils.is_default_registry(registry):

src/dstack/_internal/core/compatibility/backends.py

@@ -43,4 +43,8 @@ def _get_kubernetes_proxy_jump_excludes(
         excludes["cpu"] = True
     if all(proxy_jump.memory is None for proxy_jump in proxy_jumps):
         excludes["memory"] = True
+    if all(proxy_jump.node_selector is None for proxy_jump in proxy_jumps):
+        excludes["node_selector"] = True
+    if all(proxy_jump.tolerations is None for proxy_jump in proxy_jumps):
+        excludes["tolerations"] = True
     return excludes