This is an automated email from the ASF dual-hosted git repository.
kaxilnaik pushed a commit to branch v3-0-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-0-test by this push:
new f92f5c8fa63 [v3-0-test] Prevent CPU spike in task supervisor when
heartbeat timeout exceeded (#51023) (#51047)
f92f5c8fa63 is described below
commit f92f5c8fa63066f2d0c3857ebf11cb2e1b0e6af0
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Mon Jun 2 23:37:50 2025 +0530
[v3-0-test] Prevent CPU spike in task supervisor when heartbeat timeout
exceeded (#51023) (#51047)
(cherry picked from commit beb7b629161c0e51f7995ee630ec53820a41b6f0)
Co-authored-by: Kaxil Naik <[email protected]>
---
.../src/airflow/sdk/execution_time/supervisor.py | 4 +-
.../task_sdk/execution_time/test_supervisor.py | 67 ++++++++++++++++++++++
2 files changed, 70 insertions(+), 1 deletion(-)
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index 0328525b0c8..fcbf7126d97 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -685,7 +685,9 @@ class WatchedSubprocess:
:param expect_signal: Signal not to log if the task exits with this
code.
:returns: The process exit code, or None if it's still alive
"""
- events = self.selector.select(timeout=max_wait_time)
+ # Ensure minimum timeout to prevent CPU spike with tight loop when
timeout is 0 or negative
+ timeout = max(0.01, max_wait_time)
+ events = self.selector.select(timeout=timeout)
for key, _ in events:
# Retrieve the handler responsible for processing this file object
(e.g., stdout, stderr)
socket_handler = key.data
diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
index 1d19d19021e..7861a9d5997 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -958,6 +958,73 @@ class TestWatchedSubprocessKill:
# Validate that `_check_subprocess_exit` is called
mock_process.wait.assert_called_once_with(timeout=0)
+ def test_max_wait_time_prevents_cpu_spike(self, watched_subprocess,
mock_process, monkeypatch):
+ """Test that max_wait_time calculation prevents CPU spike when
heartbeat timeout is reached."""
+ # Mock the configuration to reproduce the CPU spike scenario
+ # Set heartbeat timeout to be very small relative to
MIN_HEARTBEAT_INTERVAL
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.HEARTBEAT_TIMEOUT",
1)
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.MIN_HEARTBEAT_INTERVAL",
10)
+
+ # Set up a scenario where the last successful heartbeat was a long
time ago
+ # This will cause the heartbeat calculation to result in a negative
value
+ mock_process._last_successful_heartbeat = time.monotonic() - 100 #
100 seconds ago
+
+ # Mock process to still be alive (not exited)
+ mock_process.wait.side_effect = psutil.TimeoutExpired(pid=12345,
seconds=0)
+
+ # Call _service_subprocess which is used in _monitor_subprocess
+ # This tests the max_wait_time calculation directly
+ watched_subprocess._service_subprocess(max_wait_time=0.005) # Very
small timeout to verify our fix
+
+ # Verify that selector.select was called with a minimum timeout of 0.01
+ # This proves our fix prevents the timeout=0 scenario that causes CPU
spike
+ watched_subprocess.selector.select.assert_called_once()
+ call_args = watched_subprocess.selector.select.call_args
+ timeout_arg = call_args[1]["timeout"] if "timeout" in call_args[1]
else call_args[0][0]
+
+ # The timeout should be at least 0.01 (our minimum), never 0
+ assert timeout_arg >= 0.01, f"Expected timeout >= 0.01, got
{timeout_arg}"
+
+ @pytest.mark.parametrize(
+ ["heartbeat_timeout", "min_interval", "heartbeat_ago",
"expected_min_timeout"],
+ [
+ # Normal case: heartbeat is recent, should use calculated value
+ pytest.param(30, 5, 5, 0.01, id="normal_heartbeat"),
+ # Edge case: heartbeat timeout exceeded, should use minimum
+ pytest.param(10, 20, 50, 0.01, id="heartbeat_timeout_exceeded"),
+ # Bug reproduction case: timeout < interval, heartbeat very old
+ pytest.param(5, 10, 100, 0.01, id="cpu_spike_scenario"),
+ ],
+ )
+ def test_max_wait_time_calculation_edge_cases(
+ self,
+ watched_subprocess,
+ mock_process,
+ monkeypatch,
+ heartbeat_timeout,
+ min_interval,
+ heartbeat_ago,
+ expected_min_timeout,
+ ):
+ """Test max_wait_time calculation in various edge case scenarios."""
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.HEARTBEAT_TIMEOUT",
heartbeat_timeout)
+
monkeypatch.setattr("airflow.sdk.execution_time.supervisor.MIN_HEARTBEAT_INTERVAL",
min_interval)
+
+ watched_subprocess._last_successful_heartbeat = time.monotonic() -
heartbeat_ago
+ mock_process.wait.side_effect = psutil.TimeoutExpired(pid=12345,
seconds=0)
+
+ # Call the method and verify timeout is never less than our minimum
+ watched_subprocess._service_subprocess(
+ max_wait_time=999
+ ) # Large value, should be overridden by calculation
+
+ # Extract the timeout that was actually used
+ watched_subprocess.selector.select.assert_called_once()
+ call_args = watched_subprocess.selector.select.call_args
+ actual_timeout = call_args[1]["timeout"] if "timeout" in call_args[1]
else call_args[0][0]
+
+ assert actual_timeout >= expected_min_timeout
+
class TestHandleRequest:
@pytest.fixture
