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 34ffab8e4a4 Fix GKEStartPodOperator ignoring execution_timeout when
deferrable (#71528)
34ffab8e4a4 is described below
commit 34ffab8e4a4232470cd9d3c06847c6960529f615
Author: Ayush Khati <[email protected]>
AuthorDate: Sun Oct 18 14:18:28 2026 +0530
Fix GKEStartPodOperator ignoring execution_timeout when deferrable (#71528)
Co-authored-by: Shahar Epstein <[email protected]>
---
.../google/cloud/operators/kubernetes_engine.py | 16 ++++++
.../google/cloud/triggers/kubernetes_engine.py | 1 +
.../cloud/operators/test_kubernetes_engine.py | 61 ++++++++++++++++++++++
.../cloud/triggers/test_kubernetes_engine.py | 1 +
4 files changed, 79 insertions(+)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py
b/providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py
index 506558b1c60..72758b7a7d0 100644
---
a/providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py
+++
b/providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py
@@ -19,6 +19,7 @@
from __future__ import annotations
+import datetime
import time
import warnings
from collections.abc import Sequence
@@ -930,6 +931,19 @@ class GKEStartPodOperator(GKEOperatorMixin,
KubernetesPodOperator):
on_finish_action = self.on_finish_action
if type(on_finish_action) is str and self.on_finish_action not in
[i.value for i in OnFinishAction]:
on_finish_action = self.on_finish_action.split(".")[-1].lower() #
type: ignore[assignment]
+
+ # Anchor on ti.start_date so the deadline survives re-deferrals.
+ trigger_kwargs = dict(self.trigger_kwargs or {})
+ defer_timeout: datetime.timedelta | None = None
+ if self.execution_timeout is not None and context is not None:
+ ti_start_date = context["ti"].start_date
+ execution_deadline = int(ti_start_date.timestamp() +
self.execution_timeout.total_seconds())
+ trigger_kwargs["_execution_deadline"] = execution_deadline
+ # Allow the trigger to emit its timeout event before the framework
backstop fires.
+ remaining = execution_deadline - time.time()
+ poll_buffer = max(60, int(self.poll_interval * 2))
+ defer_timeout = datetime.timedelta(seconds=max(remaining, 0) +
poll_buffer)
+
self.defer(
trigger=GKEStartPodTrigger(
pod_name=self.pod.metadata.name, # type: ignore[union-attr]
@@ -950,8 +964,10 @@ class GKEStartPodOperator(GKEOperatorMixin,
KubernetesPodOperator):
logging_interval=self.logging_interval,
last_log_time=last_log_time,
use_dns_endpoint=self.use_dns_endpoint,
+ trigger_kwargs=trigger_kwargs,
),
method_name="trigger_reentry",
+ timeout=defer_timeout,
)
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 833a4cfe4a2..b1638b17ffb 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
@@ -157,6 +157,7 @@ class GKEStartPodTrigger(KubernetesPodTrigger):
"impersonation_chain": self.impersonation_chain,
"logging_interval": self.logging_interval,
"last_log_time": self.last_log_time,
+ "trigger_kwargs": self.trigger_kwargs,
},
)
diff --git
a/providers/google/tests/unit/google/cloud/operators/test_kubernetes_engine.py
b/providers/google/tests/unit/google/cloud/operators/test_kubernetes_engine.py
index eb06480c971..857fee86253 100644
---
a/providers/google/tests/unit/google/cloud/operators/test_kubernetes_engine.py
+++
b/providers/google/tests/unit/google/cloud/operators/test_kubernetes_engine.py
@@ -17,12 +17,14 @@
# under the License.
from __future__ import annotations
+import datetime
import runpy
from copy import deepcopy
from unittest import mock
from unittest.mock import PropertyMock, call
import pytest
+import time_machine
from google.api_core.exceptions import AlreadyExists, FailedPrecondition,
PermissionDenied
from google.cloud.container_v1.types import Cluster, NodePool
from kubernetes.client import models as k8s
@@ -1272,10 +1274,69 @@ class TestGKEStartPodOperator:
logging_interval=None,
last_log_time=mock_last_log_time,
use_dns_endpoint=False,
+ trigger_kwargs={},
)
mock_defer.assert_called_once_with(
trigger=mock_trigger.return_value,
method_name="trigger_reentry",
+ timeout=None,
+ )
+
+ @pytest.mark.parametrize(
+ (
+ "execution_timeout",
+ "elapsed",
+ "poll_interval",
+ "trigger_kwargs",
+ "expected_timeout",
+ "pass_context",
+ ),
+ [
+ pytest.param(None, 30, 2, {"_redefer_count": 1}, None, True,
id="no-execution-timeout"),
+ pytest.param(300, 30, 2, None, 270 + 60, True,
id="remaining-plus-min-buffer"),
+ pytest.param(300, 30, 45, None, 270 + 90, True,
id="buffer-scales-with-poll-interval"),
+ pytest.param(300, 600, 2, None, 60, True,
id="deadline-passed-clamps-remaining"),
+ pytest.param(120, 0, 2, {"_redefer_count": 1}, 120 + 60, True,
id="keeps-user-trigger-kwargs"),
+ pytest.param(300, 30, 2, None, None, False,
id="no-context-skips-deadline"),
+ ],
+ )
+ @mock.patch(GKE_OPERATORS_PATH.format("GKEStartPodOperator.defer"))
+
@mock.patch(GKE_OPERATORS_PATH.format("GKEClusterAuthDetails.fetch_cluster_info"))
+ @mock.patch(GKE_OPERATORS_PATH.format("GKEHook"))
+ @mock.patch(GKE_OPERATORS_PATH.format("GKEStartPodTrigger"))
+ def test_invoke_defer_method_execution_timeout(
+ self,
+ mock_trigger,
+ mock_cluster_hook,
+ mock_fetch_cluster_info,
+ mock_defer,
+ execution_timeout,
+ elapsed,
+ poll_interval,
+ trigger_kwargs,
+ expected_timeout,
+ pass_context,
+ ):
+ ti_start = datetime.datetime(2026, 1, 1, 12, 0, 0, tzinfo=datetime.UTC)
+ self.operator.pod = make_pod()
+ self.operator.poll_interval = poll_interval
+ self.operator.trigger_kwargs = deepcopy(trigger_kwargs)
+ if execution_timeout is not None:
+ self.operator.execution_timeout =
datetime.timedelta(seconds=execution_timeout)
+ mock_fetch_cluster_info.return_value = GKE_CLUSTER_URL, GKE_SSL_CA_CERT
+ ti = mock.MagicMock(spec=["start_date"])
+ ti.start_date = ti_start
+
+ with time_machine.travel(ti_start +
datetime.timedelta(seconds=elapsed), tick=False):
+ self.operator.invoke_defer_method(context={"ti": ti} if
pass_context else None)
+
+ expected_trigger_kwargs = dict(trigger_kwargs or {})
+ if execution_timeout is not None and pass_context:
+ expected_trigger_kwargs["_execution_deadline"] =
int(ti_start.timestamp()) + execution_timeout
+ assert mock_trigger.call_args.kwargs["trigger_kwargs"] ==
expected_trigger_kwargs
+ assert self.operator.trigger_kwargs == trigger_kwargs
+ assert mock_defer.call_args.kwargs["timeout"] == (
+ None if expected_timeout is None else
datetime.timedelta(seconds=expected_timeout)
)
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 a577dccf0db..fbcd0f5c9de 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
@@ -144,6 +144,7 @@ class TestGKEStartPodTrigger:
"last_log_time": None,
"logging_interval": None,
"use_dns_endpoint": False,
+ "trigger_kwargs": {},
}
def test_serialize_round_trip_preserves_on_finish_action(self):
