From e25dccc6b63ca3dacf01596ee0fd757d5854a633 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Fri, 30 Oct 2026 14:59:47 +0000 Subject: [PATCH] Kubernetes: manage the jump pod with a ReplicaSet The jump pod was a bare pod, reused for as long as it existed. Once evicted, for example due to node pressure, it stayed in the Failed phase and was never recreated, so every job in the project failed until someone deleted the pod by hand. The jump pod is now managed by a ReplicaSet, which replaces the pod after an eviction or a node removal. Not a Deployment: its ReplicaSet selects pods by the pod-template-hash label and thus cannot adopt the existing pod, and rollouts and rollbacks are not needed anyway. * run_job() creates the ReplicaSet if it doesn't exist. The ReplicaSet adopts a running legacy pod, a finished one is deleted. A pod created before #3891 has different labels and is left as is until it's gone. * A service created before #3891 selects pods by the old labels. Once the old pod is gone, the service is recreated with the new selector and the same NodePort, so that ssh_proxy of running jobs stays valid. * update_provisioning_data() lists jump pods by label and picks a ready one. It falls back to the name for a pod created before #3891 and fails right away if that pod is finished, as it's never replaced. * Creating the namespace, the ReplicaSet, and the service ignores 409 to tolerate concurrent run_job() calls. Notes: * The Role now needs get/create on replicasets and list on pods. Without the update, Kubernetes provisioning fails with 403. * Once the legacy pod is gone, servers before this commit recreate a bare pod and the service on each run_job(). The ReplicaSet adopts and deletes the pod, and the NodePort changes if proxy_jump.port is not set. Downgrading requires deleting the ReplicaSet. Part-of: https://github.com/dstackai/dstack/issues/4407 --- .../kubernetes/dstack-backend-role.yaml | 5 +- .../core/backends/kubernetes/compute.py | 320 ++++++++++------ .../core/backends/kubernetes/resources.py | 22 +- .../core/backends/kubernetes/utils.py | 2 +- .../core/backends/kubernetes/test_compute.py | 350 +++++++++++++++++- 5 files changed, 571 insertions(+), 128 deletions(-) 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()