feat(kubernetes): retry disrupted Job pods

README.md

@@ -168,6 +168,17 @@ Currently the following launchers are supported:
 * Local Kubernetes using local storage via HostPath volumes
 * Google Cloud Kubernetes Engine using Google Cloud Storage
 
+Kubernetes Job tasks can opt into bounded replacement of infrastructure-disrupted pods:
+
+```yaml
+annotations:
+  tangleml.com/launchers/kubernetes/job/disruption_retries: "1"
+```
+
+The launcher retries pods marked by Kubernetes with `DisruptionTarget=True`; a non-zero exit from the
+main container still fails immediately. The default is `0` (no retries). Only enable this for
+restart-safe components because a disrupted attempt may have produced partial external side effects.
+
 More launchers may be added in the future.
 
 ### Credits

cloud_pipelines_backend/launchers/kubernetes_launchers.py

@@ -57,6 +57,10 @@
 SECURITY_CONTEXT_CAPABILITY_IPC_LOCK_ANNOTATION_KEY = (
     "tangleml.com/launchers/kubernetes/security_context.capability.IPC_LOCK"
 )
+JOB_DISRUPTION_RETRIES_ANNOTATION_KEY = (
+    "tangleml.com/launchers/kubernetes/job/disruption_retries"
+)
+_JOB_MAX_DISRUPTION_RETRIES = 10
 
 
 # Multi-node constants
@@ -996,6 +1000,70 @@ def cleanup(self):
         self._delete_pod()
 
 
