This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 634d6bd7242 Enforce execution_timeout from the task supervisor side
(#73806)
634d6bd7242 is described below
commit 634d6bd7242b44b7012e1a3b36afe197a6cc95b7
Author: Andrew Chang <[email protected]>
AuthorDate: Wed Oct 7 19:43:38 2026 +0800
Enforce execution_timeout from the task supervisor side (#73806)
The timeout is only enforced by a SIGALRM handler inside the task
process. When that process is stuck in native code, deadlocked on a
lock inherited across fork, or corrupted by SIGSEGV, the handler never
runs and the task stays running forever. The task process now tells the
supervisor its timeout before starting execute(), and the supervisor
sends SIGTERM and then SIGKILL if the task has not reported a terminal
state within a grace period. The in-process handler stays as the first
responder because on_kill, the retry policy and callbacks can only run
there.
---
airflow-core/docs/core-concepts/tasks.rst | 5 +-
airflow-core/newsfragments/73806.improvement.rst | 1 +
.../src/airflow/config_templates/config.yml | 4 +-
.../tests/unit/dag_processing/test_processor.py | 1 +
airflow-core/tests/unit/jobs/test_triggerer_job.py | 1 +
.../pkg/execution/genmodels/discriminators.gen.go | 4 +
go-sdk/pkg/execution/genmodels/models.gen.go | 18 ++-
go-sdk/schema/supervisor-schema.json | 21 +++
task-sdk/src/airflow/sdk/execution_time/comms.py | 8 ++
.../airflow/sdk/execution_time/schema/schema.json | 21 +++
.../src/airflow/sdk/execution_time/supervisor.py | 50 ++++++-
.../src/airflow/sdk/execution_time/task_runner.py | 2 +
.../task_sdk/execution_time/test_supervisor.py | 154 ++++++++++++++++++++-
.../task_sdk/execution_time/test_task_runner.py | 12 +-
14 files changed, 290 insertions(+), 12 deletions(-)
diff --git a/airflow-core/docs/core-concepts/tasks.rst
b/airflow-core/docs/core-concepts/tasks.rst
index fd193333bab..0e125fda92d 100644
--- a/airflow-core/docs/core-concepts/tasks.rst
+++ b/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
diff --git a/airflow-core/newsfragments/73806.improvement.rst
b/airflow-core/newsfragments/73806.improvement.rst
new file mode 100644
index 00000000000..0c50d09b67d
--- /dev/null
+++ b/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.
diff --git a/airflow-core/src/airflow/config_templates/config.yml
b/airflow-core/src/airflow/config_templates/config.yml
index c761c3c0e14..25cffae10c3 100644
--- a/airflow-core/src/airflow/config_templates/config.yml
+++ b/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: ~
diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py
b/airflow-core/tests/unit/dag_processing/test_processor.py
index 9363937693f..76a5ea85cde 100644
--- a/airflow-core/tests/unit/dag_processing/test_processor.py
+++ b/airflow-core/tests/unit/dag_processing/test_processor.py
@@ -2315,6 +2315,7 @@ class TestDagProcessingMessageTypes:
"UpdateHITLDetail",
"GetHITLDetailResponse",
"SetRenderedMapIndex",
+ "SetExecutionTimeout",
# AIP-103 task/asset store — Dag processor has no task execution
context.
"GetTaskStateStore",
"SetTaskStateStore",
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py
b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index 67a5fb522ad..9347bd32dbc 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -3068,6 +3068,7 @@ class TestTriggererMessageTypes:
"ResendLoggingFD",
"CreateHITLDetailPayload",
"SetRenderedMapIndex",
+ "SetExecutionTimeout",
"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 09da3118f55..8736a6e2c97 100644
--- a/go-sdk/pkg/execution/genmodels/discriminators.gen.go
+++ b/go-sdk/pkg/execution/genmodels/discriminators.gen.go
@@ -90,6 +90,7 @@ const (
TypeSentFDs = "SentFDs"
TypeSetAssetStateStoreByName = "SetAssetStateStoreByName"
TypeSetAssetStateStoreByURI = "SetAssetStateStoreByUri"
+ TypeSetExecutionTimeout = "SetExecutionTimeout"
TypeSetRenderedFields = "SetRenderedFields"
TypeSetRenderedMapIndex = "SetRenderedMapIndex"
TypeSetTaskStateStore = "SetTaskStateStore"
@@ -320,6 +321,9 @@ func EnsureType(m any) any {
case SetAssetStateStoreByURI:
b.Type = TypeSetAssetStateStoreByURI
return b
+ case SetExecutionTimeout:
+ b.Type = TypeSetExecutionTimeout
+ return b
case SetRenderedFields:
b.Type = TypeSetRenderedFields
return b
diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go
b/go-sdk/pkg/execution/genmodels/models.gen.go
index dee470082f1..8904ad408f8 100644
--- a/go-sdk/pkg/execution/genmodels/models.gen.go
+++ b/go-sdk/pkg/execution/genmodels/models.gen.go
@@ -1497,6 +1497,16 @@ type SetAssetStateStoreByURI struct {
Value JsonValue `msgpack:"value"`
}
+// Tell the supervisor the task's “execution_timeout“ so it can enforce it from
+// outside the task process.
+type SetExecutionTimeout struct {
+ // TimeoutSeconds corresponds to the JSON schema field
"timeout_seconds".
+ TimeoutSeconds float64 `msgpack:"timeout_seconds"`
+
+ // Type corresponds to the JSON schema field "type".
+ Type string `msgpack:"type,omitempty"`
+}
+
// Payload for setting RTIF for a task instance.
type SetRenderedFields struct {
// RenderedFields corresponds to the JSON schema field
"rendered_fields".
@@ -1882,6 +1892,10 @@ type UpdateHITLDetail struct {
Type string `msgpack:"type,omitempty"`
}
+type VersionData map[string]interface{}
+
+type Warnings []interface{}
+
// Variable schema for responses with fields that are needed for Runtime.
type VariableResponse struct {
// Key corresponds to the JSON schema field "key".
@@ -1891,10 +1905,6 @@ type VariableResponse struct {
Value interface{} `msgpack:"value"`
}
-type Warnings []interface{}
-
-type VersionData map[string]interface{}
-
type ValidateInletsAndOutlets struct {
// TIID corresponds to the JSON schema field "ti_id".
TIID string `msgpack:"ti_id"`
diff --git a/go-sdk/schema/supervisor-schema.json
b/go-sdk/schema/supervisor-schema.json
index d6ee95927ef..8161c498827 100644
--- a/go-sdk/schema/supervisor-schema.json
+++ b/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": {
diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py
b/task-sdk/src/airflow/sdk/execution_time/comms.py
index ff151a01777..6472863fbd8 100644
--- a/task-sdk/src/airflow/sdk/execution_time/comms.py
+++ b/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 @@ ToSupervisor = Annotated[
| RetryTask
| SetAssetStateStoreByName
| SetAssetStateStoreByUri
+ | SetExecutionTimeout
| 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 d6ee95927ef..8161c498827 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
+++ b/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": {
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index 372d12bc621..2be985c0424 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -24,6 +24,7 @@ import contextlib
import functools
import io
import logging
+import math
import os
import pkgutil
import selectors
@@ -129,6 +130,7 @@ from airflow.sdk.execution_time.comms import (
SentFDs,
SetAssetStateStoreByName,
SetAssetStateStoreByUri,
+ SetExecutionTimeout,
SetRenderedFields,
SetRenderedMapIndex,
SetTaskStateStore,
@@ -294,6 +296,10 @@ SOCKET_CLEANUP_TIMEOUT: float = conf.getfloat("workers",
"socket_cleanup_timeout
# 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 @@ class ActivitySubprocess(WatchedSubprocess):
# 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 @@ class ActivitySubprocess(WatchedSubprocess):
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 @@ class ActivitySubprocess(WatchedSubprocess):
)
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 @@ class ActivitySubprocess(WatchedSubprocess):
# 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 @@ class ActivitySubprocess(WatchedSubprocess):
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 @@ class ActivitySubprocess(WatchedSubprocess):
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),
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 a251f11d341..9c4a999c27a 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -112,6 +112,7 @@ from airflow.sdk.execution_time.comms import (
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)
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 fac24c99c58..deba6630554 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -138,6 +138,7 @@ from airflow.sdk.execution_time.comms import (
SentFDs,
SetAssetStateStoreByName,
SetAssetStateStoreByUri,
+ SetExecutionTimeout,
SetRenderedFields,
SetRenderedMapIndex,
SetTaskStateStore,
@@ -1561,6 +1562,153 @@ class TestWatchedSubprocess:
mock_kill.assert_not_called()
mock_logger.warning.assert_not_called()
+ def test_server_terminated_task_gets_killed_task_cleanup_time(self,
mocker):
+
mocker.patch("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
42.0)
+
mocker.patch("airflow.sdk.execution_time.supervisor.MIN_HEARTBEAT_INTERVAL", 0)
+ mock_kill =
mocker.patch("airflow.sdk.execution_time.supervisor.WatchedSubprocess.kill")
+ client = mocker.Mock()
+ client.task_instances.heartbeat.side_effect =
ServerResponseError.from_response(
+ httpx.Response(
+ 409,
+ request=httpx.Request("PUT", "http://server/heartbeat";),
+ json={"detail": {"reason": "not_running", "current_state":
"failed"}},
+ )
+ )
+ proc = ActivitySubprocess(
+ process_log=mocker.MagicMock(),
+ id=TI_ID,
+ pid=12345,
+ stdin=mocker.Mock(),
+ process=mocker.Mock(),
+ client=client,
+ )
+
+ proc._send_heartbeat_if_needed()
+
+ mock_kill.assert_called_once_with(signal.SIGTERM, force=True,
escalation_delay=42.0)
+
+ def test_execution_timeout_kills_task_that_does_not_stop(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)
+
+ monotonic.return_value += 34
+ proc._handle_execution_timeout_if_needed()
+ mock_kill.assert_not_called()
+
+ monotonic.return_value += 1
+ 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",
+ timeout_seconds=30,
+ grace_period_seconds=5.0,
+ )
+
+ proc._handle_execution_timeout_if_needed()
+ mock_kill.assert_called_once()
+
+ @pytest.mark.parametrize(
+ ("terminal_state", "exit_code"),
+ [
+ pytest.param(TaskInstanceState.FAILED, None,
id="terminal_state_reported"),
+ pytest.param(None, 0, id="process_exited"),
+ ],
+ )
+ def test_execution_timeout_not_enforced_once_task_is_over(self, mocker,
terminal_state, exit_code):
+
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._terminal_state = terminal_state
+ proc._exit_code = exit_code
+
+ 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(
+ ("stops_on_sigterm", "expected_exit_code"),
+ [
+ pytest.param(True, 0, id="task_reports_state_on_sigterm"),
+ pytest.param(False, -signal.SIGKILL, id="task_ignores_sigterm"),
+ ],
+ )
+ def test_execution_timeout_enforced_by_supervisor(
+ self, stops_on_sigterm, expected_exit_code, mocker, captured_logs,
client_with_ti_start
+ ):
+
mocker.patch("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
0.3)
+ # Far longer than the test may take: the monitor loop has to wake up
for the timeout on its own
+
mocker.patch("airflow.sdk.execution_time.supervisor.MIN_HEARTBEAT_INTERVAL", 30)
+
+ def subprocess_main():
+ import signal
+
+ comms = CommsDecoder()
+ comms._get_response()
+
+ def _on_term(signum, frame):
+ if stops_on_sigterm:
+ comms.send(TaskState(state=TaskInstanceState.FAILED,
end_date=timezone.utcnow()))
+ exit(0)
+ print("Ignoring SIGTERM", file=sys.stderr)
+
+ signal.signal(signal.SIGTERM, _on_term)
+ comms.send(SetExecutionTimeout(timeout_seconds=0.1))
+ sleep(30)
+ exit(5)
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=TI_ID,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ queue="default",
+ ),
+ client=client_with_ti_start,
+ target=subprocess_main,
+ )
+
+ started = time.monotonic()
+ assert proc.wait() == expected_exit_code
+ assert time.monotonic() - started < 10
+ assert proc.final_state == TaskInstanceState.FAILED
+
+ assert {
+ "event": "Task did not stop before execution_timeout elapsed;
terminating process",
+ "level": "error",
+ "timeout_seconds": 0.1,
+ "grace_period_seconds": 0.3,
+ "logger": "task",
+ "timestamp": mocker.ANY,
+ "loc": mocker.ANY,
+ } in captured_logs
+
@pytest.mark.parametrize(
("signal_to_raise", "log_pattern", "level"),
(
@@ -2451,6 +2599,10 @@ REQUEST_TEST_CASES = [
),
test_id="set_rtif",
),
+ RequestTestCase(
+ message=SetExecutionTimeout(timeout_seconds=30.0),
+ test_id="set_execution_timeout",
+ ),
RequestTestCase(
message=SetRenderedMapIndex(rendered_map_index="Label: task_1"),
client_mock=ClientMock(
@@ -3820,7 +3972,7 @@ class TestHandleRequest:
)
observed_at_kill = []
- def terminate(self, signal_to_send, force):
+ def terminate(self, signal_to_send, force, escalation_delay):
observed_at_kill.append((self._terminal_state,
self._pending_terminal_state_msg))
if arrival == "during_kill":
self._handle_request(msg, structlog.get_logger(), req_id=2)
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 5d6fad9903a..d13ca9225c9 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
@@ -147,6 +147,7 @@ from airflow.sdk.execution_time.comms import (
RetryTask,
SetAssetStateStoreByName,
SetAssetStateStoreByUri,
+ SetExecutionTimeout,
SetRenderedFields,
SetTaskStateStore,
SetXCom,
@@ -1756,7 +1757,7 @@ def test_run_task_timeout(time_machine,
create_runtime_ti, mock_supervisor_comms
mock_supervisor_comms.send.assert_called_with(TaskState(state=TaskInstanceState.FAILED,
end_date=instant))
-def test_execution_timeout(create_runtime_ti):
+def test_execution_timeout(create_runtime_ti, mock_supervisor_comms):
def sleep_and_catch_other_exceptions():
with contextlib.suppress(Exception):
# Catching Exception should NOT catch AirflowTaskTimeout
@@ -6375,7 +6376,7 @@ class TestRunExecuteCallable:
task.execution_timeout = execution_timeout
return task
- def test_runs_in_isolated_context_with_safeguard_tracker_set(self):
+ def test_runs_in_isolated_context_with_safeguard_tracker_set(self,
mock_supervisor_comms):
"""The callable runs in an internal context copy that has the
safeguard tracker set and does not leak."""
var = contextvars.ContextVar("marker")
var.set("outer")
@@ -6397,8 +6398,9 @@ class TestRunExecuteCallable:
# The .set was confined to the copy, so the tracker never leaked to
the caller's context.
assert ExecutorSafeguard.tracker.get(None) is not task
task.on_kill.assert_not_called()
+ mock_supervisor_comms.send.assert_not_called()
- def test_applies_execution_timeout(self):
+ def test_applies_execution_timeout(self, mock_supervisor_comms):
"""When a timeout is set and the callable overruns, AirflowTaskTimeout
is raised and on_kill is called."""
task = self._make_task(execution_timeout=timedelta(milliseconds=10))
@@ -6409,8 +6411,9 @@ class TestRunExecuteCallable:
_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))
- def test_fast_fails_when_timeout_already_elapsed(self):
+ 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."""
task = self._make_task(execution_timeout=timedelta(seconds=-1))
execute = mock.MagicMock()
@@ -6420,6 +6423,7 @@ class TestRunExecuteCallable:
execute.assert_not_called()
task.on_kill.assert_called_once()
+ mock_supervisor_comms.send.assert_not_called()
def test_emits_task_execute_span_at_detail_level_2(self):
"""At detail level 2, running the callable produces a recorded
``task.execute`` span."""
