diff --git a/README.md b/README.md index 1101f17..3e2112c 100644 --- a/README.md +++ b/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 diff --git a/cloud_pipelines_backend/launchers/kubernetes_launchers.py b/cloud_pipelines_backend/launchers/kubernetes_launchers.py index 9b1a661..47e4fcd 100644 --- a/cloud_pipelines_backend/launchers/kubernetes_launchers.py +++ b/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,6 +1266,20 @@ 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, @@ -1205,20 +1287,7 @@ def launch_container_task( # 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) diff --git a/tests/test_kubernetes_job_disruption_retries.py b/tests/test_kubernetes_job_disruption_retries.py new file mode 100644 index 0000000..68aff3c --- /dev/null +++ b/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 + }, + )