Enforce execution_timeout during task deferrals

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."""

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)
 

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()

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

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)