This is an automated email from the ASF dual-hosted git repository.
shahar1 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 b95498cd566 Restore the on_operator_resuming callback in
KubernetesPodOperator (#76035)
b95498cd566 is described below
commit b95498cd566e0a08a408f0eb1c89db5ddba520e3
Author: Henry Chen <[email protected]>
AuthorDate: Sun Nov 1 03:21:35 2026 +0800
Restore the on_operator_resuming callback in KubernetesPodOperator (#76035)
This callback is documented as supported, but it has had no call site since
the triggerer took over reading pod logs: the invocation lived inside the
`event["status"] != "running"` guard that the same change removed. Users who
override it lost the only hook that receives the raw trigger event, and
because the base implementation is a no-op there is nothing to indicate it
stopped running.
The unit test meant to cover it was left asserting `on_pod_cleanup` instead,
so nothing detected the gap.
Invoking it ahead of the re-defer branch keeps every re-entry observable,
including the `_redefer_count` retries taken when the trigger reports an
error for a pod that is still alive -- the case where the trigger event is
the only record of what went wrong.
---
providers/cncf/kubernetes/docs/operators.rst | 7 +-
.../airflow/providers/cncf/kubernetes/callbacks.py | 7 +-
.../providers/cncf/kubernetes/operators/pod.py | 13 ++++
.../unit/cncf/kubernetes/operators/test_pod.py | 84 ++++++++++++++++++----
4 files changed, 97 insertions(+), 14 deletions(-)
diff --git a/providers/cncf/kubernetes/docs/operators.rst
b/providers/cncf/kubernetes/docs/operators.rst
index 9c491a275ca..70a8996379e 100644
--- a/providers/cncf/kubernetes/docs/operators.rst
+++ b/providers/cncf/kubernetes/docs/operators.rst
@@ -375,14 +375,19 @@ parameter.
The following callbacks are supported:
* on_sync_client_creation: called after creating the sync client
+* on_pod_manifest_created: called after the pod manifest is built but before
the pod is created
* on_pod_creation: called after creating the pod
* on_pod_starting: called after the pod starts
* on_pod_completion: called when the pod completes
+* on_pod_teardown: called after the completion callbacks but before the pod is
deleted
* on_pod_cleanup: called after cleaning/deleting the pod
* on_operator_resuming: when resuming the task from deferred state
* progress_callback: called on each line of containers logs
-Currently, the callbacks methods are not called in the async mode, this
support will be added in the future.
+In deferrable mode the operator invokes ``on_pod_manifest_created``,
``on_pod_creation``,
+``on_operator_resuming``, ``on_pod_completion``, ``on_pod_teardown`` and
``on_pod_cleanup``.
+``on_pod_starting`` is not invoked, because the triggerer -- not the operator
-- observes pod
+startup. The ``mode`` argument is currently always ``sync``.
Example:
~~~~~~~~
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/callbacks.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/callbacks.py
index f78e83fabd4..fe77c8c0e03 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/callbacks.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/callbacks.py
@@ -181,7 +181,12 @@ class KubernetesPodOperatorCallback:
**kwargs,
) -> None:
"""
- Invoke this callback when resuming the `KubernetesPodOperator` from
deferred state.
+ Invoke this callback when the operator processes a trigger event.
+
+ This fires on every re-entry, including re-entries that defer again,
and
+ including the synchronous short-circuit taken when the pod already
reached
+ a terminal state before deferral. It does not fire when the pod was
garbage
+ collected before re-entry, since there is no pod to pass.
:param pod: the current state of the pod.
:param event: the returned event from the Trigger.
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 44a5d09d359..f9716532765 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
@@ -1181,6 +1181,19 @@ class KubernetesPodOperator(BaseOperator):
# keep this for any subclass override that returns None instead.
raise PodNotFoundException("Could not find pod after resuming from
deferral")
+ # Must stay ahead of the re-defer branch below so every resume is
observable,
+ # including the ones that defer again (``_redefer_count``) -- those
carry the
+ # only record of why the trigger reported an error for a still-live
pod.
+ for callback in self.callbacks:
+ callback.on_operator_resuming(
+ pod=self.pod,
+ event=event,
+ client=self.client,
+ mode=ExecutionMode.SYNC,
+ context=context,
+ operator=self,
+ )
+
follow = self.logging_interval is None
last_log_time = event.get("last_log_time")
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 3c75746df89..5f3219001b4 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
@@ -2497,24 +2497,28 @@ class TestKubernetesPodOperator:
)
context = create_context(k)
- k.trigger_reentry(
- context=context,
- event={
- "status": "success",
- "message": TEST_SUCCESS_MESSAGE,
- "name": TEST_NAME,
- "namespace": TEST_NAMESPACE,
- },
- )
+ # ``on_operator_resuming`` sees the pod re-fetched on re-entry, which
is a
+ # different object from the one ``await_pod_completion`` later returns.
+ resumed_pod_mock = MagicMock()
+ k.hook.get_pod.return_value = resumed_pod_mock
+
+ callback_event = {
+ "status": "success",
+ "message": TEST_SUCCESS_MESSAGE,
+ "name": TEST_NAME,
+ "namespace": TEST_NAMESPACE,
+ }
+ k.trigger_reentry(context=context, event=callback_event)
# check on_operator_resuming callback
- mock_callbacks.on_pod_cleanup.assert_called_once()
- assert mock_callbacks.on_pod_cleanup.call_args.kwargs == {
+ mock_callbacks.on_operator_resuming.assert_called_once()
+ assert mock_callbacks.on_operator_resuming.call_args.kwargs == {
"client": k.client,
"mode": ExecutionMode.SYNC,
- "pod": remote_pod_mock,
+ "pod": resumed_pod_mock,
"operator": k,
"context": context,
+ "event": callback_event,
}
# check on_pod_cleanup callback
@@ -3979,6 +3983,62 @@ class TestKubernetesPodOperatorAsync:
mock_invoke_defer.assert_called_once()
assert op.trigger_kwargs["_redefer_count"] == 1
+ @patch(KUB_OP_PATH.format("invoke_defer_method"))
+ @patch(HOOK_CLASS)
+ def test_on_operator_resuming_fires_before_redefer(self, mocked_hook,
mock_invoke_defer):
+ """The resuming callback fires on re-entries that defer again, not
only terminal ones."""
+ from airflow.providers.cncf.kubernetes.callbacks import ExecutionMode
+
+ from unit.cncf.kubernetes.test_callbacks import (
+ MockKubernetesPodOperatorCallback,
+ MockWrapper,
+ )
+
+ MockWrapper.reset()
+ mock_callbacks = MockWrapper.mock_callbacks
+ mock_invoke_defer.side_effect = TaskDeferred(trigger=MagicMock(),
method_name="trigger_reentry")
+
+ base_container = MagicMock()
+ base_container.name = "base"
+ base_container.state.running = MagicMock()
+ base_container.state.terminated = None
+ base_container.state.waiting = None
+
+ remote_pod = MagicMock()
+ remote_pod.status.phase = "Running"
+ remote_pod.status.container_statuses = [base_container]
+ mocked_hook.return_value.get_pod.return_value = remote_pod
+
+ op = KubernetesPodOperator(
+ task_id="test_task",
+ name="test-pod",
+ deferrable=True,
+ callbacks=MockKubernetesPodOperatorCallback,
+ )
+ context = create_context(op)
+ event = {
+ "status": "error",
+ "message": "some transient error",
+ "name": TEST_NAME,
+ "namespace": TEST_NAMESPACE,
+ }
+
+ with pytest.raises(TaskDeferred):
+ op.trigger_reentry(context, event)
+
+ mock_callbacks.on_operator_resuming.assert_called_once()
+ assert mock_callbacks.on_operator_resuming.call_args.kwargs == {
+ "client": op.client,
+ "mode": ExecutionMode.SYNC,
+ "pod": remote_pod,
+ "operator": op,
+ "context": context,
+ "event": event,
+ }
+ # The task has not finished yet, so the terminal callbacks must stay
untouched.
+ mock_callbacks.on_pod_completion.assert_not_called()
+ mock_callbacks.on_pod_cleanup.assert_not_called()
+
@patch(KUB_OP_PATH.format("invoke_defer_method"))
@patch(HOOK_CLASS)
def test_redefer_when_pod_pending(self, mocked_hook, mock_invoke_defer):
