diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/index.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/index.rst index ea8d4b8448d10..f2ad01e921ca9 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/index.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/index.rst @@ -224,13 +224,13 @@ release that shipped it. See the legend below the tables for what each mark mean - Java - TypeScript * - Min. Airflow version + - 3.5.0 - 3.4.0 - - 3.4.0 - - 3.4.0 + - 3.5.0 * - Current Supervisor Schema Version + - 2027-02-28 - 2026-10-30 - - 2026-10-30 - - 2026-10-30 + - 2027-02-28 .. list-table:: Common to both authoring modes :header-rows: 1 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 c90d64da1ec3f..4a98f3f205e70 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 @@ -556,13 +556,15 @@ class TIRunContext(BaseModel): ``None`` for regular tasks and for stub tasks that declare no parameters. """ - execution_timeout: float | None = None - """ - The ``execution_timeout`` in seconds of a task a language SDK runs, from its serialized Dag. - - The supervisor enforces the timeout with this value. A Python task runs its own timer in the task - runner instead, and gets no value here. ``None`` when no timeout is set. - """ + execution_timeout: float | None = Field( + default=None, + description=( + "The execution_timeout in seconds set on a task a language SDK runs, from its serialized Dag. " + "Unset when the task sets none. Zero means no timeout and a negative value has already " + "elapsed, as for a Python task. A Python task runs its own timer in the task runner instead, " + "and gets no value here." + ), + ) 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 9473accb162a8..05fbe062d5aef 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 @@ -384,36 +384,38 @@ def ti_run( if first_task_reschedule_start_date is not None: context.first_task_reschedule_start_date = first_task_reschedule_start_date - # Avoid loading the serialized Dag on the normal Python task path: only a lang-SDK - # (foreign-runtime) task needs its definition here, and ti.is_stub carries the row's own copy - # of the serialized flag. - serialized_task = get_task(dag_bag, ti, session=session) if ti.is_stub else None - - # 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(serialized_task)): - try: - context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) - except ValidationError: - log.exception( - "Serialized arg_bindings spec failed validation", - dag_id=ti.dag_id, - task_id=ti.task_id, - dag_version_id=ti.dag_version_id, - ) - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={ - "reason": "invalid_arg_bindings", - "message": "The serialized TaskFlow arg spec for this stub task is not valid.", - }, - ) + # Only a lang-SDK (foreign-runtime) task needs its definition here: a Python task parses its + # Dag file itself. ti.is_stub carries the row's own copy of the serialized flag, so a Python + # task never loads the serialized Dag. The arg_bindings version check also covers + # execution_timeout, which only exists in a later API version. + if ( + client_supports_arg_bindings() + and ti.is_stub + and (serialized_task := get_task(dag_bag, ti, session=session)) is not None + ): + # The route excludes unset fields, keeping responses without a captured TaskFlow arg spec lean. + if arg_bindings := get_arg_bindings(serialized_task): + try: + context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) + except ValidationError: + log.exception( + "Serialized arg_bindings spec failed validation", + dag_id=ti.dag_id, + task_id=ti.task_id, + dag_version_id=ti.dag_version_id, + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ + "reason": "invalid_arg_bindings", + "message": "The serialized TaskFlow arg spec for this stub task is not valid.", + }, + ) - # Resolved here for a lang-SDK task only, since a regular task no longer loads its - # definition. #75387 has a Python task send its own execution_timeout to the supervisor, - # which is where the value for that kind of task belongs. - if (execution_timeout := get_execution_timeout(serialized_task)) is not None: - context.execution_timeout = execution_timeout + # The runtime cannot read the execution_timeout set on a Python stub, so it gets it here. + # A Python task enforces its own and sends that to the supervisor. + if (execution_timeout := get_execution_timeout(serialized_task)) is not None: + context.execution_timeout = execution_timeout # Only set if they are non-null if ti.next_method: 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 15e0d6fcf8cc8..9f33341144e98 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(task: SerializedOperator | None) -> list | None: +def get_arg_bindings(task: SerializedOperator) -> list | None: """ Return the TaskFlow arg spec of *task* if it is a stub task. @@ -62,18 +62,18 @@ def get_arg_bindings(task: SerializedOperator | None) -> list | None: ``None`` here and keep the legacy ignored-args behavior; per-map-index delivery lands in a follow-up. """ - if task is None or not task.is_stub: + if not task.is_stub: return None return task.arg_bindings -def get_execution_timeout(task: SerializedOperator | None) -> float | None: +def get_execution_timeout(task: SerializedOperator) -> float | None: """ - Return the ``execution_timeout`` of *task* in seconds. + Return the ``execution_timeout`` of *task* in seconds, or ``None`` if it sets none. - Like the Python task runner, a zero timeout means no timeout, and a negative one has - already elapsed. + The value is sent as it is, so a runtime handles it the way the Python task runner does: zero + means no timeout, and a negative one has already elapsed. """ - if task is None or not task.execution_timeout: + if task.execution_timeout is None: return None - return max(task.execution_timeout.total_seconds(), 0.0) + 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 d66d483f2aacd..2c1127b0d5e81 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,7 +56,6 @@ AddArgBindingsToTIRunContext, AddCallbackRunEndpoint, AddDagRunNoteUpdateEndpoint, - AddExecutionTimeoutToTIRunContext, AddExpandedTICountToTIRunContext, AddFirstTaskRescheduleStartDateField, AddLoopContext, @@ -70,17 +69,19 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30_xcom_params import ( AddPreviousIterationToXComFilterParams, ) -from airflow.api_fastapi.execution_api.versions.v2027_02_28 import AddDagProcessorJobEndpoints +from airflow.api_fastapi.execution_api.versions.v2027_02_28 import ( + AddDagProcessorJobEndpoints, + AddExecutionTimeoutToTIRunContext, +) bundle = VersionBundle( HeadVersion(), - Version("2027-02-28", AddDagProcessorJobEndpoints), + Version("2027-02-28", AddDagProcessorJobEndpoints, AddExecutionTimeoutToTIRunContext), Version( "2026-10-30", 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 5d23b073ec51c..1f7780ee32171 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,20 +93,6 @@ def remove_arg_bindings_field(response: ResponseInfo) -> None: response.body.pop("arg_bindings", None) -class AddExecutionTimeoutToTIRunContext(VersionChange): - """Add ``execution_timeout`` so the supervisor can enforce it.""" - - description = __doc__ - instructions_to_migrate_to_previous_version = ( - schema(TIRunContext).field("execution_timeout").didnt_exist, - ) - - @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] - def remove_execution_timeout_field(response: ResponseInfo) -> None: # type: ignore[misc] - """Strip ``execution_timeout`` from the run context for older clients.""" - response.body.pop("execution_timeout", None) - - class AddCallbackRunEndpoint(VersionChange): """Add the callbacks/{callback_id}/run endpoint a worker uses to exchange its single-use callback token.""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py index b54a50236284a..3b39aa4e83a64 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py @@ -17,7 +17,9 @@ from __future__ import annotations -from cadwyn import VersionChange, endpoint +from cadwyn import ResponseInfo, VersionChange, convert_response_to_previous_version_for, endpoint, schema + +from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext class AddDagProcessorJobEndpoints(VersionChange): @@ -31,3 +33,17 @@ class AddDagProcessorJobEndpoints(VersionChange): endpoint("/jobs/{job_id}/complete", ["POST"]).didnt_exist, endpoint("/jobs/{job_id}/parse-token", ["POST"]).didnt_exist, ) + + +class AddExecutionTimeoutToTIRunContext(VersionChange): + """Add ``execution_timeout`` so the supervisor can enforce it.""" + + description = __doc__ + instructions_to_migrate_to_previous_version = ( + schema(TIRunContext).field("execution_timeout").didnt_exist, + ) + + @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] + def remove_execution_timeout_field(response: ResponseInfo) -> None: # type: ignore[misc] + """Strip ``execution_timeout`` from the run context for older clients.""" + response.body.pop("execution_timeout", None) 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 2c8dc7e95a9c0..caa0266fc731c 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 @@ -1199,9 +1199,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(0)), 0.0, id="stub-zero-timeout"), pytest.param( - task.stub(execution_timeout=timedelta(seconds=-30)), 0.0, id="stub-negative-timeout" + task.stub(execution_timeout=timedelta(seconds=-30)), -30.0, id="stub-negative-timeout" ), # A Python task's timeout stays with the task runner, which sends it to the supervisor, # so the server resolves one only for a task whose definition it loads anyway. @@ -2153,6 +2153,7 @@ def test_ti_run_skips_queued_duration_metric_when_the_response_fails( dag_id=str(uuid4()), ) ti.queued_dttm = queued_at + ti.is_stub = True session.commit() with mock.patch( 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 ae9029f3bea85..513d15f6673dc 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,7 +17,6 @@ from __future__ import annotations -from datetime import timedelta from uuid import uuid4 import pytest @@ -54,7 +53,6 @@ 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( @@ -289,38 +287,6 @@ 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/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2027_02_28/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2027_02_28/test_task_instances.py new file mode 100644 index 0000000000000..430bcf2c5a568 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2027_02_28/test_task_instances.py @@ -0,0 +1,84 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from datetime import timedelta + +import pytest + +from airflow.sdk import task +from airflow.utils.state import State + +from tests_common.test_utils.db import clear_db_runs + +pytestmark = pytest.mark.db_test + +RUN_PATCH_BODY = { + "state": "running", + "hostname": "h", + "unixname": "u", + "pid": 1, + "start_date": "2024-09-30T12:00:00Z", +} + + +@pytest.fixture +def old_ver_client(client): + """Execution API version immediately before ``execution_timeout`` was added.""" + client.headers["Airflow-API-Version"] = "2026-10-30" + return client + + +@pytest.mark.parametrize("version", ["2026-10-30", "2027-02-28"]) +def test_execution_timeout_schema_by_version(client, version): + response = client.get(f"/execution/openapi.json?version={version}") + assert response.status_code == 200 + context = response.json()["components"]["schemas"]["TIRunContext"]["properties"] + assert ("execution_timeout" in context) == (version == "2027-02-28") + + +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 diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index c9c4ab19e089f..f56de3f3141d3 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -2380,7 +2380,7 @@ def get_type_names(union_type): "UpdateHITLDetail", "GetHITLDetailResponse", "SetRenderedMapIndex", - "StartExecutionTimeout", + "SetExecutionTimeout", # 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 ff0e03e8099e2..a614b30cc78e7 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -3359,7 +3359,7 @@ def get_type_names(union_type): "ResendLoggingFD", "CreateHITLDetailPayload", "SetRenderedMapIndex", - "StartExecutionTimeout", + "SetExecutionTimeout", "GetDag", # AIP-103 task store — triggerer has no task execution context. "GetTaskStateStore", diff --git a/go-sdk/README.md b/go-sdk/README.md index 6da0ccba329aa..8ecef83ee845f 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -117,7 +117,7 @@ prek hook regenerate it. -*Min. Airflow version: 3.4.0 · supervisor schema: 2026-10-30* +*Min. Airflow version: 3.5.0 · supervisor schema: 2027-02-28* | Dimension | Tier | Supported | Since (Airflow) | Since (SDK) | Notes | |---|---|---|---|---|---| diff --git a/go-sdk/capabilities.yaml b/go-sdk/capabilities.yaml index 95b39727d4981..1fab164e58f36 100644 --- a/go-sdk/capabilities.yaml +++ b/go-sdk/capabilities.yaml @@ -17,11 +17,11 @@ --- sdk: go -min_airflow_version: "3.4.0" +min_airflow_version: "3.5.0" # Keep in sync with SupervisorSchemaVersion in pkg/execution/messages.go, which is embedded in # bundle metadata. The render hook fails if the two disagree. -supervisor_schema_version: "2026-10-30" +supervisor_schema_version: "2027-02-28" # The task runner terminates a task with SucceedTask, RetryTask, or TaskState (failed/removed); only # an airflow.TriggerDagRun task emits TaskState skipped or DeferTask. It does not emit RescheduleTask diff --git a/go-sdk/pkg/execution/genmodels/discriminators.gen.go b/go-sdk/pkg/execution/genmodels/discriminators.gen.go index 3a81b6ce7a88a..8736a6e2c9730 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,6 +321,9 @@ 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 @@ -336,9 +339,6 @@ 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 f95ceb53d629a..3cefba3e3f9e0 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1555,6 +1555,16 @@ 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". @@ -1629,13 +1639,6 @@ 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"` @@ -1706,7 +1709,10 @@ 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". + // The execution_timeout in seconds set on a task a language SDK runs, from its + // serialized Dag. Unset when the task sets none. Zero means no timeout and a + // negative value has already elapsed, as for a Python task. A Python task runs + // its own timer in the task runner instead, and gets no value here. ExecutionTimeout interface{} `msgpack:"execution_timeout,omitempty"` // ExpandedTICount corresponds to the JSON schema field "expanded_ti_count". diff --git a/go-sdk/pkg/execution/messages.go b/go-sdk/pkg/execution/messages.go index 1dd596201258a..3f45c1c3fc185 100644 --- a/go-sdk/pkg/execution/messages.go +++ b/go-sdk/pkg/execution/messages.go @@ -32,7 +32,7 @@ import ( // reported in a bundle's airflow-metadata manifest as // sdk.supervisor_schema_version so the supervisor can down/upgrade messages to // a shape the bundle understands. -const SupervisorSchemaVersion = "2026-10-30" +const SupervisorSchemaVersion = "2027-02-28" // The message-type discriminator strings (genmodels.Type*) are generated from the // schema's "type" consts in discriminators.gen.go; outbound messages stamp the diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json index ebe8c2ab65ce8..1db287cb17db5 100644 --- a/go-sdk/schema/supervisor-schema.json +++ b/go-sdk/schema/supervisor-schema.json @@ -1,6 +1,6 @@ { "$schema": "https://json-schema.org/draft/2020-12/schema", - "api_version": "2026-10-30", + "api_version": "2027-02-28", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { "AssetAliasReferenceAssetEventDagRun": { @@ -3791,6 +3791,27 @@ "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": { @@ -3984,19 +4005,6 @@ "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": { @@ -5401,6 +5409,7 @@ } ], "default": null, + "description": "The execution_timeout in seconds set on a task a language SDK runs, from its serialized Dag. Unset when the task sets none. Zero means no timeout and a negative value has already elapsed, as for a Python task. A Python task runs its own timer in the task runner instead, and gets no value here.", "title": "Execution Timeout" }, "multi_team": { diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 648aa3bda2833..701be36e39513 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -1017,5 +1017,11 @@ 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 + execution_timeout: Annotated[ + float | None, + Field( + description="The execution_timeout in seconds set on a task a language SDK runs, from its serialized Dag. Unset when the task sets none. Zero means no timeout and a negative value has already elapsed, as for a Python task. A Python task runs its own timer in the task runner instead, and gets no value here.", + title="Execution Timeout", + ), + ] = None multi_team: Annotated[bool | None, Field(title="Multi Team")] = False diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index aa30b834ef64d..241f8733aea11 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -546,7 +546,8 @@ def _on_child_started( bundle_info=bundle_info, sentry_integration=sentry_integration, ) - # Language runtimes do not send StartExecutionTimeout, so the clock starts once they have the task. + # A language runtime may not send SetExecutionTimeout, so the clock starts once it has the task. + # If the runtime sends it, the clock restarts from that message. self._arm_execution_timeout() def _prepare_ti_context(self, ti: TaskInstance, ti_context: TIRunContext) -> TIRunContext | None: @@ -562,6 +563,7 @@ def _prepare_ti_context(self, ti: TaskInstance, ti_context: TIRunContext) -> TIR If an XCom cannot be deleted, this try fails and the task does not run. """ + self._execution_timeout_seconds = ti_context.execution_timeout for key in ti_context.xcom_keys_to_clear or (): self.process_log.debug("Clearing XCom with key", key=key) try: diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 45af5b757a3fb..f0e8336426c34 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -1136,10 +1136,11 @@ class SetRenderedMapIndex(BaseModel): type: Literal["SetRenderedMapIndex"] = "SetRenderedMapIndex" -class StartExecutionTimeout(BaseModel): - """Tell the supervisor to start enforcing the run context's ``execution_timeout`` from outside the task process.""" +class SetExecutionTimeout(BaseModel): + """Tell the supervisor the task's ``execution_timeout`` so it can enforce it from outside the task process.""" - type: Literal["StartExecutionTimeout"] = "StartExecutionTimeout" + timeout_seconds: Annotated[float, Field(gt=0)] + type: Literal["SetExecutionTimeout"] = "SetExecutionTimeout" class TriggerDagRun(TriggerDAGRunPayload): @@ -1351,7 +1352,7 @@ class GetDag(BaseModel): | SetTaskStateStore | SetXCom | SkipDownstreamTasks - | StartExecutionTimeout + | SetExecutionTimeout | 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 ebe8c2ab65ce8..1db287cb17db5 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1,6 +1,6 @@ { "$schema": "https://json-schema.org/draft/2020-12/schema", - "api_version": "2026-10-30", + "api_version": "2027-02-28", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { "AssetAliasReferenceAssetEventDagRun": { @@ -3791,6 +3791,27 @@ "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": { @@ -3984,19 +4005,6 @@ "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": { @@ -5401,6 +5409,7 @@ } ], "default": null, + "description": "The execution_timeout in seconds set on a task a language SDK runs, from its serialized Dag. Unset when the task sets none. Zero means no timeout and a negative value has already elapsed, as for a Python task. A Python task runs its own timer in the task runner instead, and gets no value here.", "title": "Execution Timeout" }, "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 27ff074bb237b..83446d7b776c8 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,16 +40,19 @@ def get_bundle() -> VersionBundle: from airflow.sdk.execution_time.schema.versions.v2026_10_30 import ( AddArgBindingsToSupervisorTIRunContext, AddDagDefinitionsToDagFileParsingResult, - AddExecutionTimeoutToSupervisorTIRunContext, AddExpandedTICountToSupervisorTIRunContext, AddFirstTaskRescheduleStartDateToSupervisorTIRunContext, AddLoopRuntime, AddRegionCoordinatesToSupervisorTaskInstance, AddRetryReasonToTaskState, ) + from airflow.sdk.execution_time.schema.versions.v2027_02_28 import ( + AddExecutionTimeoutToSupervisorTIRunContext, + ) return VersionBundle( HeadVersion(), + Version("2027-02-28", AddExecutionTimeoutToSupervisorTIRunContext), Version( "2026-10-30", AddArgBindingsToSupervisorTIRunContext, @@ -59,7 +62,6 @@ 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 ddc99b274b987..d69abf63b5f34 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,18 +133,3 @@ class AddExpandedTICountToSupervisorTIRunContext(VersionChange): instructions_to_migrate_to_previous_version = ( schema(TIRunContext).field("expanded_ti_count").didnt_exist, ) - - -class AddExecutionTimeoutToSupervisorTIRunContext(VersionChange): - """ - Add the ``execution_timeout`` field carrying the task's timeout, which the supervisor enforces. - - The supervisor-schema mirror of the execution API's ``AddExecutionTimeoutToTIRunContext``, - named apart so the two migrations are not confused. - """ - - description = __doc__ - - instructions_to_migrate_to_previous_version = ( - schema(TIRunContext).field("execution_timeout").didnt_exist, - ) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2027_02_28.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2027_02_28.py new file mode 100644 index 0000000000000..8a0e320a1e8bb --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2027_02_28.py @@ -0,0 +1,39 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Supervisor schema version 2027-02-28.""" + +from __future__ import annotations + +from cadwyn import VersionChange, schema + +from airflow.sdk.api.datamodels._generated import TIRunContext + + +class AddExecutionTimeoutToSupervisorTIRunContext(VersionChange): + """ + Add the ``execution_timeout`` field carrying the task's timeout, which the supervisor enforces. + + The supervisor-schema mirror of the execution API's ``AddExecutionTimeoutToTIRunContext``, + named apart so the two migrations are not confused. + """ + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + schema(TIRunContext).field("execution_timeout").didnt_exist, + ) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 05810f0e45326..4f9622e34fa33 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, @@ -1744,7 +1744,6 @@ 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! @@ -2020,10 +2019,10 @@ 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, ``StartExecutionTimeout`` arms the deadline right before ``execute()``, so bundle - loading and Dag parsing are not covered; Dag parsing has its own ``dagbag_import_timeout``. For a - task a language SDK runs, it is armed once the runtime has been sent ``StartupDetails``. Either - way, the process is sent SIGTERM ``killed_task_cleanup_time`` after the timeout elapses. + ``SetExecutionTimeout`` arms the deadline right before ``execute()`` with the timeout the task + process enforces itself, so both use the same value; bundle loading and Dag parsing are not covered, + as Dag parsing has its own ``dagbag_import_timeout``. The process is sent SIGTERM + ``killed_task_cleanup_time`` after the timeout elapses. """ if self._execution_timeout_due_in() > 0: return @@ -2167,17 +2166,20 @@ def _handle_skip_downstream_tasks( self.client.task_instances.skip_downstream_tasks(self.id, msg) return None, {} - def _handle_start_execution_timeout( - self, msg: StartExecutionTimeout, log: FilteringBoundLogger, req_id: int + def _handle_set_execution_timeout( + self, msg: SetExecutionTimeout, log: FilteringBoundLogger, req_id: int ) -> RequestResult: + self._execution_timeout_seconds = msg.timeout_seconds self._arm_execution_timeout() return None, {} def _arm_execution_timeout(self) -> None: - if self._execution_timeout_seconds is None: + # Zero means no timeout, as a zero timedelta does for a Python task. + if not self._execution_timeout_seconds: return + # A negative timeout has already elapsed, so the task gets only the grace period. self._execution_timeout_enforce_at = ( - time.monotonic() + self._execution_timeout_seconds + KILLED_TASK_CLEANUP_TIME + time.monotonic() + max(self._execution_timeout_seconds, 0.0) + KILLED_TASK_CLEANUP_TIME ) def _handle_set_rendered_fields( @@ -2478,7 +2480,7 @@ def _handle_clear_asset_state_store_by_uri( 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(SetExecutionTimeout, _handle_set_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 caccb6d737c37..ee6f50295c47e 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, @@ -2352,7 +2352,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(StartExecutionTimeout()) + SUPERVISOR_COMMS.send(SetExecutionTimeout(timeout_seconds=timeout_seconds)) # Run task in timeout wrapper with timeout(timeout_seconds): result = ctx.run(execute, context=context) diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index b9f16aeeb9a84..a9b7ab4736408 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -1324,10 +1324,13 @@ def test_stopping_the_task_also_stops_what_the_runtime_spawned(self, tmp_path): class TestPopenActivitySubprocessExecutionTimeout: @pytest.mark.parametrize( - ("execution_timeout", "killed"), + ("execution_timeout", "elapsed", "killed"), [ - pytest.param(30.0, True, id="with-timeout"), - pytest.param(None, False, id="without-timeout"), + pytest.param(30.0, 35, True, id="with-timeout"), + pytest.param(None, 35, False, id="without-timeout"), + pytest.param(0.0, 35, False, id="zero-means-no-timeout"), + pytest.param(-30.0, 3, False, id="negative-keeps-grace-period"), + pytest.param(-30.0, 6, True, id="negative-has-elapsed"), ], ) @patch("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME", 5.0) @@ -1335,7 +1338,7 @@ class TestPopenActivitySubprocessExecutionTimeout: @patch.object(_PopenActivitySubprocess, "send_msg", autospec=True) @patch("time.monotonic", autospec=True, return_value=1.0) def test_arms_execution_timeout_from_run_context( - self, mock_monotonic, mock_send_msg, mock_kill, execution_timeout, killed + self, mock_monotonic, mock_send_msg, mock_kill, execution_timeout, elapsed, killed ): client = MagicMock() client.task_instances.start.return_value = TIRunContext.model_construct( @@ -1359,7 +1362,7 @@ def test_arms_execution_timeout_from_run_context( sentry_integration="", ) - mock_monotonic.return_value += 35 + mock_monotonic.return_value += elapsed proc._handle_execution_timeout_if_needed() assert mock_kill.called is killed 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 f89e6b87e81f7..2e6db689c8fc0 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,20 +550,20 @@ 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"]) + @pytest.mark.parametrize("version", ["2026-06-16", "2026-10-30", "2027-02-28"]) 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: + if version == "2027-02-28": assert context["execution_timeout"] == 90.0 + else: + assert "execution_timeout" not in context def test_head_version_keeps_arg_bindings(self, real_migrator, startup_details): from airflow.sdk.api.datamodels._generated import LiteralArgBinding, XComArgBinding - out = real_migrator.downgrade(startup_details, "2026-10-30") + out = real_migrator.downgrade(startup_details, "2027-02-28") assert out.ti_context.arg_bindings is not None literal, xcom, defaulted = (a.root for a in out.ti_context.arg_bindings) assert isinstance(literal, LiteralArgBinding) 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 72a9e29c7e38b..bf20ba7e111b2 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, @@ -1711,7 +1711,7 @@ def test_execution_timeout_kills_task_that_does_not_stop(self, mocker): monotonic = mocker.patch("time.monotonic", autospec=True, return_value=1.0) proc = self._start_activity_subprocess(mocker, execution_timeout=30.0) - proc._handle_request(StartExecutionTimeout(), log=mocker.Mock(), req_id=1) + proc._handle_request(SetExecutionTimeout(timeout_seconds=30.0), log=mocker.Mock(), req_id=1) monotonic.return_value += 34 proc._handle_execution_timeout_if_needed() @@ -1763,16 +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): + def test_set_execution_timeout_uses_task_process_value(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=None) + proc = self._start_activity_subprocess(mocker, execution_timeout=3600.0) - proc._handle_request(StartExecutionTimeout(), log=mocker.Mock(), req_id=1) + proc._handle_request(SetExecutionTimeout(timeout_seconds=30.0), log=mocker.Mock(), req_id=1) - monotonic.return_value += 3600 + monotonic.return_value += 35 proc._handle_execution_timeout_if_needed() - mock_kill.assert_not_called() + mock_kill.assert_called_once_with(signal.SIGTERM, force=True, escalation_delay=5.0) @pytest.mark.parametrize( ("terminal_state", "exit_code"), @@ -1786,7 +1787,7 @@ def test_execution_timeout_not_enforced_once_task_is_over(self, mocker, terminal 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) - proc._handle_request(StartExecutionTimeout(), log=mocker.Mock(), req_id=1) + proc._handle_request(SetExecutionTimeout(timeout_seconds=30.0), log=mocker.Mock(), req_id=1) proc._terminal_state = terminal_state proc._exit_code = exit_code @@ -1810,12 +1811,6 @@ 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 @@ -1829,7 +1824,7 @@ def _on_term(signum, frame): print("Ignoring SIGTERM", file=sys.stderr) signal.signal(signal.SIGTERM, _on_term) - comms.send(StartExecutionTimeout()) + comms.send(SetExecutionTimeout(timeout_seconds=0.1)) sleep(30) exit(5) @@ -2771,8 +2766,8 @@ class RequestTestCase: test_id="set_rtif", ), RequestTestCase( - message=StartExecutionTimeout(), - test_id="start_execution_timeout", + message=SetExecutionTimeout(timeout_seconds=30.0), + test_id="set_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 ea1b4afd964b0..311b1f64407f0 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, @@ -6951,7 +6951,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(StartExecutionTimeout()) + mock_supervisor_comms.send.assert_called_once_with(SetExecutionTimeout(timeout_seconds=0.01)) 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/README.md b/ts-sdk/README.md index 490c031bfe6a6..39adbe38f23e1 100644 --- a/ts-sdk/README.md +++ b/ts-sdk/README.md @@ -100,7 +100,7 @@ Do not edit the table by hand. Update the manifest and run the `update-ts-sdk-re -*Min. Airflow version: 3.4.0 · supervisor schema: 2026-10-30* +*Min. Airflow version: 3.5.0 · supervisor schema: 2027-02-28* | Dimension | Tier | Supported | Since (Airflow) | Since (SDK) | Notes | |---|---|---|---|---|---| diff --git a/ts-sdk/capabilities.yaml b/ts-sdk/capabilities.yaml index 6fef910ad011d..a049b7f9fb79a 100644 --- a/ts-sdk/capabilities.yaml +++ b/ts-sdk/capabilities.yaml @@ -17,11 +17,11 @@ --- sdk: ts -min_airflow_version: "3.4.0" +min_airflow_version: "3.5.0" # Keep in sync with SUPERVISOR_API_VERSION in src/generated/supervisor.ts. The render hook fails # if the two disagree. -supervisor_schema_version: "2026-10-30" +supervisor_schema_version: "2027-02-28" # The runtime terminates a task with SucceedTask, RetryTask, TaskState (failed/removed), or, for a # human-in-the-loop task, AwaitInputTask; only a triggerDagRun task emits TaskState skipped or diff --git a/ts-sdk/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json index ebe8c2ab65ce8..1db287cb17db5 100644 --- a/ts-sdk/schema/supervisor-schema.json +++ b/ts-sdk/schema/supervisor-schema.json @@ -1,6 +1,6 @@ { "$schema": "https://json-schema.org/draft/2020-12/schema", - "api_version": "2026-10-30", + "api_version": "2027-02-28", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { "AssetAliasReferenceAssetEventDagRun": { @@ -3791,6 +3791,27 @@ "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": { @@ -3984,19 +4005,6 @@ "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": { @@ -5401,6 +5409,7 @@ } ], "default": null, + "description": "The execution_timeout in seconds set on a task a language SDK runs, from its serialized Dag. Unset when the task sets none. Zero means no timeout and a negative value has already elapsed, as for a Python task. A Python task runs its own timer in the task runner instead, and gets no value here.", "title": "Execution Timeout" }, "multi_team": { diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 69285544734c5..8cc93eb7bedf4 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -270,6 +270,9 @@ export type TaskId1 = string; export type Name9 = string; export type Kind1 = "literal"; export type FromDefault = boolean; +/** + * The execution_timeout in seconds set on a task a language SDK runs, from its serialized Dag. Unset when the task sets none. Zero means no timeout and a negative value has already elapsed, as for a Python task. A Python task runs its own timer in the task runner instead, and gets no value here. + */ export type ExecutionTimeout = number | null; export type MultiTeam = boolean; export type Type13 = "TaskCallbackRequest"; @@ -580,13 +583,15 @@ export type Type66 = "SetAssetStateStoreByName"; export type Uri9 = string; export type Key16 = string; export type Type67 = "SetAssetStateStoreByUri"; -export type Type68 = "SetRenderedFields"; +export type TimeoutSeconds = number; +export type Type68 = "SetExecutionTimeout"; +export type Type69 = "SetRenderedFields"; export type RenderedMapIndex3 = string; -export type Type69 = "SetRenderedMapIndex"; +export type Type70 = "SetRenderedMapIndex"; export type TiId8 = string; export type Key17 = string; export type ExpiresAt = string | null; -export type Type70 = "SetTaskStateStore"; +export type Type71 = "SetTaskStateStore"; export type Key18 = string; export type DagId21 = string; export type RunId14 = string; @@ -594,10 +599,9 @@ export type TaskId9 = string; export type MapIndex7 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; -export type Type71 = "SetXCom"; +export type Type72 = "SetXCom"; export type Tasks = (string | [unknown, unknown])[]; -export type Type72 = "SkipDownstreamTasks"; -export type Type73 = "StartExecutionTimeout"; +export type Type73 = "SkipDownstreamTasks"; export type DagRelPath = string; export type StartDate6 = string; export type SentryIntegration = string; @@ -1782,6 +1786,16 @@ 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. * @@ -1790,7 +1804,7 @@ export interface SetAssetStateStoreByUri { */ export interface SetRenderedFields { rendered_fields: RenderedFields; - type?: Type68; + type?: Type69; } export interface RenderedFields { [k: string]: JsonValue; @@ -1803,7 +1817,7 @@ export interface RenderedFields { */ export interface SetRenderedMapIndex { rendered_map_index: RenderedMapIndex3; - type?: Type69; + type?: Type70; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1814,7 +1828,7 @@ export interface SetTaskStateStore { key: Key17; value: JsonValue; expires_at: ExpiresAt; - type?: Type70; + type?: Type71; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1829,7 +1843,7 @@ export interface SetXCom { map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; - type?: Type71; + type?: Type72; } /** * Update state of downstream tasks within a task instance to 'skipped', while updating current task to success state. @@ -1839,15 +1853,6 @@ 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; } /** @@ -2048,4 +2053,4 @@ export interface XComSequenceSliceResult { * (e.g. bundle metadata) and runs the migrator accordingly. * Exposed so the SDK author / operator can confirm which schema * version their build is pinned to. */ -export const SUPERVISOR_API_VERSION = "2026-10-30" as const; +export const SUPERVISOR_API_VERSION = "2027-02-28" as const;