Add job_postprocessor hook to Kubernetes Job launchers

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

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

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