Enforce execution_timeout from the task supervisor side

airflow-core/docs/core-concepts/tasks.rst

@@ -119,7 +119,10 @@ Timeouts
 If you want a task to have a maximum runtime, set its ``execution_timeout`` attribute to a ``datetime.timedelta`` value
 that is the maximum permissible runtime. This applies to all Airflow tasks, including sensors. ``execution_timeout`` controls the
 maximum time allowed for every execution. If ``execution_timeout`` is breached, the task times out and
-``AirflowTaskTimeout`` is raised.
+``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.
 
 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/newsfragments/73806.improvement.rst

@@ -0,0 +1 @@
+``execution_timeout`` is now also enforced by the task supervisor. A task process that does not stop on its own within ``[core] killed_task_cleanup_time`` seconds after the timeout elapses is sent SIGTERM and then SIGKILL, so tasks stuck in native code or with a broken signal handler no longer run forever.

airflow-core/src/airflow/config_templates/config.yml

@@ -312,7 +312,9 @@ core:
     killed_task_cleanup_time:
       description: |
         When a task is killed forcefully, this is the amount of time in seconds that
-        it has to cleanup after it is sent a SIGTERM, before it is SIGKILLED
+        it has to cleanup after it is sent a SIGTERM, before it is SIGKILLED.
+        The task supervisor also waits this long after ``execution_timeout`` elapses before
+        it sends SIGTERM to a task process that has not stopped on its own.
       version_added: ~
       type: integer
       example: ~

airflow-core/tests/unit/dag_processing/test_processor.py

@@ -2315,6 +2315,7 @@ def get_type_names(union_type):
             "UpdateHITLDetail",
             "GetHITLDetailResponse",
             "SetRenderedMapIndex",
+            "SetExecutionTimeout",
             # AIP-103 task/asset store — Dag processor has no task execution context.
             "GetTaskStateStore",
             "SetTaskStateStore",

airflow-core/tests/unit/jobs/test_triggerer_job.py

@@ -3068,6 +3068,7 @@ def get_type_names(union_type):
             "ResendLoggingFD",
             "CreateHITLDetailPayload",
             "SetRenderedMapIndex",
+            "SetExecutionTimeout",
             "GetDag",
             # AIP-103 task store — triggerer has no task execution context.
             "GetTaskStateStore",

go-sdk/schema/supervisor-schema.json

@@ -3719,6 +3719,27 @@
       "title": "SetAssetStateStoreByUri",
       "type": "object"
     },
