You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by am...@apache.org on 2026/07/23 07:34:01 UTC
(airflow) branch main updated: Replace `reattach_on_restart` with `durable` execution in KPO (#69914)
This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 901d9190d11 Replace `reattach_on_restart` with `durable` execution in KPO (#69914)
901d9190d11 is described below
commit 901d9190d111f3f410b7bb3da278b41c1895404d
Author: Amogh Desai <am...@gmail.com>
AuthorDate: Thu Jul 23 13:03:51 2026 +0530
Replace `reattach_on_restart` with `durable` execution in KPO (#69914)
---
providers/cncf/kubernetes/docs/operators.rst | 37 +++
.../providers/cncf/kubernetes/operators/pod.py | 99 +++++++-
.../decorators/test_kubernetes_commons.py | 18 ++
.../unit/cncf/kubernetes/operators/test_pod.py | 259 ++++++++++++++++++++-
4 files changed, 402 insertions(+), 11 deletions(-)
diff --git a/providers/cncf/kubernetes/docs/operators.rst b/providers/cncf/kubernetes/docs/operators.rst
index 0eb7dcedc4b..abdd4307aba 100644
--- a/providers/cncf/kubernetes/docs/operators.rst
+++ b/providers/cncf/kubernetes/docs/operators.rst
@@ -187,6 +187,43 @@ 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) 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 -- the reconnect step uses this persisted identity instead of a label search.
+
+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 18062c71526..178f61cda9a 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 re
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.client.exceptions import ApiException
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 @@ from airflow.providers.cncf.kubernetes.utils.pod_manager import (
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 @@ alphanum_lower = string.ascii_lowercase + string.digits
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
@@ -321,7 +335,8 @@ class KubernetesPodOperator(BaseOperator):
in_cluster: bool | None = None,
cluster_context: str | None = None,
labels: dict | None = None,
- reattach_on_restart: bool = True,
+ reattach_on_restart: bool | None = None,
+ durable: bool = True,
startup_timeout_seconds: int = 120,
startup_check_interval_seconds: int = 5,
schedule_timeout_seconds: int | None = None,
@@ -380,6 +395,16 @@ class KubernetesPodOperator(BaseOperator):
log_formatter: Callable[[str, str], str] | None = None,
**kwargs,
) -> None:
+ if reattach_on_restart is not None:
+ # Kept as a real named parameter (not **kwargs) so `default_args` still applies correctly.
+ 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
@@ -409,7 +434,11 @@ class KubernetesPodOperator(BaseOperator):
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.
@@ -614,9 +643,63 @@ class KubernetesPodOperator(BaseOperator):
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
+ if (pod.metadata.labels or {}).get(self.POD_CHECKED_KEY) == "True":
+ # A prior attempt already fully processed this pod, it belongs to a previous
+ # try_number, not one to reattach to. Fall through to fresh pod creation.
+ self.log.info(
+ "Pod %s/%s from task state store was already checked by a prior attempt, "
+ "creating a fresh pod instead of reattaching.",
+ 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},
+ )
+
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
@@ -644,7 +727,9 @@ class KubernetesPodOperator(BaseOperator):
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 d6ffbd3a81f..16c4ed65c2e 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 @@ class TestKubernetesDecoratorsBase:
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(
@@ -113,6 +120,17 @@ class TestKubernetesDecoratorsBase:
self.mock_hook = mock.patch(HOOK_CLASS).start()
self.mock_hook.return_value.get_xcom_sidecar_container_security_context.return_value = None
+ 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 8163b95947b..ea5ed5a207c 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 @@ import time_machine
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 airflow.utils.types import DagRunType
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 @@ class TestKubernetesPodOperator:
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
@@ -2438,7 +2450,7 @@ class TestKubernetesPodOperator:
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)
@@ -2475,7 +2487,7 @@ class TestKubernetesPodOperator:
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)
@@ -2508,7 +2520,7 @@ class TestKubernetesPodOperator:
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)
@@ -2526,6 +2538,245 @@ class TestKubernetesPodOperator:
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"
+ 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"},
+ )
+ 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"}
+ 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"}
+ 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_reconnect_skips_already_checked_pod_and_creates_fresh(self):
+ k = KubernetesPodOperator(
+ image="ubuntu:16.04",
+ cmds=["bash", "-cx"],
+ arguments=["echo 10"],
+ task_id="task",
+ name="hello",
+ log_pod_spec_on_failure=False,
+ on_finish_action="keep_pod",
+ )
+ context = create_context(k)
+ task_state_store = MagicMock()
+ task_state_store.get.return_value = {"name": "prior-pod", "namespace": "default"}
+ context["task_state_store"] = task_state_store
+
+ stale_pod = MagicMock()
+ stale_pod.status.phase = PodPhase.FAILED
+ stale_pod.status.reason = None
+ stale_pod.metadata.name = "prior-pod"
+ stale_pod.metadata.labels = {"try_number": "1", "already_checked": "True"}
+ k.hook.get_pod = MagicMock(return_value=stale_pod)
+
+ created_pod = MagicMock()
+ created_pod.metadata.name = "new-pod"
+ created_pod.metadata.namespace = "default"
+ self.create_mock.return_value = created_pod
+
+ 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", return_value=None) as mock_find,
+ patch(f"{KPO_MODULE}.KubernetesPodOperator.process_pod_deletion") as mock_process_pod_deletion,
+ ):
+ result = k.get_or_create_pod(pod_request_obj=mock_pod_request_obj, context=context)
+
+ mock_find.assert_called_once()
+ mock_process_pod_deletion.assert_not_called()
+ 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": "new-pod", "namespace": "default"},
+ )
+ 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"}
+ 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
+
+ def test_reattach_on_restart_via_default_args_reaches_durable(self, dag_maker):
+ with pytest.warns(AirflowProviderDeprecationWarning, match="reattach_on_restart"):
+ with dag_maker(dag_id="test_reattach_default_args", default_args={"reattach_on_restart": False}):
+ k = KubernetesPodOperator(task_id="task")
+ assert k.durable is False
+ assert k.reattach_on_restart is False
+
+
class TestSuppress:
def test__suppress(self, caplog):
with _optionally_suppress(ValueError):