Enforce execution_timeout for tasks run by a language SDK

airflow-core/docs/core-concepts/tasks.rst

@@ -140,10 +140,10 @@ Timeouts
 If you want a task to have a maximum runtime, set its ``execution_timeout`` attribute to a ``datetime.timedelta`` value
 that is the maximum permissible runtime. This applies to all Airflow tasks, including sensors. ``execution_timeout`` controls the
 maximum time allowed for every execution. If ``execution_timeout`` is breached, the task times out and
-``AirflowTaskTimeout`` is raised. For tasks executed by the Python task runner, the task supervisor enforces
-the timeout from outside the task process too: if the task has not stopped within ``[core] killed_task_cleanup_time``
-seconds after ``execution_timeout`` elapses (for example because it is blocked in native code that never returns
-control to Python), the supervisor sends it SIGTERM, and SIGKILL after the same period again.
+``AirflowTaskTimeout`` is raised. The task supervisor also enforces the timeout from outside the task process:
+if the task has not stopped within ``[core] killed_task_cleanup_time`` seconds after ``execution_timeout`` elapses
+(for example because it is blocked in native code), the supervisor sends it SIGTERM, and SIGKILL after the same
+period again. Tasks run by a language SDK do not raise ``AirflowTaskTimeout`` and are stopped only this way.
 
 In addition, sensors have a ``timeout`` parameter. This only matters for sensors in ``reschedule`` mode. ``timeout`` controls the maximum
 time allowed for the sensor to succeed. If ``timeout`` is breached, ``AirflowSensorTimeout`` will be raised and the sensor fails immediately

airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py

@@ -555,6 +555,13 @@ class TIRunContext(BaseModel):
     ``None`` for regular tasks and for stub tasks that declare no parameters.
     """
 
+    execution_timeout: float | None = None
+    """
+    The task's ``execution_timeout`` in seconds, taken from its serialized Dag.
+
+    The supervisor enforces the timeout with this value. ``None`` when no timeout is set.
+    """
+
     multi_team: bool = False
     """
     Whether the deployment runs in multi-team mode.

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

@@ -92,6 +92,8 @@
 from airflow.api_fastapi.execution_api.services.task_instances import (
     client_supports_arg_bindings,
     get_arg_bindings,
+    get_execution_timeout,
+    get_task,
 )
 from airflow.api_fastapi.execution_api.versions.v2026_10_30 import IdentifyArchivedTaskStateUpdates
 from airflow.configuration import conf
@@ -380,11 +382,11 @@ def ti_run(
         if first_task_reschedule_start_date is not None:
             context.first_task_reschedule_start_date = first_task_reschedule_start_date
 
+        serialized_task = get_task(dag_bag, ti, session=session)
+
         # Only set for lang-SDK (foreign-runtime) tasks with a captured TaskFlow arg
         # spec; the route excludes unset fields, keeping regular responses lean.
-        if client_supports_arg_bindings() and (
-            arg_bindings := get_arg_bindings(dag_bag, ti, session=session)
-        ):
+        if client_supports_arg_bindings() and (arg_bindings := get_arg_bindings(serialized_task)):
             try:
                 context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings)
             except ValidationError:
@@ -402,6 +404,9 @@ def ti_run(
                     },
                 )
 
+        if (execution_timeout := get_execution_timeout(serialized_task)) is not None:
+            context.execution_timeout = execution_timeout
+
         # Only set if they are non-null
         if ti.next_method:
             context.next_method = ti.next_method

airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py

@@ -24,6 +24,7 @@
     from sqlalchemy.orm import Session
 
     from airflow.models.dagbag import DBDagBag
+    from airflow.serialization.serialized_objects import SerializedOperator
 
 
 def client_supports_arg_bindings() -> bool:
@@ -44,18 +45,35 @@ def client_supports_arg_bindings() -> bool:
     return AddArgBindingsToTIRunContext.is_applied
 
 
-def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | None:
+def get_task(dag_bag: DBDagBag, ti: Any, *, session: Session) -> SerializedOperator | None:
+    """Return the task of *ti* from its Dag version, or ``None`` if it cannot be found."""
+    if ti.dag_version_id is None:
+        return None
+    if (dag := dag_bag.get_dag(ti.dag_version_id, session=session)) is None:
+        return None
+    return dag.task_dict.get(ti.task_id)
+
+
+def get_arg_bindings(task: SerializedOperator | None) -> list | None:
     """
-    Extract the stub task's TaskFlow arg spec from its Dag version.
+    Return the TaskFlow arg spec of *task* if it is a stub task.
 
     Mapped (``.expand()``) stubs never capture a parse-time spec, so they resolve to
     ``None`` here and keep the legacy ignored-args behavior; per-map-index delivery
     lands in a follow-up.
     """
-    if ti.dag_version_id is None:
-        return None
-    if (dag := dag_bag.get_dag(ti.dag_version_id, session=session)) is None:
-        return None
-    if (task := dag.task_dict.get(ti.task_id)) is None or not task.is_stub:
+    if task is None or not task.is_stub:
         return None
     return task.arg_bindings
