ashb commented on code in PR #73806:
URL: https://github.com/apache/airflow/pull/73806#discussion_r4195494175
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -1305,6 +1306,204 @@ def test_overtime_handling(
mock_kill.assert_not_called()
mock_logger.warning.assert_not_called()
+ def test_server_terminated_task_gets_killed_task_cleanup_time(self,
mocker, monkeypatch):
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
42.0)
+
monkeypatch.setattr("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_set_execution_timeout_schedules_enforcement(self, mocker,
monkeypatch):
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
5.0)
+ mocker.patch("time.monotonic", return_value=100.0)
Review Comment:
Use freezegun instead of general mocking please.
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1882,6 +1892,8 @@ def _monitor_subprocess(self):
MIN_HEARTBEAT_INTERVAL,
),
)
+ if self._exit_code is None and (due_in :=
self._execution_timeout_due_in()) is not None:
Review Comment:
This condition feels complicated. Thoughts?
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -1305,6 +1306,204 @@ def test_overtime_handling(
mock_kill.assert_not_called()
mock_logger.warning.assert_not_called()
+ def test_server_terminated_task_gets_killed_task_cleanup_time(self,
mocker, monkeypatch):
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
42.0)
+
monkeypatch.setattr("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_set_execution_timeout_schedules_enforcement(self, mocker,
monkeypatch):
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
5.0)
+ mocker.patch("time.monotonic", return_value=100.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)
+
+ assert proc._execution_timeout_seconds == 30
+ assert proc._execution_timeout_enforce_at == 135.0
Review Comment:
These are impl details that I don't think should be asserted,
##########
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.
Review Comment:
This points out a slight disparity.
By using this same threshold for timeout overrun, this would mean that you
get n seconds after timeout, then TERM is sent, and then possibly _another_ n
seconds before the task is killed?
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -1305,6 +1306,204 @@ def test_overtime_handling(
mock_kill.assert_not_called()
mock_logger.warning.assert_not_called()
+ def test_server_terminated_task_gets_killed_task_cleanup_time(self,
mocker, monkeypatch):
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
42.0)
+
monkeypatch.setattr("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_set_execution_timeout_schedules_enforcement(self, mocker,
monkeypatch):
Review Comment:
Using mocker and monkeypatch in the same test is "okay, but just a bit odd".
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1921,6 +1934,61 @@ def _handle_process_overtime_if_needed(self):
)
self.kill(signal.SIGTERM, force=True)
+ def _execution_timeout_due_in(self) -> float | None:
+ """Seconds until the supervisor must act on an overrunning task, or
None if there is nothing to enforce."""
+ if (
+ self._execution_timeout_next_signal is None
+ or self._execution_timeout_enforce_at is None
+ or self._terminal_state
+ or self._pending_terminal_state_msg is not None
+ ):
+ return None
+ return self._execution_timeout_enforce_at - time.monotonic()
+
+ def _handle_execution_timeout_if_needed(self):
Review Comment:
I'm a little uneasy about having the same "TERM then KILL" here as it
already exists. Is there any way we could reuse that logic. That would also
make this simpler as we wouldn't need `_execution_timeout_next_signal` etc
either.
Discuss?
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1921,6 +1934,61 @@ def _handle_process_overtime_if_needed(self):
)
self.kill(signal.SIGTERM, force=True)
+ def _execution_timeout_due_in(self) -> float | None:
+ """Seconds until the supervisor must act on an overrunning task, or
None if there is nothing to enforce."""
+ if (
+ self._execution_timeout_next_signal is None
+ or self._execution_timeout_enforce_at is None
+ or self._terminal_state
+ or self._pending_terminal_state_msg is not None
+ ):
+ return None
+ 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.
+ """
+ due_in = self._execution_timeout_due_in()
+ if due_in is None or due_in > 0:
+ return
+
+ next_signal = self._execution_timeout_next_signal
+ if next_signal == signal.SIGTERM:
+ self.process_log.error(
+ "Task did not stop after execution_timeout elapsed;
terminating process",
Review Comment:
```suggestion
"Task did not stop before execution_timeout elapsed;
terminating process",
```
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -1305,6 +1306,204 @@ def test_overtime_handling(
mock_kill.assert_not_called()
mock_logger.warning.assert_not_called()
+ def test_server_terminated_task_gets_killed_task_cleanup_time(self,
mocker, monkeypatch):
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
42.0)
+
monkeypatch.setattr("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_set_execution_timeout_schedules_enforcement(self, mocker,
monkeypatch):
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
5.0)
+ mocker.patch("time.monotonic", return_value=100.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)
+
+ assert proc._execution_timeout_seconds == 30
+ assert proc._execution_timeout_enforce_at == 135.0
+ assert proc._execution_timeout_next_signal == signal.SIGTERM
+ assert proc._execution_timeout_due_in() == 35.0
+
+ @pytest.mark.parametrize(
+ ("enforce_at", "next_signal", "terminal_state",
"pending_terminal_msg", "expected_action"),
+ [
+ pytest.param(None, None, None, False, None, id="no_timeout_set"),
+ pytest.param(25.0, signal.SIGTERM, None, False, None,
id="not_due_yet"),
+ pytest.param(20.0, signal.SIGTERM, None, False, "sigterm",
id="due_sends_sigterm"),
+ pytest.param(20.0, signal.SIGKILL, None, False, "sigkill",
id="due_again_sends_sigkill"),
+ pytest.param(
+ 15.0, signal.SIGTERM, TaskInstanceState.FAILED, False, None,
id="terminal_state_reported"
+ ),
+ pytest.param(15.0, signal.SIGTERM, None, True, None,
id="terminal_state_pending_api_retry"),
+ ],
+ )
+ def test_execution_timeout_enforcement(
+ self,
+ mocker,
+ monkeypatch,
+ enforce_at,
+ next_signal,
+ terminal_state,
+ pending_terminal_msg,
+ expected_action,
+ ):
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.KILLED_TASK_CLEANUP_TIME",
5.0)
+ mocker.patch("time.monotonic", return_value=20.0)
+ mock_kill =
mocker.patch("airflow.sdk.execution_time.supervisor.WatchedSubprocess.kill")
+ mock_signal = mocker.patch(
+
"airflow.sdk.execution_time.supervisor.WatchedSubprocess._signal_subprocess"
+ )
+ proc = ActivitySubprocess(
+ process_log=mocker.MagicMock(),
+ id=TI_ID,
+ pid=12345,
+ stdin=mocker.Mock(),
+ process=mocker.Mock(),
+ client=mocker.Mock(),
+ )
+ proc._execution_timeout_seconds = 30.0
+ proc._execution_timeout_enforce_at = enforce_at
+ proc._execution_timeout_next_signal = next_signal
+ proc._terminal_state = terminal_state
+ if pending_terminal_msg:
+ proc._pending_terminal_state_msg =
SucceedTask(end_date=timezone.utcnow())
+
+ proc._handle_execution_timeout_if_needed()
+
+ if expected_action == "sigterm":
Review Comment:
I don't love the parametrize plus "if name ==" pattern that agents are prone
to. See if you can make this more data driven via parametrization cleanly
please.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
