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