diff --git a/providers/cncf/kubernetes/docs/operators.rst b/providers/cncf/kubernetes/docs/operators.rst index a3f8e417237e8..127269f5475b9 100644 --- a/providers/cncf/kubernetes/docs/operators.rst +++ b/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 ` 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 diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py index 9f59fcb66d368..6e7970cb53050 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/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: diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_commons.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_commons.py index 9a23349956f91..abad6feea3368 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_commons.py +++ b/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: ... diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py index 11483b7f0a174..869f1d6ffc430 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py @@ -31,9 +31,11 @@ from kubernetes.client import ApiClient, V1Pod, V1PodSecurityContext, V1PodStatus, models as k8s from kubernetes.client.exceptions import ApiException +from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.models import DAG, DagModel, DagRun, TaskInstance from airflow.providers.cncf.kubernetes import pod_generator from airflow.providers.cncf.kubernetes.operators.pod import ( + POD_IDENTIFIER_STATE_KEY, KubernetesPodOperator, PodEventType, _optionally_suppress, @@ -62,7 +64,7 @@ from tests_common.test_utils import db from tests_common.test_utils.dag import sync_dag_to_db from tests_common.test_utils.taskinstance import create_task_instance, get_template_context -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS if AIRFLOW_V_3_0_PLUS or AIRFLOW_V_3_1_PLUS: from airflow.models.xcom import XComModel as XCom @@ -189,6 +191,16 @@ def setup_tests(self, dag_maker): self.await_start_mock = self.await_pod_patch.start() self.await_pod_mock = self.await_pod_completion_patch.start() self._default_client_mock = self._default_client_patch.start() + + try: + from airflow.sdk.execution_time.context import TaskStateStoreAccessor + except ImportError: + # Airflow versions < 3.3, do not have task_state_store + pass + else: + patch.object(TaskStateStoreAccessor, "get", return_value=None).start() + patch.object(TaskStateStoreAccessor, "set").start() + self.dag_maker = dag_maker yield @@ -2383,7 +2395,7 @@ def test_process_duplicate_label_pods__label_patched_if_action_is_not_delete_pod name="test", task_id="task", do_xcom_push=False, - reattach_on_restart=False, + durable=False, on_finish_action=on_finish_action, ) context = create_context(k) @@ -2420,7 +2432,7 @@ def test_process_duplicate_label_pods_with_start_time_none( name="test", task_id="task", do_xcom_push=False, - reattach_on_restart=False, + durable=False, on_finish_action=on_finish_action, ) context = create_context(k) @@ -2453,7 +2465,7 @@ def test_process_duplicate_label_pods__pod_removed_if_delete_pod( name="test", task_id="task", do_xcom_push=False, - reattach_on_restart=False, + durable=False, on_finish_action=OnFinishAction.DELETE_POD, ) context = create_context(k) @@ -2471,6 +2483,194 @@ def test_process_duplicate_label_pods__pod_removed_if_delete_pod( assert result.metadata.name == pod_2.metadata.name +@pytest.mark.skipif( + not AIRFLOW_V_3_3_PLUS, reason="durable execution (task_state_store) requires Airflow 3.3+" +) +class TestKubernetesPodOperatorDurableExecution: + @pytest.fixture(autouse=True) + def setup_tests(self): + self.create_pod_patch = patch(f"{POD_MANAGER_CLASS}.create_pod") + self.create_mock = self.create_pod_patch.start() + self._default_client_patch = patch(f"{HOOK_CLASS}._get_default_client") + self._default_client_mock = self._default_client_patch.start() + + yield + + patch.stopall() + + def test_durable_fresh_submit_persists_pod_identity(self): + k = KubernetesPodOperator( + image="ubuntu:16.04", + cmds=["bash", "-cx"], + arguments=["echo 10"], + task_id="task", + name="hello", + log_pod_spec_on_failure=False, + ) + context = create_context(k) + task_state_store = MagicMock() + task_state_store.get.return_value = None + context["task_state_store"] = task_state_store + + mock_pod_request_obj = MagicMock() + mock_pod_request_obj.to_dict.return_value = {"metadata": {"name": "test-pod"}} + created_pod = MagicMock() + created_pod.metadata.name = "test-pod" + created_pod.metadata.namespace = "default" + created_pod.metadata.uid = "abc-uid" + self.create_mock.return_value = created_pod + + with patch(f"{KPO_MODULE}.KubernetesPodOperator.find_pod", return_value=None) as mock_find: + result = k.get_or_create_pod(pod_request_obj=mock_pod_request_obj, context=context) + + mock_find.assert_called_once() + self.create_mock.assert_called_once_with(pod=mock_pod_request_obj) + task_state_store.set.assert_called_once_with( + POD_IDENTIFIER_STATE_KEY, + {"name": "test-pod", "namespace": "default", "uid": "abc-uid"}, + ) + assert result == mock_pod_request_obj + + def test_durable_reconnects_directly_via_task_state_store(self): + k = KubernetesPodOperator( + image="ubuntu:16.04", + cmds=["bash", "-cx"], + arguments=["echo 10"], + task_id="task", + name="hello", + log_pod_spec_on_failure=False, + ) + context = create_context(k) + task_state_store = MagicMock() + task_state_store.get.return_value = {"name": "prior-pod", "namespace": "default", "uid": "abc"} + context["task_state_store"] = task_state_store + + running_pod = MagicMock() + running_pod.status.phase = "Running" + running_pod.status.reason = None + running_pod.metadata.name = "prior-pod" + running_pod.metadata.labels = {"try_number": "1"} + k.hook.get_pod = MagicMock(return_value=running_pod) + + mock_pod_request_obj = MagicMock() + + with patch(f"{KPO_MODULE}.KubernetesPodOperator.find_pod") as mock_find: + result = k.get_or_create_pod(pod_request_obj=mock_pod_request_obj, context=context) + + mock_find.assert_not_called() + k.hook.get_pod.assert_called_once_with("prior-pod", "default") + self.create_mock.assert_not_called() + assert result == running_pod + + @pytest.mark.parametrize("pod_phase", [PodPhase.SUCCEEDED, PodPhase.FAILED]) + @patch(f"{KPO_MODULE}.KubernetesPodOperator.process_pod_deletion") + def test_durable_reconnect_terminal_state_runs_existing_handling( + self, mock_process_pod_deletion, pod_phase + ): + k = KubernetesPodOperator( + image="ubuntu:16.04", + cmds=["bash", "-cx"], + arguments=["echo 10"], + task_id="task", + name="hello", + log_pod_spec_on_failure=False, + ) + context = create_context(k) + task_state_store = MagicMock() + task_state_store.get.return_value = {"name": "prior-pod", "namespace": "default", "uid": "abc"} + context["task_state_store"] = task_state_store + + terminal_pod = MagicMock() + terminal_pod.status.phase = pod_phase + terminal_pod.status.reason = None + terminal_pod.metadata.name = "prior-pod" + terminal_pod.metadata.labels = {"try_number": "1"} + k.hook.get_pod = MagicMock(return_value=terminal_pod) + mock_process_pod_deletion.return_value = True + + mock_pod_request_obj = MagicMock() + mock_pod_request_obj.to_dict.return_value = {"metadata": {"name": "new-pod"}} + + with patch(f"{KPO_MODULE}.KubernetesPodOperator.find_pod") as mock_find: + result = k.get_or_create_pod(pod_request_obj=mock_pod_request_obj, context=context) + + mock_find.assert_not_called() + mock_process_pod_deletion.assert_called_once_with(terminal_pod) + self.create_mock.assert_called_once_with(pod=mock_pod_request_obj) + assert result == mock_pod_request_obj + + def test_durable_true_without_state_store_fallsback(self): + """No task_state_store key in context (e.g. Airflow <3.3): unchanged label-search behavior.""" + k = KubernetesPodOperator( + image="ubuntu:16.04", + cmds=["bash", "-cx"], + arguments=["echo 10"], + task_id="task", + name="hello", + log_pod_spec_on_failure=False, + ) + context = create_context(k) + + found_pod = MagicMock() + found_pod.status.phase = "Running" + found_pod.status.reason = None + + mock_pod_request_obj = MagicMock() + + with patch(f"{KPO_MODULE}.KubernetesPodOperator.find_pod", return_value=found_pod) as mock_find: + result = k.get_or_create_pod(pod_request_obj=mock_pod_request_obj, context=context) + + mock_find.assert_called_once() + assert result == found_pod + + def test_durable_false_never_touches_task_state_store(self): + k = KubernetesPodOperator( + image="ubuntu:16.04", + cmds=["bash", "-cx"], + arguments=["echo 10"], + task_id="task", + name="hello", + log_pod_spec_on_failure=False, + durable=False, + ) + context = create_context(k) + task_state_store = MagicMock() + task_state_store.get.return_value = {"name": "prior-pod", "namespace": "default", "uid": "abc"} + context["task_state_store"] = task_state_store + + mock_pod_request_obj = MagicMock() + mock_pod_request_obj.to_dict.return_value = {"metadata": {"name": "test-pod"}} + + with patch(f"{KPO_MODULE}.KubernetesPodOperator.find_pod") as mock_find: + result = k.get_or_create_pod(pod_request_obj=mock_pod_request_obj, context=context) + + mock_find.assert_not_called() + task_state_store.get.assert_not_called() + task_state_store.set.assert_not_called() + self.create_mock.assert_called_once_with(pod=mock_pod_request_obj) + assert result == mock_pod_request_obj + + @pytest.mark.parametrize("reattach_value", [True, False]) + def test_reattach_on_restart_deprecation_maps_to_durable(self, reattach_value): + with pytest.warns(AirflowProviderDeprecationWarning, match="reattach_on_restart"): + k = KubernetesPodOperator( + task_id="task", + reattach_on_restart=reattach_value, + ) + assert k.durable is reattach_value + assert k.reattach_on_restart is reattach_value + + def test_reattach_on_restart_and_durable_conflict_reattach_on_restart_wins(self): + with pytest.warns(AirflowProviderDeprecationWarning, match="reattach_on_restart"): + k = KubernetesPodOperator( + task_id="task", + durable=False, + reattach_on_restart=True, + ) + assert k.durable is True + assert k.reattach_on_restart is True + + class TestSuppress: def test__suppress(self, caplog): with _optionally_suppress(ValueError):