diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index ad051b3e6d340..5a6ab230ca597 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -138,6 +138,7 @@ class TIDeferredStatePayload(StrictBaseModel): """ trigger_timeout: timedelta | None = None + execution_timeout: timedelta | None = None queue: str | None = None next_method: str """The name of the method on the operator to call in the worker after the trigger has fired.""" 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..a94f040f9c077 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 @@ -693,8 +693,14 @@ 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 + if ti_patch_payload.execution_timeout is not None: + ti = session.get(TI, task_instance_id) + if ti and ti.start_date: + execution_timeout_date = ti.start_date + ti_patch_payload.execution_timeout + if timeout: + timeout = min(timeout, execution_timeout_date) + else: + timeout = execution_timeout_date query = update(TI).where(TI.id == task_instance_id) 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..90add65047759 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 @@ -1512,6 +1512,8 @@ def test_ti_update_state_to_deferred( state=State.RUNNING, session=session, ) + # simulate a task that started 1 hour ago + ti.start_date = timezone.datetime(2024, 11, 22) - timedelta(hours=1) session.commit() instant = timezone.datetime(2024, 11, 22) @@ -1536,6 +1538,7 @@ def test_ti_update_state_to_deferred( }, }, "trigger_timeout": "P1D", # 1 day + "execution_timeout": "PT2H", # 2 hours "queue": "default" if queues_enabled else None, "classpath": "my-classpath", "next_method": "execute_callback", @@ -1585,8 +1588,11 @@ def test_ti_update_state_to_deferred( }, "bar": "abc", } + # execution_timeout_date = start_date + 2 hours = 2024-11-22 - 1hr + 2hrs = 2024-11-22 01:00:00 + # trigger_timeout = instant + 1 day = 2024-11-23 00:00:00 + # min is execution_timeout_date assert tis[0].trigger_timeout == timezone.make_aware( - datetime(2024, 11, 23), timezone=timezone.utc + datetime(2024, 11, 22, 1, 0, 0), timezone=timezone.utc ) t = session.scalars(select(Trigger)).all() diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 72aebeab3e49b..c2c333803192c 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -273,6 +273,7 @@ class TIDeferredStatePayload(BaseModel): classpath: Annotated[str, Field(title="Classpath")] trigger_kwargs: Annotated[dict[str, JsonValue] | str | None, Field(title="Trigger Kwargs")] = None trigger_timeout: Annotated[timedelta | None, Field(title="Trigger Timeout")] = None + execution_timeout: Annotated[timedelta | None, Field(title="Execution Timeout")] = None queue: Annotated[str | None, Field(title="Queue")] = None next_method: Annotated[str, Field(title="Next Method")] next_kwargs: Annotated[dict[str, JsonValue] | None, Field(title="Next Kwargs")] = None 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..1bad152e25efe 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1476,6 +1476,7 @@ def _defer_task( classpath=classpath, trigger_kwargs=trigger_kwargs, trigger_timeout=defer.timeout, + execution_timeout=ti.task.execution_timeout, queue=queue, next_method=defer.method_name, next_kwargs=next_kwargs, @@ -2070,6 +2071,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. @@ -2084,9 +2086,15 @@ def _run_execute_callable( ctx.run(ExecutorSafeguard.tracker.set, task) if task.execution_timeout: from airflow.sdk.execution_time.timeout import timeout + from airflow.utils.timezone import utcnow - # TODO: handle timeout in case of deferral timeout_seconds = task.execution_timeout.total_seconds() + + # If the task resumes from deferral, subtract the time already spent + if ti.start_date: + time_passed = (utcnow() - ti.start_date).total_seconds() + timeout_seconds -= time_passed + try: # It's possible we're already timed out, so fast-fail if true if timeout_seconds <= 0: @@ -2140,7 +2148,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)