[ROCm][CI] Warm up the RLHF dev server before the pause/resume timing checks (#53594)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
stefankoncarevic committed
Aug 28, 2026 at 11:07 UTC
5f213ed1592903b7bc38f173d320dac1b2769303
1 file changed
+25
-4
tests/entrypoints/serve/dev/rlhf/conftest.py
+25
-4
@@ -68,6 +68,16 @@ _DUMMY_ARGS = [
68
# ---------------------------------------------------------------------------
69
70
71
+def _warm_up(url: str) -> None:
72
+ """Put one request through the engine before tests start timing things.
73
+
74
+ /health turns green before any request has travelled the request path, and
75
+ that first pass costs seconds on a loaded machine.
76
+ """
77
+ response = gen(url, max_tokens=4, timeout=120)
78
+ assert ok(response), f"warm-up generation failed: {response}"
79
+
80
+
81
@contextmanager
82
def server(
83
extra_args=None,
@@ -118,6 +128,7 @@ def server(
128
else:
129
proc.terminate()
130
raise RuntimeError("vllm server did not start in time")
131
+ _warm_up(url)
132
yield url
133
finally:
134
proc.terminate()
@@ -216,6 +227,10 @@ def ok(resp) -> bool:
227
# ---------------------------------------------------------------------------
228
229
230
+# First-token wait for a streaming request; loaded machines need the slack.
231
+STREAM_START_TIMEOUT = 20.0
232
+
233
+
234
@dataclass
235
class StreamResult:
236
started: threading.Event = field(default_factory=threading.Event)
@@ -265,12 +280,18 @@ def start_stream(url: str, max_tokens: int) -> tuple[StreamResult, threading.Thr
280
args=(url, result, max_tokens),
281
)
282
thread.start()
268
- started = result.started.wait(timeout=10)
283
+ started = result.started.wait(timeout=STREAM_START_TIMEOUT)
284
if not started or result.done.is_set():
270
- pause(url, mode="abort")
271
- resume(url)
285
+ # Best-effort: on a stalled server these time out too, and would then
286
+ # mask the assertions below.
287
+ with contextlib.suppress(requests.RequestException):
288
+ pause(url, mode="abort")
289
+ resume(url)
290
thread.join(timeout=10)
273
- assert started, "request did not start generating"
291
+ assert started, (
292
+ f"request did not start generating within {STREAM_START_TIMEOUT}s "
293
+ f"(stream error: {result.error})"
294
+ )
295
assert not result.done.is_set(), "request completed before it could be paused"
296
return result, thread
297