anukratb commented on code in PR #75297:
URL: https://github.com/apache/airflow/pull/75297#discussion_r4337266005
##########
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:
@Andrushika thats gone too now - did the same in the new on_kill test which
now checks that both the task process and its session child are 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]
