SameerMesiah97 commented on code in PR #73806:
URL: https://github.com/apache/airflow/pull/73806#discussion_r4116927261
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -291,6 +292,10 @@ def _handle_mask_secret_request(
# like listeners after task is complete.
TASK_OVERTIME_THRESHOLD: float = conf.getfloat("core", "task_success_overtime")
+# How long a task process gets to stop itself (raise AirflowTaskTimeout, run
on_kill, report its state) after
+# execution_timeout elapses before the supervisor sends SIGTERM, and again
before it escalates to SIGKILL.
+EXECUTION_TIMEOUT_GRACE_PERIOD: float = 5.0
Review Comment:
Could you clarify the reasoning behind the fixed five-second grace period?
When `execution_timeout` raises in the task process, provider cleanup in
`finally` blocks or `on_kill()` may need to make remote API calls before the
task can report its state. If that takes longer than five seconds, the
supervisor sends `SIGTERM` and may then escalate to `SIGKILL`. Should this
grace period be configurable so tasks with longer cleanup requirements can
finish?
I am not sure if this is feasible but is it possible for grace period to be
surface from the provider code through the task runner?
##########
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:
optional: I would add a test for the `ProcessNotFound` branch. Please see
the below:
```
def test_execution_timeout_clears_escalation_if_process_is_gone(self,
mocker):
process_not_found = type("ProcessNotFound", (Exception,), {})
process = mocker.Mock()
process.ProcessNotFound = process_not_found
mocker.patch("time.monotonic", return_value=20.0)
mock_signal = mocker.patch(
"airflow.sdk.execution_time.supervisor.WatchedSubprocess._signal_subprocess",
side_effect=process_not_found(),
)
mock_kill =
mocker.patch("airflow.sdk.execution_time.supervisor.WatchedSubprocess.kill")
proc = ActivitySubprocess(
process_log=mocker.MagicMock(),
id=TI_ID,
pid=12345,
stdin=mocker.Mock(),
process=process,
client=mocker.Mock(),
)
proc._execution_timeout_seconds = 30.0
proc._execution_timeout_enforce_at = 20.0
proc._execution_timeout_next_signal = signal.SIGTERM
proc._handle_execution_timeout_if_needed()
mock_signal.assert_called_once_with(signal.SIGTERM)
mock_kill.assert_not_called()
assert proc._execution_timeout_next_signal is None
assert proc._execution_timeout_enforce_at is None
assert proc._execution_timeout_due_in() is None
```
Keep in mind that the above test might not be 100% correct.
--
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]
