From 9272320798cbfe0e726f0f9330ff2c7f76811877 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Mon, 6 Jul 2026 20:41:28 -0700 Subject: [PATCH] feat(timing): warmup progress heartbeat and empty-drain guard Co-Authored-By: Claude Fable 5 Signed-off-by: Anthony Casagrande --- docs/environment-variables.md | 1 + src/aiperf/common/environment.py | 7 +++ src/aiperf/timing/phase/runner.py | 33 +++++++++++ tests/unit/timing/phase/test_runner.py | 81 ++++++++++++++++++++++++++ 4 files changed, 122 insertions(+) diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 7b4b64ca6d..22fe44d57a 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -219,6 +219,7 @@ Service lifecycle and inter-service communication configuration. Controls timeou | `AIPERF_SERVICE_CONNECTION_PROBE_INTERVAL` | `0.1` | ≥ 0.1, ≤ 600.0 | Interval in seconds for connection probes while waiting for initial connection to the zmq message bus | | `AIPERF_SERVICE_CONNECTION_PROBE_TIMEOUT` | `90.0` | ≥ 1.0, ≤ 100000.0 | Maximum time in seconds to wait for connection probe response while waiting for initial connection to the zmq message bus | | `AIPERF_SERVICE_CREDIT_PROGRESS_REPORT_INTERVAL` | `2.0` | ≥ 1, ≤ 100000.0 | Interval in seconds between credit progress report messages | +| `AIPERF_SERVICE_WARMUP_PROGRESS_LOG_INTERVAL` | `30.0` | ≥ 0.0, ≤ 100000.0 | Interval in seconds between warmup progress heartbeat log messages. Set to 0 to disable. | | `AIPERF_SERVICE_DISABLE_UVLOOP` | `False` | — | Disable uvloop and use default asyncio event loop instead | | `AIPERF_SERVICE_HEARTBEAT_INTERVAL` | `5.0` | ≥ 1.0, ≤ 100000.0 | Interval in seconds between heartbeat messages for component services | | `AIPERF_SERVICE_PROFILE_CONFIGURE_TIMEOUT` | `600.0` | ≥ 1.0, ≤ 100000.0 | Timeout in seconds for profile configure command | diff --git a/src/aiperf/common/environment.py b/src/aiperf/common/environment.py index b530e8b4d3..ac9bd3c89c 100644 --- a/src/aiperf/common/environment.py +++ b/src/aiperf/common/environment.py @@ -872,6 +872,13 @@ class _ServiceSettings(BaseSettings): default=2.0, description="Interval in seconds between credit progress report messages", ) + WARMUP_PROGRESS_LOG_INTERVAL: float = Field( + ge=0.0, + le=100000.0, + default=30.0, + description="Interval in seconds between warmup progress heartbeat log messages. " + "Set to 0 to disable.", + ) DISABLE_UVLOOP: bool = Field( default=False, description="Disable uvloop and use default asyncio event loop instead", diff --git a/src/aiperf/timing/phase/runner.py b/src/aiperf/timing/phase/runner.py index 19dd63746f..b376d9715d 100644 --- a/src/aiperf/timing/phase/runner.py +++ b/src/aiperf/timing/phase/runner.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio +import time from collections.abc import Callable from typing import TYPE_CHECKING @@ -585,6 +586,24 @@ def _format_phase_complete(self, stats: CreditPhaseStats) -> str: parts.append("was_cancelled=True") return " | ".join(parts) + @staticmethod + def _format_warmup_progress(stats: CreditPhaseStats) -> str: + """Format a periodic warmup heartbeat for non-interactive logs.""" + returned = stats.requests_completed + stats.requests_cancelled + target = stats.final_requests_sent or stats.total_expected_requests + returned_desc = ( + f"returned={returned:,}/{target:,}" if target else f"returned={returned:,}" + ) + parts = [ + f"Phase {stats.phase} progress", + returned_desc, + f"sent={stats.requests_sent:,}", + f"in_flight={stats.in_flight_requests:,}", + f"errors={stats.request_errors:,}", + f"elapsed={stats.requests_elapsed_time:.1f}s", + ] + return " | ".join(parts) + async def _wait_for_sending_complete(self) -> None: """Wait for phase to send all credits (with timeout). @@ -665,6 +684,8 @@ async def _wait_for_returning_complete(self) -> None: f"Waiting for all cancelled credits to be returned for " f"phase {self._config.phase}. Need {need} more credits." ) + if need <= 0: + self._progress.all_credits_returned_event.set() # Wait with timeout to avoid hanging indefinitely drain_timeout = Environment.TIMING.CANCEL_DRAIN_TIMEOUT try: @@ -767,13 +788,25 @@ async def _progress_report_loop(self) -> None: Runs as a background task until the phase is complete. Publishes progress at CREDIT_PROGRESS_REPORT_INTERVAL intervals. + During warmup, also emits a throttled INFO heartbeat so headless runs + remain observable when no interactive UI consumes progress messages. """ self.debug(f"Starting progress reporting loop for phase {self._config.phase}") + warmup_log_interval = Environment.SERVICE.WARMUP_PROGRESS_LOG_INTERVAL + next_warmup_log_at = time.monotonic() + warmup_log_interval try: while True: try: stats = self._progress.create_stats(self._lifecycle) await self._phase_publisher.publish_progress(stats) + now = time.monotonic() + if ( + self._config.phase == CreditPhase.WARMUP + and warmup_log_interval > 0 + and now >= next_warmup_log_at + ): + self.info(self._format_warmup_progress(stats)) + next_warmup_log_at = now + warmup_log_interval except Exception as e: self.error( f"Error publishing progress for phase {self._config.phase}: {e!r}" diff --git a/tests/unit/timing/phase/test_runner.py b/tests/unit/timing/phase/test_runner.py index 326614a781..eaf22d1179 100644 --- a/tests/unit/timing/phase/test_runner.py +++ b/tests/unit/timing/phase/test_runner.py @@ -654,6 +654,51 @@ async def test_cleanup_runs_on_exception( await r.run(is_final_phase=True) mock_ramper.stop.assert_called_once() + async def test_timeout_skips_cancel_drain_when_all_credits_returned( + self, + conv_src: MagicMock, + pub: MagicMock, + router: MagicMock, + conc: MagicMock, + cancel: MagicMock, + cb: MagicMock, + ) -> None: + r = make_runner(cfg(dur=1.0), conv_src, pub, router, conc, cancel, cb) + stats = CreditPhaseStats( + phase=CreditPhase.PROFILING, + start_ns=1, + sent_end_ns=2, + requests_end_ns=3, + final_requests_sent=94, + requests_sent=94, + requests_completed=94, + requests_cancelled=0, + final_requests_completed=94, + final_requests_cancelled=0, + final_request_errors=0, + final_sent_sessions=1, + final_completed_sessions=1, + final_cancelled_sessions=0, + ) + r._progress.create_stats = MagicMock(return_value=stats) + # Force the post-timeout cancel branch: pretend credits are still + # outstanding at the initial check, time out the wait, then recompute + # need == 0 so the empty-drain guard sets the event without waiting. + r._progress.check_all_returned_or_cancelled = MagicMock(return_value=False) + r._lifecycle.start() + r._lifecycle.mark_sending_complete() + + with patch.object( + r, + "_wait_for_event_with_timeout", + new=AsyncMock(return_value=True), + ): + await r._wait_for_returning_complete() + + router.cancel_all_credits.assert_awaited_once() + assert r._progress.all_credits_returned_event.is_set() + r._concurrency_manager.release_stuck_slots.assert_not_called() + class TestFixedScheduleConfigCorrection: """Tests for FIXED_SCHEDULE mode config correction using actual dataset size.""" @@ -830,3 +875,39 @@ async def execute_phase(self) -> None: real_router._register_worker("worker-1") await asyncio.wait_for(run_task, timeout=5.0) assert strategy.execute_called + + +class TestWarmupProgressHeartbeat: + """Warmup emits a throttled INFO heartbeat in _progress_report_loop + (AIPERF_SERVICE_WARMUP_PROGRESS_LOG_INTERVAL, default 30s, 0 disables) so + headless warmup stays observable when no interactive UI consumes progress. + """ + + def test_format_warmup_progress_includes_returned_target_and_elapsed(self): + stats = CreditPhaseStats( + phase=CreditPhase.WARMUP, + requests_sent=40, + requests_completed=30, + requests_cancelled=2, + request_errors=1, + total_expected_requests=50, + ) + msg = PhaseRunner._format_warmup_progress(stats) + assert "Phase warmup progress" in msg + assert "returned=32/50" in msg # completed + cancelled / target + assert "sent=40" in msg + assert "errors=1" in msg + assert "elapsed=" in msg + + def test_format_warmup_progress_without_target(self): + stats = CreditPhaseStats( + phase=CreditPhase.WARMUP, + requests_sent=5, + requests_completed=3, + ) + msg = PhaseRunner._format_warmup_progress(stats) + assert "returned=3" in msg and "returned=3/" not in msg # no target + + def test_warmup_log_interval_env_field_exists(self): + # 0 disables; default is a positive heartbeat cadence. + assert Environment.SERVICE.WARMUP_PROGRESS_LOG_INTERVAL >= 0.0