From 2753883be378bac8703f2e671521cb8720e14b69 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Tue, 3 Nov 2026 14:25:25 +0000 Subject: [PATCH] Kubernetes: add proxy_jump.node_selector and tolerations The jump pod could land on any node. On some clusters, e.g., Nebius, a NodePort is reachable only on the node that runs the pod, so the jump pod is unreachable if that node has no ExternalIP. Users may also want to keep the jump pod off GPU nodes. * proxy_jump.node_selector restricts the jump pod to labeled nodes. Items are key=value, key= (an empty value), or key (any value). As nodeSelector can't match any value, all items go into a single required node affinity term. * proxy_jump.tolerations items are key[=value][:effect], the kubectl taint format with an optional effect. Without a value, any value is tolerated; without an effect, any effect. An empty key, which would tolerate all taints, is rejected. * If tolerations is not set, GPU node taints are autodetected as before, unless node_selector is set, since autodetection considers all nodes, not only the selected ones. tolerations: [] turns autodetection off. Notes: * New values don't apply to an existing jump pod. The docs describe how to apply them with kubectl. Co-Authored-By: Claude Opus 5.5 --- mkdocs/docs/concepts/backends.md | 22 +++- .../core/backends/kubernetes/compute.py | 111 ++++++++++++------ .../core/backends/kubernetes/configurator.py | 26 +++- .../core/backends/kubernetes/models.py | 21 ++++ .../core/backends/kubernetes/resources.py | 59 +++++++++- .../_internal/core/compatibility/backends.py | 4 + .../core/backends/kubernetes/test_compute.py | 104 +++++++++++++++- .../backends/kubernetes/test_configurator.py | 67 +++++++++++ .../backends/kubernetes/test_resources.py | 99 ++++++++++++++++ .../core/compatibility/test_backends.py | 13 ++ 10 files changed, 479 insertions(+), 47 deletions(-) diff --git a/mkdocs/docs/concepts/backends.md b/mkdocs/docs/concepts/backends.md index 2b071bf5a..62abf33c5 100644 --- a/mkdocs/docs/concepts/backends.md +++ b/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 set resources replicaset dstack--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 edit replicaset dstack--ssh-jump-pod + # Recreate the pod kubectl -n delete pod -l app.kubernetes.io/instance=dstack--ssh-jump-pod ``` diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index e996c5882..74922b78c 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/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 = [ diff --git a/src/dstack/_internal/core/backends/kubernetes/configurator.py b/src/dstack/_internal/core/backends/kubernetes/configurator.py index 89cab248c..55e75f8ff 100644 --- a/src/dstack/_internal/core/backends/kubernetes/configurator.py +++ b/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}" + ) diff --git a/src/dstack/_internal/core/backends/kubernetes/models.py b/src/dstack/_internal/core/backends/kubernetes/models.py index ac126f9fa..5ea9f90a7 100644 --- a/src/dstack/_internal/core/backends/kubernetes/models.py +++ b/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): diff --git a/src/dstack/_internal/core/backends/kubernetes/resources.py b/src/dstack/_internal/core/backends/kubernetes/resources.py index c08df7568..8828c0e7b 100644 --- a/src/dstack/_internal/core/backends/kubernetes/resources.py +++ b/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): diff --git a/src/dstack/_internal/core/compatibility/backends.py b/src/dstack/_internal/core/compatibility/backends.py index 725df7127..72e58d8e5 100644 --- a/src/dstack/_internal/core/compatibility/backends.py +++ b/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 diff --git a/src/tests/_internal/core/backends/kubernetes/test_compute.py b/src/tests/_internal/core/backends/kubernetes/test_compute.py index c8ffe10d0..f4e0d1b12 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_compute.py +++ b/src/tests/_internal/core/backends/kubernetes/test_compute.py @@ -127,6 +127,7 @@ def test_get_offers_without_full_offers_adjusts_to_requested_slice(): JUMP_POD_SELECTOR = {"app.kubernetes.io/instance": JUMP_POD_NAME} # Both the labels and the service selector before https://github.com/dstackai/dstack/pull/3891 PRE_3891_JUMP_POD_LABELS = {"app.kubernetes.io/name": JUMP_POD_NAME} +GPU_NODE_TAINT = client.V1Taint(key="nvidia.com/gpu", value="present", effect="NoSchedule") def _mock_read(method: Mock, obj: Optional[object]) -> None: @@ -180,13 +181,17 @@ def _make_service(selector: dict[str, str], node_port: int) -> client.V1Service: ) -def _make_node(name: str, external_ip: Optional[str] = None) -> client.V1Node: +def _make_node( + name: str, + external_ip: Optional[str] = None, + taints: Optional[list[client.V1Taint]] = None, +) -> client.V1Node: addresses: list[client.V1NodeAddress] = [] if external_ip is not None: addresses.append(client.V1NodeAddress(type="ExternalIP", address=external_ip)) return client.V1Node( metadata=client.V1ObjectMeta(name=name), - spec=client.V1NodeSpec(), + spec=client.V1NodeSpec(taints=taints), status=client.V1NodeStatus(addresses=addresses), ) @@ -199,13 +204,16 @@ def make_apis( pod: Optional[client.V1Pod] = None, replica_set: Optional[client.V1ReplicaSet] = None, service: Optional[client.V1Service] = None, + nodes: Optional[list[client.V1Node]] = None, ) -> tuple[Mock, Mock]: api = Mock(spec=client.CoreV1Api) _mock_read(api.read_namespace, client.V1Namespace() if namespace_exists else None) _mock_read(api.read_namespaced_pod, pod) _mock_read(api.read_namespaced_service, service) - # An untainted node, so the jump pod needs no tolerations. - api.list_node.return_value = client.V1NodeList(items=[_make_node("node-a")]) + if nodes is None: + # An untainted node, so the jump pod needs no tolerations. + nodes = [_make_node("node-a")] + api.list_node.return_value = client.V1NodeList(items=nodes) apps_api = Mock(spec=client.AppsV1Api) _mock_read(apps_api.read_namespaced_replica_set, replica_set) return api, apps_api @@ -217,6 +225,8 @@ def call( jump_pod_port: Optional[int] = None, cpu: str = "500m", memory: str = "1Gi", + node_selector: Optional[list[str]] = None, + tolerations: Optional[list[str]] = None, ) -> None: _create_jump_pod_service_if_not_exists( api=api, @@ -229,8 +239,15 @@ def call( authorized_keys=["ssh-ed25519 AAAA project"], cpu_request_limit=cpu, memory_request_limit=memory, + node_selector=node_selector, + tolerations=tolerations, ) + def get_pod_spec(self, apps_api: Mock) -> client.V1PodSpec: + apps_api.create_namespaced_replica_set.assert_called_once() + replica_set = apps_api.create_namespaced_replica_set.call_args.kwargs["body"] + return replica_set.spec.template.spec + def assert_replica_set_created( self, apps_api: Mock, *, cpu: str = "500m", memory: str = "1Gi" ) -> None: @@ -285,6 +302,85 @@ def test_sets_jump_pod_resources(self): self.assert_replica_set_created(apps_api, cpu="1", memory="2Gi") + def test_autodetects_tolerations_if_all_nodes_are_tainted(self): + api, apps_api = self.make_apis(nodes=[_make_node("node-a", taints=[GPU_NODE_TAINT])]) + + self.call(api, apps_api) + + pod_spec = self.get_pod_spec(apps_api) + assert pod_spec.affinity is None + assert pod_spec.tolerations == [ + client.V1Toleration(key="nvidia.com/gpu", operator="Exists", effect="NoSchedule") + ] + + def test_sets_node_affinity_without_autodetected_tolerations(self): + api, apps_api = self.make_apis(nodes=[_make_node("node-a", taints=[GPU_NODE_TAINT])]) + + self.call(api, apps_api, node_selector=["example.com/jump-host", "example.com/pool=cpu"]) + + pod_spec = self.get_pod_spec(apps_api) + assert pod_spec.affinity == client.V1Affinity( + node_affinity=client.V1NodeAffinity( + required_during_scheduling_ignored_during_execution=client.V1NodeSelector( + # A single term, so that all requirements must match + node_selector_terms=[ + client.V1NodeSelectorTerm( + match_expressions=[ + client.V1NodeSelectorRequirement( + key="example.com/jump-host", operator="Exists" + ), + client.V1NodeSelectorRequirement( + key="example.com/pool", operator="In", values=["cpu"] + ), + ] + ) + ] + ) + ) + ) + assert pod_spec.tolerations == [] + api.list_node.assert_not_called() + + def test_empty_node_selector_disables_autodetected_tolerations(self): + api, apps_api = self.make_apis(nodes=[_make_node("node-a", taints=[GPU_NODE_TAINT])]) + + self.call(api, apps_api, node_selector=[]) + + pod_spec = self.get_pod_spec(apps_api) + # An empty node selector term would match no nodes + assert pod_spec.affinity is None + assert pod_spec.tolerations == [] + api.list_node.assert_not_called() + + @pytest.mark.parametrize("node_selector", [None, ["example.com/pool=cpu"]]) + @pytest.mark.parametrize( + ["tolerations", "expected"], + [ + pytest.param( + ["example.com/pool=cpu:NoSchedule"], + [ + client.V1Toleration( + key="example.com/pool", operator="Equal", value="cpu", effect="NoSchedule" + ) + ], + id="set", + ), + pytest.param([], [], id="empty"), + ], + ) + def test_sets_tolerations_without_autodetection( + self, + node_selector: Optional[list[str]], + tolerations: list[str], + expected: list[client.V1Toleration], + ): + api, apps_api = self.make_apis(nodes=[_make_node("node-a", taints=[GPU_NODE_TAINT])]) + + self.call(api, apps_api, node_selector=node_selector, tolerations=tolerations) + + assert self.get_pod_spec(apps_api).tolerations == expected + api.list_node.assert_not_called() + def test_does_nothing_if_replica_set_and_service_exist(self): api, apps_api = self.make_apis( replica_set=client.V1ReplicaSet( diff --git a/src/tests/_internal/core/backends/kubernetes/test_configurator.py b/src/tests/_internal/core/backends/kubernetes/test_configurator.py index a8b983355..05f9be084 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_configurator.py +++ b/src/tests/_internal/core/backends/kubernetes/test_configurator.py @@ -156,6 +156,73 @@ def test_validate_config_invalid_context_proxy_jump_resources(self): ): KubernetesConfigurator().validate_config(config, default_creds_enabled=True) + @pytest.mark.usefixtures("get_clusters_mock") + def test_validate_config_valid_proxy_jump_placement(self): + config = KubernetesBackendConfigWithCreds( + kubeconfig=KubeconfigConfig(data="mocked", filename="-"), + contexts=[ + KubernetesContextConfig( + name="ctx", + proxy_jump=KubernetesProxyJumpConfig( + node_selector=["example.com/jump-host", "example.com/pool=cpu"], + tolerations=["example.com/pool", "nvidia.com/gpu=:NoSchedule"], + ), + ), + ], + ) + KubernetesConfigurator().validate_config(config, default_creds_enabled=True) + + @pytest.mark.usefixtures("get_clusters_mock") + @pytest.mark.parametrize( + ["proxy_jump", "expected_error"], + [ + pytest.param( + KubernetesProxyJumpConfig(node_selector=["example.com/pool=cpu", "=cpu"]), + "invalid proxy_jump.node_selector item '=cpu': Invalid key: Empty name", + id="node-selector-empty-key", + ), + pytest.param( + KubernetesProxyJumpConfig(tolerations=[":NoSchedule"]), + "invalid proxy_jump.tolerations item ':NoSchedule': Invalid key: Empty name", + id="tolerations-empty-key", + ), + pytest.param( + KubernetesProxyJumpConfig(tolerations=["nvidia.com/gpu:noschedule"]), + "invalid proxy_jump.tolerations item 'nvidia.com/gpu:noschedule': Invalid effect", + id="tolerations-invalid-effect", + ), + ], + ) + def test_validate_config_invalid_proxy_jump_placement( + 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_placement(self): + config = KubernetesBackendConfigWithCreds( + kubeconfig=KubeconfigConfig(data="mocked", filename="-"), + contexts=[ + KubernetesContextConfig( + name="ctx", + proxy_jump=KubernetesProxyJumpConfig(tolerations=["key=a b"]), + ), + ], + ) + with pytest.raises( + ServerClientError, + match=re.escape( + "invalid proxy_jump.tolerations item 'key=a b' in context 'ctx': Invalid value" + ), + ): + 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 bd2324d40..e4a3f3d5a 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_resources.py +++ b/src/tests/_internal/core/backends/kubernetes/test_resources.py @@ -2,6 +2,7 @@ import pytest from gpuhunt import AcceleratorVendor +from kubernetes.client import V1NodeSelectorRequirement, V1Toleration from dstack._internal.core.backends.kubernetes.resources import ( KubernetesResource, @@ -10,6 +11,8 @@ adjust_resources_by_resource_requests, get_amd_gpu_from_node_labels, get_nvidia_gpu_from_node_labels, + parse_node_selector_requirement, + parse_toleration, validate_label_key, validate_label_value, validate_quantity, @@ -221,6 +224,102 @@ def test_invalid_value(self, value: str, expected_error: str): validate_label_value(value) +class TestParseNodeSelectorRequirement: + @pytest.mark.parametrize( + ["requirement", "expected"], + [ + pytest.param( + "example.com/pool=jump", + V1NodeSelectorRequirement(key="example.com/pool", operator="In", values=["jump"]), + id="value", + ), + pytest.param( + "node-role.kubernetes.io/worker=", + V1NodeSelectorRequirement( + key="node-role.kubernetes.io/worker", operator="In", values=[""] + ), + id="empty-value", + ), + pytest.param( + "example.com/pool", + V1NodeSelectorRequirement(key="example.com/pool", operator="Exists"), + id="any-value", + ), + ], + ) + def test_valid(self, requirement: str, expected: V1NodeSelectorRequirement): + assert parse_node_selector_requirement(requirement) == expected + + @pytest.mark.parametrize( + ["requirement", "expected_error"], + [ + pytest.param("", "Invalid key: Empty name", id="empty"), + pytest.param("=jump", "Invalid key: Empty name", id="empty-key"), + pytest.param("Example.com/pool=jump", "Invalid key: Invalid prefix", id="invalid-key"), + pytest.param("pool=a=b", "Invalid value", id="invalid-value"), + ], + ) + def test_invalid(self, requirement: str, expected_error: str): + with pytest.raises(ValueError, match=expected_error): + parse_node_selector_requirement(requirement) + + +class TestParseToleration: + @pytest.mark.parametrize( + ["toleration", "expected"], + [ + pytest.param( + "nvidia.com/gpu", + V1Toleration(key="nvidia.com/gpu", operator="Exists"), + id="any-value-any-effect", + ), + pytest.param( + "nvidia.com/gpu:NoSchedule", + V1Toleration(key="nvidia.com/gpu", operator="Exists", effect="NoSchedule"), + id="any-value", + ), + pytest.param( + "nvidia.com/gpu=present", + V1Toleration(key="nvidia.com/gpu", operator="Equal", value="present"), + id="any-effect", + ), + pytest.param( + "nvidia.com/gpu=present:NoExecute", + V1Toleration( + key="nvidia.com/gpu", operator="Equal", value="present", effect="NoExecute" + ), + id="value-and-effect", + ), + pytest.param( + "dedicated=:PreferNoSchedule", + V1Toleration( + key="dedicated", operator="Equal", value="", effect="PreferNoSchedule" + ), + id="empty-value", + ), + ], + ) + def test_valid(self, toleration: str, expected: V1Toleration): + assert parse_toleration(toleration) == expected + + @pytest.mark.parametrize( + ["toleration", "expected_error"], + [ + # An empty key with the Exists operator would tolerate all taints + pytest.param("", "Invalid key: Empty name", id="empty"), + pytest.param(":NoSchedule", "Invalid key: Empty name", id="empty-key"), + pytest.param("Example.com/key", "Invalid key: Invalid prefix", id="invalid-key"), + pytest.param("key=a b", "Invalid value", id="invalid-value"), + pytest.param("key:", "Invalid effect", id="empty-effect"), + pytest.param("key:noschedule", "Invalid effect", id="lowercase-effect"), + pytest.param("key:NoSchedule-", "Invalid effect", id="taint-removal"), + ], + ) + def test_invalid(self, toleration: str, expected_error: str): + with pytest.raises(ValueError, match=expected_error): + parse_toleration(toleration) + + class TestValidateQuantity: @pytest.mark.parametrize( "quantity", diff --git a/src/tests/_internal/core/compatibility/test_backends.py b/src/tests/_internal/core/compatibility/test_backends.py index f360cadbe..5ce7038ee 100644 --- a/src/tests/_internal/core/compatibility/test_backends.py +++ b/src/tests/_internal/core/compatibility/test_backends.py @@ -65,6 +65,19 @@ def test_keeps_set_kubernetes_proxy_jump_resources(self): assert _request_body(config)["proxy_jump"] == {"hostname": None, "port": None, "cpu": "1"} + def test_keeps_set_kubernetes_proxy_jump_placement(self): + config = KubernetesBackendConfigWithCreds( + kubeconfig=KubeconfigConfig(data="mocked"), + proxy_jump=KubernetesProxyJumpConfig(node_selector=["pool=cpu"], tolerations=[]), + ) + + assert _request_body(config)["proxy_jump"] == { + "hostname": None, + "port": None, + "node_selector": ["pool=cpu"], + "tolerations": [], + } + def test_excludes_unset_kubernetes_context_proxy_jump_resources(self): config = KubernetesBackendConfigWithCreds( kubeconfig=KubeconfigConfig(data="mocked"),