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/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 d7d779b15cf5f..209e6846a8aa5 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1533,16 +1533,6 @@ type SetAssetStateStoreByURI struct { Value JsonValue `msgpack:"value"` } -// Tell the supervisor the task's “execution_timeout“ so it can enforce it from -// outside the task process. -type SetExecutionTimeout struct { - // TimeoutSeconds corresponds to the JSON schema field "timeout_seconds". - TimeoutSeconds float64 `msgpack:"timeout_seconds"` - - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` -} - // Payload for setting RTIF for a task instance. type SetRenderedFields struct { // RenderedFields corresponds to the JSON schema field "rendered_fields". @@ -1578,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 7e09a736de793..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 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": { @@ -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 7c957027ba792..873a273899d3d 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -1129,11 +1129,16 @@ class SetRenderedMapIndex(BaseModel): type: Literal["SetRenderedMapIndex"] = "SetRenderedMapIndex" -class SetExecutionTimeout(BaseModel): - """Tell the supervisor the task's ``execution_timeout`` so it can enforce it from outside the task process.""" +class SetTimeout(BaseModel): + """ + 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)] - type: Literal["SetExecutionTimeout"] = "SetExecutionTimeout" + timeout_seconds: Annotated[float, Field(ge=0)] + type: Literal["SetTimeout"] = "SetTimeout" class TriggerDagRun(TriggerDAGRunPayload): @@ -1340,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 7e09a736de793..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 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": { @@ -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 baffc8c9d0f1b..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 @@ -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. + +``SetTimeout.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..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,30 +1946,31 @@ 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 ``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 ``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 execution_timeout elapsed; terminating process", - timeout_seconds=self._execution_timeout_seconds, + "Task did not stop before its timeout elapsed; terminating process", + timeout_seconds=self._timeout_seconds, grace_period_seconds=KILLED_TASK_CLEANUP_TIME, ) self.kill(signal.SIGTERM, force=True, escalation_delay=KILLED_TASK_CLEANUP_TIME) @@ -2106,11 +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: - self._execution_timeout_seconds = msg.timeout_seconds - self._execution_timeout_enforce_at = time.monotonic() + msg.timeout_seconds + KILLED_TASK_CLEANUP_TIME + def _handle_set_timeout(self, msg: SetTimeout, log: FilteringBoundLogger, req_id: int) -> RequestResult: + if msg.timeout_seconds == 0: + self._timeout_seconds = None + self._timeout_enforce_at = None + return None, {} + 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( @@ -2407,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 652225b6673aa..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, @@ -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(SetTimeout(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(SetTimeout(timeout_seconds=0)) dag_file_parse_ms = int((time.monotonic() - dag_file_parse_start) * 1000) if TYPE_CHECKING: assert what.ti.dag_id @@ -2264,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/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() 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..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,24 +1610,45 @@ 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 execution_timeout elapsed; terminating process", + "Task did not stop before its timeout elapsed; terminating process", timeout_seconds=30, 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): + 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(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_timeout_if_needed() + + mock_kill.assert_not_called() + proc.process_log.error.assert_not_called() + @pytest.mark.parametrize( ("terminal_state", "exit_code"), [ @@ -1647,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() @@ -1684,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) @@ -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, @@ -2626,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 8d4ade1a3401b..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 @@ -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 @@ -151,9 +152,9 @@ RetryTask, SetAssetStateStoreByName, SetAssetStateStoreByUri, - SetExecutionTimeout, SetRenderedFields, SetTaskStateStore, + SetTimeout, SetXCom, SkipDownstreamTasks, StartupDetails, @@ -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(SetTimeout(timeout_seconds=12.5)), + mock.call(SetTimeout(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, @@ -6589,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 7e09a736de793..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 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": { @@ -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 31196f868a5cb..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,16 +1757,6 @@ export interface SetAssetStateStoreByUri { value: JsonValue; type?: Type67; } -/** - * Tell the supervisor the task's ``execution_timeout`` so it can enforce it from outside the task process. - * - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "SetExecutionTimeout". - */ -export interface SetExecutionTimeout { - timeout_seconds: TimeoutSeconds; - type?: Type68; -} /** * Payload for setting RTIF for a task instance. * @@ -1775,7 +1765,7 @@ export interface SetExecutionTimeout { */ export interface SetRenderedFields { rendered_fields: RenderedFields; - type?: Type69; + type?: Type68; } export interface RenderedFields { [k: string]: JsonValue; @@ -1788,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 @@ -1799,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; } /**