test: cover per-phase drain timeout in BenchmarkSession#402
Conversation
|
MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅ |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces configurable per-phase in-flight drain timeouts via a new DrainConfig schema, replacing the hard-coded 240-second timeout to better support slow generative workloads. The reviewer identified a potential TypeError and rounding issue in _drain_inflight where formatting timeout_s (which can be None or a sub-second float) with %.0f in the error log could fail or display misleadingly as 0 s. Suggestions were provided to safely format the timeout string and update the corresponding unit test assertion.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| logger.info( | ||
| "Draining %d in-flight responses (timeout: %s)...", | ||
| phase_issuer.inflight, | ||
| f"{timeout_s:.0f} s" if timeout_s is not None else "none", | ||
| ) | ||
| self._drain_event.clear() | ||
| try: | ||
| await asyncio.wait_for(self._drain_event.wait(), timeout=240.0) | ||
| await asyncio.wait_for(self._drain_event.wait(), timeout=timeout_s) | ||
| except TimeoutError: | ||
| logger.error( | ||
| "Drain timed out after 240 s with %d responses still in flight; " | ||
| "Drain timed out after %.0f s with %d responses still in flight; " | ||
| "proceeding to next phase.", | ||
| timeout_s, | ||
| phase_issuer.inflight, | ||
| ) |
There was a problem hiding this comment.
The timeout_s parameter is typed as float | None. Formatting it directly with %.0f in logger.error can trigger a TypeError at runtime if TimeoutError is raised while timeout_s is None (or cause static type checkers like mypy to fail). Additionally, using :.0f or %.0f rounds sub-second timeouts (like 0.1 s in tests) to 0 s, which is misleading.
We can define a helper string timeout_str to format the timeout safely and precisely for both log messages.
| logger.info( | |
| "Draining %d in-flight responses (timeout: %s)...", | |
| phase_issuer.inflight, | |
| f"{timeout_s:.0f} s" if timeout_s is not None else "none", | |
| ) | |
| self._drain_event.clear() | |
| try: | |
| await asyncio.wait_for(self._drain_event.wait(), timeout=240.0) | |
| await asyncio.wait_for(self._drain_event.wait(), timeout=timeout_s) | |
| except TimeoutError: | |
| logger.error( | |
| "Drain timed out after 240 s with %d responses still in flight; " | |
| "Drain timed out after %.0f s with %d responses still in flight; " | |
| "proceeding to next phase.", | |
| timeout_s, | |
| phase_issuer.inflight, | |
| ) | |
| timeout_str = f"{timeout_s} s" if timeout_s is not None else "none" | |
| logger.info( | |
| "Draining %d in-flight responses (timeout: %s)...", | |
| phase_issuer.inflight, | |
| timeout_str, | |
| ) | |
| self._drain_event.clear() | |
| try: | |
| await asyncio.wait_for(self._drain_event.wait(), timeout=timeout_s) | |
| except TimeoutError: | |
| logger.error( | |
| "Drain timed out after %s with %d responses still in flight; " | |
| "proceeding to next phase.", | |
| timeout_str, | |
| phase_issuer.inflight, | |
| ) |
|
|
||
| result = await asyncio.wait_for(session.run(phases), timeout=10.0) | ||
| assert result is not None | ||
| assert "Drain timed out after 0 s" in caplog.text or "Drain timed out" in caplog.text |
There was a problem hiding this comment.
Since we updated the log formatting to preserve decimal values for sub-second timeouts (avoiding rounding 0.1 s to 0 s), we can simplify and make this assertion more precise.
| assert "Drain timed out after 0 s" in caplog.text or "Drain timed out" in caplog.text | |
| assert "Drain timed out after 0.1 s" in caplog.text |
wu6u3tw
force-pushed
the
wan22-drain-timeouts
branch
from
c143e68 to
36b43f1
Compare
|
Applied both review suggestions: drain log lines now format the timeout via a shared timeout_str (no %.0f on a possibly-None float, and sub-second values like 0.1 s print exactly), and the test asserts the precise "Drain timed out after 0.1 s" message. Also trimmed the inline comments/docstrings. 101 session+config tests pass. |
| help="Warmup drain timeout in seconds (None = wait indefinitely)", | ||
| ), | ||
| ] = Field( | ||
| warmup_timeout_s: float | None = Field( |
There was a problem hiding this comment.
I would in general avoid adding more fine-granularity flag to control timeout at each phase, because it adds burden to the whole system and makes config even more difficult to understand/write.
@arekay-nv @viraatc do you think we should just set an arbitrary large timeout and fail out if the timeout is hit? The silent timeout has caused us lots of trouble in tuning because it will pretend successful run while some samples are actually dropped.
The per-phase drain timeout feature (settings.drain, PhaseConfig.drain_timeout) landed via mlcommons#330 and mlcommons#337 but had no session-level regression test. Add one: a phase whose samples never complete must exit the drain wait after drain_timeout seconds instead of hanging.
wu6u3tw
force-pushed
the
wan22-drain-timeouts
branch
from
36b43f1 to
a31fe03
Compare
wu6u3tw
changed the title
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #402 +/- ##
=======================================
Coverage ? 80.09%
=======================================
Files ? 129
Lines ? 17276
Branches ? 0
=======================================
Hits ? 13838
Misses ? 3438
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
3 participants
Summary
The original version of this PR made per-phase in-flight drain timeouts configurable. That feature has since landed on main via #330 and #337 (
settings.drain/DrainConfig,PhaseConfig.drain_timeout, unbounded accuracy drain), so this PR is reduced to the one piece main still lacks: a session-level regression test.test_drain_respects_per_phase_timeout: a phase whose in-flight samples never complete must exit the drain wait afterdrain_timeoutseconds instead of hanging.Test plan
uv run pytest tests/unit/load_generator/test_async_session.py -k drain(3 passed)pre-commit runon the changed file (ruff, mypy, license header all pass)