Andrushika commented on code in PR #73806:
URL: https://github.com/apache/airflow/pull/73806#discussion_r4204227303
##########
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:
Yes, and this is on purpose. So basically the first n seconds is to give
enough time for users' customized `on_kill()` and failure callback (it could
take long). The second n seconds is the time between SIGTERM and SIGKILL.
And Yeah, it could be confusing since these two share the same config. I
originally just wanted the "first timeout" to be configurable to the users,
that's why I am reusing this existing config instead of adding a new one.
We could separate it into another config if you prefer?
(Or, keep only timeout -> SIGTERM configurable, and fix SIGTERM -> SIGKILL
to the default escalation delay)
##########
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:
Removed
##########
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:
Changed all of them into mocker
##########
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:
Yes you're right... Fixed it. I'll take it more carefully next time
##########
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:
Fixed. The check is gone now that the deadline is just one more entry in the
`min(...)` for the select timeout.
##########
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:
Fixed!
##########
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:
Whoops sorry I didn't notice there is the same logic... I reuse `kill()`
now, and the next-signal field is gone.
--
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]
