This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new 832602523ea [v3-4-test] Enforce execution_timeout for tasks run by a
language SDK (#74869) (#75376)
832602523ea is described below
commit 832602523ea2d7cdd3574e5fa8014f308ef1d287
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Thu Oct 22 15:59:55 2026 +0800
[v3-4-test] Enforce execution_timeout for tasks run by a language SDK
(#74869) (#75376)
* Extract a helper that looks up the stub task
* Add execution_timeout to TIRunContext
* Add execution_timeout to the supervisor schema
* Arm the supervisor execution timeout from the run context
* Document that the supervisor timeout covers language SDK tasks
* Treat a negative execution_timeout as elapsed for language SDK tasks
The Python task runner fails a task with a negative execution_timeout
right away. Sent as-is, a negative value would instead move the
supervisor's SIGTERM deadline to an arbitrary earlier point, and
treating it as no timeout would let the task run forever. Clamping it
to zero makes the supervisor treat the timeout as already elapsed,
which is the closest the outside layer can get to the Python behavior.
* Send execution_timeout in the run context for all tasks
* Start the execution timeout for language SDK tasks in the coordinator
* Replace SetExecutionTimeout with StartExecutionTimeout
* Look up the task once in ti_run
* Fill the default execution_timeout of Lang-SDK tasks from the config
* Arm the execution timeout only in _arm_execution_timeout
* Store execution_timeout at startup and arm it in _on_child_started
* Align docstrings with the execution timeout changes
* Rename the task variable in ti_run to fix mypy
(cherry picked from commit 597aff10d71ca8e02de97d891e4f758c878bb1f7)
Co-authored-by: ChickenBenny <[email protected]>
---
airflow-core/docs/core-concepts/tasks.rst | 8 +-
.../execution_api/datamodels/taskinstance.py | 7 +
.../execution_api/routes/task_instances.py | 11 +-
.../execution_api/services/task_instances.py | 32 ++++-
.../api_fastapi/execution_api/versions/__init__.py | 2 +
.../execution_api/versions/v2026_10_30.py | 14 ++
.../airflow/serialization/serialized_objects.py | 9 +-
.../versions/head/test_task_instances.py | 47 ++++++
.../versions/v2026_10_30/test_task_instances.py | 34 +++++
.../tests/unit/dag_processing/test_processor.py | 2 +-
airflow-core/tests/unit/jobs/test_triggerer_job.py | 2 +-
.../unit/serialization/test_dag_serialization.py | 18 +++
.../pkg/execution/genmodels/discriminators.gen.go | 8 +-
go-sdk/pkg/execution/genmodels/models.gen.go | 159 ++++++++++-----------
go-sdk/schema/supervisor-schema.json | 46 +++---
.../src/airflow/sdk/api/datamodels/_generated.py | 1 +
.../src/airflow/sdk/coordinators/_subprocess.py | 17 +++
task-sdk/src/airflow/sdk/execution_time/comms.py | 9 +-
.../airflow/sdk/execution_time/schema/schema.json | 46 +++---
.../sdk/execution_time/schema/versions/__init__.py | 2 +
.../execution_time/schema/versions/v2026_10_30.py | 15 ++
.../src/airflow/sdk/execution_time/supervisor.py | 33 +++--
.../src/airflow/sdk/execution_time/task_runner.py | 4 +-
.../tests/task_sdk/coordinators/test_subprocess.py | 42 ++++++
.../execution_time/schema/test_migrator.py | 10 ++
.../task_sdk/execution_time/test_supervisor.py | 84 ++++++++---
.../task_sdk/execution_time/test_task_runner.py | 4 +-
ts-sdk/schema/supervisor-schema.json | 46 +++---
ts-sdk/src/generated/supervisor.ts | 42 +++---
29 files changed, 523 insertions(+), 231 deletions(-)
diff --git a/airflow-core/docs/core-concepts/tasks.rst
b/airflow-core/docs/core-concepts/tasks.rst
index f57de76baa8..d97c6410b41 100644
--- a/airflow-core/docs/core-concepts/tasks.rst
+++ b/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
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 8a1deb203fd..db998c4ad54 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
@@ -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.
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 ff3ffbc021b..aa0fdcd0495 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
@@ -89,6 +89,8 @@ from airflow.api_fastapi.execution_api.security import (
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
@@ -377,11 +379,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:
@@ -399,6 +401,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
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py
b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py
index ab67aaa9814..15e0d6fcf8c 100644
---
a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py
+++
b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py
@@ -24,6 +24,7 @@ if TYPE_CHECKING:
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)
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
index f67687e6e23..bbb02749163 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
@@ -56,6 +56,7 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30
import (
AddArgBindingsToTIRunContext,
AddCallbackRunEndpoint,
AddDagRunNoteUpdateEndpoint,
+ AddExecutionTimeoutToTIRunContext,
AddExpandedTICountToTIRunContext,
AddFirstTaskRescheduleStartDateField,
AddLoopContext,
@@ -77,6 +78,7 @@ bundle = VersionBundle(
AddArgBindingsToTIRunContext,
AddCallbackRunEndpoint,
AddDagRunNoteUpdateEndpoint,
+ AddExecutionTimeoutToTIRunContext,
AddExpandedTICountToTIRunContext,
AddLoopContext,
AddLoopDecisionField,
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
index 22b34e5115d..e4e02b765c4 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
@@ -93,6 +93,20 @@ class
AddArgBindingsToTIRunContext(VersionChangeWithSideEffects):
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."""
diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py
b/airflow-core/src/airflow/serialization/serialized_objects.py
index f5fd8c118b9..73fc26e0114 100644
--- a/airflow-core/src/airflow/serialization/serialized_objects.py
+++ b/airflow-core/src/airflow/serialization/serialized_objects.py
@@ -2156,9 +2156,9 @@ class DagSerialization(BaseSerialization):
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 @@ class DagSerialization(BaseSerialization):
):
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:
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 8b26dd49b40..299a00d9030 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
@@ -1123,6 +1123,53 @@ class TestTIRunState:
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
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py
index 3b4fcd392aa..9328d0a9830 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py
+++
b/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 @@ class TestArgBindingsFieldBackwardCompat:
]
+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."""
diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py
b/airflow-core/tests/unit/dag_processing/test_processor.py
index 0271a5e07fb..154cb7640ce 100644
--- a/airflow-core/tests/unit/dag_processing/test_processor.py
+++ b/airflow-core/tests/unit/dag_processing/test_processor.py
@@ -2371,7 +2371,7 @@ class TestDagProcessingMessageTypes:
"UpdateHITLDetail",
"GetHITLDetailResponse",
"SetRenderedMapIndex",
- "SetExecutionTimeout",
+ "StartExecutionTimeout",
# AIP-103 task/asset store — Dag processor has no task execution
context.
"GetTaskStateStore",
"SetTaskStateStore",
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py
b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index bbf01eba0f1..6cd5c17a226 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -3313,7 +3313,7 @@ class TestTriggererMessageTypes:
"ResendLoggingFD",
"CreateHITLDetailPayload",
"SetRenderedMapIndex",
- "SetExecutionTimeout",
+ "StartExecutionTimeout",
"GetDag",
# AIP-103 task store — triggerer has no task execution context.
"GetTaskStateStore",
diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py
b/airflow-core/tests/unit/serialization/test_dag_serialization.py
index d40f36ebf6d..0caec6fa2ba 100644
--- a/airflow-core/tests/unit/serialization/test_dag_serialization.py
+++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py
@@ -5466,6 +5466,24 @@ class TestFillConfigDefaults:
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}
diff --git a/go-sdk/pkg/execution/genmodels/discriminators.gen.go
b/go-sdk/pkg/execution/genmodels/discriminators.gen.go
index 8736a6e2c97..3a81b6ce7a8 100644
--- a/go-sdk/pkg/execution/genmodels/discriminators.gen.go
+++ b/go-sdk/pkg/execution/genmodels/discriminators.gen.go
@@ -90,12 +90,12 @@ const (
TypeSentFDs = "SentFDs"
TypeSetAssetStateStoreByName = "SetAssetStateStoreByName"
TypeSetAssetStateStoreByURI = "SetAssetStateStoreByUri"
- TypeSetExecutionTimeout = "SetExecutionTimeout"
TypeSetRenderedFields = "SetRenderedFields"
TypeSetRenderedMapIndex = "SetRenderedMapIndex"
TypeSetTaskStateStore = "SetTaskStateStore"
TypeSetXCom = "SetXCom"
TypeSkipDownstreamTasks = "SkipDownstreamTasks"
+ TypeStartExecutionTimeout = "StartExecutionTimeout"
TypeStartupDetails = "StartupDetails"
TypeSucceedTask = "SucceedTask"
TypeTICount = "TICount"
@@ -321,9 +321,6 @@ func EnsureType(m any) any {
case SetAssetStateStoreByURI:
b.Type = TypeSetAssetStateStoreByURI
return b
- case SetExecutionTimeout:
- b.Type = TypeSetExecutionTimeout
- return b
case SetRenderedFields:
b.Type = TypeSetRenderedFields
return b
@@ -339,6 +336,9 @@ func EnsureType(m any) any {
case SkipDownstreamTasks:
b.Type = TypeSkipDownstreamTasks
return b
+ case StartExecutionTimeout:
+ b.Type = TypeStartExecutionTimeout
+ return b
case StartupDetails:
b.Type = TypeStartupDetails
return b
diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go
b/go-sdk/pkg/execution/genmodels/models.gen.go
index fbf36031cea..f95ceb53d62 100644
--- a/go-sdk/pkg/execution/genmodels/models.gen.go
+++ b/go-sdk/pkg/execution/genmodels/models.gen.go
@@ -1555,16 +1555,6 @@ type SetAssetStateStoreByURI struct {
Value JsonValue `msgpack:"value"`
}
-// Tell the supervisor the task's “execution_timeout“ so it can enforce it from
-// outside the task process.
-type SetExecutionTimeout struct {
- // TimeoutSeconds corresponds to the JSON schema field
"timeout_seconds".
- TimeoutSeconds float64 `msgpack:"timeout_seconds"`
-
- // Type corresponds to the JSON schema field "type".
- Type string `msgpack:"type,omitempty"`
-}
-
// Payload for setting RTIF for a task instance.
type SetRenderedFields struct {
// RenderedFields corresponds to the JSON schema field
"rendered_fields".
@@ -1639,6 +1629,13 @@ type SkipDownstreamTasks struct {
Type string `msgpack:"type,omitempty"`
}
+// Tell the supervisor to start enforcing the run context's “execution_timeout“
+// from outside the task process.
+type StartExecutionTimeout struct {
+ // Type corresponds to the JSON schema field "type".
+ Type string `msgpack:"type,omitempty"`
+}
+
type StartupDetails struct {
// BundleInfo corresponds to the JSON schema field "bundle_info".
BundleInfo BundleInfo `msgpack:"bundle_info"`
@@ -1709,6 +1706,9 @@ type TIRunContext struct {
// DagRun corresponds to the JSON schema field "dag_run".
DagRun DagRun `msgpack:"dag_run"`
+ // ExecutionTimeout corresponds to the JSON schema field
"execution_timeout".
+ ExecutionTimeout interface{} `msgpack:"execution_timeout,omitempty"`
+
// ExpandedTICount corresponds to the JSON schema field
"expanded_ti_count".
ExpandedTICount *int `msgpack:"expanded_ti_count,omitempty"`
@@ -1750,21 +1750,53 @@ type TIRunContext struct {
type TaskArgBinding interface{}
-type TriggerKwargs map[string]interface{}
+type TaskBreadcrumbsResult struct {
+ // Breadcrumbs corresponds to the JSON schema field "breadcrumbs".
+ Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem
`msgpack:"breadcrumbs"`
-type Warnings []interface{}
+ // Type corresponds to the JSON schema field "type".
+ Type string `msgpack:"type,omitempty"`
+}
-type TaskIds []string
+type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{}
-// Variable schema for responses with fields that are needed for Runtime.
-type VariableResponse struct {
- // Key corresponds to the JSON schema field "key".
- Key string `msgpack:"key"`
+// Task callback status information.
+//
+// A Class with information about the success/failure TI callback to be
executed.
+// Currently, only failure
+// callbacks when tasks are externally killed or experience heartbeat timeouts
are
+// run via DagFileProcessorProcess.
+type TaskCallbackRequest struct {
+ // BundleName corresponds to the JSON schema field "bundle_name".
+ BundleName string `msgpack:"bundle_name"`
- // Value corresponds to the JSON schema field "value".
- Value interface{} `msgpack:"value"`
+ // BundleVersion corresponds to the JSON schema field "bundle_version".
+ BundleVersion interface{} `msgpack:"bundle_version"`
+
+ // ContextFromServer corresponds to the JSON schema field
"context_from_server".
+ ContextFromServer *TIRunContext
`msgpack:"context_from_server,omitempty"`
+
+ // Filepath corresponds to the JSON schema field "filepath".
+ Filepath string `msgpack:"filepath"`
+
+ // Msg corresponds to the JSON schema field "msg".
+ Msg interface{} `msgpack:"msg,omitempty"`
+
+ // TaskCallbackType corresponds to the JSON schema field
"task_callback_type".
+ TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"`
+
+ // TI corresponds to the JSON schema field "ti".
+ TI TaskInstance `msgpack:"ti"`
+
+ // Type corresponds to the JSON schema field "type".
+ Type string `msgpack:"type,omitempty"`
+
+ // VersionData corresponds to the JSON schema field "version_data".
+ VersionData *VersionData `msgpack:"version_data,omitempty"`
}
+type TaskIds []string
+
// Schema for TaskInstance model with minimal required fields needed for
Runtime.
type TaskInstance struct {
// ContextCarrier corresponds to the JSON schema field
"context_carrier".
@@ -1804,70 +1836,24 @@ type TaskInstance struct {
TryNumber int `msgpack:"try_number"`
}
+type TaskInstanceState string
+
const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input"
const TaskInstanceStateDeferred TaskInstanceState = "deferred"
-const TaskInstanceStateSkipped TaskInstanceState = "skipped"
-const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed"
-const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule"
-const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry"
const TaskInstanceStateFailed TaskInstanceState = "failed"
+const TaskInstanceStateQueued TaskInstanceState = "queued"
+const TaskInstanceStateRemoved TaskInstanceState = "removed"
const TaskInstanceStateRestarting TaskInstanceState = "restarting"
-const TaskInstanceStateSuccess TaskInstanceState = "success"
const TaskInstanceStateRunning TaskInstanceState = "running"
-const TaskInstanceStateQueued TaskInstanceState = "queued"
-
const TaskInstanceStateScheduled TaskInstanceState = "scheduled"
-const TaskInstanceStateRemoved TaskInstanceState = "removed"
-
-type TaskInstanceState string
+const TaskInstanceStateSkipped TaskInstanceState = "skipped"
+const TaskInstanceStateSuccess TaskInstanceState = "success"
+const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule"
+const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry"
+const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed"
type TaskOutlets []AssetProfile
-type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{}
-
-type TaskBreadcrumbsResult struct {
- // Breadcrumbs corresponds to the JSON schema field "breadcrumbs".
- Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem
`msgpack:"breadcrumbs"`
-
- // Type corresponds to the JSON schema field "type".
- Type string `msgpack:"type,omitempty"`
-}
-
-// Task callback status information.
-//
-// A Class with information about the success/failure TI callback to be
executed.
-// Currently, only failure
-// callbacks when tasks are externally killed or experience heartbeat timeouts
are
-// run via DagFileProcessorProcess.
-type TaskCallbackRequest struct {
- // BundleName corresponds to the JSON schema field "bundle_name".
- BundleName string `msgpack:"bundle_name"`
-
- // BundleVersion corresponds to the JSON schema field "bundle_version".
- BundleVersion interface{} `msgpack:"bundle_version"`
-
- // ContextFromServer corresponds to the JSON schema field
"context_from_server".
- ContextFromServer *TIRunContext
`msgpack:"context_from_server,omitempty"`
-
- // Filepath corresponds to the JSON schema field "filepath".
- Filepath string `msgpack:"filepath"`
-
- // Msg corresponds to the JSON schema field "msg".
- Msg interface{} `msgpack:"msg,omitempty"`
-
- // TaskCallbackType corresponds to the JSON schema field
"task_callback_type".
- TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"`
-
- // TI corresponds to the JSON schema field "ti".
- TI TaskInstance `msgpack:"ti"`
-
- // Type corresponds to the JSON schema field "type".
- Type string `msgpack:"type,omitempty"`
-
- // VersionData corresponds to the JSON schema field "version_data".
- VersionData *VersionData `msgpack:"version_data,omitempty"`
-}
-
// Response containing the first reschedule date for a task instance.
type TaskRescheduleStartDate struct {
// StartDate corresponds to the JSON schema field "start_date".
@@ -1877,12 +1863,6 @@ type TaskRescheduleStartDate struct {
Type string `msgpack:"type,omitempty"`
}
-type TaskStateState string
-
-const TaskStateStateFailed TaskStateState = "failed"
-const TaskStateStateSkipped TaskStateState = "skipped"
-const TaskStateStateRemoved TaskStateState = "removed"
-
// Update a task's state.
//
// If a process exits without sending one of these the state will be derived
from
@@ -1909,6 +1889,12 @@ type TaskState struct {
Type string `msgpack:"type,omitempty"`
}
+type TaskStateState string
+
+const TaskStateStateFailed TaskStateState = "failed"
+const TaskStateStateRemoved TaskStateState = "removed"
+const TaskStateStateSkipped TaskStateState = "skipped"
+
// Response to GetTaskStateStore; wraps the generated API response for
supervisor
// to worker comms.
type TaskStateStoreResult struct {
@@ -1958,6 +1944,19 @@ type TriggerDagRun struct {
Type string `msgpack:"type,omitempty"`
}
+type TriggerKwargs map[string]interface{}
+
+type Warnings []interface{}
+
+// Variable schema for responses with fields that are needed for Runtime.
+type VariableResponse struct {
+ // Key corresponds to the JSON schema field "key".
+ Key string `msgpack:"key"`
+
+ // Value corresponds to the JSON schema field "value".
+ Value interface{} `msgpack:"value"`
+}
+
type UpdateDagRunNote struct {
// Note corresponds to the JSON schema field "note".
Note interface{} `msgpack:"note"`
diff --git a/go-sdk/schema/supervisor-schema.json
b/go-sdk/schema/supervisor-schema.json
index e739708d12d..ebe8c2ab65c 100644
--- a/go-sdk/schema/supervisor-schema.json
+++ b/go-sdk/schema/supervisor-schema.json
@@ -3791,27 +3791,6 @@
"title": "SetAssetStateStoreByUri",
"type": "object"
},
- "SetExecutionTimeout": {
- "description": "Tell the supervisor the task's ``execution_timeout`` so
it can enforce it from outside the task process.",
- "properties": {
- "timeout_seconds": {
- "exclusiveMinimum": 0,
- "title": "Timeout Seconds",
- "type": "number"
- },
- "type": {
- "const": "SetExecutionTimeout",
- "default": "SetExecutionTimeout",
- "title": "Type",
- "type": "string"
- }
- },
- "required": [
- "timeout_seconds"
- ],
- "title": "SetExecutionTimeout",
- "type": "object"
- },
"SetRenderedFields": {
"description": "Payload for setting RTIF for a task instance.",
"properties": {
@@ -4005,6 +3984,19 @@
"title": "SkipDownstreamTasks",
"type": "object"
},
+ "StartExecutionTimeout": {
+ "description": "Tell the supervisor to start enforcing the run context's
``execution_timeout`` from outside the task process.",
+ "properties": {
+ "type": {
+ "const": "StartExecutionTimeout",
+ "default": "StartExecutionTimeout",
+ "title": "Type",
+ "type": "string"
+ }
+ },
+ "title": "StartExecutionTimeout",
+ "type": "object"
+ },
"StartupDetails": {
"properties": {
"ti": {
@@ -5399,6 +5391,18 @@
"default": null,
"title": "Arg Bindings"
},
+ "execution_timeout": {
+ "anyOf": [
+ {
+ "type": "number"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Execution Timeout"
+ },
"multi_team": {
"default": false,
"title": "Multi Team",
diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
index 72d3cc9a33b..b10bc03f363 100644
--- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
+++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
@@ -887,4 +887,5 @@ class TIRunContext(BaseModel):
should_retry: Annotated[bool | None, Field(title="Should Retry")] = False
start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] =
None
arg_bindings: Annotated[list[TaskArgBinding] | None, Field(title="Arg
Bindings")] = None
+ execution_timeout: Annotated[float | None, Field(title="Execution
Timeout")] = None
multi_team: Annotated[bool | None, Field(title="Multi Team")] = False
diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
index 0bd84e147be..e102b424a73 100644
--- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
+++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
@@ -408,6 +408,23 @@ class _PopenActivitySubprocess(ActivitySubprocess):
return self
+ def _on_child_started(
+ self,
+ *,
+ ti: TaskInstance,
+ dag_rel_path: str | os.PathLike[str],
+ bundle_info,
+ sentry_integration: str,
+ ) -> None:
+ super()._on_child_started(
+ ti=ti,
+ dag_rel_path=dag_rel_path,
+ bundle_info=bundle_info,
+ sentry_integration=sentry_integration,
+ )
+ # Language runtimes do not send StartExecutionTimeout, so the clock
starts once they have the task.
+ self._arm_execution_timeout()
+
def _prepare_ti_context(self, ti: TaskInstance, ti_context: TIRunContext)
-> TIRunContext | None:
"""
Delete the XComs in ``ti_context.xcom_keys_to_clear`` so the launched
runtime does not have to.
diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py
b/task-sdk/src/airflow/sdk/execution_time/comms.py
index bd5781f5dfe..45af5b757a3 100644
--- a/task-sdk/src/airflow/sdk/execution_time/comms.py
+++ b/task-sdk/src/airflow/sdk/execution_time/comms.py
@@ -1136,11 +1136,10 @@ class SetRenderedMapIndex(BaseModel):
type: Literal["SetRenderedMapIndex"] = "SetRenderedMapIndex"
-class SetExecutionTimeout(BaseModel):
- """Tell the supervisor the task's ``execution_timeout`` so it can enforce
it from outside the task process."""
+class StartExecutionTimeout(BaseModel):
+ """Tell the supervisor to start enforcing the run context's
``execution_timeout`` from outside the task process."""
- timeout_seconds: Annotated[float, Field(gt=0)]
- type: Literal["SetExecutionTimeout"] = "SetExecutionTimeout"
+ type: Literal["StartExecutionTimeout"] = "StartExecutionTimeout"
class TriggerDagRun(TriggerDAGRunPayload):
@@ -1347,12 +1346,12 @@ ToSupervisor = Annotated[
| RetryTask
| SetAssetStateStoreByName
| SetAssetStateStoreByUri
- | SetExecutionTimeout
| SetRenderedFields
| SetRenderedMapIndex
| SetTaskStateStore
| SetXCom
| SkipDownstreamTasks
+ | StartExecutionTimeout
| SucceedTask
| ValidateInletsAndOutlets
| TaskState
diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
index e739708d12d..ebe8c2ab65c 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
@@ -3791,27 +3791,6 @@
"title": "SetAssetStateStoreByUri",
"type": "object"
},
- "SetExecutionTimeout": {
- "description": "Tell the supervisor the task's ``execution_timeout`` so
it can enforce it from outside the task process.",
- "properties": {
- "timeout_seconds": {
- "exclusiveMinimum": 0,
- "title": "Timeout Seconds",
- "type": "number"
- },
- "type": {
- "const": "SetExecutionTimeout",
- "default": "SetExecutionTimeout",
- "title": "Type",
- "type": "string"
- }
- },
- "required": [
- "timeout_seconds"
- ],
- "title": "SetExecutionTimeout",
- "type": "object"
- },
"SetRenderedFields": {
"description": "Payload for setting RTIF for a task instance.",
"properties": {
@@ -4005,6 +3984,19 @@
"title": "SkipDownstreamTasks",
"type": "object"
},
+ "StartExecutionTimeout": {
+ "description": "Tell the supervisor to start enforcing the run context's
``execution_timeout`` from outside the task process.",
+ "properties": {
+ "type": {
+ "const": "StartExecutionTimeout",
+ "default": "StartExecutionTimeout",
+ "title": "Type",
+ "type": "string"
+ }
+ },
+ "title": "StartExecutionTimeout",
+ "type": "object"
+ },
"StartupDetails": {
"properties": {
"ti": {
@@ -5399,6 +5391,18 @@
"default": null,
"title": "Arg Bindings"
},
+ "execution_timeout": {
+ "anyOf": [
+ {
+ "type": "number"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Execution Timeout"
+ },
"multi_team": {
"default": false,
"title": "Multi Team",
diff --git
a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
index 40da65651ae..27ff074bb23 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py
@@ -40,6 +40,7 @@ def get_bundle() -> VersionBundle:
from airflow.sdk.execution_time.schema.versions.v2026_10_30 import (
AddArgBindingsToSupervisorTIRunContext,
AddDagDefinitionsToDagFileParsingResult,
+ AddExecutionTimeoutToSupervisorTIRunContext,
AddExpandedTICountToSupervisorTIRunContext,
AddFirstTaskRescheduleStartDateToSupervisorTIRunContext,
AddLoopRuntime,
@@ -58,6 +59,7 @@ def get_bundle() -> VersionBundle:
AddRegionCoordinatesToSupervisorTaskInstance,
AddFirstTaskRescheduleStartDateToSupervisorTIRunContext,
AddExpandedTICountToSupervisorTIRunContext,
+ AddExecutionTimeoutToSupervisorTIRunContext,
),
Version("2026-06-16"),
)
diff --git
a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
index d69abf63b5f..ddc99b274b9 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
@@ -133,3 +133,18 @@ class
AddExpandedTICountToSupervisorTIRunContext(VersionChange):
instructions_to_migrate_to_previous_version = (
schema(TIRunContext).field("expanded_ti_count").didnt_exist,
)
+
+
+class AddExecutionTimeoutToSupervisorTIRunContext(VersionChange):
+ """
+ Add the ``execution_timeout`` field carrying the task's timeout, which the
supervisor enforces.
+
+ The supervisor-schema mirror of the execution API's
``AddExecutionTimeoutToTIRunContext``,
+ named apart so the two migrations are not confused.
+ """
+
+ description = __doc__
+
+ instructions_to_migrate_to_previous_version = (
+ schema(TIRunContext).field("execution_timeout").didnt_exist,
+ )
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index 388075858bd..656efa07b13 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -130,12 +130,12 @@ from airflow.sdk.execution_time.comms import (
SentFDs,
SetAssetStateStoreByName,
SetAssetStateStoreByUri,
- SetExecutionTimeout,
SetRenderedFields,
SetRenderedMapIndex,
SetTaskStateStore,
SetXCom,
SkipDownstreamTasks,
+ StartExecutionTimeout,
StartupDetails,
SucceedTask,
TaskBreadcrumbsResult,
@@ -1726,6 +1726,7 @@ class ActivitySubprocess(WatchedSubprocess):
# tell us "no, stop!" for any reason)
ti_context = self.client.task_instances.start(ti.id, self.pid,
datetime.now(tz=UTC))
self._should_retry = ti_context.should_retry
+ self._execution_timeout_seconds = ti_context.execution_timeout
self._last_successful_heartbeat = time.monotonic()
except Exception as e:
# On any error kill that subprocess!
@@ -1996,13 +1997,15 @@ class ActivitySubprocess(WatchedSubprocess):
"""
Enforce ``execution_timeout`` from outside the task process.
- The task process raises ``AirflowTaskTimeout`` itself when the timeout
elapses, which is the only
- place ``on_kill``, the retry policy, callbacks and listeners can run.
That needs a process that can
- still run Python signal handlers; one that has not reported a terminal
state within the grace period
- is stuck (native code, SIGSEGV) and is sent SIGTERM, then SIGKILL.
+ A Python task process raises ``AirflowTaskTimeout`` itself when the
timeout elapses, which is the
+ only place ``on_kill``, the retry policy, callbacks and listeners can
run. That needs a process that
+ can still run Python signal handlers; one that has not reported a
terminal state within the grace
+ period is stuck (native code, SIGSEGV) and is sent SIGTERM, then
SIGKILL.
- The deadline is armed by ``SetExecutionTimeout``, which the task
process sends right before
- ``execute()``. Anything before that (bundle load, Dag parsing) is not
covered here.
+ For a Python task, ``StartExecutionTimeout`` arms the deadline right
before ``execute()``, so bundle
+ loading and Dag parsing are not covered; Dag parsing has its own
``dagbag_import_timeout``. For a
+ task a language SDK runs, it is armed once the runtime has been sent
``StartupDetails``. Either
+ way, the process is sent SIGTERM ``killed_task_cleanup_time`` after
the timeout elapses.
"""
if self._execution_timeout_due_in() > 0:
return
@@ -2146,13 +2149,19 @@ class ActivitySubprocess(WatchedSubprocess):
self.client.task_instances.skip_downstream_tasks(self.id, msg)
return None, {}
- def _handle_set_execution_timeout(
- self, msg: SetExecutionTimeout, log: FilteringBoundLogger, req_id: int
+ def _handle_start_execution_timeout(
+ self, msg: StartExecutionTimeout, log: FilteringBoundLogger, req_id:
int
) -> RequestResult:
- self._execution_timeout_seconds = msg.timeout_seconds
- self._execution_timeout_enforce_at = time.monotonic() +
msg.timeout_seconds + KILLED_TASK_CLEANUP_TIME
+ self._arm_execution_timeout()
return None, {}
+ def _arm_execution_timeout(self) -> None:
+ if self._execution_timeout_seconds is None:
+ return
+ self._execution_timeout_enforce_at = (
+ time.monotonic() + self._execution_timeout_seconds +
KILLED_TASK_CLEANUP_TIME
+ )
+
def _handle_set_rendered_fields(
self, msg: SetRenderedFields, log: FilteringBoundLogger, req_id: int
) -> RequestResult:
@@ -2447,11 +2456,11 @@ class ActivitySubprocess(WatchedSubprocess):
register_request_method(RetryTask, _handle_task_state),
register_request_method(SetAssetStateStoreByName,
_handle_set_asset_state_store_by_name),
register_request_method(SetAssetStateStoreByUri,
_handle_set_asset_state_store_by_uri),
- register_request_method(SetExecutionTimeout,
_handle_set_execution_timeout),
register_request_method(SetRenderedFields,
_handle_set_rendered_fields),
register_request_method(SetRenderedMapIndex,
_handle_set_rendered_map_index),
register_request_method(SetTaskStateStore,
_handle_set_task_state_store),
register_request_method(SkipDownstreamTasks,
_handle_skip_downstream_tasks),
+ register_request_method(StartExecutionTimeout,
_handle_start_execution_timeout),
register_request_method(SucceedTask, _handle_finished_task),
register_request_method(TaskState, _handle_task_state),
register_request_method(TriggerDagRun,
_handle_trigger_dag_run),
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 ee78312e95e..6651a7785e9 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -119,10 +119,10 @@ from airflow.sdk.execution_time.comms import (
ResendLoggingFD,
RetryTask,
SentFDs,
- SetExecutionTimeout,
SetRenderedFields,
SetRenderedMapIndex,
SkipDownstreamTasks,
+ StartExecutionTimeout,
StartupDetails,
SucceedTask,
TaskBreadcrumbsResult,
@@ -2343,7 +2343,7 @@ def _run_execute_callable(
# It's possible we're already timed out, so fast-fail if true
if timeout_seconds <= 0:
raise AirflowTaskTimeout()
-
SUPERVISOR_COMMS.send(SetExecutionTimeout(timeout_seconds=timeout_seconds))
+ SUPERVISOR_COMMS.send(StartExecutionTimeout())
# Run task in timeout wrapper
with timeout(timeout_seconds):
result = ctx.run(execute, context=context)
diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py
b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py
index 63eafc6d1ab..f2e2622255f 100644
--- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py
+++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py
@@ -1150,6 +1150,48 @@ class TestPopenActivitySubprocessClearXComs:
} in captured_logs
+class TestPopenActivitySubprocessExecutionTimeout:
+ @pytest.mark.parametrize(
+ ("execution_timeout", "killed"),
+ [
+ pytest.param(30.0, True, id="with-timeout"),
+ pytest.param(None, False, id="without-timeout"),
+ ],
+ )
+ @patch("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
5.0)
+ @patch("airflow.sdk.execution_time.supervisor.WatchedSubprocess.kill",
autospec=True)
+ @patch.object(_PopenActivitySubprocess, "send_msg", autospec=True)
+ @patch("time.monotonic", autospec=True, return_value=1.0)
+ def test_arms_execution_timeout_from_run_context(
+ self, mock_monotonic, mock_send_msg, mock_kill, execution_timeout,
killed
+ ):
+ client = MagicMock()
+ client.task_instances.start.return_value =
TIRunContext.model_construct(
+ should_retry=False, execution_timeout=execution_timeout,
xcom_keys_to_clear=[]
+ )
+ proc = _PopenActivitySubprocess(
+ process_log=MagicMock(),
+ id=uuid7(),
+ pid=12345,
+ stdin=MagicMock(spec=socket.socket),
+ process=MagicMock(),
+ client=client,
+ comm_server=MagicMock(spec=socket.socket),
+ logs_server=MagicMock(spec=socket.socket),
+ )
+
+ proc._on_child_started(
+ ti=_make_ti(),
+ dag_rel_path="d.py",
+ bundle_info=BundleInfo(name="my-bundle", version=None),
+ sentry_integration="",
+ )
+
+ mock_monotonic.return_value += 35
+ proc._handle_execution_timeout_if_needed()
+ assert mock_kill.called is killed
+
+
class TestResolveArtifactBundle:
"""Execute-time resolution of the Dag bundle holding the artifacts."""
diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
index 6e86037fe05..f89e6b87e81 100644
--- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
+++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
@@ -550,6 +550,16 @@ class TestRealBundleArgBindingsDowngrade:
assert ti["region_id"] == startup_details.ti.region_id
assert ti["region_index"] == 7
+ @pytest.mark.parametrize("version", ["2026-06-16", "2026-10-30"])
+ def test_execution_timeout_follows_supervisor_version(self, real_migrator,
startup_details, version):
+ startup_details.ti_context.execution_timeout = 90.0
+ context = real_migrator.downgrade(startup_details,
version).model_dump()["ti_context"]
+
+ if version == "2026-06-16":
+ assert "execution_timeout" not in context
+ else:
+ assert context["execution_timeout"] == 90.0
+
def test_head_version_keeps_arg_bindings(self, real_migrator,
startup_details):
from airflow.sdk.api.datamodels._generated import LiteralArgBinding,
XComArgBinding
diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
index 12a76d25a1b..72a9e29c7e3 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -72,6 +72,7 @@ from airflow.sdk.api.datamodels._generated import (
PreviousTIResponse,
TaskInstance,
TaskInstanceState,
+ TIRunContext,
)
from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType,
TaskAlreadyRunningError
@@ -140,12 +141,12 @@ from airflow.sdk.execution_time.comms import (
SentFDs,
SetAssetStateStoreByName,
SetAssetStateStoreByUri,
- SetExecutionTimeout,
SetRenderedFields,
SetRenderedMapIndex,
SetTaskStateStore,
SetXCom,
SkipDownstreamTasks,
+ StartExecutionTimeout,
StartupDetails,
SucceedTask,
TaskBreadcrumbsResult,
@@ -1708,16 +1709,9 @@ class TestWatchedSubprocess:
mocker.patch("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
5.0)
mock_kill =
mocker.patch("airflow.sdk.execution_time.supervisor.WatchedSubprocess.kill")
monotonic = mocker.patch("time.monotonic", autospec=True,
return_value=1.0)
- proc = ActivitySubprocess(
- process_log=mocker.MagicMock(),
- id=TI_ID,
- pid=12345,
- stdin=mocker.Mock(),
- process=mocker.Mock(),
- client=mocker.Mock(),
- )
+ proc = self._start_activity_subprocess(mocker, execution_timeout=30.0)
- proc._handle_request(SetExecutionTimeout(timeout_seconds=30),
log=mocker.Mock(), req_id=1)
+ proc._handle_request(StartExecutionTimeout(), log=mocker.Mock(),
req_id=1)
monotonic.return_value += 34
proc._handle_execution_timeout_if_needed()
@@ -1728,13 +1722,58 @@ class TestWatchedSubprocess:
mock_kill.assert_called_once_with(signal.SIGTERM, force=True,
escalation_delay=5.0)
proc.process_log.error.assert_called_once_with(
"Task did not stop before execution_timeout elapsed; terminating
process",
- timeout_seconds=30,
+ timeout_seconds=30.0,
grace_period_seconds=5.0,
)
proc._handle_execution_timeout_if_needed()
mock_kill.assert_called_once()
+ @staticmethod
+ def _start_activity_subprocess(mocker, *, execution_timeout):
+ client = mocker.Mock()
+ client.task_instances.start.return_value =
TIRunContext.model_construct(
+ should_retry=False, execution_timeout=execution_timeout
+ )
+ proc = ActivitySubprocess(
+ process_log=mocker.MagicMock(),
+ id=TI_ID,
+ pid=12345,
+ stdin=mocker.Mock(),
+ process=mocker.Mock(),
+ client=client,
+ )
+ mocker.patch.object(ActivitySubprocess, "send_msg", autospec=True)
+ proc._on_child_started(
+ ti=TaskInstance(
+ id=TI_ID, task_id="b", dag_id="c", run_id="d", try_number=1,
dag_version_id=uuid7()
+ ),
+ dag_rel_path="d.py",
+ bundle_info=BundleInfo(name="my-bundle", version=None),
+ sentry_integration="",
+ )
+ return proc
+
+ def test_python_task_does_not_arm_execution_timeout_at_startup(self,
mocker):
+ mock_kill =
mocker.patch("airflow.sdk.execution_time.supervisor.WatchedSubprocess.kill")
+ monotonic = mocker.patch("time.monotonic", autospec=True,
return_value=1.0)
+ proc = self._start_activity_subprocess(mocker, execution_timeout=30.0)
+
+ monotonic.return_value += 3600
+ proc._handle_execution_timeout_if_needed()
+ mock_kill.assert_not_called()
+
+ def test_start_execution_timeout_without_timeout_never_kills_task(self,
mocker):
+ mock_kill =
mocker.patch("airflow.sdk.execution_time.supervisor.WatchedSubprocess.kill")
+ monotonic = mocker.patch("time.monotonic", autospec=True,
return_value=1.0)
+ proc = self._start_activity_subprocess(mocker, execution_timeout=None)
+
+ proc._handle_request(StartExecutionTimeout(), log=mocker.Mock(),
req_id=1)
+
+ monotonic.return_value += 3600
+ proc._handle_execution_timeout_if_needed()
+ mock_kill.assert_not_called()
+
@pytest.mark.parametrize(
("terminal_state", "exit_code"),
[
@@ -1746,15 +1785,8 @@ class TestWatchedSubprocess:
mocker.patch("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
5.0)
mock_kill =
mocker.patch("airflow.sdk.execution_time.supervisor.WatchedSubprocess.kill")
monotonic = mocker.patch("time.monotonic", autospec=True,
return_value=1.0)
- proc = ActivitySubprocess(
- process_log=mocker.MagicMock(),
- id=TI_ID,
- pid=12345,
- stdin=mocker.Mock(),
- process=mocker.Mock(),
- client=mocker.Mock(),
- )
- proc._handle_request(SetExecutionTimeout(timeout_seconds=30),
log=mocker.Mock(), req_id=1)
+ proc = self._start_activity_subprocess(mocker, execution_timeout=30.0)
+ proc._handle_request(StartExecutionTimeout(), log=mocker.Mock(),
req_id=1)
proc._terminal_state = terminal_state
proc._exit_code = exit_code
@@ -1778,6 +1810,12 @@ class TestWatchedSubprocess:
# Far longer than the test may take: the monitor loop has to wake up
for the timeout on its own
mocker.patch("airflow.sdk.execution_time.supervisor.MIN_HEARTBEAT_INTERVAL", 30)
+ client_with_ti_start.task_instances.start.return_value = (
+ client_with_ti_start.task_instances.start.return_value.model_copy(
+ update={"execution_timeout": 0.1}
+ )
+ )
+
def subprocess_main():
import signal
@@ -1791,7 +1829,7 @@ class TestWatchedSubprocess:
print("Ignoring SIGTERM", file=sys.stderr)
signal.signal(signal.SIGTERM, _on_term)
- comms.send(SetExecutionTimeout(timeout_seconds=0.1))
+ comms.send(StartExecutionTimeout())
sleep(30)
exit(5)
@@ -2733,8 +2771,8 @@ REQUEST_TEST_CASES = [
test_id="set_rtif",
),
RequestTestCase(
- message=SetExecutionTimeout(timeout_seconds=30.0),
- test_id="set_execution_timeout",
+ message=StartExecutionTimeout(),
+ test_id="start_execution_timeout",
),
RequestTestCase(
message=SetRenderedMapIndex(rendered_map_index="Label: task_1"),
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 14e1ec4801e..bef7f005c56 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
@@ -151,11 +151,11 @@ from airflow.sdk.execution_time.comms import (
RetryTask,
SetAssetStateStoreByName,
SetAssetStateStoreByUri,
- SetExecutionTimeout,
SetRenderedFields,
SetTaskStateStore,
SetXCom,
SkipDownstreamTasks,
+ StartExecutionTimeout,
StartupDetails,
SucceedTask,
TaskRescheduleStartDate,
@@ -6887,7 +6887,7 @@ class TestRunExecuteCallable:
_run_execute_callable(context={}, execute=execute, task=task)
task.on_kill.assert_called_once()
-
mock_supervisor_comms.send.assert_called_once_with(SetExecutionTimeout(timeout_seconds=0.01))
+
mock_supervisor_comms.send.assert_called_once_with(StartExecutionTimeout())
def test_fast_fails_when_timeout_already_elapsed(self,
mock_supervisor_comms):
"""A non-positive timeout fast-fails before running the callable and
still calls on_kill."""
diff --git a/ts-sdk/schema/supervisor-schema.json
b/ts-sdk/schema/supervisor-schema.json
index e739708d12d..ebe8c2ab65c 100644
--- a/ts-sdk/schema/supervisor-schema.json
+++ b/ts-sdk/schema/supervisor-schema.json
@@ -3791,27 +3791,6 @@
"title": "SetAssetStateStoreByUri",
"type": "object"
},
- "SetExecutionTimeout": {
- "description": "Tell the supervisor the task's ``execution_timeout`` so
it can enforce it from outside the task process.",
- "properties": {
- "timeout_seconds": {
- "exclusiveMinimum": 0,
- "title": "Timeout Seconds",
- "type": "number"
- },
- "type": {
- "const": "SetExecutionTimeout",
- "default": "SetExecutionTimeout",
- "title": "Type",
- "type": "string"
- }
- },
- "required": [
- "timeout_seconds"
- ],
- "title": "SetExecutionTimeout",
- "type": "object"
- },
"SetRenderedFields": {
"description": "Payload for setting RTIF for a task instance.",
"properties": {
@@ -4005,6 +3984,19 @@
"title": "SkipDownstreamTasks",
"type": "object"
},
+ "StartExecutionTimeout": {
+ "description": "Tell the supervisor to start enforcing the run context's
``execution_timeout`` from outside the task process.",
+ "properties": {
+ "type": {
+ "const": "StartExecutionTimeout",
+ "default": "StartExecutionTimeout",
+ "title": "Type",
+ "type": "string"
+ }
+ },
+ "title": "StartExecutionTimeout",
+ "type": "object"
+ },
"StartupDetails": {
"properties": {
"ti": {
@@ -5399,6 +5391,18 @@
"default": null,
"title": "Arg Bindings"
},
+ "execution_timeout": {
+ "anyOf": [
+ {
+ "type": "number"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Execution Timeout"
+ },
"multi_team": {
"default": false,
"title": "Multi Team",
diff --git a/ts-sdk/src/generated/supervisor.ts
b/ts-sdk/src/generated/supervisor.ts
index 75a05dde12c..69285544734 100644
--- a/ts-sdk/src/generated/supervisor.ts
+++ b/ts-sdk/src/generated/supervisor.ts
@@ -270,6 +270,7 @@ export type TaskId1 = string;
export type Name9 = string;
export type Kind1 = "literal";
export type FromDefault = boolean;
+export type ExecutionTimeout = number | null;
export type MultiTeam = boolean;
export type Type13 = "TaskCallbackRequest";
export type Filepath2 = string;
@@ -579,15 +580,13 @@ export type Type66 = "SetAssetStateStoreByName";
export type Uri9 = string;
export type Key16 = string;
export type Type67 = "SetAssetStateStoreByUri";
-export type TimeoutSeconds = number;
-export type Type68 = "SetExecutionTimeout";
-export type Type69 = "SetRenderedFields";
+export type Type68 = "SetRenderedFields";
export type RenderedMapIndex3 = string;
-export type Type70 = "SetRenderedMapIndex";
+export type Type69 = "SetRenderedMapIndex";
export type TiId8 = string;
export type Key17 = string;
export type ExpiresAt = string | null;
-export type Type71 = "SetTaskStateStore";
+export type Type70 = "SetTaskStateStore";
export type Key18 = string;
export type DagId21 = string;
export type RunId14 = string;
@@ -595,9 +594,10 @@ export type TaskId9 = string;
export type MapIndex7 = number | null;
export type DagResult1 = boolean;
export type MappedLength = number | null;
-export type Type72 = "SetXCom";
+export type Type71 = "SetXCom";
export type Tasks = (string | [unknown, unknown])[];
-export type Type73 = "SkipDownstreamTasks";
+export type Type72 = "SkipDownstreamTasks";
+export type Type73 = "StartExecutionTimeout";
export type DagRelPath = string;
export type StartDate6 = string;
export type SentryIntegration = string;
@@ -1068,6 +1068,7 @@ export interface TIRunContext {
should_retry?: ShouldRetry;
start_date?: StartDate2;
arg_bindings?: ArgBindings;
+ execution_timeout?: ExecutionTimeout;
multi_team?: MultiTeam;
}
/**
@@ -1781,16 +1782,6 @@ export interface SetAssetStateStoreByUri {
value: JsonValue;
type?: Type67;
}
-/**
- * Tell the supervisor the task's ``execution_timeout`` so it can enforce it
from outside the task process.
- *
- * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema
- * via the `definition` "SetExecutionTimeout".
- */
-export interface SetExecutionTimeout {
- timeout_seconds: TimeoutSeconds;
- type?: Type68;
-}
/**
* Payload for setting RTIF for a task instance.
*
@@ -1799,7 +1790,7 @@ export interface SetExecutionTimeout {
*/
export interface SetRenderedFields {
rendered_fields: RenderedFields;
- type?: Type69;
+ type?: Type68;
}
export interface RenderedFields {
[k: string]: JsonValue;
@@ -1812,7 +1803,7 @@ export interface RenderedFields {
*/
export interface SetRenderedMapIndex {
rendered_map_index: RenderedMapIndex3;
- type?: Type70;
+ type?: Type69;
}
/**
* This interface was referenced by `SupervisorWireSchema`'s JSON-Schema
@@ -1823,7 +1814,7 @@ export interface SetTaskStateStore {
key: Key17;
value: JsonValue;
expires_at: ExpiresAt;
- type?: Type71;
+ type?: Type70;
}
/**
* This interface was referenced by `SupervisorWireSchema`'s JSON-Schema
@@ -1838,7 +1829,7 @@ export interface SetXCom {
map_index?: MapIndex7;
dag_result?: DagResult1;
mapped_length?: MappedLength;
- type?: Type72;
+ type?: Type71;
}
/**
* Update state of downstream tasks within a task instance to 'skipped', while
updating current task to success state.
@@ -1848,6 +1839,15 @@ export interface SetXCom {
*/
export interface SkipDownstreamTasks {
tasks: Tasks;
+ type?: Type72;
+}
+/**
+ * Tell the supervisor to start enforcing the run context's
``execution_timeout`` from outside the task process.
+ *
+ * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema
+ * via the `definition` "StartExecutionTimeout".
+ */
+export interface StartExecutionTimeout {
type?: Type73;
}
/**
