Replace reattach_on_restart with durable execution in KPO

providers/cncf/kubernetes/docs/operators.rst

@@ -187,6 +187,45 @@ for debugging), set ``on_kill_action="keep_pod"``:
 The ``termination_grace_period`` parameter is also respected during cleanup, giving the
 pod time to shut down gracefully before being forcefully terminated.
 
+Durable execution
+^^^^^^^^^^^^^^^^^
+
+If the worker running ``KubernetesPodOperator`` dies while the pod is still running (e.g. the
+worker is preempted or crashes) and the task is retried, the operator can reattach to the pod
+that is already running instead of creating a duplicate. This is controlled by the ``durable``
+parameter, which defaults to ``True``.
+
+On Airflow 3.3+, ``durable=True`` persists the running pod's identity (name, namespace, uid) to
+:doc:`task state store <apache-airflow:core-concepts/task-state-store>` before the operator starts
+waiting on it. On retry, the operator reads this identity back and reconnects directly to that
+specific pod -- no label search, and no possibility of the ambiguity failure
+(``FoundMoreThanOnePodFailure``) that a label search can hit when more than one matching pod
+exists.
+
+If no identity has been persisted yet to the task state store - either because this is the first attempt, or because the
+worker crashed in the narrow window after the pod was created but before its identity could be
+persisted, the operator falls back to the same label search ``reattach_on_restart`` has always
+used, so a running pod from a prior attempt is still found and reattached to rather than
+duplicated. Once an identity is persisted, subsequent retries skip the label search entirely.
+
+To always create a fresh pod on retry rather than reattaching, set ``durable=False``:
+
+.. code-block:: python
+
+    k = KubernetesPodOperator(
+        task_id="task",
+        image="my-image:latest",
+        durable=False,
+    )
+
+Durable execution requires Airflow 3.3 or newer, since it relies on the task state store. On
+earlier Airflow versions ``durable=True`` (the default) falls back to the same label-search
+reattach behavior this operator has always used.
+
+``durable`` supersedes the deprecated ``reattach_on_restart`` parameter -- passing
+``reattach_on_restart`` still works but emits ``AirflowProviderDeprecationWarning`` (on Airflow
+3.3+) and maps its value onto ``durable``.
+
 How does XCom work?
 ^^^^^^^^^^^^^^^^^^^
 The :class:`~airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator` handles

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py

@@ -29,6 +29,7 @@
 import shlex
 import string
 import time
+import warnings
 from collections.abc import Callable, Container, Iterable, Mapping, Sequence
 from contextlib import AbstractContextManager, suppress
 from enum import Enum
@@ -43,6 +44,7 @@
 from kubernetes.stream import stream
 from urllib3.exceptions import HTTPError
 
+from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.cncf.kubernetes import pod_generator
 from airflow.providers.cncf.kubernetes.backcompat.backwards_compat_converters import (
     convert_affinity,
@@ -84,7 +86,7 @@
     PodPhase,
     detect_pod_terminate_early_issues,
 )
-from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_1_PLUS
+from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS
 from airflow.providers.common.compat.sdk import XCOM_RETURN_KEY, AirflowSkipException, conf
 
 if AIRFLOW_V_3_1_PLUS:
@@ -113,6 +115,10 @@
 
 KUBE_CONFIG_ENV_VAR = "KUBECONFIG"
 
+# Key used to persist/retrieve a submitted pod's identity in task_state_store across retries.
+# Renaming this on a deployed operator breaks in flight retries because the old key is already stored.
+POD_IDENTIFIER_STATE_KEY = "pod_identifier"
+
 
 class PodEventType(Enum):
     """Type of Events emitted by kubernetes pod."""
@@ -170,8 +176,16 @@ class KubernetesPodOperator(BaseOperator):
     :param in_cluster: run kubernetes client with in_cluster configuration.
     :param cluster_context: context that points to kubernetes cluster.
         Ignored when in_cluster is True. If None, current-context is used. (templated)
