Skip to content

Honor execution timeout across task deferrals #69734

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
GautamSharma99 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Honor execution timeout across task deferrals #69734

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -677,6 +677,15 @@ def _create_ti_state_update_query_and_update_state(
if ti_patch_payload.trigger_timeout is not None:
timeout = timezone.utcnow() + ti_patch_payload.trigger_timeout

ti = session.get(TI, task_instance_id)
if ti is not None and ti.start_date is not None:
dag = dag_bag.get_dag_for_run(dag_run=ti.dag_run, session=session)

@kaxil kaxil Jul 21, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This deserializes the full SerializedDAG on every deferral just to read one task attribute (execution_timeout). ti_update_state with state=deferred is a hot path -- every deferrable-sensor poke and every deferrable-operator await hits it.

The sibling helper _handle_fail_fast_for_dag in this same file deliberately avoids exactly this:

# Check fail_fast from DagModel (simple column lookup) - early exit if False
# This avoids loading 5-50 MB SerializedDAG in 99% of cases
fail_fast = session.scalar(select(DagModel.fail_fast).where(DagModel.dag_id == dag_id))
if not fail_fast:
    return
# Only load SerializedDAG when fail_fast=True (rare case ~1%)
ser_dag = dag_bag.get_dag_for_run(dag_run=dr, session=session)

That guard only works because fail_fast is a DagModel column. execution_timeout is task-level and lives only in the serialized blob, so there's no cheap column to gate on -- every deferral here pays the full deserialization unconditionally, even for the tasks (the majority) that set no execution_timeout.

The worker already has the value with no DB hit: _run_execute_callable reads task.execution_timeout, and the worker already ships trigger_timeout=defer.timeout in the defer payload (task_runner.py:1494 -> TIDeferredStatePayload.trigger_timeout). Compute the execution deadline (or the remaining budget) on the worker at defer time and send it in the payload the way trigger_timeout already flows; the server then does min(...) with zero deserialization.

Tradeoff: a new payload field needs an Execution API version bump (Cadwyn), while the server-side load doesn't touch the wire. But a 5-50 MB deserialization on every defer isn't worth saving a versioned field, and it regresses the fail-fast optimization this file just added. The extra session.get(TI, ...) full-ORM reload (the route only selected the columns it needed) folds into the same fix.

if dag is not None:
with contextlib.suppress(TaskNotFound):
if execution_timeout := dag.get_task(ti.task_id).execution_timeout:
execution_deadline = ti.start_date + execution_timeout
timeout = min(timeout, execution_deadline) if timeout else execution_deadline

trigger_kwargs = ti_patch_payload.trigger_kwargs
if not isinstance(trigger_kwargs, str):
# If it's passed as a string, assume the client encrypted it, otherwise assume it doesn't need to
Expand All @@ -693,9 +702,6 @@ def _create_ti_state_update_query_and_update_state(
session.add(trigger_row)
session.flush()

# TODO: HANDLE execution timeout later as it requires a call to the DB
# either get it from the serialised DAG or get it from the API

query = update(TI).where(TI.id == task_instance_id)

# Store next_kwargs directly (already serialized by worker)
Expand Down
36 changes: 35 additions & 1 deletion airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

from __future__ import annotations

from datetime import datetime
from datetime import datetime, timedelta
from typing import TYPE_CHECKING
from unittest import mock
from uuid import UUID, uuid4
Expand Down Expand Up @@ -1602,6 +1602,40 @@ def test_ti_update_state_to_deferred(
else:
assert t[0].queue is None

@pytest.mark.parametrize(
("trigger_timeout", "expected_timeout"),
[
(None, timedelta(hours=1)),
("PT30M", timedelta(minutes=30)),
("PT2H", timedelta(hours=1)),
],
)
def test_ti_update_state_to_deferred_respects_execution_and_trigger_timeouts(
self, client, session, create_task_instance, time_machine, trigger_timeout, expected_timeout
):
instant = timezone.datetime(2024, 11, 22)
task = EmptyOperator(task_id="test_deferred_execution_timeout", execution_timeout=timedelta(hours=2))
ti = create_task_instance(task=task, state=State.RUNNING, session=session)
ti.start_date = instant - timedelta(hours=1)
session.commit()
time_machine.move_to(instant, tick=False)

payload = {
"state": "deferred",
"trigger_kwargs": {},
"classpath": "my-classpath",
"next_method": "execute_callback",
"next_kwargs": {},
}
if trigger_timeout is not None:
payload["trigger_timeout"] = trigger_timeout

response = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)

assert response.status_code == 204
session.refresh(ti)
assert ti.trigger_timeout == instant + expected_timeout

@staticmethod
def _defer_ti_in_team_bundle(client, session, create_task_instance):
"""Map the TI's Dag to a bundle/team, then defer it via the Execution API."""
Expand Down
6 changes: 4 additions & 2 deletions task-sdk/src/airflow/sdk/execution_time/task_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2070,6 +2070,7 @@ def _run_execute_callable(
context: Context,
execute: Callable[..., Any] | functools.partial[Any],
task: BaseOperator,
ti: RuntimeTaskInstance,
) -> Any:
"""
Run the task's execute callable, applying the execution timeout if one is set.
Expand All @@ -2085,8 +2086,9 @@ def _run_execute_callable(
if task.execution_timeout:
from airflow.sdk.execution_time.timeout import timeout

# TODO: handle timeout in case of deferral
timeout_seconds = task.execution_timeout.total_seconds()
if ti._ti_context_from_server and ti._ti_context_from_server.next_method:
timeout_seconds -= (datetime.now(tz=timezone.utc) - ti.start_date).total_seconds()
try:
# It's possible we're already timed out, so fast-fail if true
if timeout_seconds <= 0:
Expand Down Expand Up @@ -2140,7 +2142,7 @@ def _execute_task(context: Context, ti: RuntimeTaskInstance, log: Logger):

log.info("::endgroup::")

result = _run_execute_callable(context, execute, task)
result = _run_execute_callable(context, execute, task, ti)

if (post_execute_hook := task._post_execute_hook) is not None:
create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context, result)
Expand Down
52 changes: 46 additions & 6 deletions task-sdk/tests/task_sdk/execution_time/test_task_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5700,6 +5700,17 @@ def _make_task(execution_timeout=None):
task.execution_timeout = execution_timeout
return task

@staticmethod
def _make_ti(*, start_date=None, next_method=None):
ti = mock.MagicMock(spec=RuntimeTaskInstance)
ti.start_date = start_date or datetime(2026, 7, 11, tzinfo=dt_timezone.utc)
if next_method is None:
ti._ti_context_from_server = None
else:
ti._ti_context_from_server = mock.MagicMock(spec=TIRunContext)
ti._ti_context_from_server.next_method = next_method
return ti

def test_runs_in_isolated_context_with_safeguard_tracker_set(self):
"""The callable runs in an internal context copy that has the safeguard tracker set and does not leak."""
var = contextvars.ContextVar("marker")
Expand All @@ -5712,7 +5723,7 @@ def execute(context):
seen["tracker"] = ExecutorSafeguard.tracker.get(None)
return context["value"] * 2

result = _run_execute_callable(context={"value": 21}, execute=execute, task=task)
result = _run_execute_callable(context={"value": 21}, execute=execute, task=task, ti=self._make_ti())

assert result == 42
# The safeguard tracker is set to the task inside the copy used to run execute.
Expand All @@ -5731,7 +5742,7 @@ def execute(context):
time.sleep(0.2)

with pytest.raises(AirflowTaskTimeout):
_run_execute_callable(context={}, execute=execute, task=task)
_run_execute_callable(context={}, execute=execute, task=task, ti=self._make_ti())

task.on_kill.assert_called_once()

Expand All @@ -5741,7 +5752,32 @@ def test_fast_fails_when_timeout_already_elapsed(self):
execute = mock.MagicMock()

with pytest.raises(AirflowTaskTimeout):
_run_execute_callable(context={}, execute=execute, task=task)
_run_execute_callable(context={}, execute=execute, task=task, ti=self._make_ti())

execute.assert_not_called()
task.on_kill.assert_called_once()

@mock.patch("airflow.sdk.execution_time.timeout.timeout", autospec=True)
def test_applies_remaining_execution_timeout_after_deferral(self, mock_timeout, time_machine):
now = datetime(2026, 7, 11, 12, tzinfo=dt_timezone.utc)
time_machine.move_to(now, tick=False)
task = self._make_task(execution_timeout=timedelta(minutes=5))
ti = self._make_ti(start_date=now - timedelta(minutes=2), next_method="execute_complete")

result = _run_execute_callable(context={}, execute=lambda context: "ok", task=task, ti=ti)

assert result == "ok"
mock_timeout.assert_called_once_with(180)

def test_fast_fails_when_resumed_execution_timeout_elapsed(self, time_machine):
now = datetime(2026, 7, 11, 12, tzinfo=dt_timezone.utc)
time_machine.move_to(now, tick=False)
task = self._make_task(execution_timeout=timedelta(minutes=5))
ti = self._make_ti(start_date=now - timedelta(minutes=6), next_method="execute_complete")
execute = mock.MagicMock()

with pytest.raises(AirflowTaskTimeout):
_run_execute_callable(context={}, execute=execute, task=task, ti=ti)

execute.assert_not_called()
task.on_kill.assert_called_once()
Expand All @@ -5759,7 +5795,9 @@ def test_emits_task_execute_span_at_detail_level_2(self):

with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t):
with t.start_as_current_span("parent", context=parent_ctx):
result = _run_execute_callable(context={}, execute=lambda context: "ok", task=task)
result = _run_execute_callable(
context={}, execute=lambda context: "ok", task=task, ti=self._make_ti()
)

assert result == "ok"
names = [s.name for s in exporter.get_finished_spans()]
Expand Down Expand Up @@ -5787,7 +5825,7 @@ def execute(context):

with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t):
with t.start_as_current_span("parent", context=parent_ctx):
result = _run_execute_callable(context={}, execute=execute, task=task)
result = _run_execute_callable(context={}, execute=execute, task=task, ti=self._make_ti())

assert result == "ok"
spans = {s.name: s for s in exporter.get_finished_spans()}
Expand All @@ -5806,7 +5844,9 @@ def test_no_task_execute_span_at_detail_level_1(self):

with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t):
with t.start_as_current_span("parent", context=parent_ctx):
result = _run_execute_callable(context={}, execute=lambda context: "ok", task=task)
result = _run_execute_callable(
context={}, execute=lambda context: "ok", task=task, ti=self._make_ti()
)

assert result == "ok"
names = [s.name for s in exporter.get_finished_spans()]
Expand Down