+def _configure_job_disruption_retries(
+    job_spec: k8s_client_lib.V1JobSpec,
+    *,
+    annotations: dict[str, Any] | None,
+) -> None:
+    """Configure bounded retries for Kubernetes-marked pod disruptions.
+
+    The policy is opt-in and leaves existing Jobs at zero retries. With retries
+    enabled, ``DisruptionTarget`` pods consume the finite retry budget and are
+    replaced, while a non-zero exit from the main container fails its index
+    immediately. The numeric backoff is also a finite guard for unclassified
+    infrastructure failures that do not expose a main-container exit code.
+    """
+    disruption_retries_value = (annotations or {}).get(
+        JOB_DISRUPTION_RETRIES_ANNOTATION_KEY, "0"
+    )
+    try:
+        disruption_retries = int(disruption_retries_value)
+    except (TypeError, ValueError) as ex:
+        raise interfaces.LauncherError(
+            f"Invalid {JOB_DISRUPTION_RETRIES_ANNOTATION_KEY}={disruption_retries_value!r}; "
+            f"expected an integer between 0 and {_JOB_MAX_DISRUPTION_RETRIES}."
+        ) from ex
+    if str(disruption_retries) != str(disruption_retries_value).strip():
+        raise interfaces.LauncherError(
+            f"Invalid {JOB_DISRUPTION_RETRIES_ANNOTATION_KEY}={disruption_retries_value!r}; "
+            f"expected an integer between 0 and {_JOB_MAX_DISRUPTION_RETRIES}."
+        )
+    if not 0 <= disruption_retries <= _JOB_MAX_DISRUPTION_RETRIES:
+        raise interfaces.LauncherError(
+            f"Invalid {JOB_DISRUPTION_RETRIES_ANNOTATION_KEY}={disruption_retries!r}; "
+            f"expected an integer between 0 and {_JOB_MAX_DISRUPTION_RETRIES}."
+        )
+
+    job_spec.backoff_limit_per_index = disruption_retries
+    if not disruption_retries:
+        return
+
+    job_spec.pod_failure_policy = k8s_client_lib.V1PodFailurePolicy(
+        rules=[
+            k8s_client_lib.V1PodFailurePolicyRule(
+                action="Count",
+                on_pod_conditions=[
+                    k8s_client_lib.V1PodFailurePolicyOnPodConditionsPattern(
+                        type="DisruptionTarget",
+                        status="True",
+                    )
+                ],
+            ),
+            k8s_client_lib.V1PodFailurePolicyRule(
+                action="FailIndex",
+                on_exit_codes=k8s_client_lib.V1PodFailurePolicyOnExitCodesRequirement(
+                    container_name=_MAIN_CONTAINER_NAME,
+                    operator="NotIn",
+                    values=[0],
+                ),
+            ),
+        ]
+    )
+    # Required when podFailurePolicy is configured. Waiting for a fully failed
+    # pod also makes the old/new attempt boundary unambiguous.
+    job_spec.pod_replacement_policy = "Failed"
+
+
 class _KubernetesJobLauncher(
     _KubernetesContainerLauncherBase,
     interfaces.ContainerTaskLauncher["LaunchedKubernetesJob"],
@@ -1198,27 +1266,28 @@ def launch_container_task(
             # This requires the service name to be known.
             pod.spec.subdomain = explicit_service_name
 
+        job_spec = k8s_client_lib.V1JobSpec(
+            template=k8s_client_lib.V1PodTemplateSpec(
+                metadata=pod.metadata,
+                spec=pod.spec,
+            ),
+            # Let's always use Indexed Jobs. There are no downsides.
+            completion_mode="Indexed",
+            # Without explicit max_failed_indexes=0, the job waits for all pods to end and then succeeds ("Complete") despite pod failures!
+            max_failed_indexes=0,
+            completions=num_nodes,
+            parallelism=num_nodes,
+        )
+        _configure_job_disruption_retries(job_spec, annotations=annotations)
+
         job = k8s_client_lib.V1Job(
             metadata=k8s_client_lib.V1ObjectMeta(
                 name=explicit_job_name,
                 namespace=namespace,
                 # annotations=self._pod_annotations,
                 # labels=self._pod_labels,
             ),
-            spec=k8s_client_lib.V1JobSpec(
-                template=k8s_client_lib.V1PodTemplateSpec(
-                    metadata=pod.metadata,
-                    spec=pod.spec,
-                ),
-                # Let's always use Indexed Jobs. There are no downsides.
-                completion_mode="Indexed",
-                # backoff_limit=0,
-                backoff_limit_per_index=0,
-                # Without explicit max_failed_indexes=0, the job waits for all pods to end and then succeeds ("Complete") despite pod failures!
-                max_failed_indexes=0,
-                completions=num_nodes,
-                parallelism=num_nodes,
-            ),
+            spec=job_spec,
         )
 
         job = self._transform_job_before_launching(job=job, annotations=annotations)

tests/test_kubernetes_job_disruption_retries.py

@@ -0,0 +1,91 @@
+"""Tests for opt-in retries of infrastructure-disrupted Kubernetes Job pods."""
+
+import pytest
+from kubernetes import client as k8s_client_lib
+
+from cloud_pipelines_backend.launchers import interfaces
+from cloud_pipelines_backend.launchers import kubernetes_launchers
+
+
+def _job_spec() -> k8s_client_lib.V1JobSpec:
+    return k8s_client_lib.V1JobSpec(
+        completion_mode="Indexed",
+        completions=1,
+        max_failed_indexes=0,
+        parallelism=1,
+        template=k8s_client_lib.V1PodTemplateSpec(
+            spec=k8s_client_lib.V1PodSpec(
+                containers=[k8s_client_lib.V1Container(name="main", image="python")],
+                restart_policy="Never",
+            )
+        ),
+    )
+
+
+def test_default_keeps_zero_retry_job_contract() -> None:
+    spec = _job_spec()
+
+    kubernetes_launchers._configure_job_disruption_retries(spec, annotations=None)
+
+    assert spec.backoff_limit_per_index == 0
+    assert spec.pod_failure_policy is None
+    assert spec.pod_replacement_policy is None
+
+
+def test_opt_in_counts_disruptions_but_fails_user_code_immediately() -> None:
+    spec = _job_spec()
+
+    kubernetes_launchers._configure_job_disruption_retries(
+        spec,
+        annotations={kubernetes_launchers.JOB_DISRUPTION_RETRIES_ANNOTATION_KEY: "1"},
+    )
+
+    assert spec.backoff_limit_per_index == 1
+    assert spec.pod_replacement_policy == "Failed"
+    assert spec.pod_failure_policy is not None
+    disruption_rule, user_code_rule = spec.pod_failure_policy.rules
+
+    assert disruption_rule.action == "Count"
+    assert disruption_rule.on_exit_codes is None
+    assert len(disruption_rule.on_pod_conditions) == 1
+    disruption_condition = disruption_rule.on_pod_conditions[0]
+    assert disruption_condition.type == "DisruptionTarget"
+    assert disruption_condition.status == "True"
+
+    assert user_code_rule.action == "FailIndex"
+    assert user_code_rule.on_pod_conditions is None
+    assert user_code_rule.on_exit_codes.container_name == "main"
+    assert user_code_rule.on_exit_codes.operator == "NotIn"
+    assert user_code_rule.on_exit_codes.values == [0]
+
+    serialized = kubernetes_launchers._kubernetes_serialize(spec)
+    assert serialized["backoffLimitPerIndex"] == 1
+    assert serialized["maxFailedIndexes"] == 0
+    assert serialized["podReplacementPolicy"] == "Failed"
+    assert serialized["podFailurePolicy"]["rules"] == [
+        {
+            "action": "Count",
+            "onPodConditions": [{"status": "True", "type": "DisruptionTarget"}],
+        },
+        {
+            "action": "FailIndex",
+            "onExitCodes": {
+                "containerName": "main",
+                "operator": "NotIn",
+                "values": [0],
+            },
+        },
+    ]
+
+
+@pytest.mark.parametrize("value", ["-1", "11", "1.0", "true", ""])
+def test_invalid_disruption_retry_count_fails_closed(value: str) -> None:
+    with pytest.raises(
+        interfaces.LauncherError, match="expected an integer between 0 and 10"
+    ):
+        kubernetes_launchers._configure_job_disruption_retries(
+            _job_spec(),
+            annotations={
+                kubernetes_launchers.JOB_DISRUPTION_RETRIES_ANNOTATION_KEY: value
+            },
+        )