Add --min-completed-minutes to cleanup-pods to prevent KPO race condition

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/cli/definition.py

@@ -68,6 +68,16 @@
     ),
 )
 
+ARG_MIN_COMPLETED_MINUTES = Arg(
+    ("--min-completed-minutes",),
+    default=1,
+    type=positive_int(allow_zero=True),
+    help=(
+        "Minimum age in minutes of a completed (Succeeded/Failed/Evicted) pod before it is deleted. "
+        "Default is 1. Set to 0 to delete immediately."
+    ),
+)
+
 ARG_TEAM = Arg(
     ("--team",),
     default=None,
@@ -84,7 +94,7 @@
             "in evicted/failed/succeeded/pending states"
         ),
         func=lazy_load_command("airflow.providers.cncf.kubernetes.cli.kubernetes_command.cleanup_pods"),
-        args=(ARG_NAMESPACE, ARG_MIN_PENDING_MINUTES, ARG_VERBOSE),
+        args=(ARG_NAMESPACE, ARG_MIN_PENDING_MINUTES, ARG_MIN_COMPLETED_MINUTES, ARG_VERBOSE),
     ),
     ActionCommand(
         name="generate-dag-yaml",

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/cli/kubernetes_command.py

@@ -119,6 +119,31 @@ def generate_pod_yaml(args):
     print(f"YAML output can be found at {yaml_output_path}")
 
 
+def _get_pod_completion_time(pod):
+    """
+    Return the time the pod's last container finished.
+
+    Scans both ``container_statuses`` and ``init_container_statuses`` and returns the
+    latest ``finished_at`` timestamp. Falls back to the latest condition
+    ``last_transition_time`` (which is updated at the terminal transition), and finally
+    to ``creation_timestamp`` as a last resort.
+    """
+    statuses = [*(pod.status.container_statuses or []), *(pod.status.init_container_statuses or [])]
+    times = [
+        s.state.terminated.finished_at
+        for s in statuses
+        if s.state and s.state.terminated and s.state.terminated.finished_at
+    ]
+    if times:
+        return max(times)
+    condition_times = [
+        c.last_transition_time for c in (pod.status.conditions or []) if c.last_transition_time
+    ]
+    if condition_times:
+        return max(condition_times)
+    return pod.metadata.creation_timestamp
+
+
 @cli_utils.action_cli(check_db=False)
 @providers_configuration_loaded
 def cleanup_pods(args):
@@ -130,6 +155,8 @@ def cleanup_pods(args):
     if min_pending_minutes < 5:
         min_pending_minutes = 5
 
+    min_completed_minutes = args.min_completed_minutes
+
     # https://kubernetes.io/docs/concepts/workloads/pods/pod-lifecycle/
     # All Containers in the Pod have terminated in success, and will not be restarted.
     pod_succeeded = "succeeded"
@@ -173,16 +200,20 @@ def cleanup_pods(args):
             pod_restart_policy = pod.spec.restart_policy.lower()
             current_time = datetime.now(pod.metadata.creation_timestamp.tzinfo)
 
-            if (
+            is_terminal = (
                 pod_phase == pod_succeeded
                 or (pod_phase == pod_failed and pod_restart_policy == pod_restart_policy_never)
                 or (pod_reason == pod_reason_evicted)
-                or (
-                    pod_phase == pod_pending
-                    and current_time - pod.metadata.creation_timestamp
-                    > timedelta(minutes=min_pending_minutes)
-                )
-            ):
+            )
+            is_terminal_old_enough = is_terminal and (
+                min_completed_minutes == 0
+                or current_time - _get_pod_completion_time(pod) > timedelta(minutes=min_completed_minutes)
+            )
+            is_pending_too_long = (
+                pod_phase == pod_pending
+                and current_time - pod.metadata.creation_timestamp > timedelta(minutes=min_pending_minutes)
+            )
+            if is_terminal_old_enough or is_pending_too_long:
                 print(
                     f'Deleting pod "{pod_name}" phase "{pod_phase}" and reason "{pod_reason}", '
                     f'restart policy "{pod_restart_policy}"'

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/cli/test_kubernetes_command.py

@@ -18,12 +18,15 @@
 
 import importlib
 import os
+from datetime import timedelta
 from unittest import mock
 from unittest.mock import MagicMock, call
 
 import kubernetes
 import pytest
+import time_machine
 from dateutil.parser import parse
+from kubernetes.client import models as k8s
 
 from airflow.cli import cli_parser
 from airflow.executors import executor_loader
@@ -34,6 +37,47 @@
 
 pytestmark = pytest.mark.db_test
 
+NOW = parse("2024-01-01T13:15:17Z")
+
+
+def make_container_status(state):
+    return k8s.V1ContainerStatus(
+        name="base", ready=False, restart_count=0, image="img", image_id="id", state=state
+    )
+
+
+def make_terminated_status(finished_at):
+    return make_container_status(
+        k8s.V1ContainerState(terminated=k8s.V1ContainerStateTerminated(exit_code=0, finished_at=finished_at))
+    )
+
+
+def make_terminal_pod(
+    name,
+    phase,
+    reason=None,
+    restart_policy="Never",
+    finished_at=None,
+    init_finished_at=None,
+    condition_time=None,
+):
+    conditions = (
+        [k8s.V1PodCondition(type="Ready", status="False", last_transition_time=condition_time)]
+        if condition_time
+        else []
+    )
+    return k8s.V1Pod(
+        metadata=k8s.V1ObjectMeta(name=name, creation_timestamp=NOW - timedelta(hours=1)),
+        spec=k8s.V1PodSpec(containers=[], restart_policy=restart_policy),
+        status=k8s.V1PodStatus(
+            phase=phase,
+            reason=reason,
+            container_statuses=[make_terminated_status(finished_at)] if finished_at else [],
+            init_container_statuses=[make_terminated_status(init_finished_at)] if init_finished_at else [],
+            conditions=conditions,
+        ),
+    )
+
 
 class TestGenerateDagYamlCommand:
     @classmethod
@@ -283,3 +327,139 @@ def test_list_pod_with_continue_token(self, load_incluster_config, list_namespac
         list_namespaced_pod.assert_has_calls(calls)
         delete_pod.assert_called_with("dummy", "awesome-namespace")
         load_incluster_config.assert_called_once()
+
+    @pytest.mark.parametrize(
+        ("pod_kwargs", "min_completed_minutes", "expect_deleted"),
+        [
+            pytest.param(
+                {"phase": "Succeeded", "finished_at": NOW - timedelta(seconds=3)},
+                1,
+                False,
+                id="succeeded-just-finished-kept",
+            ),
+            pytest.param(
+                {"phase": "Succeeded", "finished_at": NOW - timedelta(minutes=5)},
+                1,
+                True,
+                id="succeeded-old-enough-deleted",
+            ),
+            pytest.param(
+                {"phase": "Succeeded", "finished_at": NOW - timedelta(seconds=3)},
+                0,
+                True,
+                id="zero-disables-guard",
+            ),
+            pytest.param(
+                {"phase": "Failed", "finished_at": NOW - timedelta(seconds=3)},
+                1,
+                False,
+                id="failed-just-finished-kept",
+            ),
+            pytest.param(
+                {"phase": "Failed", "init_finished_at": NOW - timedelta(seconds=3)},
+                1,
+                False,
+                id="init-container-failed-just-finished-kept",
+            ),
+            pytest.param(
+                {"phase": "Failed", "reason": "Evicted", "condition_time": NOW - timedelta(seconds=3)},
+                1,
+                False,
+                id="evicted-before-containers-started-kept",
+            ),
+            pytest.param(
+                {"phase": "Failed", "reason": "Evicted", "condition_time": NOW - timedelta(minutes=5)},
+                1,
+                True,
+                id="evicted-old-enough-deleted",
+            ),
+        ],
+    )
+    @time_machine.travel(NOW, tick=False)
+    @mock.patch("airflow.providers.cncf.kubernetes.cli.kubernetes_command._delete_pod")
+    @mock.patch("kubernetes.client.CoreV1Api.list_namespaced_pod")
+    @mock.patch("kubernetes.config.load_incluster_config")
+    def test_cleanup_pods_min_completed_minutes(
+        self,
+        load_incluster_config,
+        list_namespaced_pod,
+        delete_pod,
+        pod_kwargs,
+        min_completed_minutes,
+        expect_deleted,
+    ):
+        pods = MagicMock()
+        pods.metadata._continue = None
+        pods.items = [make_terminal_pod("run-o1sxc2on", **pod_kwargs)]
+        list_namespaced_pod.return_value = pods
+        kubernetes_command.cleanup_pods(
+            self.parser.parse_args(
+                [
+                    "kubernetes",
+                    "cleanup-pods",
+                    "--namespace",
+                    "awesome-namespace",
+                    "--min-completed-minutes",
+                    str(min_completed_minutes),
+                ]
+            )
+        )
+        if expect_deleted:
+            delete_pod.assert_called_once_with("run-o1sxc2on", "awesome-namespace")
+        else:
+            delete_pod.assert_not_called()
+
+
+class TestGetPodCompletionTime:
+    T1 = parse("2024-01-01T10:00:00Z")
+    T2 = parse("2024-01-01T10:05:00Z")
+    CREATED_AT = parse("2024-01-01T09:00:00Z")  # earlier than T1/T2
+
+    def _pod(self, container_statuses=None, init_container_statuses=None, conditions=None):
+        return k8s.V1Pod(
+            metadata=k8s.V1ObjectMeta(name="run-o1sxc2on", creation_timestamp=self.CREATED_AT),
+            status=k8s.V1PodStatus(
+                container_statuses=container_statuses,
+                init_container_statuses=init_container_statuses,
+                conditions=conditions,
+            ),
+        )
+
+    def test_single_main_container(self):
+        pod = self._pod(container_statuses=[make_terminated_status(self.T1)])
+        assert kubernetes_command._get_pod_completion_time(pod) == self.T1
+
+    def test_single_init_container_no_main(self):
+        pod = self._pod(container_statuses=[], init_container_statuses=[make_terminated_status(self.T1)])
+        assert kubernetes_command._get_pod_completion_time(pod) == self.T1
+
+    def test_main_and_init_returns_max(self):
+        pod = self._pod(
+            container_statuses=[make_terminated_status(self.T1)],
+            init_container_statuses=[make_terminated_status(self.T2)],
+        )
+        assert kubernetes_command._get_pod_completion_time(pod) == self.T2
+
+    @pytest.mark.parametrize(
+        "container_status",
+        [
+            pytest.param(make_container_status(state=None), id="no-state"),
+            pytest.param(
+                make_container_status(
+                    k8s.V1ContainerState(waiting=k8s.V1ContainerStateWaiting(reason="ContainerCreating"))
+                ),
+                id="never-terminated",
+            ),
+            pytest.param(make_terminated_status(finished_at=None), id="terminated-without-finished-at"),
+        ],
+    )
+    def test_falls_back_to_conditions(self, container_status):
+        pod = self._pod(
+            container_statuses=[container_status],
+            conditions=[k8s.V1PodCondition(type="Ready", status="False", last_transition_time=self.T1)],
+        )
+        assert kubernetes_command._get_pod_completion_time(pod) == self.T1
+
+    def test_no_containers_no_conditions_falls_back_to_creation_timestamp(self):
+        pod = self._pod(container_statuses=[], init_container_statuses=[], conditions=[])
+        assert kubernetes_command._get_pod_completion_time(pod) == self.CREATED_AT