diff --git a/cloud_pipelines_backend/launchers/google_kubernetes_launchers.py b/cloud_pipelines_backend/launchers/google_kubernetes_launchers.py index 6bbb551..8bbc454 100644 --- a/cloud_pipelines_backend/launchers/google_kubernetes_launchers.py +++ b/cloud_pipelines_backend/launchers/google_kubernetes_launchers.py @@ -26,6 +26,7 @@ def __init__( pod_labels: dict[str, str] | None = None, pod_annotations: dict[str, str] | None = None, pod_postprocessor: kubernetes_launchers.PodPostProcessor | None = None, + job_postprocessor: kubernetes_launchers.JobPostProcessor | None = None, ): pod_postprocessors = [ kubernetes_launchers._google_kubernetes_engine_accelerator_pod_postprocessor @@ -44,6 +45,7 @@ def __init__( pod_labels=pod_labels, pod_annotations={"gke-gcsfuse/volumes": "true"} | (pod_annotations or {}), pod_postprocessor=final_pod_postporocessor, + job_postprocessor=job_postprocessor, _storage_provider=google_cloud_storage.GoogleCloudStorageProvider( gcs_client ), @@ -67,6 +69,7 @@ def __init__( pod_labels: dict[str, str] | None = None, pod_annotations: dict[str, str] | None = None, pod_postprocessor: kubernetes_launchers.PodPostProcessor | None = None, + job_postprocessor: kubernetes_launchers.JobPostProcessor | None = None, ): pod_postprocessors = [ kubernetes_launchers._google_kubernetes_engine_accelerator_pod_postprocessor @@ -85,6 +88,7 @@ def __init__( pod_labels=pod_labels, pod_annotations={"gke-gcsfuse/volumes": "true"} | (pod_annotations or {}), pod_postprocessor=final_pod_postporocessor, + job_postprocessor=job_postprocessor, _storage_provider=google_cloud_storage.GoogleCloudStorageProvider( gcs_client ), diff --git a/cloud_pipelines_backend/launchers/kubernetes_launchers.py b/cloud_pipelines_backend/launchers/kubernetes_launchers.py index 25cd267..edbb266 100644 --- a/cloud_pipelines_backend/launchers/kubernetes_launchers.py +++ b/cloud_pipelines_backend/launchers/kubernetes_launchers.py @@ -152,6 +152,12 @@ def __call__( ) -> k8s_client_lib.V1Pod: ... +class JobPostProcessor(typing.Protocol): + def __call__( + self, *, job: k8s_client_lib.V1Job, annotations: dict[str, str] | None = None + ) -> k8s_client_lib.V1Job: ... + + class _KubernetesContainerLauncherBase: """Launcher that launches container using Kubernetes""" @@ -999,6 +1005,7 @@ def __init__( pod_labels: dict[str, str] | None = None, pod_annotations: dict[str, str] | None = None, pod_postprocessor: PodPostProcessor | None = None, + job_postprocessor: JobPostProcessor | None = None, _storage_provider: storage_provider_interfaces.StorageProvider, _create_volume_and_volume_mount: typing.Callable[ [str, str, str, bool], @@ -1017,6 +1024,7 @@ def __init__( pod_postprocessor=pod_postprocessor, _create_volume_and_volume_mount=_create_volume_and_volume_mount, ) + self._job_postprocessor = job_postprocessor def launch_container_task( self, @@ -1242,7 +1250,8 @@ def launch_container_task( def _transform_job_before_launching( self, *, job: k8s_client_lib.V1Job, annotations: dict[str, str] | None = None ) -> k8s_client_lib.V1Job: - del annotations + if self._job_postprocessor: + job = self._job_postprocessor(job=job, annotations=annotations) return job def get_refreshed_launched_container_from_dict( @@ -1634,6 +1643,7 @@ def __init__( pod_labels: dict[str, str] | None = None, pod_annotations: dict[str, str] | None = None, pod_postprocessor: PodPostProcessor | None = None, + job_postprocessor: JobPostProcessor | None = None, _storage_provider: storage_provider_interfaces.StorageProvider, _create_volume_and_volume_mount: typing.Callable[ [str, str, str, bool], @@ -1660,6 +1670,7 @@ def __init__( pod_labels=pod_labels, pod_annotations=pod_annotations, pod_postprocessor=pod_postprocessor, + job_postprocessor=job_postprocessor, _storage_provider=_storage_provider, _create_volume_and_volume_mount=_create_volume_and_volume_mount, ) diff --git a/tests/test_kubernetes_launchers.py b/tests/test_kubernetes_launchers.py new file mode 100644 index 0000000..9e17d03 --- /dev/null +++ b/tests/test_kubernetes_launchers.py @@ -0,0 +1,83 @@ +"""Tests for cloud_pipelines_backend.launchers.kubernetes_launchers. + +These stay offline: they exercise the job_postprocessor plumbing on the base +launcher classes directly, without a real Kubernetes API client or the GKE +launchers (which need google.cloud.storage). +""" + +from __future__ import annotations + +from kubernetes import client as k8s_client_lib + +from cloud_pipelines_backend.launchers import kubernetes_launchers + + +def _job() -> k8s_client_lib.V1Job: + return k8s_client_lib.V1Job( + spec=k8s_client_lib.V1JobSpec(template=k8s_client_lib.V1PodTemplateSpec()) + ) + + +def _job_launcher_with_postprocessor( + job_postprocessor: kubernetes_launchers.JobPostProcessor | None, +) -> kubernetes_launchers._KubernetesJobLauncher: + # Bypass __init__ (needs an api_client / storage provider); we only exercise + # the transform hook, which reads self._job_postprocessor. + launcher = object.__new__(kubernetes_launchers._KubernetesJobLauncher) + launcher._job_postprocessor = job_postprocessor + return launcher + + +def test_transform_job_before_launching_applies_job_postprocessor(): + def set_ttl(*, job: k8s_client_lib.V1Job, annotations=None) -> k8s_client_lib.V1Job: + job.spec.ttl_seconds_after_finished = 604800 + return job + + launcher = _job_launcher_with_postprocessor(set_ttl) + + result = launcher._transform_job_before_launching(job=_job(), annotations={}) + + assert result.spec.ttl_seconds_after_finished == 604800 + + +def test_transform_job_before_launching_is_noop_without_postprocessor(): + launcher = _job_launcher_with_postprocessor(None) + job = _job() + + result = launcher._transform_job_before_launching(job=job, annotations={}) + + assert result is job + assert result.spec.ttl_seconds_after_finished is None + + +def test_pod_or_job_launcher_forwards_job_postprocessor(monkeypatch): + captured: dict = {} + + class _StubPodLauncher: + def __init__(self, **kwargs): + pass + + class _StubJobLauncher: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr( + kubernetes_launchers, "_KubernetesPodLauncher", _StubPodLauncher + ) + monkeypatch.setattr( + kubernetes_launchers, "_KubernetesJobLauncher", _StubJobLauncher + ) + + def job_postprocessor( + *, job: k8s_client_lib.V1Job, annotations=None + ) -> k8s_client_lib.V1Job: + return job + + kubernetes_launchers._KubernetesPodOrJobLauncher( + api_client=None, + job_postprocessor=job_postprocessor, + _storage_provider=None, + _create_volume_and_volume_mount=None, + ) + + assert captured["job_postprocessor"] is job_postprocessor