-    :param reattach_on_restart: if the worker dies while the pod is running, reattach and monitor
-        during the next try. If False, always create a new pod for each try.
+    :param reattach_on_restart: deprecated for Airflow 3.3+, use ``durable`` instead. If the worker dies
+        while the pod is running, reattach and monitor during the next try. If False, always create a
+        new pod for each try.
+    :param durable: if the worker dies while the pod is running, reattach and monitor during the next
+        try instead of creating a duplicate pod. If False, always create a new pod for each try.
+        Supersedes ``reattach_on_restart``; on Airflow 3.3+ the reconnection uses a persisted pod
+        identity in task state store instead of a label search, removing the ambiguity failure a label
+        search can hit when more than one matching pod exists. Defaults to ``True``. On Airflow
+        versions below 3.3, ``durable`` still works but falls back to the same label-search reattach
+        behavior as ``reattach_on_restart``, since task state store is unavailable.
     :param labels: labels to apply to the Pod. (templated)
     :param startup_timeout_seconds: timeout in seconds to startup the pod after pod was scheduled.
     :param startup_check_interval_seconds: interval in seconds to check if the pod has already started
@@ -319,7 +333,7 @@ def __init__(
         in_cluster: bool | None = None,
         cluster_context: str | None = None,
         labels: dict | None = None,
-        reattach_on_restart: bool = True,
+        durable: bool = True,
         startup_timeout_seconds: int = 120,
         startup_check_interval_seconds: int = 5,
         schedule_timeout_seconds: int | None = None,
@@ -377,6 +391,18 @@ def __init__(
         log_formatter: Callable[[str, str], str] | None = None,
         **kwargs,
     ) -> None:
+        if "reattach_on_restart" in kwargs:
+            # Dropped from the named signature entirely so the warning fires on any explicit use,
+            # `True` or `False`.
+            reattach_on_restart = kwargs.pop("reattach_on_restart")
+            if AIRFLOW_V_3_3_PLUS:
+                warnings.warn(
+                    "`reattach_on_restart` is deprecated and will be removed in a future release. "
+                    "Use `durable` instead.",
+                    AirflowProviderDeprecationWarning,
+                    stacklevel=2,
+                )
+            durable = reattach_on_restart
         super().__init__(**kwargs)
         self.kubernetes_conn_id = kubernetes_conn_id
         self.do_xcom_push = do_xcom_push
@@ -406,7 +432,11 @@ def __init__(
         self.secrets = secrets or []
         self.in_cluster = in_cluster
         self.cluster_context = cluster_context
-        self.reattach_on_restart = reattach_on_restart
+        self.durable = durable
+        # `reattach_on_restart` is kept as an ordinary attribute (not a deprecated property) so
+        # internal reads of it do not themselves emit a deprecation warning; it always mirrors
+        # `self.durable`.
+        self.reattach_on_restart = durable
         self.get_logs = get_logs
         # Fallback to the class variable BASE_CONTAINER_NAME here instead of via default argument value
         # in the init method signature, to be compatible with subclasses overloading the class variable value.
@@ -610,9 +640,53 @@ def log_matching_pod(self, pod: k8s.V1Pod, context: Context) -> None:
         self.log.info("`try_number` of task_instance: %s", context["ti"].try_number)
         self.log.info("`try_number` of pod: %s", pod.metadata.labels["try_number"])
 
+    def _get_pod_from_task_state_store(self, context: Context) -> k8s.V1Pod | None:
+        task_state_store = context.get("task_state_store")
+        if task_state_store is None:
+            return None
+        stored = task_state_store.get(POD_IDENTIFIER_STATE_KEY)
+        if not isinstance(stored, dict):
+            return None
+        name, namespace = stored.get("name"), stored.get("namespace")
+        if not isinstance(name, str) or not isinstance(namespace, str):
+            self.log.warning(
+                "Pod identity stored in task state store is malformed (%s), falling back to label search.",
+                stored,
+            )
+            return None
+        try:
+            pod = self.hook.get_pod(name, namespace)
+        except ApiException as e:
+            if e.status != 404:
+                raise
+            self.log.warning(
+                "Pod %s/%s from task state store no longer exists, falling back to label search.",
+                namespace,
+                name,
+            )
+            return None
+        self.log.info(
+            "Reconnecting to pod %s/%s via identity persisted in task state store.", namespace, name
+        )
+        self.log_matching_pod(pod=pod, context=context)
+        return pod
+
+    def _persist_pod_identity_to_task_state_store(self, context: Context, pod: k8s.V1Pod) -> None:
+        task_state_store = context.get("task_state_store")
+        if task_state_store is None:
+            return
+        task_state_store.set(
+            POD_IDENTIFIER_STATE_KEY,
+            {"name": pod.metadata.name, "namespace": pod.metadata.namespace, "uid": pod.metadata.uid},
+        )
+
     def get_or_create_pod(self, pod_request_obj: k8s.V1Pod, context: Context) -> k8s.V1Pod:
         if self.reattach_on_restart:
-            pod = self.find_pod(pod_request_obj.metadata.namespace, context=context)
+            pod = self._get_pod_from_task_state_store(context)
+            if pod is None:
+                pod = self.find_pod(pod_request_obj.metadata.namespace, context=context)
+                if pod is not None:
+                    self._persist_pod_identity_to_task_state_store(context, pod)
             if pod:
                 # If pod is terminated then delete the pod an create a new as not possible to get xcom
                 pod_phase = pod.status.phase if pod.status and pod.status.phase else None
@@ -640,7 +714,9 @@ def get_or_create_pod(self, pod_request_obj: k8s.V1Pod, context: Context) -> k8s
                 self.log.info("Deleted pod to handle rerun and create new pod!")
 
         self.log.debug("Starting pod:\n%s", yaml.safe_dump(pod_request_obj.to_dict()))
-        self.pod_manager.create_pod(pod=pod_request_obj)
+        created_pod = self.pod_manager.create_pod(pod=pod_request_obj)
+        if self.reattach_on_restart:
+            self._persist_pod_identity_to_task_state_store(context, created_pod)
         return pod_request_obj
 
     def await_pod_start(self, pod: k8s.V1Pod) -> None:

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_commons.py

@@ -90,6 +90,13 @@ def setup(self, dag_maker):
         self.dag = self.dag_maker.dag
 
         self.mock_create_pod = mock.patch(f"{POD_MANAGER_CLASS}.create_pod").start()
+        self.mock_create_pod.return_value = mock.MagicMock(
+            **{
+                "metadata.name": "test-pod",
+                "metadata.namespace": "default",
+                "metadata.uid": "test-uid",
+            }
+        )
         self.mock_await_pod_start = mock.patch(f"{POD_MANAGER_CLASS}.await_pod_start").start()
         self.mock_watch_pod_events = mock.patch(f"{POD_MANAGER_CLASS}.watch_pod_events").start()
         self.mock_await_xcom_sidecar_container_start = mock.patch(
@@ -112,6 +119,17 @@ def setup(self, dag_maker):
         )
         self.mock_hook = mock.patch(HOOK_CLASS).start()
 
+        try:
+            from airflow.sdk.execution_time.context import TaskStateStoreAccessor
+        except ImportError:
+            # Airflow versions without task_state_store (e.g. pre-3.3, or pre-3.0 where
+            # airflow.sdk doesn't exist at all) don't define this class -- nothing to patch,
+            # context.get("task_state_store") already returns None on those versions.
+            pass
+        else:
+            mock.patch.object(TaskStateStoreAccessor, "get", return_value=None).start()
+            mock.patch.object(TaskStateStoreAccessor, "set").start()
+
         # Without this patch each time pod manager would try to extract logs from the pod
         # and log an error about it's inability to get containers for the log
         # {pod_manager.py:572} ERROR - Could not retrieve containers for the pod: ...