Andrushika commented on code in PR #73806:
URL: https://github.com/apache/airflow/pull/73806#discussion_r4125066916
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -1303,6 +1304,150 @@ def test_overtime_handling(
mock_kill.assert_not_called()
mock_logger.warning.assert_not_called()
+ def test_set_execution_timeout_schedules_enforcement(self, mocker,
monkeypatch):
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.EXECUTION_TIMEOUT_GRACE_PERIOD",
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.EXECUTION_TIMEOUT_GRACE_PERIOD",
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":
+ mock_signal.assert_called_once_with(signal.SIGTERM)
+ mock_kill.assert_not_called()
+ assert proc._execution_timeout_next_signal == signal.SIGKILL
+ assert proc._execution_timeout_enforce_at == 25.0
+ proc.process_log.error.assert_called_once_with(
+ "Task did not stop after execution_timeout elapsed;
terminating process",
+ timeout_seconds=30.0,
+ grace_period_seconds=5.0,
+ )
+ elif expected_action == "sigkill":
+ mock_kill.assert_called_once_with(signal.SIGKILL)
+ mock_signal.assert_not_called()
+ assert proc._execution_timeout_next_signal is None
+ assert proc._execution_timeout_enforce_at is None
+ proc.process_log.error.assert_called_once_with(
+ "Task process did not exit after SIGTERM; killing it",
timeout_seconds=30.0
+ )
+ else:
+ mock_signal.assert_not_called()
+ 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, monkeypatch,
captured_logs, client_with_ti_start
+ ):
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.EXECUTION_TIMEOUT_GRACE_PERIOD",
0.3)
+ # Far longer than the test may take: the monitor loop has to wake up
for the timeout on its own
+
monkeypatch.setattr("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
+
+ events = [m["event"] for m in captured_logs]
+ assert "Task did not stop after execution_timeout elapsed; terminating
process" in events
+ assert ("Task process did not exit after SIGTERM; killing it" in
events) is not stops_on_sigterm
+
Review Comment:
Added, thanks.
--
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]
