diff --git a/mkdocs/snippets/kubernetes/dstack-backend-role.yaml b/mkdocs/snippets/kubernetes/dstack-backend-role.yaml index 7cc37fbc16..c2c9ca5145 100644 --- a/mkdocs/snippets/kubernetes/dstack-backend-role.yaml +++ b/mkdocs/snippets/kubernetes/dstack-backend-role.yaml @@ -6,7 +6,7 @@ metadata: rules: - apiGroups: [""] resources: ["pods"] - verbs: ["get", "watch", "create", "delete"] + verbs: ["get", "list", "watch", "create", "delete"] - apiGroups: [""] resources: ["services"] verbs: ["get", "create", "delete"] @@ -16,3 +16,6 @@ rules: - apiGroups: [""] resources: ["persistentvolumeclaims"] verbs: ["get", "create", "delete"] + - apiGroups: ["apps"] + resources: ["replicasets"] + verbs: ["get", "create"] diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index 4ea57a5608..c9cd0e7f9f 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -64,9 +64,9 @@ get_amd_gpu_from_node_labels, get_instance_offer_from_node, get_instance_offers, - get_node_labels, - get_node_name, get_nvidia_gpu_from_node_labels, + get_object_labels, + get_object_name, is_hard_taint, is_taint_tolerated, parse_quantity, @@ -238,6 +238,7 @@ def run_job( jump_pod_service_name = _get_pod_service_name(jump_pod_name) _create_jump_pod_service_if_not_exists( api=api, + apps_api=client.AppsV1Api(cluster.api_client), namespace=namespace, project_name=run.project_name, jump_pod_name=jump_pod_name, @@ -804,7 +805,7 @@ def _get_nvidia_gpu_node_affinity( ) -> client.V1NodeAffinity: matching_gpu_label_values: set[str] = set() for node in nodes: - labels = get_node_labels(node) + labels = get_object_labels(node) gpu = get_nvidia_gpu_from_node_labels(labels) if gpu is not None and gpu_matches_gpu_spec(gpu, gpu_spec): matching_gpu_label_values.add(labels[NVIDIA_GPU_PRODUCT_LABEL]) @@ -837,7 +838,7 @@ def _get_amd_gpu_node_affinity( ) -> client.V1NodeAffinity: matching_device_ids: set[int] = set() for node in nodes: - labels = get_node_labels(node) + labels = get_object_labels(node) gpu = get_amd_gpu_from_node_labels(labels) if gpu is not None and gpu_matches_gpu_spec(gpu, gpu_spec): matching_device_ids.update(AMD_GPU_NAME_TO_DEVICE_IDS[gpu.name]) @@ -868,6 +869,7 @@ def _offer_modifier( def _create_jump_pod_service_if_not_exists( api: client.CoreV1Api, + apps_api: client.AppsV1Api, namespace: str, project_name: str, jump_pod_name: str, @@ -882,8 +884,9 @@ def _create_jump_pod_service_if_not_exists( ) labels = filter_invalid_labels(base_labels) - service: Optional[client.V1Service] = None pod: Optional[client.V1Pod] = None + replica_set: Optional[client.V1ReplicaSet] = None + service: Optional[client.V1Service] = None _namespace = call_api_method( api.read_namespace, expected=404, @@ -895,116 +898,165 @@ def _create_jump_pod_service_if_not_exists( name=namespace, ), ) - api.create_namespace(body=_namespace) + # 409 if created concurrently, e.g., by run_job() of another job. + call_api_method( + api.create_namespace, + expected=409, + body=_namespace, + ) else: - service = call_api_method( - api.read_namespaced_service, + pod = call_api_method( + api.read_namespaced_pod, expected=404, - name=jump_pod_service_name, + name=jump_pod_name, namespace=namespace, ) - pod = call_api_method( - api.read_namespaced_pod, + replica_set = call_api_method( + apps_api.read_namespaced_replica_set, expected=404, name=jump_pod_name, namespace=namespace, ) + service = call_api_method( + api.read_namespaced_service, + expected=404, + name=jump_pod_service_name, + namespace=namespace, + ) - # The service may exist without the pod if the node on which the jump pod was running - # has been deleted. - if service is not None and pod is not None: - return - - call_api_method( - api.delete_namespaced_pod, - expected=404, - namespace=namespace, - name=jump_pod_name, - ) - # 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) + should_create_replica_set = replica_set is None + service_selector = _build_service_selector_from_labels(base_labels) + if pod is not None: + # A legacy pod created before the ReplicaSet setup was introduced. + pod_phase = PodPhase(get_or_error(get_or_error(pod.status).phase)) + if pod_phase.is_finished(): + # Most probably evicted -- e.g., due to node pressure. Do cleanup. + call_api_method( + api.delete_namespaced_pod, + expected=404, + namespace=namespace, + name=jump_pod_name, ) - if not tolerations: - logger.warning("No appropriate node found, the jump pod may never be scheduled") - commands = _get_jump_pod_commands(authorized_keys) - pod = client.V1Pod( - metadata=client.V1ObjectMeta( - name=jump_pod_name, - labels=labels, - ), - spec=client.V1PodSpec( - containers=[ - client.V1Container( - name=f"{jump_pod_name}-container", - image=JUMP_POD_IMAGE, - command=["/bin/sh"], - args=["-c", " && ".join(commands)], - ports=[ - client.V1ContainerPort( - container_port=JUMP_POD_SSH_PORT, - ) - ], - ) - ], - restart_policy=RestartPolicy.ALWAYS, - tolerations=tolerations, - ), - ) - api.create_namespaced_pod( - namespace=namespace, - body=pod, - ) - call_api_method( - api.delete_namespaced_service, - expected=404, - namespace=namespace, - name=jump_pod_service_name, - ) - service = client.V1Service( - metadata=client.V1ObjectMeta( - name=jump_pod_service_name, - labels=labels, - ), - spec=client.V1ServiceSpec( - type="NodePort", - selector=_build_service_selector_from_labels(base_labels), - ports=[ - client.V1ServicePort( - port=JUMP_POD_SSH_PORT, - target_port=JUMP_POD_SSH_PORT, - node_port=jump_pod_port, + elif get_object_labels(pod).get("app.kubernetes.io/managed-by") != "dstack": + # A legacy pod created before https://github.com/dstackai/dstack/pull/3891 + # It uses different labels -- we don't create a ReplicaSet for it to avoid propagating + # the old label schema indefinitely but still recreate the service if it's gone using + # the old selector. + should_create_replica_set = False + service_selector = {"app.kubernetes.io/name": jump_pod_name} + # 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) ) - ], - ), - ) - api.create_namespaced_service( - namespace=namespace, - body=service, - ) + if not tolerations: + logger.warning("No appropriate node found, the jump pod may never be scheduled") + commands = _get_jump_pod_commands(authorized_keys) + replica_set = client.V1ReplicaSet( + metadata=client.V1ObjectMeta( + name=jump_pod_name, + labels=labels, + ), + spec=client.V1ReplicaSetSpec( + replicas=1, + selector=_build_label_selector_from_labels(base_labels), + template=client.V1PodTemplateSpec( + metadata=client.V1ObjectMeta( + labels=labels, + ), + spec=client.V1PodSpec( + containers=[ + client.V1Container( + name=f"{jump_pod_name}-container", + image=JUMP_POD_IMAGE, + command=["/bin/sh"], + args=["-c", " && ".join(commands)], + ports=[ + client.V1ContainerPort( + container_port=JUMP_POD_SSH_PORT, + ) + ], + ) + ], + restart_policy=RestartPolicy.ALWAYS, + tolerations=tolerations, + ), + ), + ), + ) + call_api_method( + apps_api.create_namespaced_replica_set, + expected=409, + namespace=namespace, + body=replica_set, + ) + + node_port = jump_pod_port + if service is not None: + service_spec = get_or_error(service.spec) + if service_spec.selector != service_selector: + # A service created before https://github.com/dstackai/dstack/pull/3891 selects pods + # by the old labels, which the ReplicaSet pods lack. Recreate it with the same NodePort + # so that ssh_proxy of running jobs stays valid. + if service_spec.ports: + node_port = service_spec.ports[0].node_port + call_api_method( + api.delete_namespaced_service, + expected=404, + namespace=namespace, + name=jump_pod_service_name, + ) + service = None + + if service is None: + service = client.V1Service( + metadata=client.V1ObjectMeta( + name=jump_pod_service_name, + labels=labels, + ), + spec=client.V1ServiceSpec( + type="NodePort", + selector=service_selector, + ports=[ + client.V1ServicePort( + port=JUMP_POD_SSH_PORT, + target_port=JUMP_POD_SSH_PORT, + node_port=node_port, + ) + ], + ), + ) + call_api_method( + api.create_namespaced_service, + expected=409, + namespace=namespace, + body=service, + ) def _check_and_configure_jump_pod_service( @@ -1016,15 +1068,32 @@ def _check_and_configure_jump_pod_service( project_ssh_private_key: str, authorized_keys: list[str], ) -> Optional[SSHConnectionParams]: - jump_pod = api.read_namespaced_pod( + # Matches ReplicaSet pods and a legacy pod named `jump_pod_name`, whether adopted by the + # ReplicaSet or not. + jump_pods = api.list_namespaced_pod( namespace=namespace, - name=jump_pod_name, - ) - jump_pod_phase = PodPhase(get_or_error(get_or_error(jump_pod.status).phase)) - if jump_pod_phase.is_finished(): - raise ProvisioningError(f"Jump pod {jump_pod_name} is unexpectedly finished") - if not jump_pod_phase.is_running(): - logger.debug("Jump pod %s is not running yet", jump_pod_name) + label_selector=_build_query_label_selector_from_unique_name(jump_pod_name), + ).items + if not jump_pods: + # A legacy pod created before https://github.com/dstackai/dstack/pull/3891 lacks the label. + legacy_jump_pod = call_api_method( + api.read_namespaced_pod, + expected=404, + namespace=namespace, + name=jump_pod_name, + ) + if legacy_jump_pod is not None: + # It's never adopted by the ReplicaSet and thus never replaced, no point in waiting. + jump_pod_phase = PodPhase(get_or_error(get_or_error(legacy_jump_pod.status).phase)) + if jump_pod_phase.is_finished(): + raise ProvisioningError(f"Jump pod {jump_pod_name} is unexpectedly finished") + jump_pods = [legacy_jump_pod] + # No pods at all is not an error -- the ReplicaSet may not have created a pod yet. + for jump_pod in jump_pods: + if _is_pod_ready(jump_pod): + break + else: + logger.debug("Jump pod %s is not ready yet", jump_pod_name) return None if jump_pod_hostname is None: @@ -1037,7 +1106,7 @@ def _check_and_configure_jump_pod_service( if node_address.type == "ExternalIP" ] if node_external_ips: - if get_node_name(node) == jump_pod_node_name: + if get_object_name(node) == jump_pod_node_name: jump_pod_hostname = node_external_ips[0] break cluster_external_ips.extend(node_external_ips) @@ -1386,6 +1455,16 @@ def _get_unscheduled_pod_reason_message( return None, None +def _is_pod_ready(pod: client.V1Pod) -> bool: + if pod.metadata is None: + return False + if pod.metadata.deletion_timestamp is not None: + return False + if pod.status is None: + return False + return any(c.type == "Ready" and c.status == "True" for c in (pod.status.conditions or [])) + + def _wait_for_load_balancer_address( api: client.CoreV1Api, namespace: str, @@ -1490,9 +1569,20 @@ def _run_ssh_command( return proc.returncode, proc.stdout +_SELECTOR_LABEL_KEY = "app.kubernetes.io/instance" + + def _build_service_selector_from_labels(labels: dict[str, str]) -> dict[str, str]: - label_key = "app.kubernetes.io/instance" - return {label_key: labels[label_key]} + return {_SELECTOR_LABEL_KEY: labels[_SELECTOR_LABEL_KEY]} + + +def _build_label_selector_from_labels(labels: dict[str, str]) -> client.V1LabelSelector: + return client.V1LabelSelector(match_labels={_SELECTOR_LABEL_KEY: labels[_SELECTOR_LABEL_KEY]}) + + +def _build_query_label_selector_from_unique_name(unique_name: str) -> str: + # See build_base_labels() + return f"{_SELECTOR_LABEL_KEY}={unique_name}" def _get_pod_service_name(pod_name: str) -> str: diff --git a/src/dstack/_internal/core/backends/kubernetes/resources.py b/src/dstack/_internal/core/backends/kubernetes/resources.py index b049eccc8c..37aa36ef6b 100644 --- a/src/dstack/_internal/core/backends/kubernetes/resources.py +++ b/src/dstack/_internal/core/backends/kubernetes/resources.py @@ -5,11 +5,11 @@ from collections.abc import Mapping from decimal import Decimal from enum import Enum -from typing import Callable, Literal, Optional, Union, cast, get_args +from typing import Callable, Literal, Optional, Protocol, Union, cast, get_args import gpuhunt from gpuhunt import KNOWN_AMD_GPUS, KNOWN_NVIDIA_GPUS, AcceleratorVendor -from kubernetes.client import CoreV1Api, V1Node, V1Taint +from kubernetes.client import CoreV1Api, V1Node, V1ObjectMeta, V1Taint from kubernetes.utils.quantity import parse_quantity as _parse_quantity from typing_extensions import Self @@ -349,14 +349,18 @@ def format_memory(memory: Memory) -> str: return f"{float(memory)}Gi" -def get_node_name(node: V1Node) -> Optional[str]: - if (metadata := node.metadata) is None: +class ObjectWithMeta(Protocol): + metadata: Optional[V1ObjectMeta] + + +def get_object_name(obj: ObjectWithMeta) -> Optional[str]: + if (metadata := obj.metadata) is None: return None return metadata.name -def get_node_labels(node: V1Node) -> dict[str, str]: - if (metadata := node.metadata) is None: +def get_object_labels(obj: ObjectWithMeta) -> dict[str, str]: + if (metadata := obj.metadata) is None: return {} if (labels := metadata.labels) is None: return {} @@ -386,7 +390,7 @@ def get_instance_offers( nodes_allocated_resources = _get_nodes_allocated_resources(api) offers: list[InstanceOfferWithAvailability] = [] for node in api.list_node().items: - if (node_name := get_node_name(node)) is None: + if (node_name := get_object_name(node)) is None: continue offer = _get_instance_offer_from_node( node=node, @@ -406,7 +410,7 @@ def get_instance_offers( def get_instance_offer_from_node( node: V1Node, region: str ) -> Optional[InstanceOfferWithAvailability]: - node_name = get_node_name(node) + node_name = get_object_name(node) if node_name is None: return None return _get_instance_offer_from_node( @@ -513,7 +517,7 @@ def _get_instance_offer_from_node( def _get_gpus_from_node( node: V1Node, node_name: str, node_resources: KubernetesResources ) -> list[Gpu]: - labels = get_node_labels(node) + labels = get_object_labels(node) for gpu_resource, gpu_getter in ( (NVIDIA_GPU_RESOURCE, get_nvidia_gpu_from_node_labels), (AMD_GPU_RESOURCE, get_amd_gpu_from_node_labels), diff --git a/src/dstack/_internal/core/backends/kubernetes/utils.py b/src/dstack/_internal/core/backends/kubernetes/utils.py index b6be7d0e7a..7907df7341 100644 --- a/src/dstack/_internal/core/backends/kubernetes/utils.py +++ b/src/dstack/_internal/core/backends/kubernetes/utils.py @@ -207,7 +207,7 @@ def call_api_method( If you don't expect any error status code, just call the method directly. Args: - method: the `CoreV1Api` bound method. + method: the Kubernetes API bound method, e.g., of `CoreV1Api`. expected: Expected error statuses, e.g., 404. args: positional arguments of the method. kwargs: keyword arguments of the method. diff --git a/src/tests/_internal/core/backends/kubernetes/test_compute.py b/src/tests/_internal/core/backends/kubernetes/test_compute.py index 5e7ba550d5..3d1a7a9064 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_compute.py +++ b/src/tests/_internal/core/backends/kubernetes/test_compute.py @@ -1,7 +1,20 @@ -from unittest.mock import MagicMock, patch +from collections.abc import Generator +from datetime import datetime, timezone +from typing import Optional +from unittest.mock import MagicMock, Mock, patch -from dstack._internal.core.backends.kubernetes.compute import KubernetesCompute +import pytest +from kubernetes import client +from kubernetes.client.exceptions import ApiException + +from dstack._internal.core.backends.kubernetes.compute import ( + KubernetesCompute, + _check_and_configure_jump_pod_service, + _create_jump_pod_service_if_not_exists, +) from dstack._internal.core.backends.kubernetes.models import KubeconfigConfig, KubernetesConfig +from dstack._internal.core.backends.kubernetes.resources import build_base_labels +from dstack._internal.core.errors import ProvisioningError from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.instances import ( Disk, @@ -10,6 +23,7 @@ InstanceOfferWithAvailability, InstanceType, Resources, + SSHConnectionParams, ) from dstack._internal.core.models.resources import ResourcesSpec from dstack._internal.core.models.runs import Requirements @@ -98,3 +112,335 @@ def test_get_offers_without_full_offers_adjusts_to_requested_slice(): assert adjusted_resources.memory_mib == 8 * 1024 assert len(adjusted_resources.gpus) == 1 assert adjusted_resources.disk.size_mib == 100 * 1024 + + +NAMESPACE = "dstack" +JUMP_POD_NAME = "dstack-main-ssh-jump-pod" +JUMP_POD_SERVICE_NAME = "dstack-main-ssh-jump-pod-service" +JUMP_POD_LABELS = build_base_labels( + component="ssh-proxy", + unique_name=JUMP_POD_NAME, + project="main", +) +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} + + +def _mock_read(method: Mock, obj: Optional[object]) -> None: + if obj is None: + method.side_effect = ApiException(status=404) + else: + method.return_value = obj + + +def _make_pod( + *, + name: str = JUMP_POD_NAME, + labels: Optional[dict[str, str]] = None, + phase: str = "Running", + ready: bool = True, + node_name: Optional[str] = "node-a", + terminating: bool = False, +) -> client.V1Pod: + return client.V1Pod( + metadata=client.V1ObjectMeta( + name=name, + labels=labels if labels is not None else JUMP_POD_LABELS, + deletion_timestamp=( + datetime(2026, 1, 1, tzinfo=timezone.utc) if terminating else None + ), + ), + spec=client.V1PodSpec(containers=[], node_name=node_name), + status=client.V1PodStatus( + phase=phase, + conditions=[ + client.V1PodCondition(type="Ready", status="True" if ready else "False"), + ], + ), + ) + + +def _make_service(selector: dict[str, str], node_port: int) -> client.V1Service: + return client.V1Service( + metadata=client.V1ObjectMeta(name=JUMP_POD_SERVICE_NAME), + spec=client.V1ServiceSpec( + type="NodePort", + selector=selector, + ports=[client.V1ServicePort(port=22, target_port=22, node_port=node_port)], + ), + ) + + +def _make_node(name: str, external_ip: Optional[str] = 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(), + status=client.V1NodeStatus(addresses=addresses), + ) + + +class TestCreateJumpPodServiceIfNotExists: + def make_apis( + self, + *, + namespace_exists: bool = True, + pod: Optional[client.V1Pod] = None, + replica_set: Optional[client.V1ReplicaSet] = None, + service: Optional[client.V1Service] = 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")]) + apps_api = Mock(spec=client.AppsV1Api) + _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: + _create_jump_pod_service_if_not_exists( + api=api, + apps_api=apps_api, + namespace=NAMESPACE, + project_name="main", + jump_pod_name=JUMP_POD_NAME, + jump_pod_service_name=JUMP_POD_SERVICE_NAME, + jump_pod_port=jump_pod_port, + authorized_keys=["ssh-ed25519 AAAA project"], + ) + + def assert_replica_set_created(self, apps_api: Mock) -> 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 + + def assert_service_created( + self, api: Mock, *, selector: dict[str, str], node_port: Optional[int] + ) -> None: + api.create_namespaced_service.assert_called_once() + service = api.create_namespaced_service.call_args.kwargs["body"] + assert service.metadata.name == JUMP_POD_SERVICE_NAME + assert service.spec.selector == selector + assert service.spec.ports[0].node_port == node_port + + def test_creates_namespace_replica_set_and_service(self): + api, apps_api = self.make_apis(namespace_exists=False) + + self.call(api, apps_api, jump_pod_port=30022) + + api.create_namespace.assert_called_once() + api.read_namespaced_pod.assert_not_called() + self.assert_replica_set_created(apps_api) + self.assert_service_created(api, selector=JUMP_POD_SELECTOR, node_port=30022) + + def test_creates_replica_set_and_service_if_not_found(self): + api, apps_api = self.make_apis() + + self.call(api, apps_api) + + api.create_namespace.assert_not_called() + self.assert_replica_set_created(apps_api) + self.assert_service_created(api, selector=JUMP_POD_SELECTOR, node_port=None) + + def test_does_nothing_if_replica_set_and_service_exist(self): + api, apps_api = self.make_apis( + replica_set=client.V1ReplicaSet( + spec=client.V1ReplicaSetSpec( + selector=client.V1LabelSelector(match_labels=JUMP_POD_SELECTOR), + ), + ), + service=_make_service(JUMP_POD_SELECTOR, node_port=30022), + ) + + self.call(api, apps_api) + + apps_api.create_namespaced_replica_set.assert_not_called() + api.delete_namespaced_service.assert_not_called() + api.create_namespaced_service.assert_not_called() + + def test_creates_replica_set_to_adopt_running_legacy_pod(self): + api, apps_api = self.make_apis( + pod=_make_pod(), + service=_make_service(JUMP_POD_SELECTOR, node_port=30022), + ) + + self.call(api, apps_api) + + api.delete_namespaced_pod.assert_not_called() + self.assert_replica_set_created(apps_api) + api.create_namespaced_service.assert_not_called() + + def test_deletes_finished_legacy_pod(self): + api, apps_api = self.make_apis( + pod=_make_pod(phase="Failed", ready=False), + service=_make_service(JUMP_POD_SELECTOR, node_port=30022), + ) + + self.call(api, apps_api) + + api.delete_namespaced_pod.assert_called_once_with(namespace=NAMESPACE, name=JUMP_POD_NAME) + self.assert_replica_set_created(apps_api) + api.create_namespaced_service.assert_not_called() + + def test_keeps_running_pre_3891_pod_and_its_service(self): + api, apps_api = self.make_apis( + pod=_make_pod(labels=PRE_3891_JUMP_POD_LABELS), + service=_make_service(PRE_3891_JUMP_POD_LABELS, node_port=30022), + ) + + self.call(api, apps_api) + + api.delete_namespaced_pod.assert_not_called() + apps_api.create_namespaced_replica_set.assert_not_called() + api.delete_namespaced_service.assert_not_called() + api.create_namespaced_service.assert_not_called() + + def test_creates_service_with_old_selector_for_running_pre_3891_pod(self): + api, apps_api = self.make_apis(pod=_make_pod(labels=PRE_3891_JUMP_POD_LABELS)) + + self.call(api, apps_api, jump_pod_port=30022) + + apps_api.create_namespaced_replica_set.assert_not_called() + self.assert_service_created(api, selector=PRE_3891_JUMP_POD_LABELS, node_port=30022) + + # The NodePort of the existing service wins even if proxy_jump.port is set + @pytest.mark.parametrize("jump_pod_port", [None, 30022]) + def test_recreates_pre_3891_service_with_same_node_port(self, jump_pod_port: Optional[int]): + api, apps_api = self.make_apis( + pod=_make_pod(labels=PRE_3891_JUMP_POD_LABELS, phase="Failed", ready=False), + service=_make_service(PRE_3891_JUMP_POD_LABELS, node_port=31234), + ) + + self.call(api, apps_api, jump_pod_port=jump_pod_port) + + api.delete_namespaced_pod.assert_called_once_with(namespace=NAMESPACE, name=JUMP_POD_NAME) + self.assert_replica_set_created(apps_api) + api.delete_namespaced_service.assert_called_once_with( + namespace=NAMESPACE, name=JUMP_POD_SERVICE_NAME + ) + self.assert_service_created(api, selector=JUMP_POD_SELECTOR, node_port=31234) + # The other way round, the create would fail with an ignored 409 and leave no service + method_names = [name for name, _, _ in api.method_calls] + assert method_names.index("delete_namespaced_service") < method_names.index( + "create_namespaced_service" + ) + + def test_ignores_conflicts_on_create(self): + api, apps_api = self.make_apis(namespace_exists=False) + api.create_namespace.side_effect = ApiException(status=409) + apps_api.create_namespaced_replica_set.side_effect = ApiException(status=409) + api.create_namespaced_service.side_effect = ApiException(status=409) + + self.call(api, apps_api) + + apps_api.create_namespaced_replica_set.assert_called_once() + api.create_namespaced_service.assert_called_once() + + +class TestCheckAndConfigureJumpPodService: + @pytest.fixture(autouse=True) + def run_ssh_command_mock(self) -> Generator[Mock, None, None]: + with patch( + "dstack._internal.core.backends.kubernetes.compute._run_ssh_command", + return_value=(0, b""), + ) as run_ssh_command_mock: + yield run_ssh_command_mock + + def make_api( + self, *, pods: list[client.V1Pod], pre_3891_pod: Optional[client.V1Pod] = None + ) -> Mock: + api = Mock(spec=client.CoreV1Api) + api.list_namespaced_pod.return_value = client.V1PodList(items=pods) + _mock_read(api.read_namespaced_pod, pre_3891_pod) + api.list_node.return_value = client.V1NodeList( + items=[ + _make_node("node-a", external_ip="203.0.113.1"), + _make_node("node-b", external_ip="203.0.113.2"), + ] + ) + api.read_namespaced_service.return_value = _make_service( + JUMP_POD_SELECTOR, node_port=30022 + ) + return api + + def call(self, api: Mock) -> Optional[SSHConnectionParams]: + return _check_and_configure_jump_pod_service( + api=api, + namespace=NAMESPACE, + jump_pod_name=JUMP_POD_NAME, + jump_pod_service_name=JUMP_POD_SERVICE_NAME, + jump_pod_hostname=None, + project_ssh_private_key="private-key", + authorized_keys=[], + ) + + def test_returns_ssh_proxy_on_node_of_ready_pod(self): + api = self.make_api( + pods=[ + _make_pod(name="evicted", phase="Failed", ready=False, node_name="node-a"), + _make_pod(name="ready", node_name="node-b"), + ] + ) + + ssh_proxy = self.call(api) + + assert ssh_proxy == SSHConnectionParams( + hostname="203.0.113.2", port=30022, username="root" + ) + api.list_namespaced_pod.assert_called_once_with( + namespace=NAMESPACE, label_selector=f"app.kubernetes.io/instance={JUMP_POD_NAME}" + ) + api.read_namespaced_pod.assert_not_called() + + @pytest.mark.parametrize( + "pod", + [ + pytest.param(_make_pod(phase="Pending", ready=False, node_name=None), id="pending"), + pytest.param(_make_pod(phase="Failed", ready=False), id="evicted"), + 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): + 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() + + def test_waits_if_no_pods(self, run_ssh_command_mock: Mock): + # E.g., the ReplicaSet has not created a pod yet. + api = self.make_api(pods=[]) + + assert self.call(api) is None + api.read_namespaced_pod.assert_called_once_with(namespace=NAMESPACE, name=JUMP_POD_NAME) + run_ssh_command_mock.assert_not_called() + + def test_returns_ssh_proxy_if_pre_3891_pod_is_ready(self): + api = self.make_api( + pods=[], + pre_3891_pod=_make_pod(labels=PRE_3891_JUMP_POD_LABELS, node_name="node-a"), + ) + + ssh_proxy = self.call(api) + + assert ssh_proxy == SSHConnectionParams( + hostname="203.0.113.1", port=30022, username="root" + ) + + def test_raises_if_pre_3891_pod_is_finished(self, run_ssh_command_mock: Mock): + api = self.make_api( + pods=[], + pre_3891_pod=_make_pod(labels=PRE_3891_JUMP_POD_LABELS, phase="Failed", ready=False), + ) + + with pytest.raises(ProvisioningError, match="unexpectedly finished"): + self.call(api) + run_ssh_command_mock.assert_not_called()