diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index c1bac7960234d..321a45803dca0 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -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) + 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 @@ -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) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 542ce7eaaf15b..0d7a2b96c0cbe 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -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 @@ -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.""" diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 94363b984b7f6..4afad4cc54b1c 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -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. @@ -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: @@ -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) diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 9d15ac40c6967..585e6592461c3 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -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") @@ -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. @@ -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() @@ -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() @@ -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()] @@ -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()} @@ -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()]