+    "SetExecutionTimeout": {
+      "description": "Tell the supervisor the task's ``execution_timeout`` so it can enforce it from outside the task process.",
+      "properties": {
+        "timeout_seconds": {
+          "exclusiveMinimum": 0,
+          "title": "Timeout Seconds",
+          "type": "number"
+        },
+        "type": {
+          "const": "SetExecutionTimeout",
+          "default": "SetExecutionTimeout",
+          "title": "Type",
+          "type": "string"
+        }
+      },
+      "required": [
+        "timeout_seconds"
+      ],
+      "title": "SetExecutionTimeout",
+      "type": "object"
+    },
     "SetRenderedFields": {
       "description": "Payload for setting RTIF for a task instance.",
       "properties": {

task-sdk/src/airflow/sdk/execution_time/comms.py

@@ -1125,6 +1125,13 @@ 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."""
+
+    timeout_seconds: Annotated[float, Field(gt=0)]
+    type: Literal["SetExecutionTimeout"] = "SetExecutionTimeout"
+
+
 class TriggerDagRun(TriggerDAGRunPayload):
     dag_id: str
     run_id: Annotated[str, Field(title="Dag Run Id")]
@@ -1329,6 +1336,7 @@ class GetDag(BaseModel):
     | RetryTask
     | SetAssetStateStoreByName
     | SetAssetStateStoreByUri
+    | SetExecutionTimeout
     | SetRenderedFields
     | SetRenderedMapIndex
     | SetTaskStateStore

task-sdk/src/airflow/sdk/execution_time/schema/schema.json

@@ -3719,6 +3719,27 @@
       "title": "SetAssetStateStoreByUri",
       "type": "object"
     },
+    "SetExecutionTimeout": {
+      "description": "Tell the supervisor the task's ``execution_timeout`` so it can enforce it from outside the task process.",
+      "properties": {
+        "timeout_seconds": {
+          "exclusiveMinimum": 0,
+          "title": "Timeout Seconds",
+          "type": "number"
+        },
+        "type": {
+          "const": "SetExecutionTimeout",
+          "default": "SetExecutionTimeout",
+          "title": "Type",
+          "type": "string"
+        }
+      },
+      "required": [
+        "timeout_seconds"
+      ],
+      "title": "SetExecutionTimeout",
+      "type": "object"
+    },
     "SetRenderedFields": {
       "description": "Payload for setting RTIF for a task instance.",
       "properties": {

task-sdk/src/airflow/sdk/execution_time/supervisor.py

@@ -24,6 +24,7 @@
 import functools
 import io
 import logging
+import math
 import os
 import pkgutil
 import selectors
@@ -129,6 +130,7 @@
     SentFDs,
     SetAssetStateStoreByName,
     SetAssetStateStoreByUri,
+    SetExecutionTimeout,
     SetRenderedFields,
     SetRenderedMapIndex,
     SetTaskStateStore,
@@ -294,6 +296,10 @@ def _handle_mask_secret_request(
 # like listeners after task is complete.
 TASK_OVERTIME_THRESHOLD: float = conf.getfloat("core", "task_success_overtime")
 
+# How long a task process gets to clean up when the supervisor stops it: after the server says it should no
+# longer run, SIGTERM to SIGKILL; after execution_timeout elapses, both before SIGTERM and before SIGKILL.
+KILLED_TASK_CLEANUP_TIME: float = conf.getfloat("core", "killed_task_cleanup_time")
+
 SERVER_TERMINATED = TerminalStateNonSuccess.SERVER_TERMINATED.value
 
 # These are the task instance states that require some additional information to transition into.
@@ -1637,6 +1643,10 @@ 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)
+    """Monotonic time at which the supervisor stops a task that has not reported a terminal state."""
+
     decoder: ClassVar[TypeAdapter[ToSupervisor]] = TypeAdapter(ToSupervisor)
 
     ti: RuntimeTI | None = None
@@ -1893,6 +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(),
                 ),
             )
             # Block until events are ready or the timeout is reached
@@ -1920,6 +1931,7 @@ def _monitor_subprocess(self):
                 self._send_heartbeat_if_needed()
 
                 self._handle_process_overtime_if_needed()
+                self._handle_execution_timeout_if_needed()
 
     def _handle_process_overtime_if_needed(self):
         """Handle termination of auxiliary processes if the task exceeds the configured overtime."""
@@ -1934,6 +1946,34 @@ def _handle_process_overtime_if_needed(self):
             )
             self.kill(signal.SIGTERM, force=True)
 
+    def _execution_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:
+            return math.inf
+        return self._execution_timeout_enforce_at - time.monotonic()
+
+    def _handle_execution_timeout_if_needed(self):
+        """
+        Enforce ``execution_timeout`` from outside the task process.
+
+        The task process raises ``AirflowTaskTimeout`` itself when the timeout elapses, which is the only
+        place ``on_kill``, the retry policy, callbacks and listeners can run. That needs a process that can
+        still run Python signal handlers; one that has not reported a terminal state within the grace period
+        is stuck (native code, SIGSEGV) and is sent SIGTERM, then SIGKILL.
+
+        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.
+        """
+        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",
+            timeout_seconds=self._execution_timeout_seconds,
+            grace_period_seconds=KILLED_TASK_CLEANUP_TIME,
+        )
+        self.kill(signal.SIGTERM, force=True, escalation_delay=KILLED_TASK_CLEANUP_TIME)
+
     def _send_heartbeat_if_needed(self):
         """Send a heartbeat to the client if heartbeat interval has passed."""
         # Respect the minimum interval between heartbeat attempts
@@ -1968,7 +2008,7 @@ def _send_heartbeat_if_needed(self):
                 # kill() drains worker messages while waiting for the process to exit.
                 self._terminal_state = SERVER_TERMINATED
                 self._pending_terminal_state_msg = None
-                self.kill(signal.SIGTERM, force=True)
+                self.kill(signal.SIGTERM, force=True, escalation_delay=KILLED_TASK_CLEANUP_TIME)
                 self.process_log.error("Task killed!")
             else:
                 # If we get any other error, we'll just log it and try again next time
@@ -2066,6 +2106,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
+        return None, {}
+
     def _handle_set_rendered_fields(
         self, msg: SetRenderedFields, log: FilteringBoundLogger, req_id: int
     ) -> RequestResult:
@@ -2357,6 +2404,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(SetRenderedFields, _handle_set_rendered_fields),
                 register_request_method(SetRenderedMapIndex, _handle_set_rendered_map_index),
                 register_request_method(SetTaskStateStore, _handle_set_task_state_store),

task-sdk/src/airflow/sdk/execution_time/task_runner.py

@@ -112,6 +112,7 @@
     ResendLoggingFD,
     RetryTask,
     SentFDs,
+    SetExecutionTimeout,
     SetRenderedFields,
     SetRenderedMapIndex,
     SkipDownstreamTasks,
@@ -2209,6 +2210,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))
             # Run task in timeout wrapper
             with timeout(timeout_seconds):
                 result = ctx.run(execute, context=context)