Andrushika commented on code in PR #75297:
URL: https://github.com/apache/airflow/pull/75297#discussion_r4336764775
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -2231,6 +2231,95 @@ def _handler(sig, frame):
expected_logs.extend(({"event": "Process exited", "logger":
"supervisor"},))
assert logs == expected_logs
+ def test_wait_terminates_subprocess_when_monitor_raises(self,
watched_subprocess, mock_process, mocker):
+ monitor_error = OSError(28, "No space left on device")
+ mocker.patch.object(
+ ActivitySubprocess, "_monitor_subprocess", autospec=True,
side_effect=monitor_error
+ )
+ cleanup_sockets = mocker.patch.object(ActivitySubprocess,
"cleanup_sockets_after_kill", autospec=True)
+ signal_subprocess = mocker.patch.object(ActivitySubprocess,
"_signal_subprocess", autospec=True)
+ upload_logs = mocker.patch.object(ActivitySubprocess, "_upload_logs",
autospec=True)
+ mock_process.wait.return_value = -signal.SIGTERM
+
+ with pytest.raises(OSError, match="No space left on device") as raised:
+ watched_subprocess.wait()
+
+ assert raised.value is monitor_error
+ cleanup_sockets.assert_called_once_with(watched_subprocess)
+ signal_subprocess.assert_called_once_with(watched_subprocess,
signal.SIGTERM)
+
mock_process.wait.assert_called_once_with(timeout=supervisor.KILLED_TASK_CLEANUP_TIME)
+ watched_subprocess.selector.close.assert_called_once_with()
+ upload_logs.assert_not_called()
+
+ def test_wait_escalates_subprocess_after_monitor_error(self,
watched_subprocess, mock_process, mocker):
+ mocker.patch.object(
+ ActivitySubprocess, "_monitor_subprocess", autospec=True,
side_effect=RuntimeError("boom")
+ )
+ signal_subprocess = mocker.patch.object(ActivitySubprocess,
"_signal_subprocess", autospec=True)
+ mock_process.wait.side_effect = [
+ psutil.TimeoutExpired(pid=watched_subprocess.pid, seconds=0),
+ -signal.SIGKILL,
+ ]
+
+ with pytest.raises(RuntimeError, match="boom"):
+ watched_subprocess.wait()
+
+ assert signal_subprocess.call_args_list == [
+ mocker.call(watched_subprocess, signal.SIGTERM),
+ mocker.call(watched_subprocess, signal.SIGKILL),
+ ]
+ assert mock_process.wait.call_args_list == [
+ mocker.call(timeout=supervisor.KILLED_TASK_CLEANUP_TIME),
+ mocker.call(timeout=supervisor.KILLED_TASK_CLEANUP_TIME),
+ ]
+ assert watched_subprocess._exit_code == -signal.SIGKILL
+
+ def test_monitor_log_write_error_does_not_leave_subprocess_running(
+ self, monkeypatch, client_with_ti_start
+ ):
+ def subprocess_main():
+ CommsDecoder()._get_response()
+ while True:
+ print("fill log", flush=True)
+ sleep(0.01)
+
+ def fake_reconfigure(logger, *args, **kwargs):
+ return logger
+
+ monkeypatch.setattr(supervisor, "KILLED_TASK_CLEANUP_TIME", 0.2)
+ monkeypatch.setattr(supervisor, "reconfigure_logger", fake_reconfigure)
+ logger = mock.Mock(spec=FilteringBoundLogger)
+ logger.log.side_effect = OSError(28, "No space left on device")
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=uuid7(),
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ queue="default",
+ ),
+ client=client_with_ti_start,
+ logger=logger,
+ target=subprocess_main,
+ )
+
+ try:
+ with pytest.raises(OSError, match="No space left on device"):
+ proc.wait()
+
+ assert proc._exit_code is not None
Review Comment:
Nit: `assert not psutil.pid_exists(proc.pid)` below already checks the
effect, so this assert on `_exit_code` can go.
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1485,6 +1485,34 @@ def _check_subprocess_exit(
self.process_log.warning("Process exited abnormally",
exit_code=self._exit_code)
return self._exit_code
+ def _terminate_subprocess_after_monitor_error(self) -> None:
+ if self._exit_code is not None:
+ return
+
+ try:
+ self.cleanup_sockets_after_kill()
+ except Exception:
+ log.exception("Failed to clean up subprocess sockets after
supervisor loop error", pid=self.pid)
+
+ for sig in (signal.SIGTERM, signal.SIGKILL):
Review Comment:
Let's use the existing `self.kill()` here
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -2231,6 +2231,95 @@ def _handler(sig, frame):
expected_logs.extend(({"event": "Process exited", "logger":
"supervisor"},))
assert logs == expected_logs
+ def test_wait_terminates_subprocess_when_monitor_raises(self,
watched_subprocess, mock_process, mocker):
+ monitor_error = OSError(28, "No space left on device")
+ mocker.patch.object(
+ ActivitySubprocess, "_monitor_subprocess", autospec=True,
side_effect=monitor_error
+ )
+ cleanup_sockets = mocker.patch.object(ActivitySubprocess,
"cleanup_sockets_after_kill", autospec=True)
+ signal_subprocess = mocker.patch.object(ActivitySubprocess,
"_signal_subprocess", autospec=True)
+ upload_logs = mocker.patch.object(ActivitySubprocess, "_upload_logs",
autospec=True)
+ mock_process.wait.return_value = -signal.SIGTERM
+
+ with pytest.raises(OSError, match="No space left on device") as raised:
+ watched_subprocess.wait()
+
+ assert raised.value is monitor_error
+ cleanup_sockets.assert_called_once_with(watched_subprocess)
+ signal_subprocess.assert_called_once_with(watched_subprocess,
signal.SIGTERM)
+
mock_process.wait.assert_called_once_with(timeout=supervisor.KILLED_TASK_CLEANUP_TIME)
+ watched_subprocess.selector.close.assert_called_once_with()
+ upload_logs.assert_not_called()
+
+ def test_wait_escalates_subprocess_after_monitor_error(self,
watched_subprocess, mock_process, mocker):
Review Comment:
Nit: if this path reuses `kill()`, this test can go, since
`test_kill_escalation_path` already covers this.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -2231,6 +2231,95 @@ def _handler(sig, frame):
expected_logs.extend(({"event": "Process exited", "logger":
"supervisor"},))
assert logs == expected_logs
+ def test_wait_terminates_subprocess_when_monitor_raises(self,
watched_subprocess, mock_process, mocker):
+ monitor_error = OSError(28, "No space left on device")
+ mocker.patch.object(
+ ActivitySubprocess, "_monitor_subprocess", autospec=True,
side_effect=monitor_error
+ )
+ cleanup_sockets = mocker.patch.object(ActivitySubprocess,
"cleanup_sockets_after_kill", autospec=True)
+ signal_subprocess = mocker.patch.object(ActivitySubprocess,
"_signal_subprocess", autospec=True)
+ upload_logs = mocker.patch.object(ActivitySubprocess, "_upload_logs",
autospec=True)
+ mock_process.wait.return_value = -signal.SIGTERM
+
+ with pytest.raises(OSError, match="No space left on device") as raised:
+ watched_subprocess.wait()
+
+ assert raised.value is monitor_error
+ cleanup_sockets.assert_called_once_with(watched_subprocess)
+ signal_subprocess.assert_called_once_with(watched_subprocess,
signal.SIGTERM)
+
mock_process.wait.assert_called_once_with(timeout=supervisor.KILLED_TASK_CLEANUP_TIME)
+ watched_subprocess.selector.close.assert_called_once_with()
+ upload_logs.assert_not_called()
Review Comment:
We might not want to assert this. It locks in that remote logs are not
uploaded when the loop fails, and the comment in `wait()` says a failure is
when the logs matter most.
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1485,6 +1485,34 @@ def _check_subprocess_exit(
self.process_log.warning("Process exited abnormally",
exit_code=self._exit_code)
return self._exit_code
+ def _terminate_subprocess_after_monitor_error(self) -> None:
+ if self._exit_code is not None:
+ return
+
+ try:
+ self.cleanup_sockets_after_kill()
Review Comment:
We might need to send SIGTERM before closing the sockets here. IIUC,
`on_kill()` then runs with the log socket already closed, and
`SubprocessHook.send_sigterm()` logs before it calls `os.killpg`, so if that
log call fails the bash group from #75275 is never signalled.
--
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]
