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 841d51e3636 Enforce execution_timeout in deferrable
DataprocSubmitJobOperator (#72030)
841d51e3636 is described below
commit 841d51e36367eef41ef4b6096a100dc34062ad6d
Author: Noel <[email protected]>
AuthorDate: Sat Oct 31 13:02:54 2026 -0300
Enforce execution_timeout in deferrable DataprocSubmitJobOperator (#72030)
* Enforce execution_timeout for deferred Dataproc submit jobs
DataprocSubmitJobOperator ignored execution_timeout once it deferred,
so a stuck job could run forever. The operator now derives an
absolute deadline from ti.start_date + execution_timeout, passes it
to DataprocSubmitTrigger as execution_deadline, and also passes
timeout= to defer() as a framework-level backstop.
DataprocSubmitTrigger cancels the Dataproc job once the deadline
passes (unless cancel_on_kill is False) and emits a TIMED_OUT event;
execute_complete raises AirflowException on that event.
Add tests for both the trigger and the operator, and record the new
AirflowException raise in the known-exceptions allowlist.
* Address review: KPO-style deadline handling for Dataproc trigger
Check the execution deadline before polling the job, so a deferral that
wakes past the deadline cancels the job without first waiting on
get_job() and its retries. Pad the framework defer timeout by
max(60, 2 * polling_interval_seconds), matching KubernetesPodOperator,
so the trigger's own timeout handling wins the race.
Raise AirflowTaskTimeout instead of AirflowException on the TIMED_OUT
event, and freeze time in the operator test to assert the exact defer
timeout.
* Bound Dataproc trigger polling by the execution deadline
A get_job() call started just before the deadline could keep retrying
past it, and the polling sleep could overshoot it by a full interval,
letting the framework defer timeout fire first and orphan the job.
Wrap get_job() in asyncio.wait_for() with the remaining time and cap
the sleep at the deadline.
---------
Co-authored-by: Noel <[email protected]>
Co-authored-by: Shahar Epstein <[email protected]>
---
.../providers/google/cloud/operators/dataproc.py | 24 +++++-
.../providers/google/cloud/triggers/dataproc.py | 39 ++++++++-
.../unit/google/cloud/operators/test_dataproc.py | 45 ++++++++++
.../unit/google/cloud/triggers/test_dataproc.py | 95 ++++++++++++++++++++++
4 files changed, 198 insertions(+), 5 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py
b/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py
index 88059d97f45..3f2e0e06be3 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py
@@ -36,7 +36,7 @@ from google.api_core.retry import Retry,
exponential_sleep_generator
from google.cloud.dataproc_v1 import Batch, Cluster, ClusterStatus, JobStatus
from airflow.exceptions import AirflowProviderDeprecationWarning
-from airflow.providers.common.compat.sdk import AirflowException, conf,
timezone
+from airflow.providers.common.compat.sdk import AirflowException,
AirflowTaskTimeout, conf, timezone
from airflow.providers.google.cloud.hooks.dataproc import (
DataprocHook,
DataProcJobBuilder,
@@ -1919,7 +1919,9 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
:param asynchronous: Flag to return after submitting the job to the
Dataproc API.
This is useful for submitting long-running jobs and
waiting on them asynchronously using the DataprocJobSensor
- :param deferrable: Run operator in the deferrable mode
+ :param deferrable: Run operator in the deferrable mode.
``execution_timeout`` is enforced
+ while the task is deferred: on expiry the Dataproc job is cancelled
(unless
+ ``cancel_on_kill`` is False) and the task fails.
:param polling_interval_seconds: time in seconds between polling for job
completion.
The value is considered only when running in deferrable mode. Must be
greater than 0.
:param cancel_on_kill: Flag which indicates whether cancel the hook's job
or not, when on_kill is called
@@ -2054,6 +2056,18 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
raise AirflowException(f"Job failed:\n{job}")
if state == JobStatus.State.CANCELLED:
raise AirflowException(f"Job was cancelled:\n{job}")
+ execution_deadline = None
+ trigger_timeout = None
+ if self.execution_timeout:
+ # Anchor on ti.start_date so the deadline stays stable across
+ # re-deferrals, matching non-deferrable execution_timeout
semantics.
+ start_date = context["ti"].start_date
+ anchor = start_date.timestamp() if start_date else time.time()
+ execution_deadline = anchor +
self.execution_timeout.total_seconds()
+ # Framework backstop only: slack of at least a minute (two
poll intervals) so the
+ # trigger's own deadline handling (which cancels the job)
fires first.
+ poll_buffer = max(60, 2 * self.polling_interval_seconds)
+ trigger_timeout = timedelta(seconds=max(execution_deadline -
time.time(), 0) + poll_buffer)
self.defer(
trigger=DataprocSubmitTrigger(
job_id=self.job_id,
@@ -2063,8 +2077,10 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
impersonation_chain=self.impersonation_chain,
polling_interval_seconds=self.polling_interval_seconds,
cancel_on_kill=self.cancel_on_kill,
+ execution_deadline=execution_deadline,
),
method_name="execute_complete",
+ timeout=trigger_timeout,
)
elif not self.asynchronous:
self.log.info("Waiting for job %s to complete", new_job_id)
@@ -2083,6 +2099,10 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
"""
job_state = event["job_state"]
job_id = event["job_id"]
+ if job_state == "TIMED_OUT":
+ raise AirflowTaskTimeout(
+ event.get("message") or f"Job {job_id} exceeded the task's
execution_timeout while deferred."
+ )
job = event["job"]
if job_state == JobStatus.State.ERROR.name: # type: ignore
raise AirflowException(f"Job {job_id} failed:\n{job}")
diff --git
a/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
b/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
index 665e336c0b0..cd877516028 100644
--- a/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
+++ b/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
@@ -104,10 +104,14 @@ class DataprocSubmitTrigger(DataprocBaseTrigger):
Service Account Token Creator IAM role to the directly preceding
identity, with first
account from the list granting this role to the originating account
(templated).
:param polling_interval_seconds: polling period in seconds to check for
the status
+ :param execution_deadline: Optional absolute timestamp (in seconds since
the epoch) after
+ which the task is considered timed out. On expiry the trigger cancels
the Dataproc job
+ (unless ``cancel_on_kill`` is False) and emits a ``TIMED_OUT`` event.
"""
- def __init__(self, job_id: str, **kwargs):
+ def __init__(self, job_id: str, execution_deadline: float | None = None,
**kwargs):
self.job_id = job_id
+ self.execution_deadline = execution_deadline
super().__init__(**kwargs)
def serialize(self):
@@ -121,6 +125,7 @@ class DataprocSubmitTrigger(DataprocBaseTrigger):
"impersonation_chain": self.impersonation_chain,
"polling_interval_seconds": self.polling_interval_seconds,
"cancel_on_kill": self.cancel_on_kill,
+ "execution_deadline": self.execution_deadline,
},
)
@@ -205,12 +210,40 @@ class DataprocSubmitTrigger(DataprocBaseTrigger):
try:
while True:
- job = await hook.get_job(project_id=self.project_id,
region=self.region, job_id=self.job_id)
+ remaining = None if self.execution_deadline is None else
self.execution_deadline - time.time()
+ if remaining is not None and remaining <= 0:
+ if self.cancel_on_kill:
+ self.log.info("Job %s exceeded execution_timeout;
cancelling it.", self.job_id)
+ try:
+ await
sync_to_async(self.get_sync_hook().cancel_job)(
+ job_id=self.job_id,
project_id=self.project_id, region=self.region
+ )
+ except Exception:
+ self.log.exception("Failed to cancel job %s after
execution timeout", self.job_id)
+ yield TriggerEvent(
+ {
+ "job_id": self.job_id,
+ "job_state": "TIMED_OUT",
+ "message": f"Job {self.job_id} exceeded the task's
execution_timeout "
+ "while deferred.",
+ }
+ )
+ return
+ try:
+ job = await asyncio.wait_for(
+ hook.get_job(project_id=self.project_id,
region=self.region, job_id=self.job_id),
+ timeout=remaining,
+ )
+ except TimeoutError:
+ continue
state = job.status.state
self.log.info("Dataproc job: %s is in state: %s", self.job_id,
state)
if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED,
JobStatus.State.ERROR):
break
- await asyncio.sleep(self.polling_interval_seconds)
+ sleep_seconds = self.polling_interval_seconds
+ if self.execution_deadline is not None:
+ sleep_seconds = min(sleep_seconds,
max(self.execution_deadline - time.time(), 0))
+ await asyncio.sleep(sleep_seconds)
yield TriggerEvent(
{"job_id": self.job_id, "job_state":
JobStatus.State(state).name, "job": Job.to_dict(job)}
)
diff --git
a/providers/google/tests/unit/google/cloud/operators/test_dataproc.py
b/providers/google/tests/unit/google/cloud/operators/test_dataproc.py
index 29facc7c967..3b465fd29fb 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_dataproc.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_dataproc.py
@@ -23,6 +23,7 @@ from unittest import mock
from unittest.mock import ANY, MagicMock, Mock, call
import pytest
+import time_machine
from google.api_core.exceptions import AlreadyExists, NotFound
from google.api_core.retry import Retry
from google.api_core.retry_async import AsyncRetry
@@ -1872,6 +1873,50 @@ class TestDataprocSubmitJobOperator(DataprocJobTestBase):
assert isinstance(exc.value.trigger, DataprocSubmitTrigger)
assert exc.value.method_name == GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
+ @time_machine.travel(dt.datetime(2026, 1, 1, 0, 10, tzinfo=dt.UTC),
tick=False)
+ @mock.patch(DATAPROC_PATH.format("DataprocHook"))
+ @mock.patch(DATAPROC_TRIGGERS_PATH.format("DataprocAsyncHook"))
+ def test_execute_deferrable_with_execution_timeout(self,
mock_trigger_hook, mock_hook):
+ mock_hook.return_value.submit_job.return_value.reference.job_id =
TEST_JOB_ID
+
+ op = DataprocSubmitJobOperator(
+ task_id=TASK_ID,
+ region=GCP_REGION,
+ project_id=GCP_PROJECT,
+ job={},
+ gcp_conn_id=GCP_CONN_ID,
+ deferrable=True,
+ execution_timeout=dt.timedelta(hours=1),
+ )
+ start_date = dt.datetime(2026, 1, 1, tzinfo=dt.UTC)
+ self.mock_ti.start_date = start_date
+
+ with pytest.raises(TaskDeferred) as exc:
+ op.execute(self.mock_context)
+
+ expected_deadline = start_date.timestamp() + 3600
+ assert exc.value.trigger.execution_deadline == expected_deadline
+ # 50 minutes remain, plus the 60 second minimum buffer.
+ assert exc.value.timeout == dt.timedelta(minutes=50, seconds=60)
+
+ @mock.patch(DATAPROC_PATH.format("DataprocHook"))
+ def test_execute_complete_timed_out_raises(self, mock_hook):
+ op = DataprocSubmitJobOperator(
+ task_id=TASK_ID,
+ region=GCP_REGION,
+ project_id=GCP_PROJECT,
+ job={},
+ gcp_conn_id=GCP_CONN_ID,
+ deferrable=True,
+ )
+ event = {
+ "job_id": TEST_JOB_ID,
+ "job_state": "TIMED_OUT",
+ "message": f"Job {TEST_JOB_ID} exceeded the task's
execution_timeout while deferred.",
+ }
+ with pytest.raises(AirflowTaskTimeout, match="execution_timeout"):
+ op.execute_complete(context=self.mock_context, event=event)
+
@mock.patch(DATAPROC_PATH.format("DataprocHook"))
@mock.patch(DATAPROC_PATH.format("DataprocSubmitJobOperator.defer"))
@mock.patch(DATAPROC_PATH.format("DataprocHook.submit_job"))
diff --git a/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
b/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
index d469cdf478c..946e455d2f6 100644
--- a/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
+++ b/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
@@ -18,6 +18,7 @@ from __future__ import annotations
import asyncio
import logging
+import time
from asyncio import CancelledError, Future, sleep
from types import SimpleNamespace
from unittest import mock
@@ -726,6 +727,7 @@ class TestDataprocSubmitTrigger:
"polling_interval_seconds": TEST_POLL_INTERVAL,
"cancel_on_kill": True,
"impersonation_chain": None,
+ "execution_deadline": None,
}
@pytest.mark.asyncio
@@ -761,6 +763,99 @@ class TestDataprocSubmitTrigger:
)
assert event.payload == expected_event.payload
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitTrigger.get_sync_hook")
+
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitTrigger.get_async_hook")
+ async def test_submit_trigger_run_execution_deadline_cancels_job(
+ self, mock_get_async_hook, mock_get_sync_hook
+ ):
+ """Test the trigger cancels a still-running job once the execution
deadline passes."""
+ trigger = DataprocSubmitTrigger(
+ job_id=TEST_JOB_ID,
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ polling_interval_seconds=TEST_POLL_INTERVAL,
+ execution_deadline=time.time() - 1,
+ )
+ mock_job = Job(status=JobStatus(state=JobStatus.State.RUNNING))
+ mock_get_async_hook.return_value.get_job_client = mock.AsyncMock()
+ mock_get_async_hook.return_value.get_job =
mock.AsyncMock(return_value=mock_job)
+ mock_get_sync_hook.return_value.cancel_job = mock.MagicMock()
+
+ event = await trigger.run().asend(None)
+
+ mock_get_async_hook.return_value.get_job.assert_not_awaited()
+ mock_get_sync_hook.return_value.cancel_job.assert_called_once_with(
+ job_id=TEST_JOB_ID,
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ )
+ assert event.payload["job_id"] == TEST_JOB_ID
+ assert event.payload["job_state"] == "TIMED_OUT"
+
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitTrigger.get_sync_hook")
+
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitTrigger.get_async_hook")
+ async def
test_submit_trigger_run_execution_deadline_respects_cancel_on_kill_false(
+ self, mock_get_async_hook, mock_get_sync_hook
+ ):
+ """Test the deadline event is emitted without cancelling when
cancel_on_kill is False."""
+ trigger = DataprocSubmitTrigger(
+ job_id=TEST_JOB_ID,
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ polling_interval_seconds=TEST_POLL_INTERVAL,
+ cancel_on_kill=False,
+ execution_deadline=time.time() - 1,
+ )
+ mock_job = Job(status=JobStatus(state=JobStatus.State.RUNNING))
+ mock_get_async_hook.return_value.get_job_client = mock.AsyncMock()
+ mock_get_async_hook.return_value.get_job =
mock.AsyncMock(return_value=mock_job)
+
+ event = await trigger.run().asend(None)
+
+ mock_get_async_hook.return_value.get_job.assert_not_awaited()
+ mock_get_sync_hook.return_value.cancel_job.assert_not_called()
+ assert event.payload["job_state"] == "TIMED_OUT"
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("hanging_call", ["get_job", "sleep"])
+
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitTrigger.get_sync_hook")
+
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitTrigger.get_async_hook")
+ async def test_submit_trigger_run_execution_deadline_bounds_poll_and_sleep(
+ self, mock_get_async_hook, mock_get_sync_hook, hanging_call
+ ):
+ """Test a slow get_job or a long polling sleep cannot run past the
execution deadline."""
+ trigger = DataprocSubmitTrigger(
+ job_id=TEST_JOB_ID,
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ polling_interval_seconds=60,
+ execution_deadline=time.time() + 0.2,
+ )
+ running_job = Job(status=JobStatus(state=JobStatus.State.RUNNING))
+
+ async def slow_get_job(**kwargs):
+ await asyncio.sleep(60)
+
+ mock_get_async_hook.return_value.get_job_client = mock.AsyncMock()
+ mock_get_async_hook.return_value.get_job = mock.AsyncMock(
+ side_effect=slow_get_job if hanging_call == "get_job" else None,
return_value=running_job
+ )
+ mock_get_sync_hook.return_value.cancel_job = mock.MagicMock()
+
+ event = await asyncio.wait_for(trigger.run().asend(None), timeout=5)
+
+ mock_get_sync_hook.return_value.cancel_job.assert_called_once_with(
+ job_id=TEST_JOB_ID,
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ )
+ assert event.payload["job_state"] == "TIMED_OUT"
+
@pytest.mark.asyncio
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocSubmitTrigger.get_sync_hook")
async def test_on_kill_cancels_job(self, mock_get_sync_hook,
submit_trigger):
