| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared fixtures and helpers for the RL lifecycle test suite. |
| 4 | |
| 5 | All test modules under this directory import from here to avoid duplication. |
| 6 | |
| 7 | RFC: https://github.com/vllm-project/vllm/issues/45585 |
| 8 | PR: https://github.com/vllm-project/vllm/pull/45586 |
| 9 | """ |
| 10 | |
| 11 | import contextlib |
| 12 | import json |
| 13 | import os |
| 14 | import subprocess |
| 15 | import sys |
| 16 | import threading |
| 17 | import time |
| 18 | from contextlib import contextmanager |
| 19 | from dataclasses import dataclass, field |
| 20 | from typing import Any |
| 21 | |
| 22 | import requests |
| 23 | |
| 24 | # --------------------------------------------------------------------------- |
| 25 | # Model / server defaults |
| 26 | # --------------------------------------------------------------------------- |
| 27 | |
| 28 | |
| 29 | MODEL_NAME = os.environ.get("VLLM_TEST_MODEL", "Qwen/Qwen3-0.6B") |
| 30 | |
| 31 | |
| 32 | _BASE_ARGS = [ |
| 33 | "--dtype", |
| 34 | "bfloat16", |
| 35 | "--max-model-len", |
| 36 | "2048", |
| 37 | "--max-num-seqs", |
| 38 | "32", |
| 39 | "--gpu-memory-utilization", |
| 40 | "0.75", |
| 41 | "--enable-sleep-mode", |
| 42 | "--enforce-eager", |
| 43 | ] |
| 44 | |
| 45 | |
| 46 | # Lightweight args for state-machine / protocol tests that don't need real |
| 47 | # weights (avoids spending time downloading a 1B checkpoint in T0 tests). |
| 48 | _DUMMY_ARGS = [ |
| 49 | "--dtype", |
| 50 | "bfloat16", |
| 51 | "--max-model-len", |
| 52 | "128", |
| 53 | "--max-num-seqs", |
| 54 | "8", |
| 55 | "--gpu-memory-utilization", |
| 56 | "0.5", |
| 57 | "--enable-sleep-mode", |
| 58 | "--enforce-eager", |
| 59 | "--load-format", |
| 60 | "dummy", |
| 61 | ] |
| 62 | |
| 63 | |
| 64 | # --------------------------------------------------------------------------- |
| 65 | # Server harness |
| 66 | # --------------------------------------------------------------------------- |
| 67 | |
| 68 | |
| 69 | def _warm_up(url: str) -> None: |
| 70 | """Put one request through the engine before tests start timing things. |
| 71 | |
| 72 | /health turns green before any request has travelled the request path, and |
| 73 | that first pass costs seconds on a loaded machine. |
| 74 | """ |
| 75 | response = gen(url, max_tokens=4, timeout=120) |
| 76 | assert ok(response), f"warm-up generation failed: {response}" |
| 77 | |
| 78 | |
| 79 | @contextmanager |
| 80 | def server( |
| 81 | extra_args=None, |
| 82 | port: int = 8770, |
| 83 | timeout: float = 180.0, |
| 84 | dummy_weights: bool = False, |
| 85 | ): |
| 86 | """Launch a vLLM server with the dev router; yield its base URL. |
| 87 | |
| 88 | Args: |
| 89 | extra_args: Additional CLI flags appended after the base args. |
| 90 | port: HTTP port to bind (caller is responsible for uniqueness). |
| 91 | timeout: Seconds to wait for /health before giving up. |
| 92 | dummy_weights: If True, use --load-format dummy (fast, no real weights). |
| 93 | |
| 94 | """ |
| 95 | env = {**os.environ, "VLLM_SERVER_DEV_MODE": "1"} |
| 96 | base = _DUMMY_ARGS if dummy_weights else _BASE_ARGS |
| 97 | cmd = [ |
| 98 | sys.executable, |
| 99 | "-m", |
| 100 | "vllm.entrypoints.openai.api_server", |
| 101 | "--model", |
| 102 | MODEL_NAME, |
| 103 | "--port", |
| 104 | str(port), |
| 105 | "--served-model-name", |
| 106 | "m", |
| 107 | *(base + (extra_args or [])), |
| 108 | ] |
| 109 | proc = subprocess.Popen( |
| 110 | cmd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE |
| 111 | ) |
| 112 | url = f"http://localhost:{port}" |
| 113 | try: |
| 114 | deadline = time.time() + timeout |
| 115 | while time.time() < deadline: |
| 116 | if proc.poll() is not None: |
| 117 | err = ( |
| 118 | proc.stderr.read(4000).decode(errors="replace") |
| 119 | if proc.stderr |
| 120 | else "" |
| 121 | ) |
| 122 | raise RuntimeError(f"vllm server exited during startup:\n{err}") |
| 123 | with contextlib.suppress(Exception): |
| 124 | if requests.get(f"{url}/health", timeout=3).status_code == 200: |
| 125 | break |
| 126 | time.sleep(1) |
| 127 | else: |
| 128 | proc.terminate() |
| 129 | raise RuntimeError("vllm server did not start in time") |
| 130 | _warm_up(url) |
| 131 | yield url |
| 132 | finally: |
| 133 | proc.terminate() |
| 134 | with contextlib.suppress(subprocess.TimeoutExpired): |
| 135 | proc.wait(timeout=10) |
| 136 | if proc.poll() is None: |
| 137 | proc.kill() |
| 138 | |
| 139 | |
| 140 | # --------------------------------------------------------------------------- |
| 141 | # HTTP helpers — generation |
| 142 | # --------------------------------------------------------------------------- |
| 143 | |
| 144 | |
| 145 | def gen(url, prompt="The capital of France is", max_tokens=8, timeout=30): |
| 146 | """Fire a /v1/completions request; return JSON or None on any error.""" |
| 147 | try: |
| 148 | r = requests.post( |
| 149 | f"{url}/v1/completions", |
| 150 | json={ |
| 151 | "model": "m", |
| 152 | "prompt": prompt, |
| 153 | "max_tokens": max_tokens, |
| 154 | "temperature": 0, |
| 155 | }, |
| 156 | timeout=timeout, |
| 157 | ) |
| 158 | return r.json() |
| 159 | except Exception: |
| 160 | return None |
| 161 | |
| 162 | |
| 163 | def ok(resp) -> bool: |
| 164 | """True iff resp is a successful completion (has choices, no error key).""" |
| 165 | return ( |
| 166 | resp is not None |
| 167 | and "choices" in resp |
| 168 | and bool(resp["choices"]) |
| 169 | and "error" not in resp |
| 170 | ) |
| 171 | |
| 172 | |
| 173 | # --------------------------------------------------------------------------- |
| 174 | # HTTP helpers — stream generation |
| 175 | # --------------------------------------------------------------------------- |
| 176 | |
| 177 | |
| 178 | # First-token wait for a streaming request; loaded machines need the slack. |
| 179 | STREAM_START_TIMEOUT = 20.0 |
| 180 | |
| 181 | |
| 182 | @dataclass |
| 183 | class StreamResult: |
| 184 | started: threading.Event = field(default_factory=threading.Event) |
| 185 | done: threading.Event = field(default_factory=threading.Event) |
| 186 | chunks: list[dict[str, Any]] = field(default_factory=list) |
| 187 | finish_reason: str | None = None |
| 188 | error: Exception | None = None |
| 189 | |
| 190 | |
| 191 | def stream_completion(url: str, result: StreamResult, max_tokens: int) -> None: |
| 192 | try: |
| 193 | with requests.post( |
| 194 | f"{url}/v1/completions", |
| 195 | json={ |
| 196 | "model": "m", |
| 197 | "prompt": "Count upward slowly: one, two, three,", |
| 198 | "max_tokens": max_tokens, |
| 199 | "temperature": 0, |
| 200 | "ignore_eos": True, |
| 201 | "stream": True, |
| 202 | }, |
| 203 | stream=True, |
| 204 | timeout=(5, 60), |
| 205 | ) as response: |
| 206 | response.raise_for_status() |
| 207 | for line in response.iter_lines(decode_unicode=True): |
| 208 | if not line or line == "data: [DONE]": |
| 209 | continue |
| 210 | assert line.startswith("data: ") |
| 211 | chunk = json.loads(line.removeprefix("data: ")) |
| 212 | result.chunks.append(chunk) |
| 213 | choice = chunk["choices"][0] |
| 214 | if choice.get("text"): |
| 215 | result.started.set() |
| 216 | if choice.get("finish_reason") is not None: |
| 217 | result.finish_reason = choice["finish_reason"] |
| 218 | except Exception as error: |
| 219 | result.error = error |
| 220 | finally: |
| 221 | result.done.set() |
| 222 | |
| 223 | |
| 224 | def start_stream(url: str, max_tokens: int) -> tuple[StreamResult, threading.Thread]: |
| 225 | result = StreamResult() |
| 226 | thread = threading.Thread( |
| 227 | target=stream_completion, |
| 228 | args=(url, result, max_tokens), |
| 229 | ) |
| 230 | thread.start() |
| 231 | started = result.started.wait(timeout=STREAM_START_TIMEOUT) |
| 232 | if not started or result.done.is_set(): |
| 233 | # Best-effort: on a stalled server these time out too, and would then |
| 234 | # mask the assertions below. |
| 235 | with contextlib.suppress(requests.RequestException): |
| 236 | pause(url, mode="abort") |
| 237 | resume(url) |
| 238 | thread.join(timeout=10) |
| 239 | assert started, ( |
| 240 | f"request did not start generating within {STREAM_START_TIMEOUT}s " |
| 241 | f"(stream error: {result.error})" |
| 242 | ) |
| 243 | assert not result.done.is_set(), "request completed before it could be paused" |
| 244 | return result, thread |
| 245 | |
| 246 | |
| 247 | # --------------------------------------------------------------------------- |
| 248 | # HTTP helpers — pause / resume |
| 249 | # --------------------------------------------------------------------------- |
| 250 | |
| 251 | |
| 252 | def pause(url, mode="abort", clear_cache=True): |
| 253 | return requests.post( |
| 254 | f"{url}/pause", |
| 255 | params={"mode": mode, "clear_cache": clear_cache}, |
| 256 | timeout=15, |
| 257 | ).status_code |
| 258 | |
| 259 | |
| 260 | def resume(url): |
| 261 | return requests.post(f"{url}/resume", timeout=10).status_code |
| 262 | |
| 263 | |
| 264 | def completion_with_cache_details(url: str, prompt: str) -> dict[str, Any]: |
| 265 | response = requests.post( |
| 266 | f"{url}/v1/completions", |
| 267 | json={ |
| 268 | "model": "m", |
| 269 | "prompt": prompt, |
| 270 | "max_tokens": 8, |
| 271 | "temperature": 0, |
| 272 | "logprobs": 1, |
| 273 | }, |
| 274 | timeout=30, |
| 275 | ) |
| 276 | response.raise_for_status() |
| 277 | return response.json() |
| 278 | |
| 279 | |
| 280 | def golden_output(response: dict[str, Any]) -> dict[str, Any]: |
| 281 | choice = response["choices"][0] |
| 282 | usage = response["usage"] |
| 283 | return { |
| 284 | "text": choice["text"], |
| 285 | "finish_reason": choice["finish_reason"], |
| 286 | "tokens": choice["logprobs"]["tokens"], |
| 287 | "prompt_tokens": usage["prompt_tokens"], |
| 288 | "completion_tokens": usage["completion_tokens"], |
| 289 | } |
| 290 | |
| 291 | |
| 292 | def cached_tokens(response: dict[str, Any]) -> int: |
| 293 | return response["usage"]["prompt_tokens_details"]["cached_tokens"] |
| 294 | |
| 295 | |
| 296 | # --------------------------------------------------------------------------- |
| 297 | # HTTP helpers — sleep / wake |
| 298 | # --------------------------------------------------------------------------- |
| 299 | |
| 300 | |
| 301 | def sleep(url, level=1, mode="abort"): |
| 302 | return requests.post( |
| 303 | f"{url}/sleep", params={"level": level, "mode": mode}, timeout=15 |
| 304 | ).status_code |
| 305 | |
| 306 | |
| 307 | def wake(url, tags=None): |
| 308 | params = {"tags": tags} if tags else {} |
| 309 | return requests.post(f"{url}/wake_up", params=params, timeout=20).status_code |
| 310 | |
| 311 | |
| 312 | def is_sleeping(url) -> bool: |
| 313 | return requests.get(f"{url}/is_sleeping", timeout=5).json()["is_sleeping"] |
| 314 | |
| 315 | |
| 316 | def is_paused(url) -> bool: |
| 317 | return requests.get(f"{url}/is_paused", timeout=5).json()["is_paused"] |
| 318 | |
| 319 | |
| 320 | def health(url) -> int: |
| 321 | try: |
| 322 | return requests.get(f"{url}/health", timeout=5).status_code |
| 323 | except Exception: |
| 324 | return 0 |
| 325 | |
| 326 | |
| 327 | # --------------------------------------------------------------------------- |
| 328 | # HTTP helpers — weight transfer |
| 329 | # --------------------------------------------------------------------------- |
| 330 | |
| 331 | |
| 332 | def start_weight_update(url, is_checkpoint_format=True): |
| 333 | return requests.post( |
| 334 | f"{url}/start_weight_update", |
| 335 | json={"is_checkpoint_format": is_checkpoint_format}, |
| 336 | timeout=10, |
| 337 | ) |
| 338 | |
| 339 | |
| 340 | def finish_weight_update(url): |
| 341 | return requests.post(f"{url}/finish_weight_update", timeout=10) |
| 342 | |
| 343 | |
| 344 | def get_world_size(url, include_dp=True): |
| 345 | return requests.get( |
| 346 | f"{url}/get_world_size", |
| 347 | params={"include_dp": include_dp}, |
| 348 | timeout=5, |
| 349 | ) |