Enforce the Dag import timeout from the task supervisor side

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

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",

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",

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": {

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

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": {

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

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),