+
+
+def get_execution_timeout(task: SerializedOperator | None) -> float | None:
+    """
+    Return the ``execution_timeout`` of *task* in seconds.
+
+    Like the Python task runner, a zero timeout means no timeout, and a negative one has
+    already elapsed.
+    """
+    if task is None or not task.execution_timeout:
+        return None
+    return max(task.execution_timeout.total_seconds(), 0.0)

airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py

@@ -56,6 +56,7 @@
     AddArgBindingsToTIRunContext,
     AddCallbackRunEndpoint,
     AddDagRunNoteUpdateEndpoint,
+    AddExecutionTimeoutToTIRunContext,
     AddExpandedTICountToTIRunContext,
     AddFirstTaskRescheduleStartDateField,
     AddLoopContext,
@@ -79,6 +80,7 @@
         AddArgBindingsToTIRunContext,
         AddCallbackRunEndpoint,
         AddDagRunNoteUpdateEndpoint,
+        AddExecutionTimeoutToTIRunContext,
         AddExpandedTICountToTIRunContext,
         AddLoopContext,
         AddLoopDecisionField,

airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py

@@ -93,6 +93,20 @@ def remove_arg_bindings_field(response: ResponseInfo) -> None:
         response.body.pop("arg_bindings", None)
 
 
+class AddExecutionTimeoutToTIRunContext(VersionChange):
+    """Add ``execution_timeout`` so the supervisor can enforce it."""
+
+    description = __doc__
+    instructions_to_migrate_to_previous_version = (
+        schema(TIRunContext).field("execution_timeout").didnt_exist,
+    )
+
+    @convert_response_to_previous_version_for(TIRunContext)  # type: ignore[arg-type]
+    def remove_execution_timeout_field(response: ResponseInfo) -> None:  # type: ignore[misc]
+        """Strip ``execution_timeout`` from the run context for older clients."""
+        response.body.pop("execution_timeout", None)
+
+
 class AddCallbackRunEndpoint(VersionChange):
     """Add the callbacks/{callback_id}/run endpoint a worker uses to exchange its single-use callback token."""
 

airflow-core/src/airflow/serialization/serialized_objects.py

@@ -2156,9 +2156,9 @@ def fill_config_defaults(cls, serialized_obj: dict[str, Any]) -> None:
         Fill in the Dag settings a serialized Dag leaves unset from the Airflow config, as a Python Dag does.
 
         A Lang-SDK runtime cannot read the Airflow config, so it leaves ``max_active_tasks``,
