This is an automated email from the ASF dual-hosted git repository.
potiuk 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 74d20c9b09a Rediscover succeeded pods in GKEJobTrigger to handle
Kubernetes Job retries with XCom (#63915)
74d20c9b09a is described below
commit 74d20c9b09a5bc265a7fa55579e095accd2941b5
Author: Pranay Kumar Karvi <[email protected]>
AuthorDate: Tue Apr 28 05:02:15 2026 +0530
Rediscover succeeded pods in GKEJobTrigger to handle Kubernetes Job retries
with XCom (#63915)
* fix: re-discover succeeded pods for GKEJobTrigger XCom extraction to
handle job retries
* fix: apply ruff formatting to list comprehension
* fix: apply review feedback to GKEJobTrigger xcom retry test
* fix: use dict[str, Any] for kwargs in list_namespaced_event call
* fix: restore accidentally removed blank line in get_pod_events
* fix: poll for succeeded pods in GKEJobTrigger instead of one-shot fallback
---
.../google/cloud/hooks/kubernetes_engine.py | 20 +++++++
.../google/cloud/triggers/kubernetes_engine.py | 32 ++++++++++-
.../cloud/triggers/test_kubernetes_engine.py | 64 ++++++++++++++++++++++
3 files changed, 114 insertions(+), 2 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py
b/providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py
index e9e17f1b1bf..9ec2052c45d 100644
---
a/providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py
+++
b/providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py
@@ -507,6 +507,26 @@ class GKEKubernetesAsyncHook(GoogleBaseAsyncHook,
AsyncKubernetesHook):
if kube_client is not None:
await kube_client.close()
+ async def list_pods(
+ self,
+ namespace: str,
+ label_selector: str,
+ ) -> list:
+ """
+ List pods in the given namespace matching the label selector.
+
+ :param namespace: Kubernetes namespace.
+ :param label_selector: Label selector to filter pods (e.g.
``job-name=my-job``).
+ :return: List of V1Pod objects.
+ """
+ async with self.get_conn() as connection:
+ v1_api = async_client.CoreV1Api(connection)
+ response = await v1_api.list_namespaced_pod(
+ namespace=namespace,
+ label_selector=label_selector,
+ )
+ return list(response.items) if response.items else []
+
async def _load_config(self) -> async_client.ApiClient:
configuration = self._get_config()
token = await self.get_token()
diff --git
a/providers/google/src/airflow/providers/google/cloud/triggers/kubernetes_engine.py
b/providers/google/src/airflow/providers/google/cloud/triggers/kubernetes_engine.py
index 6867b96b747..6da2c5ea2ad 100644
---
a/providers/google/src/airflow/providers/google/cloud/triggers/kubernetes_engine.py
+++
b/providers/google/src/airflow/providers/google/cloud/triggers/kubernetes_engine.py
@@ -18,6 +18,7 @@
from __future__ import annotations
import asyncio
+import time
import warnings
from collections.abc import AsyncIterator, Sequence
from functools import cached_property
@@ -333,9 +334,36 @@ class GKEJobTrigger(BaseTrigger):
f"package
{kubernetes_provider_name}=={kubernetes_provider_version} which doesn't "
f"support this feature. Please upgrade it to version
higher than or equal to {min_version}."
)
+ # Re-discover pods for this job to handle retry scenarios
+ # where the original pod failed and a new pod was created
+ label_selector = f"job-name={self.job_name}"
+ deadline = time.monotonic() + self.poll_interval * 10
+ succeeded_pods = []
+ current_pods = []
+ while time.monotonic() < deadline:
+ current_pods = await self.hook.list_pods(
+ namespace=self.pod_namespace,
+ label_selector=label_selector,
+ )
+ succeeded_pods = [p for p in current_pods if p.status and
p.status.phase == "Succeeded"]
+ if succeeded_pods:
+ break
+ self.log.info(
+ "No succeeded pods found yet for job %s, retrying in %s
seconds...",
+ self.job_name,
+ self.poll_interval,
+ )
+ await asyncio.sleep(self.poll_interval)
+ if not succeeded_pods:
+ self.log.warning(
+ "No succeeded pods found for job %s after timeout, falling
back to original pod list",
+ self.job_name,
+ )
+ succeeded_pods = [p for p in current_pods if p.metadata.name
in self.pod_names]
+
xcom_results = []
- for pod_name in self.pod_names:
- pod = await self.hook.get_pod(name=pod_name,
namespace=self.pod_namespace)
+ for pod in succeeded_pods:
+ pod_name = pod.metadata.name
await self.hook.wait_until_container_complete(
name=pod_name,
namespace=self.pod_namespace,
diff --git
a/providers/google/tests/unit/google/cloud/triggers/test_kubernetes_engine.py
b/providers/google/tests/unit/google/cloud/triggers/test_kubernetes_engine.py
index 506d2f8191c..eee26d9fa7b 100644
---
a/providers/google/tests/unit/google/cloud/triggers/test_kubernetes_engine.py
+++
b/providers/google/tests/unit/google/cloud/triggers/test_kubernetes_engine.py
@@ -34,6 +34,7 @@ except ImportError:
from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import ( #
type: ignore[no-redef]
ContainerState,
)
+from airflow.providers.cncf.kubernetes.utils.xcom_sidecar import PodDefaults
from airflow.providers.google.cloud.triggers.kubernetes_engine import (
GKEJobTrigger,
GKEOperationTrigger,
@@ -589,3 +590,66 @@ class TestGKEStartJobTrigger:
impersonation_chain=IMPERSONATION_CHAIN,
)
assert hook_actual == hook_expected
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{GKE_TRIGGERS_PATH}.ProvidersManager")
+ @mock.patch(f"{TRIGGER_GKE_JOB_PATH}.pod_manager",
new_callable=mock.PropertyMock)
+ @mock.patch(f"{TRIGGER_GKE_JOB_PATH}.hook")
+ async def
test_run_do_xcom_push_uses_succeeded_retry_pod_not_original_failed_pod(
+ self, mock_hook, mock_pod_manager_prop, mock_providers_manager,
job_trigger
+ ):
+ """When a job has a failed pod and a succeeded retry pod, use the
succeeded pod for XCom."""
+ mock_providers_manager.return_value.providers = {
+ "apache-airflow-providers-cncf-kubernetes": mock.MagicMock(
+ data={"package-name":
"apache-airflow-providers-cncf-kubernetes"},
+ version="8.4.1",
+ )
+ }
+
+ failed_pod = mock.MagicMock()
+ failed_pod.metadata.name = "original-failed-pod"
+ failed_pod.status.phase = "Failed"
+
+ succeeded_pod = mock.MagicMock()
+ succeeded_pod.metadata.name = "retry-succeeded-pod"
+ succeeded_pod.status.phase = "Succeeded"
+
+ mock_hook.list_pods = mock.AsyncMock(return_value=[failed_pod,
succeeded_pod])
+
+ mock_hook.wait_until_container_complete = mock.AsyncMock()
+ mock_hook.wait_until_container_started = mock.AsyncMock()
+ mock_hook.wait_until_job_complete = mock.AsyncMock()
+
+ mock_job = mock.MagicMock()
+ mock_job.metadata.name = JOB_NAME
+ mock_job.metadata.namespace = NAMESPACE
+ mock_job.to_dict.return_value = {"job": "dict"}
+
+ mock_hook.wait_until_job_complete.return_value = mock_job
+ mock_hook.is_job_failed = mock.MagicMock(return_value=False)
+
+ mock_pod_manager = mock.MagicMock()
+ mock_pod_manager.extract_xcom.return_value = {"xcom": "result"}
+ mock_pod_manager_prop.return_value = mock_pod_manager
+
+ job_trigger.do_xcom_push = True
+ event = await job_trigger.run().asend(None)
+
+ mock_hook.list_pods.assert_awaited_once_with(
+ namespace=NAMESPACE,
+ label_selector=f"job-name={JOB_NAME}",
+ )
+ mock_pod_manager.extract_xcom.assert_called_once_with(succeeded_pod)
+ assert event.payload["xcom_result"] == [{"xcom": "result"}]
+ mock_hook.wait_until_container_complete.assert_called_once_with(
+ name="retry-succeeded-pod",
+ namespace=NAMESPACE,
+ container_name=BASE_CONTAINER_NAME,
+ poll_interval=POLL_INTERVAL,
+ )
+ mock_hook.wait_until_container_started.assert_called_once_with(
+ name="retry-succeeded-pod",
+ namespace=NAMESPACE,
+ container_name=PodDefaults.SIDECAR_CONTAINER_NAME,
+ poll_interval=POLL_INTERVAL,
+ )
