From a68f177f48131f77129a26d153df8332479e7b77 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Fri, 16 Oct 2026 23:25:32 +0800 Subject: [PATCH 01/15] Extract a helper that looks up the stub task --- .../execution_api/services/task_instances.py | 23 +++++++++++++++---- 1 file changed, 18 insertions(+), 5 deletions(-) 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 ab67aaa9814f0..a08c63b53a2cd 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 @@ 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,13 +45,12 @@ 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_stub_task(dag_bag: DBDagBag, ti: Any, *, session: Session) -> SerializedOperator | None: """ - Extract the stub task's TaskFlow arg spec from its Dag version. + Return the task of *ti* from its Dag version if a language SDK runs it, else ``None``. - 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. + Both ``@task.stub`` tasks and the tasks of a Dag defined natively in a language SDK + serialize with ``is_stub=True``. """ if ti.dag_version_id is None: return None @@ -58,4 +58,17 @@ def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | return None if (task := dag.task_dict.get(ti.task_id)) is None or not task.is_stub: return None + return task + + +def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | None: + """ + Extract the stub task's TaskFlow arg spec from its Dag version. + + 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 (task := get_stub_task(dag_bag, ti, session=session)) is None: + return None return task.arg_bindings From d3d535b34ac316b0ada618d0601b7ebae454a4bd Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Fri, 16 Oct 2026 23:35:16 +0800 Subject: [PATCH 02/15] Add execution_timeout to TIRunContext --- .../execution_api/datamodels/taskinstance.py | 10 ++++++ .../execution_api/routes/task_instances.py | 6 ++++ .../execution_api/services/task_instances.py | 12 +++++++ .../execution_api/versions/__init__.py | 2 ++ .../execution_api/versions/v2026_10_30.py | 14 ++++++++ .../versions/head/test_task_instances.py | 29 ++++++++++++++++ .../v2026_10_30/test_task_instances.py | 34 +++++++++++++++++++ .../airflow/sdk/api/datamodels/_generated.py | 1 + 8 files changed, 108 insertions(+) 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 8a1deb203fd89..0bf68e7629c2b 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,16 @@ 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, for tasks run by a language SDK. + + Covers ``@task.stub`` tasks and the tasks of a Dag defined natively in a language SDK. + The supervisor arms its execution timeout deadline from it, as a fallback for runtimes + that do not send ``SetExecutionTimeout``. ``None`` for Python tasks, which report their + own timeout, and 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 226ec01b91edf..38d2c1f8383e2 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 @@ -92,6 +92,7 @@ from airflow.api_fastapi.execution_api.services.task_instances import ( client_supports_arg_bindings, get_arg_bindings, + get_language_sdk_execution_timeout, ) from airflow.api_fastapi.execution_api.versions.v2026_10_30 import IdentifyArchivedTaskStateUpdates from airflow.configuration import conf @@ -402,6 +403,11 @@ def ti_run( }, ) + if ( + execution_timeout := get_language_sdk_execution_timeout(dag_bag, ti, session=session) + ) 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 a08c63b53a2cd..3d706b6b0c77e 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 @@ -72,3 +72,15 @@ def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | if (task := get_stub_task(dag_bag, ti, session=session)) is None: return None return task.arg_bindings + + +def get_language_sdk_execution_timeout(dag_bag: DBDagBag, ti: Any, *, session: Session) -> float | None: + """ + Return the ``execution_timeout`` in seconds of a task that a language SDK runs. + + A Python task parses its Dag file and reports its own timeout, so it gets ``None`` here. + A zero timeout also returns ``None``, as the Python task runner treats it as no timeout. + """ + if (task := get_stub_task(dag_bag, ti, session=session)) is None or not task.execution_timeout: + return None + return task.execution_timeout.total_seconds() 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 e85afc58b9bdd..d66d483f2aacd 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 @@ AddArgBindingsToTIRunContext, AddCallbackRunEndpoint, AddDagRunNoteUpdateEndpoint, + AddExecutionTimeoutToTIRunContext, AddExpandedTICountToTIRunContext, AddFirstTaskRescheduleStartDateField, AddLoopContext, @@ -79,6 +80,7 @@ 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 1f7780ee32171..c642624114bb2 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 @@ def remove_arg_bindings_field(response: ResponseInfo) -> None: response.body.pop("arg_bindings", None) +class AddExecutionTimeoutToTIRunContext(VersionChange): + """Add ``execution_timeout`` for tasks that a language SDK runs, 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/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 d241600d281da..8a3caae3ad289 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 @@ -1151,6 +1151,35 @@ 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(execution_timeout=timedelta(seconds=10)), None, id="python-task"), + ], + ) + def test_ti_run_returns_execution_timeout_only_for_language_sdk_task( + 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_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 513d15f6673dc..ae9029f3bea85 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 @@ 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.""" diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 6c4b2536fc44c..6bef12629d956 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -1017,4 +1017,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 From 5ea481fefe5774b485bd8094153a3205b46402a1 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Fri, 16 Oct 2026 23:43:58 +0800 Subject: [PATCH 03/15] Add execution_timeout to the supervisor schema --- go-sdk/pkg/execution/genmodels/models.gen.go | 142 +++++++++--------- go-sdk/schema/supervisor-schema.json | 12 ++ .../sdk/execution_time/schema/schema.json | 12 ++ .../schema/versions/__init__.py | 2 + .../schema/versions/v2026_10_30.py | 15 ++ .../execution_time/schema/test_migrator.py | 10 ++ ts-sdk/schema/supervisor-schema.json | 12 ++ ts-sdk/src/generated/supervisor.ts | 2 + 8 files changed, 137 insertions(+), 70 deletions(-) diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index fbf36031cea99..9d94e14942365 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1709,6 +1709,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 +1753,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 +1839,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 +1866,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 +1892,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 +1947,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 e739708d12dfc..db0185ca2fbd9 100644 --- a/go-sdk/schema/supervisor-schema.json +++ b/go-sdk/schema/supervisor-schema.json @@ -5399,6 +5399,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/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index e739708d12dfc..db0185ca2fbd9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -5399,6 +5399,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 40da65651ae83..27ff074bb237b 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 d69abf63b5f34..3b669458f1706 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 timeout of a task that a language SDK runs. + + 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/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index 6e86037fe0550..f89e6b87e81f7 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 @@ def test_task_coordinates_follow_supervisor_version(self, real_migrator, startup 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/ts-sdk/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json index e739708d12dfc..db0185ca2fbd9 100644 --- a/ts-sdk/schema/supervisor-schema.json +++ b/ts-sdk/schema/supervisor-schema.json @@ -5399,6 +5399,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 75a05dde12cd6..7f436a9b5c326 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; @@ -1068,6 +1069,7 @@ export interface TIRunContext { should_retry?: ShouldRetry; start_date?: StartDate2; arg_bindings?: ArgBindings; + execution_timeout?: ExecutionTimeout; multi_team?: MultiTeam; } /** From e1d0380dae47348d88b72bb51704d5fbea3d9707 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Fri, 16 Oct 2026 23:52:17 +0800 Subject: [PATCH 04/15] Arm the supervisor execution timeout from the run context --- .../airflow/sdk/execution_time/supervisor.py | 25 ++++++--- .../execution_time/test_supervisor.py | 54 +++++++++++++++++++ 2 files changed, 71 insertions(+), 8 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 388075858bd23..1083e383a2541 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1740,6 +1740,9 @@ def _on_child_started( return ti_context = prepared + if ti_context.execution_timeout is not None: + self._arm_execution_timeout(ti_context.execution_timeout) + # ti_context.start_date is only populated by the server when resuming from a deferral (to preserve the # original start_date rather than using the resume time). We fall back to now() otherwise. This ensures # that `context["ti"].start_date` always reflects the *first* start time. See TIRunContext.start_date @@ -1996,13 +1999,16 @@ def _handle_execution_timeout_if_needed(self): """ 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, ``SetExecutionTimeout`` arms the deadline right before ``execute()``, so bundle + loading and Dag parsing are not covered. For a task a language SDK runs, the ``execution_timeout`` + the server sends in the run context arms it when the runtime starts, and a runtime that also sends + ``SetExecutionTimeout`` re-arms it right before running the task. Either way, the process is sent + SIGTERM ``killed_task_cleanup_time`` after the timeout elapses. """ if self._execution_timeout_due_in() > 0: return @@ -2149,10 +2155,13 @@ def _handle_skip_downstream_tasks( def _handle_set_execution_timeout( self, msg: SetExecutionTimeout, 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(msg.timeout_seconds) return None, {} + def _arm_execution_timeout(self, timeout_seconds: float) -> None: + self._execution_timeout_seconds = timeout_seconds + self._execution_timeout_enforce_at = time.monotonic() + timeout_seconds + KILLED_TASK_CLEANUP_TIME + def _handle_set_rendered_fields( self, msg: SetRenderedFields, log: FilteringBoundLogger, req_id: int ) -> RequestResult: 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 12a76d25a1be4..e1d5c4e1216d8 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 @@ PreviousTIResponse, TaskInstance, TaskInstanceState, + TIRunContext, ) from airflow.sdk.definitions._internal.loop import create_loop from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType, TaskAlreadyRunningError @@ -1735,6 +1736,59 @@ def test_execution_timeout_kills_task_that_does_not_stop(self, mocker): 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_execution_timeout_from_server_kills_task_that_does_not_stop(self, mocker): + 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 = self._start_activity_subprocess(mocker, execution_timeout=30.0) + + monotonic.return_value += 34 + proc._handle_execution_timeout_if_needed() + mock_kill.assert_not_called() + + monotonic.return_value += 1 + proc._handle_execution_timeout_if_needed() + 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.0, + grace_period_seconds=5.0, + ) + + def test_no_execution_timeout_from_server_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) + + monotonic.return_value += 3600 + proc._handle_execution_timeout_if_needed() + mock_kill.assert_not_called() + @pytest.mark.parametrize( ("terminal_state", "exit_code"), [ From 1642f5cd939b6fdc0e118ea83653ae04b203e0f1 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Sat, 17 Oct 2026 21:49:20 +0800 Subject: [PATCH 05/15] Document that the supervisor timeout covers language SDK tasks --- airflow-core/docs/core-concepts/tasks.rst | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/airflow-core/docs/core-concepts/tasks.rst b/airflow-core/docs/core-concepts/tasks.rst index f57de76baa839..d97c6410b41d6 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 From abf692aa35db4dd5d74984ea2d98729365140aa5 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Sat, 17 Oct 2026 22:14:45 +0800 Subject: [PATCH 06/15] 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. --- .../api_fastapi/execution_api/services/task_instances.py | 6 +++--- .../execution_api/versions/head/test_task_instances.py | 3 +++ 2 files changed, 6 insertions(+), 3 deletions(-) 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 3d706b6b0c77e..9f1e2b5991866 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 @@ -78,9 +78,9 @@ def get_language_sdk_execution_timeout(dag_bag: DBDagBag, ti: Any, *, session: S """ Return the ``execution_timeout`` in seconds of a task that a language SDK runs. - A Python task parses its Dag file and reports its own timeout, so it gets ``None`` here. - A zero timeout also returns ``None``, as the Python task runner treats it as no timeout. + Python tasks report their own timeout, so they get ``None``. Like the Python task + runner, a zero timeout means no timeout, and a negative one has already elapsed. """ if (task := get_stub_task(dag_bag, ti, session=session)) is None or not task.execution_timeout: return None - return task.execution_timeout.total_seconds() + return max(task.execution_timeout.total_seconds(), 0.0) 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 8a3caae3ad289..eb4811ece4b52 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 @@ -1157,6 +1157,9 @@ def transform(country: str): ... 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)), None, id="python-task"), ], ) From 511024bad41bcd50c12cf2f41a68163cdc9a08c6 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Mon, 19 Oct 2026 09:22:24 +0000 Subject: [PATCH 07/15] Send execution_timeout in the run context for all tasks --- .../execution_api/datamodels/taskinstance.py | 7 ++---- .../execution_api/routes/task_instances.py | 6 ++--- .../execution_api/services/task_instances.py | 25 +++++++------------ .../execution_api/versions/v2026_10_30.py | 2 +- .../versions/head/test_task_instances.py | 23 ++++++++++++++--- 5 files changed, 33 insertions(+), 30 deletions(-) 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 0bf68e7629c2b..db998c4ad5484 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 @@ -557,12 +557,9 @@ class TIRunContext(BaseModel): execution_timeout: float | None = None """ - The task's ``execution_timeout`` in seconds, for tasks run by a language SDK. + The task's ``execution_timeout`` in seconds, taken from its serialized Dag. - Covers ``@task.stub`` tasks and the tasks of a Dag defined natively in a language SDK. - The supervisor arms its execution timeout deadline from it, as a fallback for runtimes - that do not send ``SetExecutionTimeout``. ``None`` for Python tasks, which report their - own timeout, and when no timeout is set. + The supervisor enforces the timeout with this value. ``None`` when no timeout is set. """ multi_team: bool = False 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 38d2c1f8383e2..62c6df9b14ada 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 @@ -92,7 +92,7 @@ from airflow.api_fastapi.execution_api.services.task_instances import ( client_supports_arg_bindings, get_arg_bindings, - get_language_sdk_execution_timeout, + get_execution_timeout, ) from airflow.api_fastapi.execution_api.versions.v2026_10_30 import IdentifyArchivedTaskStateUpdates from airflow.configuration import conf @@ -403,9 +403,7 @@ def ti_run( }, ) - if ( - execution_timeout := get_language_sdk_execution_timeout(dag_bag, ti, session=session) - ) is not None: + if (execution_timeout := get_execution_timeout(dag_bag, ti, session=session)) is not None: context.execution_timeout = execution_timeout # Only set if they are non-null 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 9f1e2b5991866..7e1b0526fe0ff 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 @@ -45,20 +45,13 @@ def client_supports_arg_bindings() -> bool: return AddArgBindingsToTIRunContext.is_applied -def get_stub_task(dag_bag: DBDagBag, ti: Any, *, session: Session) -> SerializedOperator | None: - """ - Return the task of *ti* from its Dag version if a language SDK runs it, else ``None``. - - Both ``@task.stub`` tasks and the tasks of a Dag defined natively in a language SDK - serialize with ``is_stub=True``. - """ +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 - if (task := dag.task_dict.get(ti.task_id)) is None or not task.is_stub: - return None - return task + return dag.task_dict.get(ti.task_id) def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | None: @@ -69,18 +62,18 @@ def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | ``None`` here and keep the legacy ignored-args behavior; per-map-index delivery lands in a follow-up. """ - if (task := get_stub_task(dag_bag, ti, session=session)) is None: + if (task := get_task(dag_bag, ti, session=session)) is None or not task.is_stub: return None return task.arg_bindings -def get_language_sdk_execution_timeout(dag_bag: DBDagBag, ti: Any, *, session: Session) -> float | None: +def get_execution_timeout(dag_bag: DBDagBag, ti: Any, *, session: Session) -> float | None: """ - Return the ``execution_timeout`` in seconds of a task that a language SDK runs. + Return the ``execution_timeout`` in seconds of *ti*'s task from its Dag version. - Python tasks report their own timeout, so they get ``None``. Like the Python task - runner, a zero timeout means no timeout, and a negative one has already elapsed. + Like the Python task runner, a zero timeout means no timeout, and a negative one has + already elapsed. """ - if (task := get_stub_task(dag_bag, ti, session=session)) is None or not task.execution_timeout: + if (task := get_task(dag_bag, ti, session=session)) 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/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py index c642624114bb2..5d23b073ec51c 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 @@ -94,7 +94,7 @@ def remove_arg_bindings_field(response: ResponseInfo) -> None: class AddExecutionTimeoutToTIRunContext(VersionChange): - """Add ``execution_timeout`` for tasks that a language SDK runs, so the supervisor can enforce it.""" + """Add ``execution_timeout`` so the supervisor can enforce it.""" description = __doc__ instructions_to_migrate_to_previous_version = ( 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 eb4811ece4b52..1b320c462c439 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 @@ -1160,12 +1160,10 @@ def transform(country: str): ... pytest.param( task.stub(execution_timeout=timedelta(seconds=-30)), 0.0, id="stub-negative-timeout" ), - pytest.param(task(execution_timeout=timedelta(seconds=10)), None, id="python-task"), + pytest.param(task(execution_timeout=timedelta(seconds=10)), 10.0, id="python-task"), ], ) - def test_ti_run_returns_execution_timeout_only_for_language_sdk_task( - self, client, dag_maker, decorator, expected - ): + def test_ti_run_returns_execution_timeout(self, client, dag_maker, decorator, expected): with dag_maker("test_execution_timeout_dag", serialized=True): @decorator @@ -1183,6 +1181,23 @@ def sleeper(): ... 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 From 3523dc42c7b3f9bac5ccbb46ff53ee58cbcb2c82 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Mon, 19 Oct 2026 09:57:17 +0000 Subject: [PATCH 08/15] Start the execution timeout for language SDK tasks in the coordinator --- .../airflow/sdk/coordinators/_subprocess.py | 5 +++ .../airflow/sdk/execution_time/supervisor.py | 11 +++--- .../task_sdk/coordinators/test_subprocess.py | 35 +++++++++++++++++++ .../execution_time/test_supervisor.py | 21 +---------- 4 files changed, 45 insertions(+), 27 deletions(-) diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index 0bd84e147bea0..a358f34c447e2 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -420,6 +420,9 @@ def _prepare_ti_context(self, ti: TaskInstance, ti_context: TIRunContext) -> TIR backend's storage. If an XCom cannot be deleted, this try fails and the task does not run. + + Once the XComs are deleted, this also arms the execution timeout from the run context, since a + language SDK runtime does not tell the supervisor when the task starts running. """ for key in ti_context.xcom_keys_to_clear or (): self.process_log.debug("Clearing XCom with key", key=key) @@ -438,6 +441,8 @@ def _prepare_ti_context(self, ti: TaskInstance, ti_context: TIRunContext) -> TIR ) self.kill(signal.SIGKILL) return None + if ti_context.execution_timeout is not None: + self._arm_execution_timeout(ti_context.execution_timeout) # The supervisor has already deleted the listed XComs, so it sends the runtime an empty list. # If a runtime also deletes these XComs itself, the empty list keeps the runtime from deleting # them a second time. diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 1083e383a2541..8cd4d206f9231 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1740,9 +1740,6 @@ def _on_child_started( return ti_context = prepared - if ti_context.execution_timeout is not None: - self._arm_execution_timeout(ti_context.execution_timeout) - # ti_context.start_date is only populated by the server when resuming from a deferral (to preserve the # original start_date rather than using the resume time). We fall back to now() otherwise. This ensures # that `context["ti"].start_date` always reflects the *first* start time. See TIRunContext.start_date @@ -2005,10 +2002,10 @@ def _handle_execution_timeout_if_needed(self): period is stuck (native code, SIGSEGV) and is sent SIGTERM, then SIGKILL. For a Python task, ``SetExecutionTimeout`` arms the deadline right before ``execute()``, so bundle - loading and Dag parsing are not covered. For a task a language SDK runs, the ``execution_timeout`` - the server sends in the run context arms it when the runtime starts, and a runtime that also sends - ``SetExecutionTimeout`` re-arms it right before running the task. Either way, the process is sent - SIGTERM ``killed_task_cleanup_time`` after the timeout elapses. + loading and Dag parsing are not covered; Dag parsing has its own ``dagbag_import_timeout``. For a + task a language SDK runs, the coordinator arms it from the run context once the runtime has + started. Either way, the process is sent SIGTERM ``killed_task_cleanup_time`` after the timeout + elapses. """ if self._execution_timeout_due_in() > 0: return diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 63eafc6d1ab07..9daae4232a06b 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -1150,6 +1150,41 @@ def handle_request(request: httpx2.Request) -> httpx2.Response: } 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("time.monotonic", autospec=True, return_value=1.0) + def test_arms_execution_timeout_from_run_context( + self, mock_monotonic, mock_kill, execution_timeout, killed + ): + proc = _PopenActivitySubprocess( + process_log=MagicMock(), + id=uuid7(), + pid=12345, + stdin=MagicMock(spec=socket.socket), + process=MagicMock(), + client=MagicMock(), + comm_server=MagicMock(spec=socket.socket), + logs_server=MagicMock(spec=socket.socket), + ) + + proc._prepare_ti_context( + _make_ti(), + TIRunContext.model_construct(execution_timeout=execution_timeout, xcom_keys_to_clear=[]), + ) + + 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/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index e1d5c4e1216d8..2790e95eb9ebd 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -1761,30 +1761,11 @@ def _start_activity_subprocess(mocker, *, execution_timeout): ) return proc - def test_execution_timeout_from_server_kills_task_that_does_not_stop(self, mocker): - mocker.patch("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME", 5.0) + 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 += 34 - proc._handle_execution_timeout_if_needed() - mock_kill.assert_not_called() - - monotonic.return_value += 1 - proc._handle_execution_timeout_if_needed() - 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.0, - grace_period_seconds=5.0, - ) - - def test_no_execution_timeout_from_server_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) - monotonic.return_value += 3600 proc._handle_execution_timeout_if_needed() mock_kill.assert_not_called() From ed640230401fbbbc000c69783c146a44bd95f41a Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Mon, 19 Oct 2026 10:13:12 +0000 Subject: [PATCH 09/15] Replace SetExecutionTimeout with StartExecutionTimeout --- .../unit/dag_processing/test_processor.py | 2 +- .../tests/unit/jobs/test_triggerer_job.py | 2 +- .../execution/genmodels/discriminators.gen.go | 8 +-- go-sdk/pkg/execution/genmodels/models.gen.go | 17 +++---- go-sdk/schema/supervisor-schema.json | 34 +++++-------- .../src/airflow/sdk/execution_time/comms.py | 9 ++-- .../sdk/execution_time/schema/schema.json | 34 +++++-------- .../airflow/sdk/execution_time/supervisor.py | 14 +++--- .../airflow/sdk/execution_time/task_runner.py | 4 +- .../execution_time/test_supervisor.py | 49 ++++++++++--------- .../execution_time/test_task_runner.py | 4 +- ts-sdk/schema/supervisor-schema.json | 34 +++++-------- ts-sdk/src/generated/supervisor.ts | 40 +++++++-------- 13 files changed, 113 insertions(+), 138 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 0271a5e07fbe9..154cb7640ce6d 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 @@ 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", diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index bbf01eba0f109..6cd5c17a22612 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 @@ def get_type_names(union_type): "ResendLoggingFD", "CreateHITLDetailPayload", "SetRenderedMapIndex", - "SetExecutionTimeout", + "StartExecutionTimeout", "GetDag", # AIP-103 task store — triggerer has no task execution context. "GetTaskStateStore", diff --git a/go-sdk/pkg/execution/genmodels/discriminators.gen.go b/go-sdk/pkg/execution/genmodels/discriminators.gen.go index 8736a6e2c9730..3a81b6ce7a88a 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 9d94e14942365..f95ceb53d629a 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"` diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json index db0185ca2fbd9..ebe8c2ab65ce8 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": { diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index bd5781f5dfe3a..45af5b757a3fb 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 @@ class GetDag(BaseModel): | 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 db0185ca2fbd9..ebe8c2ab65ce8 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": { diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 8cd4d206f9231..a7ee5caec6f62 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 @@ SentFDs, SetAssetStateStoreByName, SetAssetStateStoreByUri, - SetExecutionTimeout, SetRenderedFields, SetRenderedMapIndex, SetTaskStateStore, SetXCom, SkipDownstreamTasks, + StartExecutionTimeout, StartupDetails, SucceedTask, TaskBreadcrumbsResult, @@ -1726,6 +1726,7 @@ def _on_child_started( # 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! @@ -2001,7 +2002,7 @@ def _handle_execution_timeout_if_needed(self): 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. - For a Python task, ``SetExecutionTimeout`` arms the deadline right before ``execute()``, so bundle + 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, the coordinator arms it from the run context once the runtime has started. Either way, the process is sent SIGTERM ``killed_task_cleanup_time`` after the timeout @@ -2149,10 +2150,11 @@ def _handle_skip_downstream_tasks( 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._arm_execution_timeout(msg.timeout_seconds) + if self._execution_timeout_seconds is not None: + self._arm_execution_timeout(self._execution_timeout_seconds) return None, {} def _arm_execution_timeout(self, timeout_seconds: float) -> None: @@ -2453,11 +2455,11 @@ def _handle_clear_asset_state_store_by_uri( 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 ee78312e95ed2..6651a7785e939 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 @@ 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/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 2790e95eb9ebd..72a9e29c7e38b 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -141,12 +141,12 @@ SentFDs, SetAssetStateStoreByName, SetAssetStateStoreByUri, - SetExecutionTimeout, SetRenderedFields, SetRenderedMapIndex, SetTaskStateStore, SetXCom, SkipDownstreamTasks, + StartExecutionTimeout, StartupDetails, SucceedTask, TaskBreadcrumbsResult, @@ -1709,16 +1709,9 @@ def test_execution_timeout_kills_task_that_does_not_stop(self, mocker): 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() @@ -1729,7 +1722,7 @@ def test_execution_timeout_kills_task_that_does_not_stop(self, mocker): 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, ) @@ -1770,6 +1763,17 @@ def test_python_task_does_not_arm_execution_timeout_at_startup(self, mocker): 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"), [ @@ -1781,15 +1785,8 @@ def test_execution_timeout_not_enforced_once_task_is_over(self, mocker, terminal 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 @@ -1813,6 +1810,12 @@ def test_execution_timeout_enforced_by_supervisor( # 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 @@ -1826,7 +1829,7 @@ def _on_term(signum, frame): 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) @@ -2768,8 +2771,8 @@ class RequestTestCase: 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 14e1ec4801e4e..bef7f005c5619 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 @@ RetryTask, SetAssetStateStoreByName, SetAssetStateStoreByUri, - SetExecutionTimeout, SetRenderedFields, SetTaskStateStore, SetXCom, SkipDownstreamTasks, + StartExecutionTimeout, StartupDetails, SucceedTask, TaskRescheduleStartDate, @@ -6887,7 +6887,7 @@ def execute(context): _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 db0185ca2fbd9..ebe8c2ab65ce8 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": { diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 7f436a9b5c326..69285544734c5 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -580,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; @@ -596,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; @@ -1783,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. * @@ -1801,7 +1790,7 @@ export interface SetExecutionTimeout { */ export interface SetRenderedFields { rendered_fields: RenderedFields; - type?: Type69; + type?: Type68; } export interface RenderedFields { [k: string]: JsonValue; @@ -1814,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 @@ -1825,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 @@ -1840,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. @@ -1850,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; } /** From 911f9d2bb2c67c2fd9d202c5fd7f0590affe4f3c Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Tue, 20 Oct 2026 19:51:49 +0800 Subject: [PATCH 10/15] Look up the task once in ti_run --- .../api_fastapi/execution_api/routes/task_instances.py | 8 ++++---- .../execution_api/services/task_instances.py | 10 +++++----- 2 files changed, 9 insertions(+), 9 deletions(-) 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 62c6df9b14ada..a46dd0ed4d07a 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 @@ -93,6 +93,7 @@ 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 @@ -383,9 +384,8 @@ def ti_run( # 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) - ): + task = get_task(dag_bag, ti, session=session) + if client_supports_arg_bindings() and (arg_bindings := get_arg_bindings(task)): try: context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) except ValidationError: @@ -403,7 +403,7 @@ def ti_run( }, ) - if (execution_timeout := get_execution_timeout(dag_bag, ti, session=session)) is not None: + if (execution_timeout := get_execution_timeout(task)) is not None: context.execution_timeout = execution_timeout # Only set if they are non-null 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 7e1b0526fe0ff..bc7ce2a36ede8 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 @@ -54,7 +54,7 @@ def get_task(dag_bag: DBDagBag, ti: Any, *, session: Session) -> SerializedOpera return dag.task_dict.get(ti.task_id) -def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | None: +def get_arg_bindings(task: SerializedOperator | None) -> list | None: """ Extract the stub task's TaskFlow arg spec from its Dag version. @@ -62,18 +62,18 @@ def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | ``None`` here and keep the legacy ignored-args behavior; per-map-index delivery lands in a follow-up. """ - if (task := get_task(dag_bag, ti, session=session)) 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(dag_bag: DBDagBag, ti: Any, *, session: Session) -> float | None: +def get_execution_timeout(task: SerializedOperator | None) -> float | None: """ - Return the ``execution_timeout`` in seconds of *ti*'s task from its Dag version. + 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 := get_task(dag_bag, ti, session=session)) is None or not task.execution_timeout: + if task is None or not task.execution_timeout: return None return max(task.execution_timeout.total_seconds(), 0.0) From 49386f9368fc8411e9f133797cac75da9a9db3da Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Tue, 20 Oct 2026 22:13:16 +0800 Subject: [PATCH 11/15] Fill the default execution_timeout of Lang-SDK tasks from the config --- .../serialization/serialized_objects.py | 9 ++++++--- .../serialization/test_dag_serialization.py | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py index f5fd8c118b9b3..73fc26e01147a 100644 --- a/airflow-core/src/airflow/serialization/serialized_objects.py +++ b/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: diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index d40f36ebf6d33..0caec6fa2ba80 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 @@ 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} From fad04260f5302028082c902d8317785b853c4b22 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Tue, 20 Oct 2026 22:17:15 +0800 Subject: [PATCH 12/15] Arm the execution timeout only in _arm_execution_timeout --- task-sdk/src/airflow/sdk/execution_time/supervisor.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index a7ee5caec6f62..07ed4b3c503a6 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1664,6 +1664,7 @@ class ActivitySubprocess(WatchedSubprocess): _task_end_time_monotonic: float | None = attrs.field(default=None, init=False) _rendered_map_index: str | None = attrs.field(default=None, init=False) + _run_context_execution_timeout: float | None = attrs.field(default=None, init=False) _execution_timeout_seconds: float | None = attrs.field(default=None, init=False) _execution_timeout_enforce_at: float | None = attrs.field(default=None, init=False) """Monotonic time at which the supervisor stops a task that has not reported a terminal state.""" @@ -1726,7 +1727,7 @@ def _on_child_started( # 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._run_context_execution_timeout = ti_context.execution_timeout self._last_successful_heartbeat = time.monotonic() except Exception as e: # On any error kill that subprocess! @@ -2153,8 +2154,8 @@ def _handle_skip_downstream_tasks( def _handle_start_execution_timeout( self, msg: StartExecutionTimeout, log: FilteringBoundLogger, req_id: int ) -> RequestResult: - if self._execution_timeout_seconds is not None: - self._arm_execution_timeout(self._execution_timeout_seconds) + if self._run_context_execution_timeout is not None: + self._arm_execution_timeout(self._run_context_execution_timeout) return None, {} def _arm_execution_timeout(self, timeout_seconds: float) -> None: From 30f73dff5337c6376d07295fcef9ada908e5b666 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Wed, 21 Oct 2026 09:17:47 +0000 Subject: [PATCH 13/15] Store execution_timeout at startup and arm it in _on_child_started --- .../airflow/sdk/coordinators/_subprocess.py | 22 ++++++++++++++----- .../airflow/sdk/execution_time/supervisor.py | 15 +++++++------ .../task_sdk/coordinators/test_subprocess.py | 17 +++++++++----- 3 files changed, 37 insertions(+), 17 deletions(-) diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index a358f34c447e2..e102b424a7350 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -408,6 +408,23 @@ def start( # type: ignore[override] 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. @@ -420,9 +437,6 @@ def _prepare_ti_context(self, ti: TaskInstance, ti_context: TIRunContext) -> TIR backend's storage. If an XCom cannot be deleted, this try fails and the task does not run. - - Once the XComs are deleted, this also arms the execution timeout from the run context, since a - language SDK runtime does not tell the supervisor when the task starts running. """ for key in ti_context.xcom_keys_to_clear or (): self.process_log.debug("Clearing XCom with key", key=key) @@ -441,8 +455,6 @@ def _prepare_ti_context(self, ti: TaskInstance, ti_context: TIRunContext) -> TIR ) self.kill(signal.SIGKILL) return None - if ti_context.execution_timeout is not None: - self._arm_execution_timeout(ti_context.execution_timeout) # The supervisor has already deleted the listed XComs, so it sends the runtime an empty list. # If a runtime also deletes these XComs itself, the empty list keeps the runtime from deleting # them a second time. diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 07ed4b3c503a6..d5a31040322c4 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1664,7 +1664,6 @@ class ActivitySubprocess(WatchedSubprocess): _task_end_time_monotonic: float | None = attrs.field(default=None, init=False) _rendered_map_index: str | None = attrs.field(default=None, init=False) - _run_context_execution_timeout: float | None = attrs.field(default=None, init=False) _execution_timeout_seconds: float | None = attrs.field(default=None, init=False) _execution_timeout_enforce_at: float | None = attrs.field(default=None, init=False) """Monotonic time at which the supervisor stops a task that has not reported a terminal state.""" @@ -1727,7 +1726,7 @@ def _on_child_started( # 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._run_context_execution_timeout = ti_context.execution_timeout + self._execution_timeout_seconds = ti_context.execution_timeout self._last_successful_heartbeat = time.monotonic() except Exception as e: # On any error kill that subprocess! @@ -2154,13 +2153,15 @@ def _handle_skip_downstream_tasks( def _handle_start_execution_timeout( self, msg: StartExecutionTimeout, log: FilteringBoundLogger, req_id: int ) -> RequestResult: - if self._run_context_execution_timeout is not None: - self._arm_execution_timeout(self._run_context_execution_timeout) + self._arm_execution_timeout() return None, {} - def _arm_execution_timeout(self, timeout_seconds: float) -> None: - self._execution_timeout_seconds = timeout_seconds - self._execution_timeout_enforce_at = time.monotonic() + timeout_seconds + KILLED_TASK_CLEANUP_TIME + 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 diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 9daae4232a06b..f2e2622255feb 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -1160,24 +1160,31 @@ class TestPopenActivitySubprocessExecutionTimeout: ) @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_kill, execution_timeout, killed + 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=MagicMock(), + client=client, comm_server=MagicMock(spec=socket.socket), logs_server=MagicMock(spec=socket.socket), ) - proc._prepare_ti_context( - _make_ti(), - TIRunContext.model_construct(execution_timeout=execution_timeout, xcom_keys_to_clear=[]), + 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 From ef13d374aad713e56761e0462b06b5778b6bbc07 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Wed, 21 Oct 2026 09:21:40 +0000 Subject: [PATCH 14/15] Align docstrings with the execution timeout changes --- .../api_fastapi/execution_api/routes/task_instances.py | 3 ++- .../api_fastapi/execution_api/services/task_instances.py | 2 +- .../sdk/execution_time/schema/versions/v2026_10_30.py | 2 +- task-sdk/src/airflow/sdk/execution_time/supervisor.py | 5 ++--- 4 files changed, 6 insertions(+), 6 deletions(-) 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 a46dd0ed4d07a..ee10e7defacbe 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 @@ -382,9 +382,10 @@ def ti_run( if first_task_reschedule_start_date is not None: context.first_task_reschedule_start_date = first_task_reschedule_start_date + 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. - task = get_task(dag_bag, ti, session=session) if client_supports_arg_bindings() and (arg_bindings := get_arg_bindings(task)): try: context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) 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 bc7ce2a36ede8..15e0d6fcf8cc8 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 @@ -56,7 +56,7 @@ def get_task(dag_bag: DBDagBag, ti: Any, *, session: Session) -> SerializedOpera 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 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 3b669458f1706..ddc99b274b987 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 @@ -137,7 +137,7 @@ class AddExpandedTICountToSupervisorTIRunContext(VersionChange): class AddExecutionTimeoutToSupervisorTIRunContext(VersionChange): """ - Add the ``execution_timeout`` field carrying the timeout of a task that a language SDK runs. + 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. diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index d5a31040322c4..656efa07b13de 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -2004,9 +2004,8 @@ def _handle_execution_timeout_if_needed(self): 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, the coordinator arms it from the run context once the runtime has - started. Either way, the process is sent SIGTERM ``killed_task_cleanup_time`` after the timeout - elapses. + 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 From 287202eade160885254d99d8142e6b6c5d138084 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Thu, 22 Oct 2026 05:31:58 +0000 Subject: [PATCH 15/15] Rename the task variable in ti_run to fix mypy --- .../api_fastapi/execution_api/routes/task_instances.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 ee10e7defacbe..d092b9f62a63e 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 @@ -382,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 - task = get_task(dag_bag, ti, session=session) + 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(task)): + 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: @@ -404,7 +404,7 @@ def ti_run( }, ) - if (execution_timeout := get_execution_timeout(task)) is not None: + if (execution_timeout := get_execution_timeout(serialized_task)) is not None: context.execution_timeout = execution_timeout # Only set if they are non-null