-        ``max_active_runs``, ``max_consecutive_failed_dag_runs``, ``catchup`` and
-        ``disable_bundle_versioning`` out unless the Dag sets them. A value the Dag sets is kept.
-        *serialized_obj* is changed in place.
+        ``max_active_runs``, ``max_consecutive_failed_dag_runs``, ``catchup``,
+        ``disable_bundle_versioning`` and the tasks' default ``execution_timeout`` out unless the Dag
+        sets them. A value the Dag sets is kept. *serialized_obj* is changed in place.
         """
         dag = serialized_obj.get("dag")
         if not isinstance(dag, dict):
@@ -2178,6 +2178,9 @@ def fill_config_defaults(cls, serialized_obj: dict[str, Any]) -> None:
         ):
             if key not in dag:
                 dag[key] = get(section, option)
+        if (execution_timeout := conf.gettimedelta("core", "default_task_execution_timeout")) is not None:
+            task_defaults = serialized_obj.setdefault("client_defaults", {}).setdefault("tasks", {})
+            task_defaults.setdefault("execution_timeout", execution_timeout.total_seconds())
 
     @classmethod
     def validate_serialized_dag(cls, serialized_obj: dict[str, Any]) -> SerializedDAG:

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py

@@ -1151,6 +1151,53 @@ def transform(country: str): ...
         assert response.status_code == 200
         assert "arg_bindings" not in response.json()
 
+    @pytest.mark.parametrize(
+        ("decorator", "expected"),
+        [
+            pytest.param(task.stub(execution_timeout=timedelta(seconds=90)), 90.0, id="stub-with-timeout"),
+            pytest.param(task.stub, None, id="stub-without-timeout"),
+            pytest.param(task.stub(execution_timeout=timedelta(0)), None, id="stub-zero-timeout"),
+            pytest.param(
+                task.stub(execution_timeout=timedelta(seconds=-30)), 0.0, id="stub-negative-timeout"
+            ),
+            pytest.param(task(execution_timeout=timedelta(seconds=10)), 10.0, id="python-task"),
+        ],
+    )
+    def test_ti_run_returns_execution_timeout(self, client, dag_maker, decorator, expected):
+        with dag_maker("test_execution_timeout_dag", serialized=True):
+
+            @decorator
+            def sleeper(): ...
+
+            sleeper()
+
+        (ti,) = dag_maker.create_dagrun().get_task_instances()
+        ti.set_state(State.QUEUED)
+        dag_maker.session.flush()
+
+        response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD)
+        assert response.status_code == 200
+        body = response.json()
+        assert body.get("execution_timeout") == expected
+        assert ("execution_timeout" in body) == (expected is not None)
+
+    def test_ti_run_returns_execution_timeout_for_mapped_stub(self, client, dag_maker):
+        with dag_maker("test_mapped_stub_execution_timeout", serialized=True):
+
+            @task.stub(execution_timeout=timedelta(seconds=90))
+            def transform(country: str): ...
+
+            transform.expand(country=["uk", "fr"])
+
+        dr = dag_maker.create_dagrun()
+        ti = next(t for t in dr.get_task_instances() if t.map_index == 0)
+        ti.set_state(State.QUEUED)
+        dag_maker.session.flush()
+
+        response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD)
+        assert response.status_code == 200
+        assert response.json()["execution_timeout"] == 90.0
+
     def test_arg_bindings_adapter_rejects_unknown_kind(self):
         """The discriminated union refuses serialized specs with an unrecognised kind."""
         from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter

airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+from datetime import timedelta
 from uuid import uuid4
 
 import pytest
@@ -53,6 +54,7 @@ def test_task_coordinate_schema_by_version(client, version):
     assert ("region_index" in previous) == (version == "2026-10-30")
     context = response.json()["components"]["schemas"]["TIRunContext"]["properties"]
     assert ("loop" in context) == (version == "2026-10-30")
+    assert ("execution_timeout" in context) == (version == "2026-10-30")
 
 
 @pytest.mark.parametrize(
@@ -287,6 +289,38 @@ def test_head_version_includes_arg_bindings(self, client, stub_ti):
         ]
 
 
+class TestExecutionTimeoutBackwardCompat:
+    def setup_method(self):
+        clear_db_runs()
+
+    def teardown_method(self):
+        clear_db_runs()
+
+    @pytest.fixture
+    def stub_ti(self, dag_maker):
+        with dag_maker("test_execution_timeout_compat_dag", serialized=True):
+
+            @task.stub(execution_timeout=timedelta(seconds=90))
+            def sleeper(): ...
+
+            sleeper()
+
+        (ti,) = dag_maker.create_dagrun().get_task_instances()
+        ti.set_state(State.QUEUED)
+        dag_maker.session.flush()
+        return ti
+
+    def test_old_version_strips_execution_timeout(self, old_ver_client, stub_ti):
+        response = old_ver_client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY)
+        assert response.status_code == 200
+        assert "execution_timeout" not in response.json()
+
+    def test_head_version_includes_execution_timeout(self, client, stub_ti):
+        response = client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY)
+        assert response.status_code == 200
+        assert response.json()["execution_timeout"] == 90.0
+
+
 @pytest.fixture
 def archived_attempt(client, exec_app, monkeypatch, create_task_instance, session):
     """A running attempt, authenticated with a signed token, that has since been retried."""

airflow-core/tests/unit/dag_processing/test_processor.py

@@ -2371,7 +2371,7 @@ def get_type_names(union_type):
             "UpdateHITLDetail",
             "GetHITLDetailResponse",
             "SetRenderedMapIndex",
-            "SetExecutionTimeout",
+            "StartExecutionTimeout",
             # AIP-103 task/asset store — Dag processor has no task execution context.
             "GetTaskStateStore",
             "SetTaskStateStore",

airflow-core/tests/unit/jobs/test_triggerer_job.py

@@ -3313,7 +3313,7 @@ def get_type_names(union_type):
             "ResendLoggingFD",
             "CreateHITLDetailPayload",
             "SetRenderedMapIndex",
-            "SetExecutionTimeout",
+            "StartExecutionTimeout",
             "GetDag",
             # AIP-103 task store — triggerer has no task execution context.
             "GetTaskStateStore",

airflow-core/tests/unit/serialization/test_dag_serialization.py

@@ -5466,6 +5466,24 @@ def test_keeps_a_setting_the_dag_sets(self):
 
         assert data == before
 
+    @conf_vars({("core", "default_task_execution_timeout"): "300"})
+    def test_fills_the_default_execution_timeout_of_tasks_from_the_config(self):
+        with DAG(dag_id="native_dag", schedule=None) as dag:
+            BaseOperator(task_id="unset")
+            BaseOperator(task_id="set", execution_timeout=timedelta(seconds=60))
+            MockOperator.partial(task_id="mapped").expand(arg1=[1, 2])
+        data = DagSerialization.to_dict(dag)
+        data.pop("client_defaults", None)
+
+        DagSerialization.fill_config_defaults(data)
+
+        tasks = DagSerialization.from_dict(data).task_dict
+        assert {task_id: task.execution_timeout for task_id, task in tasks.items()} == {
+            "unset": timedelta(seconds=300),
+            "set": timedelta(seconds=60),
+            "mapped": timedelta(seconds=300),
+        }
+
     def test_leaves_a_payload_without_a_dag_to_validation(self):
         data = {"__version": 3}