[Test] Add pause/resume E2E tests (#52144)
Signed-off-by: flb_ <floatlibai@gmail.com>
flb_ committed
Aug 18, 2026 at 14:46 UTC
d785eb51ce24210ec6da82140729d30822332032
3 files changed
+582
tests/entrypoints/serve/dev/rlhf/conftest.py
new
+414
@@ -0,0 +1,414 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Shared fixtures and helpers for the RL lifecycle test suite.
5
+
6
+All test modules under this directory import from here to avoid duplication.
7
+
8
+RFC: https://github.com/vllm-project/vllm/issues/45585
9
+PR: https://github.com/vllm-project/vllm/pull/45586
10
+"""
11
+
12
+import contextlib
13
+import json
14
+import os
15
+import subprocess
16
+import sys
17
+import threading
18
+import time
19
+from collections.abc import Callable
20
+from contextlib import contextmanager
21
+from dataclasses import dataclass, field
22
+from typing import Any
23
+
24
+import requests
25
+
26
+# ---------------------------------------------------------------------------
27
+# Model / server defaults
28
+# ---------------------------------------------------------------------------
29
+
30
+
31
+MODEL_NAME = os.environ.get("VLLM_TEST_MODEL", "Qwen/Qwen3-0.6B")
32
+
33
+
34
+_BASE_ARGS = [
35
+ "--dtype",
36
+ "bfloat16",
37
+ "--max-model-len",
38
+ "2048",
39
+ "--max-num-seqs",
40
+ "32",
41
+ "--gpu-memory-utilization",
42
+ "0.75",
43
+ "--enable-sleep-mode",
44
+ "--enforce-eager",
45
+]
46
+
47
+
48
+# Lightweight args for state-machine / protocol tests that don't need real
49
+# weights (avoids spending time downloading a 1B checkpoint in T0 tests).
50
+_DUMMY_ARGS = [
51
+ "--dtype",
52
+ "bfloat16",
53
+ "--max-model-len",
54
+ "128",
55
+ "--max-num-seqs",
56
+ "8",
57
+ "--gpu-memory-utilization",
58
+ "0.5",
59
+ "--enable-sleep-mode",
60
+ "--enforce-eager",
61
+ "--load-format",
62
+ "dummy",
63
+]
64
+
65
+
66
+# ---------------------------------------------------------------------------
67
+# Server harness
68
+# ---------------------------------------------------------------------------
69
+
70
+
71
+@contextmanager
72
+def server(
73
+ extra_args=None,
74
+ port: int = 8770,
75
+ timeout: float = 180.0,
76
+ dummy_weights: bool = False,
77
+):
78
+ """Launch a vLLM server with the dev router; yield its base URL.
79
+
80
+ Args:
81
+ extra_args: Additional CLI flags appended after the base args.
82
+ port: HTTP port to bind (caller is responsible for uniqueness).
83
+ timeout: Seconds to wait for /health before giving up.
84
+ dummy_weights: If True, use --load-format dummy (fast, no real weights).
85
+ """
86
+ env = {**os.environ, "VLLM_SERVER_DEV_MODE": "1"}
87
+ base = _DUMMY_ARGS if dummy_weights else _BASE_ARGS
88
+ cmd = [
89
+ sys.executable,
90
+ "-m",
91
+ "vllm.entrypoints.openai.api_server",
92
+ "--model",
93
+ MODEL_NAME,
94
+ "--port",
95
+ str(port),
96
+ "--served-model-name",
97
+ "m",
98
+ *(base + (extra_args or [])),
99
+ ]
100
+ proc = subprocess.Popen(
101
+ cmd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE
102
+ )
103
+ url = f"http://localhost:{port}"
104
+ try:
105
+ deadline = time.time() + timeout
106
+ while time.time() < deadline:
107
+ if proc.poll() is not None:
108
+ err = (
109
+ proc.stderr.read(4000).decode(errors="replace")
110
+ if proc.stderr
111
+ else ""
112
+ )
113
+ raise RuntimeError(f"vllm server exited during startup:\n{err}")
114
+ with contextlib.suppress(Exception):
115
+ if requests.get(f"{url}/health", timeout=3).status_code == 200:
116
+ break
117
+ time.sleep(1)
118
+ else:
119
+ proc.terminate()
120
+ raise RuntimeError("vllm server did not start in time")
121
+ yield url
122
+ finally:
123
+ proc.terminate()
124
+ with contextlib.suppress(subprocess.TimeoutExpired):
125
+ proc.wait(timeout=10)
126
+ if proc.poll() is None:
127
+ proc.kill()
128
+
129
+
130
+# ---------------------------------------------------------------------------
131
+# Polling helper (200-lie workaround)
132
+# ---------------------------------------------------------------------------
133
+
134
+
135
+def poll_until(
136
+ predicate: Callable[[], bool],
137
+ timeout: float = 10.0,
138
+ interval: float = 0.5,
139
+) -> bool:
140
+ """Poll predicate() until it returns True or timeout expires.
141
+
142
+ Workaround for the vLLM sleep/wake "200-lie" — the HTTP endpoints may
143
+ return 200 before the underlying operation is complete, so callers that
144
+ need to verify state *after* an operation can use this helper instead of
145
+ assuming the 200 means completion.
146
+
147
+ Returns True if predicate became true within timeout, False otherwise.
148
+ """
149
+ deadline = time.time() + timeout
150
+ while time.time() < deadline:
151
+ try:
152
+ if predicate():
153
+ return True
154
+ except Exception:
155
+ pass
156
+ time.sleep(interval)
157
+ return False
158
+
159
+
160
+# ---------------------------------------------------------------------------
161
+# HTTP helpers — generation
162
+# ---------------------------------------------------------------------------
163
+
164
+
165
+def gen(url, prompt="The capital of France is", max_tokens=8, timeout=30):
166
+ """Fire a /v1/completions request; return JSON or None on any error."""
167
+ try:
168
+ r = requests.post(
169
+ f"{url}/v1/completions",
170
+ json={
171
+ "model": "m",
172
+ "prompt": prompt,
173
+ "max_tokens": max_tokens,
174
+ "temperature": 0,
175
+ },
176
+ timeout=timeout,
177
+ )
178
+ return r.json()
179
+ except Exception:
180
+ return None
181
+
182
+
183
+def gen_with_logprobs(
184
+ url, prompt="The capital of France is", max_tokens=8, logprobs=5, timeout=30
185
+):
186
+ """Fire a /v1/completions request with logprobs; return JSON or None."""
187
+ try:
188
+ r = requests.post(
189
+ f"{url}/v1/completions",
190
+ json={
191
+ "model": "m",
192
+ "prompt": prompt,
193
+ "max_tokens": max_tokens,
194
+ "temperature": 0,
195
+ "logprobs": logprobs,
196
+ },
197
+ timeout=timeout,
198
+ )
199
+ return r.json()
200
+ except Exception:
201
+ return None
202
+
203
+
204
+def ok(resp) -> bool:
205
+ """True iff resp is a successful completion (has choices, no error key)."""
206
+ return (
207
+ resp is not None
208
+ and "choices" in resp
209
+ and bool(resp["choices"])
210
+ and "error" not in resp
211
+ )
212
+
213
+
214
+# ---------------------------------------------------------------------------
215
+# HTTP helpers — stream generation
216
+# ---------------------------------------------------------------------------
217
+
218
+
219
+@dataclass
220
+class StreamResult:
221
+ started: threading.Event = field(default_factory=threading.Event)
222
+ done: threading.Event = field(default_factory=threading.Event)
223
+ chunks: list[dict[str, Any]] = field(default_factory=list)
224
+ finish_reason: str | None = None
225
+ error: Exception | None = None
226
+
227
+
228
+def stream_completion(url: str, result: StreamResult, max_tokens: int) -> None:
229
+ try:
230
+ with requests.post(
231
+ f"{url}/v1/completions",
232
+ json={
233
+ "model": "m",
234
+ "prompt": "Count upward slowly: one, two, three,",
235
+ "max_tokens": max_tokens,
236
+ "temperature": 0,
237
+ "ignore_eos": True,
238
+ "stream": True,
239
+ },
240
+ stream=True,
241
+ timeout=(5, 60),
242
+ ) as response:
243
+ response.raise_for_status()
244
+ for line in response.iter_lines(decode_unicode=True):
245
+ if not line or line == "data: [DONE]":
246
+ continue
247
+ assert line.startswith("data: ")
248
+ chunk = json.loads(line.removeprefix("data: "))
249
+ result.chunks.append(chunk)
250
+ choice = chunk["choices"][0]
251
+ if choice.get("text"):
252
+ result.started.set()
253
+ if choice.get("finish_reason") is not None:
254
+ result.finish_reason = choice["finish_reason"]
255
+ except Exception as error:
256
+ result.error = error
257
+ finally:
258
+ result.done.set()
259
+
260
+
261
+def start_stream(url: str, max_tokens: int) -> tuple[StreamResult, threading.Thread]:
262
+ result = StreamResult()
263
+ thread = threading.Thread(
264
+ target=stream_completion,
265
+ args=(url, result, max_tokens),
266
+ )
267
+ thread.start()
268
+ started = result.started.wait(timeout=10)
269
+ if not started or result.done.is_set():
270
+ pause(url, mode="abort")
271
+ resume(url)
272
+ thread.join(timeout=10)
273
+ assert started, "request did not start generating"
274
+ assert not result.done.is_set(), "request completed before it could be paused"
275
+ return result, thread
276
+
277
+
278
+# ---------------------------------------------------------------------------
279
+# HTTP helpers — pause / resume
280
+# ---------------------------------------------------------------------------
281
+
282
+
283
+def pause(url, mode="abort", clear_cache=True):
284
+ return requests.post(
285
+ f"{url}/pause",
286
+ params={"mode": mode, "clear_cache": clear_cache},
287
+ timeout=15,
288
+ ).status_code
289
+
290
+
291
+def resume(url):
292
+ return requests.post(f"{url}/resume", timeout=10).status_code
293
+
294
+
295
+def completion_with_cache_details(url: str, prompt: str) -> dict[str, Any]:
296
+ response = requests.post(
297
+ f"{url}/v1/completions",
298
+ json={
299
+ "model": "m",
300
+ "prompt": prompt,
301
+ "max_tokens": 8,
302
+ "temperature": 0,
303
+ "logprobs": 1,
304
+ },
305
+ timeout=30,
306
+ )
307
+ response.raise_for_status()
308
+ return response.json()
309
+
310
+
311
+def golden_output(response: dict[str, Any]) -> dict[str, Any]:
312
+ choice = response["choices"][0]
313
+ usage = response["usage"]
314
+ return {
315
+ "text": choice["text"],
316
+ "finish_reason": choice["finish_reason"],
317
+ "tokens": choice["logprobs"]["tokens"],
318
+ "prompt_tokens": usage["prompt_tokens"],
319
+ "completion_tokens": usage["completion_tokens"],
320
+ }
321
+
322
+
323
+def cached_tokens(response: dict[str, Any]) -> int:
324
+ return response["usage"]["prompt_tokens_details"]["cached_tokens"]
325
+
326
+
327
+# ---------------------------------------------------------------------------
328
+# HTTP helpers — sleep / wake
329
+# ---------------------------------------------------------------------------
330
+
331
+
332
+def sleep(url, level=1, mode="abort"):
333
+ return requests.post(
334
+ f"{url}/sleep", params={"level": level, "mode": mode}, timeout=15
335
+ ).status_code
336
+
337
+
338
+def wake(url, tags=None):
339
+ params = {"tags": tags} if tags else {}
340
+ return requests.post(f"{url}/wake_up", params=params, timeout=20).status_code
341
+
342
+
343
+def is_sleeping(url) -> bool:
344
+ return requests.get(f"{url}/is_sleeping", timeout=5).json()["is_sleeping"]
345
+
346
+
347
+def is_paused(url) -> bool:
348
+ return requests.get(f"{url}/is_paused", timeout=5).json()["is_paused"]
349
+
350
+
351
+def health(url) -> int:
352
+ try:
353
+ return requests.get(f"{url}/health", timeout=5).status_code
354
+ except Exception:
355
+ return 0
356
+
357
+
358
+# ---------------------------------------------------------------------------
359
+# HTTP helpers — weight transfer
360
+# ---------------------------------------------------------------------------
361
+
362
+
363
+def start_weight_update(url, is_checkpoint_format=True):
364
+ return requests.post(
365
+ f"{url}/start_weight_update",
366
+ json={"is_checkpoint_format": is_checkpoint_format},
367
+ timeout=10,
368
+ )
369
+
370
+
371
+def finish_weight_update(url):
372
+ return requests.post(f"{url}/finish_weight_update", timeout=10)
373
+
374
+
375
+def get_world_size(url, include_dp=True):
376
+ return requests.get(
377
+ f"{url}/get_world_size",
378
+ params={"include_dp": include_dp},
379
+ timeout=5,
380
+ )
381
+
382
+
383
+# ---------------------------------------------------------------------------
384
+# GPU / metrics helpers
385
+# ---------------------------------------------------------------------------
386
+
387
+
388
+def gpu_free_bytes(device: int = 0) -> int:
389
+ """Read GPU free bytes via subprocess to avoid import-time torch init."""
390
+ out = subprocess.check_output(
391
+ [
392
+ sys.executable,
393
+ "-c",
394
+ f"import torch; f,_=torch.accelerator.get_memory_info({device}); print(f)",
395
+ ],
396
+ timeout=10,
397
+ )
398
+ return int(out.strip())
399
+
400
+
401
+def sleep_metrics(url):
402
+ """Return (awake, weights_offloaded, discard_all) from /metrics."""
403
+ try:
404
+ from prometheus_client.parser import text_string_to_metric_families
405
+ except ImportError:
406
+ return None, None, None
407
+
408
+ r = requests.get(f"{url}/metrics", timeout=5)
409
+ vals: dict = {}
410
+ for family in text_string_to_metric_families(r.text):
411
+ if family.name == "vllm:engine_sleep_state":
412
+ for s in family.samples:
413
+ vals[s.labels.get("sleep_state", "")] = s.value
414
+ return vals.get("awake"), vals.get("weights_offloaded"), vals.get("discard_all")
tests/entrypoints/serve/dev/rlhf/state_transitions/__init__.py
tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py
new
+168
@@ -0,0 +1,168 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""End-to-end tests for the vLLM RL pause/resume lifecycle."""
4
+
5
+import os
6
+import threading
7
+from typing import Any
8
+from unittest.mock import patch
9
+
10
+import pytest
11
+import requests
12
+
13
+from tests.entrypoints.serve.dev.rlhf.conftest import (
14
+ cached_tokens,
15
+ completion_with_cache_details,
16
+ gen,
17
+ golden_output,
18
+ is_paused,
19
+ ok,
20
+ pause,
21
+ resume,
22
+ server,
23
+ start_stream,
24
+)
25
+
26
+
27
+@pytest.fixture(scope="module", params=[False, True], ids=["MRV1", "MRV2"])
28
+def use_v2(request):
29
+ return request.param
30
+
31
+
32
+@pytest.fixture(scope="module")
33
+def server_url(use_v2):
34
+ env_vars = {
35
+ "VLLM_USE_V2_MODEL_RUNNER": "1" if use_v2 else "0",
36
+ }
37
+
38
+ with (
39
+ patch.dict(os.environ, env_vars),
40
+ server(
41
+ extra_args=[
42
+ "--enable-prefix-caching",
43
+ "--enable-prompt-tokens-details",
44
+ ]
45
+ ) as url,
46
+ ):
47
+ yield url
48
+
49
+
50
+@pytest.fixture(autouse=True)
51
+def restore_unpaused_state(server_url):
52
+ assert resume(server_url) == 200
53
+ yield
54
+ assert resume(server_url) == 200
55
+
56
+
57
+class TestPauseResume:
58
+ def test_state_and_idempotency_across_cycles(self, server_url):
59
+ assert not is_paused(server_url)
60
+
61
+ assert resume(server_url) == 200
62
+ assert resume(server_url) == 200
63
+ assert not is_paused(server_url)
64
+
65
+ for _ in range(2):
66
+ for mode in ("abort", "wait", "keep"):
67
+ assert pause(server_url, mode=mode) == 200
68
+ assert pause(server_url, mode=mode) == 200
69
+ assert is_paused(server_url)
70
+
71
+ assert resume(server_url) == 200
72
+ assert resume(server_url) == 200
73
+ assert not is_paused(server_url)
74
+
75
+ def test_invalid_mode_preserves_state(self, server_url):
76
+ for paused in (False, True):
77
+ if paused:
78
+ assert pause(server_url) == 200
79
+ assert is_paused(server_url) is paused
80
+
81
+ response = requests.post(
82
+ f"{server_url}/pause",
83
+ params={"mode": "invalid"},
84
+ timeout=10,
85
+ )
86
+ assert response.status_code == 400
87
+ assert response.json()["error"]["param"] == "query.mode"
88
+ assert is_paused(server_url) is paused
89
+
90
+ assert resume(server_url) == 200
91
+
92
+ @pytest.mark.parametrize(
93
+ ("mode", "max_tokens", "inflight_finish_reason"),
94
+ [
95
+ pytest.param("abort", 256, "abort", id="abort"),
96
+ pytest.param("wait", 256, "length", id="wait"),
97
+ pytest.param("keep", 256, "length", id="keep"),
98
+ ],
99
+ )
100
+ def test_mode_request_lifecycle(
101
+ self,
102
+ server_url,
103
+ mode,
104
+ max_tokens,
105
+ inflight_finish_reason,
106
+ ):
107
+ inflight, inflight_thread = start_stream(server_url, max_tokens)
108
+ new_result: dict[str, Any] = {}
109
+ new_done = threading.Event()
110
+
111
+ def _new_request():
112
+ new_result["response"] = gen(server_url, max_tokens=4, timeout=60)
113
+ new_done.set()
114
+
115
+ new_thread = threading.Thread(target=_new_request)
116
+ try:
117
+ assert pause(server_url, mode=mode) == 200
118
+ assert is_paused(server_url)
119
+
120
+ if mode in ("abort", "wait"):
121
+ assert inflight.done.is_set()
122
+ else:
123
+ chunks_after_pause = len(inflight.chunks)
124
+ assert not inflight.done.wait(timeout=5)
125
+ assert len(inflight.chunks) == chunks_after_pause, (
126
+ "in-flight request continued generating in keep mode"
127
+ )
128
+
129
+ new_thread.start()
130
+ assert not new_done.wait(timeout=0.3), (
131
+ "new request completed while generation was paused"
132
+ )
133
+ finally:
134
+ assert resume(server_url) == 200
135
+ inflight_thread.join(timeout=30)
136
+ if new_thread.ident is not None:
137
+ new_thread.join(timeout=30)
138
+
139
+ assert not inflight_thread.is_alive()
140
+ assert inflight.error is None
141
+ assert inflight.finish_reason == inflight_finish_reason
142
+ assert not new_thread.is_alive()
143
+ assert ok(new_result.get("response"))
144
+
145
+ def test_clear_cache_preserves_output_and_controls_prefix_cache(self, server_url):
146
+ prompt = (
147
+ "Paris is the capital of France. Berlin is the capital of Germany. "
148
+ ) * 20
149
+
150
+ assert pause(server_url, clear_cache=True) == 200
151
+ assert resume(server_url) == 200
152
+ baseline = completion_with_cache_details(server_url, prompt)
153
+ warmed = completion_with_cache_details(server_url, prompt)
154
+
155
+ assert cached_tokens(baseline) == 0
156
+ assert cached_tokens(warmed) > 0
157
+
158
+ assert pause(server_url, clear_cache=False) == 200
159
+ assert resume(server_url) == 200
160
+ preserved = completion_with_cache_details(server_url, prompt)
161
+ assert golden_output(preserved) == golden_output(baseline)
162
+ assert cached_tokens(preserved) > 0
163
+
164
+ assert pause(server_url, clear_cache=True) == 200
165
+ assert resume(server_url) == 200
166
+ cleared = completion_with_cache_details(server_url, prompt)
167
+ assert golden_output(cleared) == golden_output(baseline)
168
+ assert cached_tokens(cleared) == 0