diff --git a/cloud_pipelines_backend/launchers/kubernetes_launchers.py b/cloud_pipelines_backend/launchers/kubernetes_launchers.py index 32f3e8a..33d95dc 100644 --- a/cloud_pipelines_backend/launchers/kubernetes_launchers.py +++ b/cloud_pipelines_backend/launchers/kubernetes_launchers.py @@ -1157,6 +1157,7 @@ def launch_container_task( # we should prohibit/ignore changing pod namespace in the pod post-processor. namespace = pod.metadata.namespace + service: k8s_client_lib.V1Service | None = None if enable_multi_node: main_container_spec = pod.spec.containers[0] main_container_spec.env = main_container_spec.env or [] @@ -1193,17 +1194,6 @@ def launch_container_task( }, ), ) - core_api_client = k8s_client_lib.CoreV1Api(api_client=self._api_client) - try: - _: k8s_client_lib.V1Service = core_api_client.create_namespaced_service( - namespace=namespace, - body=service, - _request_timeout=self._request_timeout, - ) - except Exception as ex: - raise interfaces.LauncherError( - f"Failed to create Kubernetes Service {explicit_service_name}: {_kubernetes_serialize(service)}" - ) from ex # Setting Pod's spec.subdomain to exact name of teh service. # This requires the service name to be known. pod.spec.subdomain = explicit_service_name @@ -1268,8 +1258,48 @@ def launch_container_task( job_name: str = created_job.metadata.name job_namespace: str = created_job.metadata.namespace + job_uid: str = created_job.metadata.uid _logger.info(f"Created Kubernetes Job {job_name} in namespace {job_namespace}") + if enable_multi_node: + assert service + # The Job's server-assigned UID is required to make the Service a + # garbage-collected dependent of the Job. + service.metadata.owner_references = [ + k8s_client_lib.V1OwnerReference( + api_version="batch/v1", + kind="Job", + name=job_name, + uid=job_uid, + ) + ] + core_api_client = k8s_client_lib.CoreV1Api(api_client=self._api_client) + try: + _: k8s_client_lib.V1Service = core_api_client.create_namespaced_service( + namespace=namespace, + body=service, + _request_timeout=self._request_timeout, + ) + except Exception as ex: + # Do not leave a runnable multi-node Job without the Service its + # Pod DNS names depend on. If the Service request succeeded but + # its response was lost, its owner reference also removes it. + try: + batch_api_client.delete_namespaced_job( + name=job_name, + namespace=job_namespace, + grace_period_seconds=0, + propagation_policy="Background", + _request_timeout=self._request_timeout, + ) + except Exception: + _logger.exception( + f"Failed to roll back Kubernetes Job {job_name} after Service creation failed." + ) + raise interfaces.LauncherError( + f"Failed to create Kubernetes Service {explicit_service_name}: {_kubernetes_serialize(service)}" + ) from ex + launched_container = LaunchedKubernetesJob( job_name=job_name, namespace=job_namespace, @@ -1687,12 +1717,34 @@ def __str__(self) -> str: def _delete_job(self): launcher = self._get_launcher() batch_api_client = k8s_client_lib.BatchV1Api(api_client=launcher._api_client) - batch_api_client.delete_namespaced_job( - name=self._job_name, - namespace=self._namespace, - grace_period_seconds=10, - propagation_policy="Foreground", - ) + try: + batch_api_client.delete_namespaced_job( + name=self._job_name, + namespace=self._namespace, + grace_period_seconds=10, + propagation_policy="Foreground", + _request_timeout=launcher._request_timeout, + ) + except kubernetes.client.exceptions.ApiException as ex: + if ex.status != http.HTTPStatus.NOT_FOUND: + raise + + # Jobs created before Services gained owner references still need + # explicit cleanup. New Services normally disappear through Kubernetes + # garbage collection; deleting by the deterministic Job name keeps this + # operation idempotent in both cases. + core_api_client = k8s_client_lib.CoreV1Api(api_client=launcher._api_client) + try: + core_api_client.delete_namespaced_service( + name=self._job_name, + namespace=self._namespace, + grace_period_seconds=0, + propagation_policy="Background", + _request_timeout=launcher._request_timeout, + ) + except kubernetes.client.exceptions.ApiException as ex: + if ex.status != http.HTTPStatus.NOT_FOUND: + raise def terminate(self): self._delete_job() diff --git a/tests/launchers/test_kubernetes_job_service_lifecycle.py b/tests/launchers/test_kubernetes_job_service_lifecycle.py new file mode 100644 index 0000000..cd91ab4 --- /dev/null +++ b/tests/launchers/test_kubernetes_job_service_lifecycle.py @@ -0,0 +1,188 @@ +from types import SimpleNamespace + +import pytest +from kubernetes import client as k8s_client_lib + +from cloud_pipelines_backend.launchers import common_annotations +from cloud_pipelines_backend.launchers import interfaces +from cloud_pipelines_backend.launchers import kubernetes_launchers + + +def _launcher(monkeypatch, *, service_error=None): + events = [] + batch = SimpleNamespace(created=[], deleted=[]) + core = SimpleNamespace(created=[], deleted=[]) + + def create_job(namespace, body, _request_timeout): + events.append("job") + body.metadata.namespace = namespace + body.metadata.uid = "job-uid" + batch.created.append((namespace, body, _request_timeout)) + return body + + def delete_job(**kwargs): + batch.deleted.append(kwargs) + + def create_service(namespace, body, _request_timeout): + events.append("service") + core.created.append((namespace, body, _request_timeout)) + if service_error: + raise service_error + return body + + def delete_service(**kwargs): + core.deleted.append(kwargs) + + batch.create_namespaced_job = create_job + batch.delete_namespaced_job = delete_job + core.create_namespaced_service = create_service + core.delete_namespaced_service = delete_service + + monkeypatch.setattr( + kubernetes_launchers.k8s_client_lib, + "BatchV1Api", + lambda api_client: batch, + ) + monkeypatch.setattr( + kubernetes_launchers.k8s_client_lib, + "CoreV1Api", + lambda api_client: core, + ) + + launcher = object.__new__(kubernetes_launchers._KubernetesJobLauncher) + launcher._api_client = SimpleNamespace( + configuration=SimpleNamespace(host="https://cluster.example") + ) + launcher._request_timeout = (3, 30) + launcher._namespace = "tangle-jobs" + launcher._service_account_name = "tangle" + launcher._pod_name_prefix = "task-" + launcher._pod_labels = {} + launcher._pod_annotations = {} + launcher._pod_postprocessor = None + launcher._choose_namespace = lambda annotations: "tangle-jobs" + launcher._prepare_kubernetes_pod = lambda **kwargs: k8s_client_lib.V1Pod( + metadata=k8s_client_lib.V1ObjectMeta(namespace="tangle-jobs"), + spec=k8s_client_lib.V1PodSpec( + containers=[k8s_client_lib.V1Container(name="main")] + ), + ) + launcher._transform_job_before_launching = lambda job, annotations: job + + return launcher, batch, core, events + + +def _launch(launcher): + return launcher.launch_container_task( + component_spec=object(), + input_arguments={}, + output_uris={}, + log_uri="memory://log", + annotations={ + common_annotations.CONTAINER_EXECUTION_ID_ANNOTATION_KEY: "execution-1", + kubernetes_launchers.MULTI_NODE_NUMBER_OF_NODES_ANNOTATION_KEY: "2", + }, + ) + + +def test_job_service_is_owned_by_the_job(monkeypatch): + launcher, _, core, events = _launcher(monkeypatch) + + launched = _launch(launcher) + + assert events == ["job", "service"] + namespace, service, timeout = core.created[0] + assert namespace == "tangle-jobs" + assert timeout == (3, 30) + assert service.metadata.name == "tangle-ce-execution-1" + assert service.spec.cluster_ip == "None" + assert service.spec.selector == {"job-name": "tangle-ce-execution-1"} + assert service.metadata.owner_references == [ + k8s_client_lib.V1OwnerReference( + api_version="batch/v1", + kind="Job", + name="tangle-ce-execution-1", + uid="job-uid", + ) + ] + assert launched._debug_job.spec.template.spec.subdomain == ("tangle-ce-execution-1") + + +def test_service_creation_failure_rolls_back_the_job(monkeypatch): + launcher, batch, _, _ = _launcher( + monkeypatch, service_error=RuntimeError("service create failed") + ) + + with pytest.raises( + interfaces.LauncherError, match="Failed to create Kubernetes Service" + ): + _launch(launcher) + + assert batch.deleted == [ + { + "name": "tangle-ce-execution-1", + "namespace": "tangle-jobs", + "grace_period_seconds": 0, + "propagation_policy": "Background", + "_request_timeout": (3, 30), + } + ] + + +def test_cleanup_removes_legacy_service_when_job_is_already_gone(monkeypatch): + batch = SimpleNamespace() + core = SimpleNamespace(deleted=[]) + + def delete_job(**kwargs): + raise k8s_client_lib.ApiException(status=404) + + def delete_service(**kwargs): + core.deleted.append(kwargs) + raise k8s_client_lib.ApiException(status=404) + + batch.delete_namespaced_job = delete_job + core.delete_namespaced_service = delete_service + monkeypatch.setattr( + kubernetes_launchers.k8s_client_lib, + "BatchV1Api", + lambda api_client: batch, + ) + monkeypatch.setattr( + kubernetes_launchers.k8s_client_lib, + "CoreV1Api", + lambda api_client: core, + ) + + launcher = SimpleNamespace( + _api_client=object(), + _request_timeout=(3, 30), + ) + launched = kubernetes_launchers.LaunchedKubernetesJob( + job_name="tangle-ce-execution-1", + namespace="tangle-jobs", + output_uris={}, + log_uri="memory://log", + debug_job=k8s_client_lib.V1Job( + metadata=k8s_client_lib.V1ObjectMeta(name="tangle-ce-execution-1"), + spec=k8s_client_lib.V1JobSpec( + template=k8s_client_lib.V1PodTemplateSpec( + spec=k8s_client_lib.V1PodSpec( + containers=[k8s_client_lib.V1Container(name="main")] + ) + ) + ), + ), + launcher=launcher, + ) + + launched.cleanup() + + assert core.deleted == [ + { + "name": "tangle-ce-execution-1", + "namespace": "tangle-jobs", + "grace_period_seconds": 0, + "propagation_policy": "Background", + "_request_timeout": (3, 30), + } + ]