Consider XCOM sidecar container during pod cleanup

kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py

@@ -965,6 +965,11 @@ def test_pod_template_file(
         )
         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)
 

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py

@@ -759,7 +759,7 @@ def execute_sync(self, context: Context):
                 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 @@ def _clean(self, event: dict[str, Any], result: dict | None, context: Context) -
         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 @@ def cleanup(
         ):
             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:

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py

@@ -741,21 +741,33 @@ def await_container_completion(self, pod: V1Pod, container_name: str, polling_ti
             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.
 
         :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):

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_commons.py

@@ -99,7 +99,16 @@ def setup(self, dag_maker):
         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

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py

@@ -266,6 +266,11 @@ def run_pod(self, operator: KubernetesPodOperator, map_index: int = -1) -> tuple
 
         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 @@ def test_pod_with_istio_delete_after_await_container_error(
         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 @@ def test_get_logs_but_not_for_base_container(
         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)
 

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py

@@ -1079,6 +1079,36 @@ def test_await_pod_completion_breaks_on_istio_container_completed(self, mock_sle
         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")