From 4809aa175a9b8fe00568ef94a073c82fb4bee1a5 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Thu, 23 Jul 2026 14:24:58 +0530 Subject: [PATCH 1/4] Mark KubernetesPodOperator and AgentOperator as durable capable --- .../src/airflow/providers/cncf/kubernetes/operators/pod.py | 6 +++++- .../tests/unit/cncf/kubernetes/operators/test_pod.py | 3 +++ .../ai/src/airflow/providers/common/ai/operators/agent.py | 4 ++++ .../common/ai/tests/unit/common/ai/operators/test_agent.py | 3 +++ 4 files changed, 15 insertions(+), 1 deletion(-) 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 178f61cda9a6a..27f130b05a987 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 @@ -34,7 +34,7 @@ from contextlib import AbstractContextManager, suppress from enum import Enum from functools import cached_property -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, ClassVar, Literal import kubernetes import pendulum @@ -283,6 +283,10 @@ class KubernetesPodOperator(BaseOperator): # !!! Changes in KubernetesPodOperator's arguments should be also reflected in !!! # - airflow-core/src/airflow/decorators/__init__.pyi (by a separate PR) + # This operator supports durable execution directly, without ResumableJobMixin -- + # it reconnects via task_state_store on retry instead of resubmitting. + _supports_durable_execution: ClassVar[bool] = True + # This field can be overloaded at the instance level via base_container_name BASE_CONTAINER_NAME = "base" ISTIO_CONTAINER_NAME = "istio-proxy" 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 ea5ed5a207ce6..d76c16efd04b1 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 @@ -2776,6 +2776,9 @@ def test_reattach_on_restart_via_default_args_reaches_durable(self, dag_maker): assert k.durable is False assert k.reattach_on_restart is False + def test_supports_durable_execution_marker(self): + assert KubernetesPodOperator._supports_durable_execution is True + class TestSuppress: def test__suppress(self, caplog): diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index f83a58f0e2aed..c71a1550b286b 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py @@ -221,6 +221,10 @@ class AgentOperator(BaseOperator, HITLReviewMixin): deserialization_allowed_class_fields: ClassVar[tuple[str, ...]] = ("output_type",) + # This operator supports durable execution directly, without ResumableJobMixin -- + # it caches step results via task_state_store for replay on retry. + _supports_durable_execution: ClassVar[bool] = True + template_fields: Sequence[str] = ( "prompt", "llm_conn_id", diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py index 2631c544138f5..44d2d0efba5cc 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py @@ -759,6 +759,9 @@ def test_cleanup_skipped_when_post_run_step_fails(self, mock_hook_cls, mock_buil storage.cleanup.assert_not_called() + def test_supports_durable_execution_marker(self): + assert AgentOperator._supports_durable_execution is True + @pytest.mark.skipif( not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0" From 4831c24c14b1fd1a48952ebd0a842d5c2e1a0acc Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Thu, 23 Jul 2026 14:53:20 +0530 Subject: [PATCH 2/4] Mark KubernetesPodOperator and AgentOperator as durable capable --- .../src/airflow/providers/cncf/kubernetes/operators/pod.py | 2 +- .../kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py | 2 +- .../ai/src/airflow/providers/common/ai/operators/agent.py | 2 +- .../common/ai/tests/unit/common/ai/operators/test_agent.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) 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 27f130b05a987..afbe07b827a08 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 @@ -285,7 +285,7 @@ class KubernetesPodOperator(BaseOperator): # This operator supports durable execution directly, without ResumableJobMixin -- # it reconnects via task_state_store on retry instead of resubmitting. - _supports_durable_execution: ClassVar[bool] = True + __supports_durable_execution: ClassVar[bool] = True # This field can be overloaded at the instance level via base_container_name BASE_CONTAINER_NAME = "base" 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 d76c16efd04b1..56c9f13d77974 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 @@ -2777,7 +2777,7 @@ def test_reattach_on_restart_via_default_args_reaches_durable(self, dag_maker): assert k.reattach_on_restart is False def test_supports_durable_execution_marker(self): - assert KubernetesPodOperator._supports_durable_execution is True + assert KubernetesPodOperator.__supports_durable_execution is True class TestSuppress: diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index c71a1550b286b..dac5e5b283077 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py @@ -223,7 +223,7 @@ class AgentOperator(BaseOperator, HITLReviewMixin): # This operator supports durable execution directly, without ResumableJobMixin -- # it caches step results via task_state_store for replay on retry. - _supports_durable_execution: ClassVar[bool] = True + __supports_durable_execution: ClassVar[bool] = True template_fields: Sequence[str] = ( "prompt", diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py index 44d2d0efba5cc..fb61086d7160c 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py @@ -760,7 +760,7 @@ def test_cleanup_skipped_when_post_run_step_fails(self, mock_hook_cls, mock_buil storage.cleanup.assert_not_called() def test_supports_durable_execution_marker(self): - assert AgentOperator._supports_durable_execution is True + assert AgentOperator.__supports_durable_execution is True @pytest.mark.skipif( From 0f3931470a5dd843167342438e512f86fd035195 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Thu, 23 Jul 2026 15:05:18 +0530 Subject: [PATCH 3/4] Add support for operators that achieve durable execution manually in registry --- dev/registry/extract_parameters.py | 26 ++++++++++++++++--- dev/registry/tests/test_extract_parameters.py | 24 +++++++++++++++++ 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/dev/registry/extract_parameters.py b/dev/registry/extract_parameters.py index 0e9eff89a4df7..e32b06bae7ac7 100644 --- a/dev/registry/extract_parameters.py +++ b/dev/registry/extract_parameters.py @@ -386,11 +386,29 @@ def load_resumable_job_mixin() -> type | None: def is_durable_capable(cls: type, resumable_mixin: type | None) -> bool: - """Return True if a class fully implements ResumableJobMixin's crash-recovery contract. - - Inheriting the mixin is not sufficient: a complete override is inert unless - execute() actually calls execute_resumable(). + """Return True if a class implements durable/crash-safe execution. + + Two ways to qualify: + 1. A class-level `__supports_durable_execution = True` + declaration (for operators that implement this directly against + task_state_store, without ResumableJobMixin -- e.g. KubernetesPodOperator, + AgentOperator). + 2. Genuinely implementing ResumableJobMixin's contract. + + The first path deliberately looks up the class prefixed attribute + (`_{ClassName}__supports_durable_execution`) rather than a fixed string. + A subclass that overrides execute() itself (e.g. SparkKubernetesOperator) + may not preserve the parent's task_state_store reconnect behavior, so the + declaration must not be inherited -- only the exact class that wrote + `__supports_durable_execution` in its own body qualifies this way. + + Inheriting the mixin alone is not sufficient for the second path: a + complete override is inert unless execute() actually calls + execute_resumable(). """ + if getattr(cls, f"_{cls.__name__}__supports_durable_execution", None) is True: + return True + if resumable_mixin is None or resumable_mixin not in cls.__mro__: return False diff --git a/dev/registry/tests/test_extract_parameters.py b/dev/registry/tests/test_extract_parameters.py index be0f608e17ed7..2cdb09532299f 100644 --- a/dev/registry/tests/test_extract_parameters.py +++ b/dev/registry/tests/test_extract_parameters.py @@ -225,6 +225,24 @@ def execute(self, context): return None +class ManuallyDurableOperator: + """Implements durable execution directly (e.g. via task_state_store), without + ResumableJobMixin -- mirrors KubernetesPodOperator/AgentOperator.""" + + __supports_durable_execution = True + + def execute(self, context): + return None + + +class ManuallyDurableSubclass(ManuallyDurableOperator): + """Overrides execute() itself -- must NOT inherit the parent's declaration, + since nothing here verifies it preserves the reconnect behavior.""" + + def execute(self, context): + return "something else entirely" + + class TestIsDurableCapable: def test_fully_implemented_and_wired_qualifies(self): assert is_durable_capable(FullyImplementedResumableOperator, FakeResumableJobMixin) is True @@ -241,6 +259,12 @@ def test_no_mixin_in_mro_disqualifies(self): def test_mixin_unavailable_disqualifies(self): assert is_durable_capable(FullyImplementedResumableOperator, None) is False + def test_manual_durable_marker_qualifies_without_mixin(self): + assert is_durable_capable(ManuallyDurableOperator, None) is True + + def test_subclass_not_redeclaring_marker_disqualifies(self): + assert is_durable_capable(ManuallyDurableSubclass, FakeResumableJobMixin) is False + # --------------------------------------------------------------------------- # Module dataclass From 55a6b8ffb3906371f0d616b7c10aeee943e4ad01 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Fri, 24 Jul 2026 11:23:14 +0530 Subject: [PATCH 4/4] fixing tests --- .../kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py | 2 +- .../common/ai/tests/unit/common/ai/operators/test_agent.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 56c9f13d77974..0e903ec644b1a 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 @@ -2777,7 +2777,7 @@ def test_reattach_on_restart_via_default_args_reaches_durable(self, dag_maker): assert k.reattach_on_restart is False def test_supports_durable_execution_marker(self): - assert KubernetesPodOperator.__supports_durable_execution is True + assert KubernetesPodOperator._KubernetesPodOperator__supports_durable_execution is True class TestSuppress: diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py index fb61086d7160c..41470724d6008 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py @@ -760,7 +760,7 @@ def test_cleanup_skipped_when_post_run_step_fails(self, mock_hook_cls, mock_buil storage.cleanup.assert_not_called() def test_supports_durable_execution_marker(self): - assert AgentOperator.__supports_durable_execution is True + assert AgentOperator._AgentOperator__supports_durable_execution is True @pytest.mark.skipif(