From d979aa9b57ca453ac93418f59e4cbda2179fb00a Mon Sep 17 00:00:00 2001 From: Andrew Chang Date: Mon, 12 Oct 2026 12:44:02 +0800 Subject: [PATCH 1/4] Extract the Dag import timeout lookup in the Python importer Move the get_dagbag_import_timeout policy lookup, its fallback when airflow settings cannot be imported, and the int/float check into a module-level function, so the task runner can resolve the same value for a Dag file before it is parsed. --- .../airflow/sdk/importers/python_importer.py | 29 ++++++++++--------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/python_importer.py b/task-sdk/src/airflow/sdk/importers/python_importer.py index e6fae5b44529a..72cf8b9807e53 100644 --- a/task-sdk/src/airflow/sdk/importers/python_importer.py +++ b/task-sdk/src/airflow/sdk/importers/python_importer.py @@ -59,6 +59,21 @@ log = logging.getLogger(__name__) +def get_dagbag_import_timeout(source_reference: str) -> float: + """Resolve the Dag import timeout for ``source_reference`` through the ``get_dagbag_import_timeout`` policy.""" + try: + from airflow import settings # noqa: SDK002 + + dagbag_import_timeout = settings.get_dagbag_import_timeout(source_reference) + except (ImportError, AttributeError): + return 30.0 + if not isinstance(dagbag_import_timeout, (int, float)): + raise AirflowConfigException( + f"Value ({dagbag_import_timeout}) from get_dagbag_import_timeout must be int or float" + ) + return dagbag_import_timeout + + class _DefinitionSourceLoader(importlib.abc.SourceLoader): """ A SourceLoader that executes a DagDefinition straight from its bytes. @@ -286,19 +301,7 @@ def parse() -> list[types.ModuleType]: ) return [] - dagbag_import_timeout: float - try: - from airflow import settings # noqa: SDK002 - - dagbag_import_timeout = settings.get_dagbag_import_timeout(repr(definition)) - except (ImportError, AttributeError): - dagbag_import_timeout = 30.0 - - if not isinstance(dagbag_import_timeout, (int, float)): - raise AirflowConfigException( - f"Value ({dagbag_import_timeout}) from get_dagbag_import_timeout must be int or float" - ) - + dagbag_import_timeout = get_dagbag_import_timeout(repr(definition)) if dagbag_import_timeout <= 0: return parse() From 4ad31d156c9a1678de2bc5173db25f108bfd63a5 Mon Sep 17 00:00:00 2001 From: Andrew Chang Date: Mon, 12 Oct 2026 12:45:11 +0800 Subject: [PATCH 2/4] Let SetExecutionTimeout(0) clear the supervisor deadline The task process needs to drop the supervisor's deadline once a phase it armed is over, so a value of 0 now means "no deadline" instead of being rejected. The supervisor's timeout log line no longer names execution_timeout, since it will also fire for the Dag parse phase. SetExecutionTimeout has not been released yet, so the in-progress supervisor schema version only gets a note instead of a VersionChange. --- go-sdk/pkg/execution/genmodels/models.gen.go | 8 ++++-- go-sdk/schema/supervisor-schema.json | 4 +-- .../src/airflow/sdk/execution_time/comms.py | 9 +++++-- .../sdk/execution_time/schema/schema.json | 4 +-- .../schema/versions/v2026_10_30.py | 3 +++ .../airflow/sdk/execution_time/supervisor.py | 13 +++++++--- .../execution_time/test_supervisor.py | 25 +++++++++++++++++-- ts-sdk/schema/supervisor-schema.json | 4 +-- ts-sdk/src/generated/supervisor.ts | 5 +++- 9 files changed, 58 insertions(+), 17 deletions(-) diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index d7d779b15cf5f..34302015653ad 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1533,8 +1533,12 @@ 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. +// Tell the supervisor how long the task process may take before it must be stopped +// from outside. +// +// Sent before parsing a plain Python Dag file with “dagbag_import_timeout“ and +// before “execute()“ +// with the task's “execution_timeout“. “0“ clears the deadline. type SetExecutionTimeout struct { // TimeoutSeconds corresponds to the JSON schema field "timeout_seconds". TimeoutSeconds float64 `msgpack:"timeout_seconds"` diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json index 7e09a736de793..88d6ec0a9072a 100644 --- a/go-sdk/schema/supervisor-schema.json +++ b/go-sdk/schema/supervisor-schema.json @@ -3801,10 +3801,10 @@ "type": "object" }, "SetExecutionTimeout": { - "description": "Tell the supervisor the task's ``execution_timeout`` so it can enforce it from outside the task process.", + "description": "Tell the supervisor how long the task process may take before it must be stopped from outside.\n\nSent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()``\nwith the task's ``execution_timeout``. ``0`` clears the deadline.", "properties": { "timeout_seconds": { - "exclusiveMinimum": 0, + "minimum": 0, "title": "Timeout Seconds", "type": "number" }, diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 7c957027ba792..b6c15ca3918b9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -1130,9 +1130,14 @@ class SetRenderedMapIndex(BaseModel): class SetExecutionTimeout(BaseModel): - """Tell the supervisor the task's ``execution_timeout`` so it can enforce it from outside the task process.""" + """ + Tell the supervisor how long the task process may take before it must be stopped from outside. + + Sent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()`` + with the task's ``execution_timeout``. ``0`` clears the deadline. + """ - timeout_seconds: Annotated[float, Field(gt=0)] + timeout_seconds: Annotated[float, Field(ge=0)] type: Literal["SetExecutionTimeout"] = "SetExecutionTimeout" 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 7e09a736de793..88d6ec0a9072a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -3801,10 +3801,10 @@ "type": "object" }, "SetExecutionTimeout": { - "description": "Tell the supervisor the task's ``execution_timeout`` so it can enforce it from outside the task process.", + "description": "Tell the supervisor how long the task process may take before it must be stopped from outside.\n\nSent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()``\nwith the task's ``execution_timeout``. ``0`` clears the deadline.", "properties": { "timeout_seconds": { - "exclusiveMinimum": 0, + "minimum": 0, "title": "Timeout Seconds", "type": "number" }, 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 baffc8c9d0f1b..6a998a7c5c0a3 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 @@ -22,6 +22,9 @@ pinned to an older version simply never sends it, so there is nothing to strip on the way down. Only *changes* to bodies that already existed at an earlier dated version require a ``VersionChange`` entry below. + +``SetExecutionTimeout.timeout_seconds`` accepts ``0`` (clear the deadline) since the same +version that introduced the body, so it needs no entry either. """ from __future__ import annotations diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index a5f164daa6d83..526496a523690 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1954,21 +1954,22 @@ def _execution_timeout_due_in(self) -> float: def _handle_execution_timeout_if_needed(self): """ - Enforce ``execution_timeout`` from outside the task process. + Enforce the task process's current timeout from outside it. 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. - 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. + The deadline is armed by ``SetExecutionTimeout``: with ``dagbag_import_timeout`` around parsing a + plain Python Dag file and with ``execution_timeout`` right before ``execute()``. Bundle loading and + zipped Dags, whose members are each imported under their own timeout, are not covered. """ if self._execution_timeout_due_in() > 0: return self._execution_timeout_enforce_at = None self.process_log.error( - "Task did not stop before execution_timeout elapsed; terminating process", + "Task did not stop before its timeout elapsed; terminating process", timeout_seconds=self._execution_timeout_seconds, grace_period_seconds=KILLED_TASK_CLEANUP_TIME, ) @@ -2109,6 +2110,10 @@ def _handle_skip_downstream_tasks( def _handle_set_execution_timeout( self, msg: SetExecutionTimeout, log: FilteringBoundLogger, req_id: int ) -> RequestResult: + if msg.timeout_seconds == 0: + self._execution_timeout_seconds = None + self._execution_timeout_enforce_at = None + return None, {} self._execution_timeout_seconds = msg.timeout_seconds self._execution_timeout_enforce_at = time.monotonic() + msg.timeout_seconds + KILLED_TASK_CLEANUP_TIME return None, {} 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 0f507a87f8e47..46434faa73ffc 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -1620,7 +1620,7 @@ def test_execution_timeout_kills_task_that_does_not_stop(self, mocker): 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", + "Task did not stop before its timeout elapsed; terminating process", timeout_seconds=30, grace_period_seconds=5.0, ) @@ -1628,6 +1628,27 @@ def test_execution_timeout_kills_task_that_does_not_stop(self, mocker): proc._handle_execution_timeout_if_needed() mock_kill.assert_called_once() + def test_set_execution_timeout_zero_clears_deadline(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._handle_request(SetExecutionTimeout(timeout_seconds=30), log=mocker.Mock(), req_id=1) + proc._handle_request(SetExecutionTimeout(timeout_seconds=0), log=mocker.Mock(), req_id=2) + + monotonic.return_value += 100 + proc._handle_execution_timeout_if_needed() + + mock_kill.assert_not_called() + proc.process_log.error.assert_not_called() + @pytest.mark.parametrize( ("terminal_state", "exit_code"), [ @@ -1710,7 +1731,7 @@ def _on_term(signum, frame): assert proc.final_state == TaskInstanceState.FAILED assert { - "event": "Task did not stop before execution_timeout elapsed; terminating process", + "event": "Task did not stop before its timeout elapsed; terminating process", "level": "error", "timeout_seconds": 0.1, "grace_period_seconds": 0.3, diff --git a/ts-sdk/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json index 7e09a736de793..88d6ec0a9072a 100644 --- a/ts-sdk/schema/supervisor-schema.json +++ b/ts-sdk/schema/supervisor-schema.json @@ -3801,10 +3801,10 @@ "type": "object" }, "SetExecutionTimeout": { - "description": "Tell the supervisor the task's ``execution_timeout`` so it can enforce it from outside the task process.", + "description": "Tell the supervisor how long the task process may take before it must be stopped from outside.\n\nSent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()``\nwith the task's ``execution_timeout``. ``0`` clears the deadline.", "properties": { "timeout_seconds": { - "exclusiveMinimum": 0, + "minimum": 0, "title": "Timeout Seconds", "type": "number" }, diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 31196f868a5cb..e70d9ba31610c 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -1758,7 +1758,10 @@ export interface SetAssetStateStoreByUri { type?: Type67; } /** - * Tell the supervisor the task's ``execution_timeout`` so it can enforce it from outside the task process. + * Tell the supervisor how long the task process may take before it must be stopped from outside. + * + * Sent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()`` + * with the task's ``execution_timeout``. ``0`` clears the deadline. * * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema * via the `definition` "SetExecutionTimeout". From c050fb573aedd679d94d0d8db437ba60993f6f0f Mon Sep 17 00:00:00 2001 From: Andrew Chang Date: Mon, 12 Oct 2026 12:45:42 +0800 Subject: [PATCH 3/4] Arm the supervisor timeout around Dag parsing in the task runner Before parsing a plain Python Dag file, the task process now sends the file's dagbag_import_timeout to the supervisor and clears it once parsing is done, so a parse stuck in native code is stopped from outside like an execute() that overruns execution_timeout. The value comes from the same get_dagbag_import_timeout policy the importer uses for the file's own alarm. Nothing is sent when the policy disables the timeout. Zipped Dags are not armed: the importer gives each member its own timeout, so a single deadline for the archive could be too short. --- airflow-core/docs/core-concepts/tasks.rst | 3 +- .../airflow/sdk/execution_time/task_runner.py | 10 ++ .../execution_time/test_task_runner.py | 115 ++++++++++++++++++ 3 files changed, 127 insertions(+), 1 deletion(-) diff --git a/airflow-core/docs/core-concepts/tasks.rst b/airflow-core/docs/core-concepts/tasks.rst index 00f095aa66e4c..15b2ca0ad232a 100644 --- a/airflow-core/docs/core-concepts/tasks.rst +++ b/airflow-core/docs/core-concepts/tasks.rst @@ -143,7 +143,8 @@ maximum time allowed for every execution. If ``execution_timeout`` is breached, ``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. +control to Python), the supervisor sends it SIGTERM, and SIGKILL after the same period again. The same +backstop covers parsing a plain Python Dag file inside the task process, using ``[core] dagbag_import_timeout``. In addition, sensors have a ``timeout`` parameter. This only matters for sensors in ``reschedule`` mode. ``timeout`` controls the maximum time allowed for the sensor to succeed. If ``timeout`` is breached, ``AirflowSensorTimeout`` will be raised and the sensor fails immediately diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 652225b6673aa..ca3624f6e94b4 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -158,6 +158,7 @@ from airflow.sdk.execution_time.sentry import Sentry from airflow.sdk.execution_time.tracing import detail_span from airflow.sdk.execution_time.xcom import XCom +from airflow.sdk.importers.python_importer import PythonDagImporter, get_dagbag_import_timeout from airflow.sdk.listener import get_listener_manager from airflow.sdk.observability.metrics import stats_utils from airflow.sdk.serde import allow_class, iter_pydantic_models @@ -1119,6 +1120,13 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: _fail_lang_sdk_task(what, importer.coordinator_classpath, log) dag_file_parse_start = time.monotonic() + import_timeout = ( + get_dagbag_import_timeout(dag_absolute_path) + if PythonDagImporter().can_handle(dag_absolute_path) + else 0 + ) + if import_timeout > 0: + SUPERVISOR_COMMS.send(SetExecutionTimeout(timeout_seconds=import_timeout)) bag = BundleDagBag( dag_folder=dag_absolute_path, safe_mode=False, @@ -1126,6 +1134,8 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: bundle_path=bundle_instance.path, bundle_name=bundle_info.name, ) + if import_timeout > 0: + SUPERVISOR_COMMS.send(SetExecutionTimeout(timeout_seconds=0)) dag_file_parse_ms = int((time.monotonic() - dag_file_parse_start) * 1000) if TYPE_CHECKING: assert what.ti.dag_id 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 8d4ade1a3401b..0acc5dbf44254 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 @@ -24,6 +24,7 @@ import os import textwrap import time +import zipfile from collections.abc import Iterable from datetime import UTC, datetime, timedelta from pathlib import Path @@ -237,6 +238,7 @@ def execute(self, context): ], ) @pytest.mark.parametrize("map_index", [-1, 3]) +@pytest.mark.usefixtures("mock_supervisor_comms") def test_parse(test_dags_dir: Path, make_ti_context, coordinates, map_index): """Test that checks parsing of a basic dag with an un-mocked parse.""" what = StartupDetails( @@ -283,7 +285,116 @@ def test_parse(test_dags_dir: Path, make_ti_context, coordinates, map_index): assert ti.map_index == map_index +@pytest.mark.parametrize( + ("policy_timeout", "expected_sends"), + [ + pytest.param( + 12.5, + [ + mock.call(SetExecutionTimeout(timeout_seconds=12.5)), + mock.call(SetExecutionTimeout(timeout_seconds=0)), + ], + id="armed_then_cleared", + ), + pytest.param(0, [], id="policy_disables_timeout"), + ], +) +@mock.patch("airflow.sdk.execution_time.task_runner.get_dagbag_import_timeout", autospec=True) +def test_parse_arms_supervisor_import_timeout( + mock_get_timeout, + test_dags_dir: Path, + make_ti_context, + mock_supervisor_comms, + policy_timeout, + expected_sends, +): + mock_get_timeout.return_value = policy_timeout + + what = StartupDetails( + ti=TaskInstance( + id=uuid7(), + task_id="a", + dag_id="super_basic", + run_id="c", + try_number=1, + dag_version_id=uuid7(), + queue="default", + ), + dag_rel_path="super_basic.py", + bundle_info=BundleInfo(name="my-bundle", version=None), + ti_context=make_ti_context(), + start_date=timezone.utcnow(), + sentry_integration="", + ) + + with patch.dict( + os.environ, + { + "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps( + [ + { + "name": "my-bundle", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": str(test_dags_dir), "refresh_interval": 1}, + } + ] + ), + }, + ): + parse(what, mock.Mock()) + + mock_get_timeout.assert_called_once_with(str(test_dags_dir / "super_basic.py")) + assert mock_supervisor_comms.send.call_args_list == expected_sends + + +@mock.patch("airflow.sdk.execution_time.task_runner.get_dagbag_import_timeout", autospec=True) +def test_parse_does_not_arm_supervisor_import_timeout_for_zipped_dag( + mock_get_timeout, test_dags_dir: Path, tmp_path: Path, make_ti_context, mock_supervisor_comms +): + mock_get_timeout.return_value = 30 + source = (test_dags_dir / "super_basic.py").read_text().replace("super_basic", "zipped_basic") + with zipfile.ZipFile(tmp_path / "zipped.zip", "w") as archive: + archive.writestr("zipped_basic.py", source) + what = StartupDetails( + ti=TaskInstance( + id=uuid7(), + task_id="a", + dag_id="zipped_basic", + run_id="c", + try_number=1, + dag_version_id=uuid7(), + queue="default", + ), + dag_rel_path="zipped.zip", + bundle_info=BundleInfo(name="my-bundle", version=None), + ti_context=make_ti_context(), + start_date=timezone.utcnow(), + sentry_integration="", + ) + + with patch.dict( + os.environ, + { + "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps( + [ + { + "name": "my-bundle", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": str(tmp_path), "refresh_interval": 1}, + } + ] + ), + }, + ): + ti = parse(what, mock.Mock()) + + assert ti.task.dag.dag_id == "zipped_basic" + mock_get_timeout.assert_not_called() + mock_supervisor_comms.send.assert_not_called() + + @mock.patch("airflow.dag_processing.dagbag.BundleDagBag") +@pytest.mark.usefixtures("mock_supervisor_comms") def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context): """Test that checks that the BundleDagBag is constructed as expected during parsing""" mock_bag_instance = mock.Mock() @@ -358,6 +469,7 @@ def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context): ), ), ) +@pytest.mark.usefixtures("mock_supervisor_comms") def test_parse_not_found(test_dags_dir: Path, make_ti_context, dag_id, task_id, expected_error): """Check for nice error messages on dag not found.""" what = StartupDetails( @@ -402,6 +514,7 @@ def test_parse_not_found(test_dags_dir: Path, make_ti_context, dag_id, task_id, log.error.assert_has_calls([expected_error]) +@pytest.mark.usefixtures("mock_supervisor_comms") def test_parse_not_found_does_not_reschedule_when_max_attempts_reached(test_dags_dir: Path, make_ti_context): """ If the startup reschedule attempt limit is reached, parsing failures should not be rescheduled @@ -883,6 +996,7 @@ def test_with_a_tracer_provider_installed_the_dag_runs_sampling_decision_holds(m assert task_code_exporter.get_finished_spans() == () +@pytest.mark.usefixtures("mock_supervisor_comms") def test_parse_module_in_bundle_root(tmp_path: Path, make_ti_context): """Check that the bundle path is added to sys.path, so Dags can import shared modules.""" tmp_path.joinpath("util.py").write_text("NAME = 'dag_name'") @@ -943,6 +1057,7 @@ def test_parse_module_in_bundle_root(tmp_path: Path, make_ti_context): @mock.patch("airflow.sdk.execution_time.task_runner.get_listener_manager") @mock.patch("airflow.dag_processing.dagbag.BundleDagBag") @mock.patch("airflow.sdk.execution_time.bundles.DagBundlesManager") +@pytest.mark.usefixtures("mock_supervisor_comms") def test_bundle_initialize_parsing_context( mock_manager_cls, mock_bag_cls, From fcdaddc357fd8d685eb8a05e80f6ae40984085d2 Mon Sep 17 00:00:00 2001 From: Andrew Chang Date: Mon, 12 Oct 2026 14:02:47 +0800 Subject: [PATCH 4/4] Rename SetExecutionTimeout to SetTimeout The message now carries the Dag import timeout as well as the task's execution_timeout, and 0 clears the deadline, so the old name described only one of its uses. It has not been released yet, so this is the last point where the name can change without a migration. --- .../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 | 28 ++++++------- go-sdk/schema/supervisor-schema.json | 42 +++++++++---------- .../src/airflow/sdk/execution_time/comms.py | 6 +-- .../sdk/execution_time/schema/schema.json | 42 +++++++++---------- .../schema/versions/v2026_10_30.py | 2 +- .../airflow/sdk/execution_time/supervisor.py | 40 +++++++++--------- .../airflow/sdk/execution_time/task_runner.py | 8 ++-- .../execution_time/test_supervisor.py | 26 ++++++------ .../execution_time/test_task_runner.py | 8 ++-- ts-sdk/schema/supervisor-schema.json | 42 +++++++++---------- ts-sdk/src/generated/supervisor.ts | 40 +++++++++--------- 14 files changed, 147 insertions(+), 149 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 629cb6b1c49b1..a1699a6dbb0cb 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -2350,7 +2350,7 @@ def get_type_names(union_type): "UpdateHITLDetail", "GetHITLDetailResponse", "SetRenderedMapIndex", - "SetExecutionTimeout", + "SetTimeout", # 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 674c3737f08a6..12889829f2b1a 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -3106,7 +3106,7 @@ def get_type_names(union_type): "ResendLoggingFD", "CreateHITLDetailPayload", "SetRenderedMapIndex", - "SetExecutionTimeout", + "SetTimeout", "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..a5193649052fd 100644 --- a/go-sdk/pkg/execution/genmodels/discriminators.gen.go +++ b/go-sdk/pkg/execution/genmodels/discriminators.gen.go @@ -90,10 +90,10 @@ const ( TypeSentFDs = "SentFDs" TypeSetAssetStateStoreByName = "SetAssetStateStoreByName" TypeSetAssetStateStoreByURI = "SetAssetStateStoreByUri" - TypeSetExecutionTimeout = "SetExecutionTimeout" TypeSetRenderedFields = "SetRenderedFields" TypeSetRenderedMapIndex = "SetRenderedMapIndex" TypeSetTaskStateStore = "SetTaskStateStore" + TypeSetTimeout = "SetTimeout" TypeSetXCom = "SetXCom" TypeSkipDownstreamTasks = "SkipDownstreamTasks" TypeStartupDetails = "StartupDetails" @@ -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 @@ -333,6 +330,9 @@ func EnsureType(m any) any { case SetTaskStateStore: b.Type = TypeSetTaskStateStore return b + case SetTimeout: + b.Type = TypeSetTimeout + return b case SetXCom: b.Type = TypeSetXCom return b diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 34302015653ad..209e6846a8aa5 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1533,20 +1533,6 @@ type SetAssetStateStoreByURI struct { Value JsonValue `msgpack:"value"` } -// Tell the supervisor how long the task process may take before it must be stopped -// from outside. -// -// Sent before parsing a plain Python Dag file with “dagbag_import_timeout“ and -// before “execute()“ -// with the task's “execution_timeout“. “0“ clears the deadline. -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". @@ -1582,6 +1568,20 @@ type SetTaskStateStore struct { Value JsonValue `msgpack:"value"` } +// Tell the supervisor how long the task process may take before it must be stopped +// from outside. +// +// Sent before parsing a plain Python Dag file with “dagbag_import_timeout“ and +// before “execute()“ +// with the task's “execution_timeout“. “0“ clears the deadline. +type SetTimeout 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"` +} + type SetXCom struct { // DagID corresponds to the JSON schema field "dag_id". DagID string `msgpack:"dag_id"` diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json index 88d6ec0a9072a..3551343a7b61e 100644 --- a/go-sdk/schema/supervisor-schema.json +++ b/go-sdk/schema/supervisor-schema.json @@ -3800,27 +3800,6 @@ "title": "SetAssetStateStoreByUri", "type": "object" }, - "SetExecutionTimeout": { - "description": "Tell the supervisor how long the task process may take before it must be stopped from outside.\n\nSent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()``\nwith the task's ``execution_timeout``. ``0`` clears the deadline.", - "properties": { - "timeout_seconds": { - "minimum": 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": { @@ -3906,6 +3885,27 @@ "title": "SetTaskStateStore", "type": "object" }, + "SetTimeout": { + "description": "Tell the supervisor how long the task process may take before it must be stopped from outside.\n\nSent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()``\nwith the task's ``execution_timeout``. ``0`` clears the deadline.", + "properties": { + "timeout_seconds": { + "minimum": 0, + "title": "Timeout Seconds", + "type": "number" + }, + "type": { + "const": "SetTimeout", + "default": "SetTimeout", + "title": "Type", + "type": "string" + } + }, + "required": [ + "timeout_seconds" + ], + "title": "SetTimeout", + "type": "object" + }, "SetXCom": { "properties": { "key": { diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index b6c15ca3918b9..873a273899d3d 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -1129,7 +1129,7 @@ class SetRenderedMapIndex(BaseModel): type: Literal["SetRenderedMapIndex"] = "SetRenderedMapIndex" -class SetExecutionTimeout(BaseModel): +class SetTimeout(BaseModel): """ Tell the supervisor how long the task process may take before it must be stopped from outside. @@ -1138,7 +1138,7 @@ class SetExecutionTimeout(BaseModel): """ timeout_seconds: Annotated[float, Field(ge=0)] - type: Literal["SetExecutionTimeout"] = "SetExecutionTimeout" + type: Literal["SetTimeout"] = "SetTimeout" class TriggerDagRun(TriggerDAGRunPayload): @@ -1345,7 +1345,7 @@ class GetDag(BaseModel): | RetryTask | SetAssetStateStoreByName | SetAssetStateStoreByUri - | SetExecutionTimeout + | SetTimeout | SetRenderedFields | SetRenderedMapIndex | SetTaskStateStore 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 88d6ec0a9072a..3551343a7b61e 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -3800,27 +3800,6 @@ "title": "SetAssetStateStoreByUri", "type": "object" }, - "SetExecutionTimeout": { - "description": "Tell the supervisor how long the task process may take before it must be stopped from outside.\n\nSent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()``\nwith the task's ``execution_timeout``. ``0`` clears the deadline.", - "properties": { - "timeout_seconds": { - "minimum": 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": { @@ -3906,6 +3885,27 @@ "title": "SetTaskStateStore", "type": "object" }, + "SetTimeout": { + "description": "Tell the supervisor how long the task process may take before it must be stopped from outside.\n\nSent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()``\nwith the task's ``execution_timeout``. ``0`` clears the deadline.", + "properties": { + "timeout_seconds": { + "minimum": 0, + "title": "Timeout Seconds", + "type": "number" + }, + "type": { + "const": "SetTimeout", + "default": "SetTimeout", + "title": "Type", + "type": "string" + } + }, + "required": [ + "timeout_seconds" + ], + "title": "SetTimeout", + "type": "object" + }, "SetXCom": { "properties": { "key": { 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 6a998a7c5c0a3..adcf486d43de3 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 @@ -23,7 +23,7 @@ way down. Only *changes* to bodies that already existed at an earlier dated version require a ``VersionChange`` entry below. -``SetExecutionTimeout.timeout_seconds`` accepts ``0`` (clear the deadline) since the same +``SetTimeout.timeout_seconds`` accepts ``0`` (clear the deadline) since the same version that introduced the body, so it needs no entry either. """ diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 526496a523690..52afd02df36bf 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -130,10 +130,10 @@ SentFDs, SetAssetStateStoreByName, SetAssetStateStoreByUri, - SetExecutionTimeout, SetRenderedFields, SetRenderedMapIndex, SetTaskStateStore, + SetTimeout, SetXCom, SkipDownstreamTasks, StartupDetails, @@ -1643,8 +1643,8 @@ 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) - _execution_timeout_seconds: float | None = attrs.field(default=None, init=False) - _execution_timeout_enforce_at: float | None = attrs.field(default=None, init=False) + _timeout_seconds: float | None = attrs.field(default=None, init=False) + _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.""" decoder: ClassVar[TypeAdapter[ToSupervisor]] = TypeAdapter(ToSupervisor) @@ -1903,7 +1903,7 @@ def _monitor_subprocess(self): # Ensure we heartbeat _at most_ 75% through the task instance heartbeat timeout time HEARTBEAT_TIMEOUT - last_heartbeat_ago * 0.75, MIN_HEARTBEAT_INTERVAL, - self._execution_timeout_due_in(), + self._timeout_due_in(), ), ) # Block until events are ready or the timeout is reached @@ -1931,7 +1931,7 @@ def _monitor_subprocess(self): self._send_heartbeat_if_needed() self._handle_process_overtime_if_needed() - self._handle_execution_timeout_if_needed() + self._handle_timeout_if_needed() def _handle_process_overtime_if_needed(self): """Handle termination of auxiliary processes if the task exceeds the configured overtime.""" @@ -1946,13 +1946,13 @@ def _handle_process_overtime_if_needed(self): ) self.kill(signal.SIGTERM, force=True) - def _execution_timeout_due_in(self) -> float: + def _timeout_due_in(self) -> float: """Seconds until the supervisor must stop an overrunning task, or infinity if there is nothing to enforce.""" - if self._execution_timeout_enforce_at is None or self._terminal_state or self._exit_code is not None: + if self._timeout_enforce_at is None or self._terminal_state or self._exit_code is not None: return math.inf - return self._execution_timeout_enforce_at - time.monotonic() + return self._timeout_enforce_at - time.monotonic() - def _handle_execution_timeout_if_needed(self): + def _handle_timeout_if_needed(self): """ Enforce the task process's current timeout from outside it. @@ -1961,16 +1961,16 @@ def _handle_execution_timeout_if_needed(self): 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``: with ``dagbag_import_timeout`` around parsing a + The deadline is armed by ``SetTimeout``: with ``dagbag_import_timeout`` around parsing a plain Python Dag file and with ``execution_timeout`` right before ``execute()``. Bundle loading and zipped Dags, whose members are each imported under their own timeout, are not covered. """ - if self._execution_timeout_due_in() > 0: + if self._timeout_due_in() > 0: return - self._execution_timeout_enforce_at = None + self._timeout_enforce_at = None self.process_log.error( "Task did not stop before its timeout elapsed; terminating process", - timeout_seconds=self._execution_timeout_seconds, + timeout_seconds=self._timeout_seconds, grace_period_seconds=KILLED_TASK_CLEANUP_TIME, ) self.kill(signal.SIGTERM, force=True, escalation_delay=KILLED_TASK_CLEANUP_TIME) @@ -2107,15 +2107,13 @@ 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 - ) -> RequestResult: + def _handle_set_timeout(self, msg: SetTimeout, log: FilteringBoundLogger, req_id: int) -> RequestResult: if msg.timeout_seconds == 0: - self._execution_timeout_seconds = None - self._execution_timeout_enforce_at = None + self._timeout_seconds = None + self._timeout_enforce_at = None return None, {} - self._execution_timeout_seconds = msg.timeout_seconds - self._execution_timeout_enforce_at = time.monotonic() + msg.timeout_seconds + KILLED_TASK_CLEANUP_TIME + self._timeout_seconds = msg.timeout_seconds + self._timeout_enforce_at = time.monotonic() + msg.timeout_seconds + KILLED_TASK_CLEANUP_TIME return None, {} def _handle_set_rendered_fields( @@ -2412,7 +2410,7 @@ 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(SetTimeout, _handle_set_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), 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 ca3624f6e94b4..4327f00deb054 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -116,9 +116,9 @@ ResendLoggingFD, RetryTask, SentFDs, - SetExecutionTimeout, SetRenderedFields, SetRenderedMapIndex, + SetTimeout, SkipDownstreamTasks, StartupDetails, SucceedTask, @@ -1126,7 +1126,7 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: else 0 ) if import_timeout > 0: - SUPERVISOR_COMMS.send(SetExecutionTimeout(timeout_seconds=import_timeout)) + SUPERVISOR_COMMS.send(SetTimeout(timeout_seconds=import_timeout)) bag = BundleDagBag( dag_folder=dag_absolute_path, safe_mode=False, @@ -1135,7 +1135,7 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: bundle_name=bundle_info.name, ) if import_timeout > 0: - SUPERVISOR_COMMS.send(SetExecutionTimeout(timeout_seconds=0)) + SUPERVISOR_COMMS.send(SetTimeout(timeout_seconds=0)) dag_file_parse_ms = int((time.monotonic() - dag_file_parse_start) * 1000) if TYPE_CHECKING: assert what.ti.dag_id @@ -2274,7 +2274,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(SetTimeout(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/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 46434faa73ffc..9198d4c6699c8 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -138,10 +138,10 @@ SentFDs, SetAssetStateStoreByName, SetAssetStateStoreByUri, - SetExecutionTimeout, SetRenderedFields, SetRenderedMapIndex, SetTaskStateStore, + SetTimeout, SetXCom, SkipDownstreamTasks, SucceedTask, @@ -1610,14 +1610,14 @@ def test_execution_timeout_kills_task_that_does_not_stop(self, mocker): client=mocker.Mock(), ) - proc._handle_request(SetExecutionTimeout(timeout_seconds=30), log=mocker.Mock(), req_id=1) + proc._handle_request(SetTimeout(timeout_seconds=30), log=mocker.Mock(), req_id=1) monotonic.return_value += 34 - proc._handle_execution_timeout_if_needed() + proc._handle_timeout_if_needed() mock_kill.assert_not_called() monotonic.return_value += 1 - proc._handle_execution_timeout_if_needed() + proc._handle_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 its timeout elapsed; terminating process", @@ -1625,7 +1625,7 @@ def test_execution_timeout_kills_task_that_does_not_stop(self, mocker): grace_period_seconds=5.0, ) - proc._handle_execution_timeout_if_needed() + proc._handle_timeout_if_needed() mock_kill.assert_called_once() def test_set_execution_timeout_zero_clears_deadline(self, mocker): @@ -1640,11 +1640,11 @@ def test_set_execution_timeout_zero_clears_deadline(self, mocker): process=mocker.Mock(), client=mocker.Mock(), ) - proc._handle_request(SetExecutionTimeout(timeout_seconds=30), log=mocker.Mock(), req_id=1) - proc._handle_request(SetExecutionTimeout(timeout_seconds=0), log=mocker.Mock(), req_id=2) + proc._handle_request(SetTimeout(timeout_seconds=30), log=mocker.Mock(), req_id=1) + proc._handle_request(SetTimeout(timeout_seconds=0), log=mocker.Mock(), req_id=2) monotonic.return_value += 100 - proc._handle_execution_timeout_if_needed() + proc._handle_timeout_if_needed() mock_kill.assert_not_called() proc.process_log.error.assert_not_called() @@ -1668,12 +1668,12 @@ def test_execution_timeout_not_enforced_once_task_is_over(self, mocker, terminal process=mocker.Mock(), client=mocker.Mock(), ) - proc._handle_request(SetExecutionTimeout(timeout_seconds=30), log=mocker.Mock(), req_id=1) + proc._handle_request(SetTimeout(timeout_seconds=30), log=mocker.Mock(), req_id=1) proc._terminal_state = terminal_state proc._exit_code = exit_code monotonic.return_value += 100 - proc._handle_execution_timeout_if_needed() + proc._handle_timeout_if_needed() mock_kill.assert_not_called() proc.process_log.error.assert_not_called() @@ -1705,7 +1705,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(SetTimeout(timeout_seconds=0.1)) sleep(30) exit(5) @@ -2647,8 +2647,8 @@ class RequestTestCase: test_id="set_rtif", ), RequestTestCase( - message=SetExecutionTimeout(timeout_seconds=30.0), - test_id="set_execution_timeout", + message=SetTimeout(timeout_seconds=30.0), + test_id="set_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 0acc5dbf44254..a35d20983fa45 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 @@ -152,9 +152,9 @@ RetryTask, SetAssetStateStoreByName, SetAssetStateStoreByUri, - SetExecutionTimeout, SetRenderedFields, SetTaskStateStore, + SetTimeout, SetXCom, SkipDownstreamTasks, StartupDetails, @@ -291,8 +291,8 @@ def test_parse(test_dags_dir: Path, make_ti_context, coordinates, map_index): pytest.param( 12.5, [ - mock.call(SetExecutionTimeout(timeout_seconds=12.5)), - mock.call(SetExecutionTimeout(timeout_seconds=0)), + mock.call(SetTimeout(timeout_seconds=12.5)), + mock.call(SetTimeout(timeout_seconds=0)), ], id="armed_then_cleared", ), @@ -6704,7 +6704,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(SetTimeout(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/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json index 88d6ec0a9072a..3551343a7b61e 100644 --- a/ts-sdk/schema/supervisor-schema.json +++ b/ts-sdk/schema/supervisor-schema.json @@ -3800,27 +3800,6 @@ "title": "SetAssetStateStoreByUri", "type": "object" }, - "SetExecutionTimeout": { - "description": "Tell the supervisor how long the task process may take before it must be stopped from outside.\n\nSent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()``\nwith the task's ``execution_timeout``. ``0`` clears the deadline.", - "properties": { - "timeout_seconds": { - "minimum": 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": { @@ -3906,6 +3885,27 @@ "title": "SetTaskStateStore", "type": "object" }, + "SetTimeout": { + "description": "Tell the supervisor how long the task process may take before it must be stopped from outside.\n\nSent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()``\nwith the task's ``execution_timeout``. ``0`` clears the deadline.", + "properties": { + "timeout_seconds": { + "minimum": 0, + "title": "Timeout Seconds", + "type": "number" + }, + "type": { + "const": "SetTimeout", + "default": "SetTimeout", + "title": "Type", + "type": "string" + } + }, + "required": [ + "timeout_seconds" + ], + "title": "SetTimeout", + "type": "object" + }, "SetXCom": { "properties": { "key": { diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index e70d9ba31610c..022ad35c645e1 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -571,15 +571,15 @@ 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 TimeoutSeconds = number; +export type Type71 = "SetTimeout"; export type Key18 = string; export type DagId21 = string; export type RunId14 = string; @@ -1757,19 +1757,6 @@ export interface SetAssetStateStoreByUri { value: JsonValue; type?: Type67; } -/** - * Tell the supervisor how long the task process may take before it must be stopped from outside. - * - * Sent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()`` - * with the task's ``execution_timeout``. ``0`` clears the deadline. - * - * 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. * @@ -1778,7 +1765,7 @@ export interface SetExecutionTimeout { */ export interface SetRenderedFields { rendered_fields: RenderedFields; - type?: Type69; + type?: Type68; } export interface RenderedFields { [k: string]: JsonValue; @@ -1791,7 +1778,7 @@ export interface RenderedFields { */ export interface SetRenderedMapIndex { rendered_map_index: RenderedMapIndex3; - type?: Type70; + type?: Type69; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1802,6 +1789,19 @@ export interface SetTaskStateStore { key: Key17; value: JsonValue; expires_at: ExpiresAt; + type?: Type70; +} +/** + * Tell the supervisor how long the task process may take before it must be stopped from outside. + * + * Sent before parsing a plain Python Dag file with ``dagbag_import_timeout`` and before ``execute()`` + * with the task's ``execution_timeout``. ``0`` clears the deadline. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "SetTimeout". + */ +export interface SetTimeout { + timeout_seconds: TimeoutSeconds; type?: Type71; } /**