This is an automated email from the ASF dual-hosted git repository.
jscheffl 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 da87130071d Consider XCOM sidecar container during pod cleanup (#64962)
da87130071d is described below
commit da87130071d54fb7a88fdeff492f5dbf0534c16a
Author: Daniel Wolf <[email protected]>
AuthorDate: Wed Apr 15 21:22:04 2026 +0200
Consider XCOM sidecar container during pod cleanup (#64962)
* Consider XCOM sidecar container during pod cleanup
* Fix failing unit tests
* Apply Copilot review findings
* Implement Copilot review finding regarding failed status
* Fix some more unit tests
* More fixes for unit tests
---
.../test_kubernetes_pod_operator.py | 5 ++
.../providers/cncf/kubernetes/operators/pod.py | 11 ++--
.../providers/cncf/kubernetes/utils/pod_manager.py | 18 +++++-
.../decorators/test_kubernetes_commons.py | 11 +++-
.../unit/cncf/kubernetes/operators/test_pod.py | 75 ++++++++++++++++++++++
.../unit/cncf/kubernetes/utils/test_pod_manager.py | 30 +++++++++
6 files changed, 141 insertions(+), 9 deletions(-)
diff --git
a/kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py
b/kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py
index 64cc3dc0757..652ee0ad954 100644
--- a/kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py
+++ b/kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py
@@ -965,6 +965,11 @@ class TestKubernetesPodOperatorSystem:
)
pod_mock = MagicMock()
pod_mock.status.phase = "Succeeded"
+ base_container_status = MagicMock()
+ base_container_status.name = "base"
+ base_container_status.state.terminated.exit_code = 0
+ pod_mock.status.container_statuses = [base_container_status]
+ pod_mock.status.init_container_statuses = None
await_pod_completion_mock.return_value = pod_mock
context = create_context(k)
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 67494995f0a..eb290a881f3 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
@@ -759,7 +759,7 @@ class KubernetesPodOperator(BaseOperator):
result = self.extract_xcom(pod=self.pod)
istio_enabled = self.is_istio_enabled(self.pod)
self.remote_pod = self.pod_manager.await_pod_completion(
- self.pod, istio_enabled, self.base_container_name
+ self.pod, istio_enabled, self.base_container_name,
self.do_xcom_push
)
finally:
pod_to_clean = self.pod or self.pod_request_obj
@@ -1025,7 +1025,7 @@ class KubernetesPodOperator(BaseOperator):
if event["status"] != "timeout":
try:
self.pod = self.pod_manager.await_pod_completion(
- self.pod, istio_enabled, self.base_container_name
+ self.pod, istio_enabled, self.base_container_name,
self.do_xcom_push
)
except ApiException as e:
if e.status == 404:
@@ -1110,9 +1110,10 @@ class KubernetesPodOperator(BaseOperator):
):
self.patch_already_checked(remote_pod, reraise=False)
- failed = (pod_phase != PodPhase.SUCCEEDED and not istio_enabled) or (
- istio_enabled and not container_is_succeeded(remote_pod,
self.base_container_name)
- )
+ if istio_enabled or self.do_xcom_push:
+ failed = not container_is_succeeded(remote_pod,
self.base_container_name)
+ else:
+ failed = pod_phase != PodPhase.SUCCEEDED
if failed:
if self.do_xcom_push and xcom_result and context:
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py
index 5b95a4c7fea..3cc0229ea61 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -741,7 +741,11 @@ class PodManager(LoggingMixin):
time.sleep(polling_time)
def await_pod_completion(
- self, pod: V1Pod, istio_enabled: bool = False, container_name: str =
"base"
+ self,
+ pod: V1Pod,
+ istio_enabled: bool = False,
+ container_name: str = "base",
+ do_xcom_push: bool = False,
) -> V1Pod:
"""
Monitor a pod and return the final state.
@@ -749,13 +753,21 @@ class PodManager(LoggingMixin):
:param istio_enabled: whether istio is enabled in the namespace
:param pod: pod spec that will be monitored
:param container_name: name of the container within the pod
- :return: tuple[State, str | None]
+ :param do_xcom_push: whether to push XComs
+ :return: V1Pod
"""
while True:
remote_pod = self.read_pod(pod)
if remote_pod.status.phase in PodPhase.terminal_states:
break
- if istio_enabled and container_is_completed(remote_pod,
container_name):
+ if (istio_enabled or do_xcom_push) and
container_is_completed(remote_pod, container_name):
+ self.log.info(
+ "Container '%s' completed but pod %s still has phase %s "
+ "(likely due to a sidecar container). Skipping waiting for
pod completion.",
+ container_name,
+ pod.metadata.name,
+ remote_pod.status.phase,
+ )
break
# abort waiting if defined issues are detected
if detect_pod_terminate_early_issues(remote_pod):
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 b062ab3a310..29f261b6989 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
@@ -99,7 +99,16 @@ class TestKubernetesDecoratorsBase:
self.mock_extract_xcom.return_value = '{"key1": "value1", "key2":
"value2"}'
self.mock_await_pod_completion =
mock.patch(f"{POD_MANAGER_CLASS}.await_pod_completion").start()
- self.mock_await_pod_completion.return_value =
mock.MagicMock(**{"status.phase": "Succeeded"})
+ base_container_status = mock.MagicMock()
+ base_container_status.name = "base"
+ base_container_status.state.terminated.exit_code = 0
+ self.mock_await_pod_completion.return_value = mock.MagicMock(
+ **{
+ "status.phase": "Succeeded",
+ "status.container_statuses": [base_container_status],
+ "status.init_container_statuses": None,
+ }
+ )
self.mock_hook = mock.patch(HOOK_CLASS).start()
# Without this patch each time pod manager would try to extract logs
from 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 412e0e08a63..5a28eb2ff6d 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
@@ -266,6 +266,11 @@ class TestKubernetesPodOperator:
remote_pod_mock = MagicMock()
remote_pod_mock.status.phase = "Succeeded"
+ base_container_status = MagicMock()
+ base_container_status.name = operator.base_container_name
+ base_container_status.state.terminated.exit_code = 0
+ remote_pod_mock.status.container_statuses = [base_container_status]
+ remote_pod_mock.status.init_container_statuses = None
self.await_pod_mock.return_value = remote_pod_mock
operator.execute(context=context)
return self.await_start_mock.call_args.kwargs["pod"], context
@@ -1055,6 +1060,71 @@ class TestKubernetesPodOperator:
else:
delete_pod_mock.assert_not_called()
+ @pytest.mark.parametrize(
+ ("base_container_exit_code", "expect_failure"),
+ [
+ pytest.param(0, False, id="base-succeeded"),
+ pytest.param(1, True, id="base-failed"),
+ ],
+ )
+ @patch(f"{POD_MANAGER_CLASS}.extract_xcom")
+ @patch(f"{POD_MANAGER_CLASS}.await_xcom_sidecar_container_start")
+ @patch(f"{POD_MANAGER_CLASS}.delete_pod")
+ @patch(f"{POD_MANAGER_CLASS}.await_pod_completion")
+ @patch(f"{KPO_MODULE}.KubernetesPodOperator.find_pod")
+ def test_cleanup_with_xcom_sidecar_uses_base_container_status(
+ self,
+ find_pod_mock,
+ await_pod_completion_mock,
+ delete_pod_mock,
+ mock_await_xcom_sidecar,
+ mock_extract_xcom,
+ base_container_exit_code,
+ expect_failure,
+ ):
+ """
+ When do_xcom_push=True, cleanup should determine success/failure based
on
+ the base container's exit status, not the pod phase. The xcom sidecar
may
+ keep the pod in Running phase after the base container completes.
+ """
+ mock_extract_xcom.return_value = "{}"
+ mock_await_xcom_sidecar.return_value = None
+
+ base_status = MagicMock()
+ base_status.name = "base"
+ base_status.state.terminated.exit_code = base_container_exit_code
+ base_status.state.terminated.message = "task failed" if
base_container_exit_code else None
+
+ xcom_sidecar_status = MagicMock()
+ xcom_sidecar_status.name = "airflow-xcom-sidecar"
+ xcom_sidecar_status.state.running = True
+ xcom_sidecar_status.state.terminated = None
+
+ # Pod is still Running because xcom sidecar is alive
+ remote_pod = MagicMock()
+ remote_pod.status.phase = "Running"
+ remote_pod.status.container_statuses = [base_status,
xcom_sidecar_status]
+ remote_pod.metadata.name = "pod-with-xcom-sidecar"
+ remote_pod.metadata.namespace = "default"
+ remote_pod.spec.containers = [MagicMock(), MagicMock()]
+
+ await_pod_completion_mock.return_value = remote_pod
+ find_pod_mock.return_value = remote_pod
+
+ k = KubernetesPodOperator(
+ task_id="task",
+ do_xcom_push=True,
+ )
+
+ context = create_context(k)
+ context["ti"].xcom_push = MagicMock()
+
+ if expect_failure:
+ with pytest.raises(AirflowException, match="task failed"):
+ k.execute(context=context)
+ else:
+ k.execute(context=context)
+
@pytest.mark.parametrize(
("task_kwargs", "should_be_deleted"),
[
@@ -1883,6 +1953,11 @@ class TestKubernetesPodOperator:
mock_extract_xcom.return_value = "{}"
remote_pod_mock = MagicMock()
remote_pod_mock.status.phase = "Succeeded"
+ base_container_status = MagicMock()
+ base_container_status.name = "base"
+ base_container_status.state.terminated.exit_code = 0
+ remote_pod_mock.status.container_statuses = [base_container_status]
+ remote_pod_mock.status.init_container_statuses = None
self.await_pod_mock.return_value = remote_pod_mock
pod, _ = self.run_pod(k)
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py
index b5c57c49167..c48a6608992 100644
---
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py
+++
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py
@@ -1079,6 +1079,36 @@ class TestPodManager:
assert result is running2
assert mock_sleep.call_count == 1
+ @mock.patch("time.sleep")
+ def
test_await_pod_completion_breaks_on_xcom_sidecar_container_completed(self,
mock_sleep, pod_factory):
+ """When do_xcom_push=True and base container has completed, stop
waiting even if pod is still Running."""
+ running1 = pod_factory(pod_phase=PodPhase.RUNNING,
container_name="base", terminated=False)
+ running2 = pod_factory(pod_phase=PodPhase.RUNNING,
container_name="base", terminated=True)
+
+ self.pod_manager.read_pod = mock.MagicMock(side_effect=[running1,
running2])
+
+ result = self.pod_manager.await_pod_completion(
+ pod=mock.MagicMock(), istio_enabled=False, container_name="base",
do_xcom_push=True
+ )
+
+ assert result is running2
+ assert mock_sleep.call_count == 1
+
+ @mock.patch("time.sleep")
+ def test_await_pod_completion_waits_for_pod_phase_without_sidecars(self,
mock_sleep, pod_factory):
+ """Without istio or xcom sidecar, await_pod_completion waits for
terminal pod phase."""
+ running1 = pod_factory(pod_phase=PodPhase.RUNNING,
container_name="base", terminated=True)
+ succeeded = pod_factory(pod_phase=PodPhase.SUCCEEDED,
container_name="base", terminated=True)
+
+ self.pod_manager.read_pod = mock.MagicMock(side_effect=[running1,
succeeded])
+
+ result = self.pod_manager.await_pod_completion(
+ pod=mock.MagicMock(), istio_enabled=False, container_name="base",
do_xcom_push=False
+ )
+
+ assert result is succeeded
+ assert mock_sleep.call_count == 1
+
@mock.patch("time.sleep")
def test_await_pod_completion_breaks_on_early_termination_issue(self,
mock_sleep, pod_factory):
running1 = pod_factory(pod_phase=PodPhase.PENDING,
container_name="base")
