| """Bounded, aggregate-only internal warm-up; importing this module sends nothing.""" |
| from __future__ import annotations |
|
|
| import argparse |
| from concurrent.futures import ThreadPoolExecutor |
| import json |
| import math |
| import re |
| from threading import Barrier |
| import time |
| from typing import Callable, cast |
| from urllib.request import Request, urlopen |
| import uuid |
|
|
| BUDGET = 32 |
| SHAPES = (1, 2, 3, 4) |
| TIMEOUT = 60 |
| |
| |
| DEADLINE_S = 90.0 |
| Reader = Callable[[str, dict[str, object] | None], str] |
|
|
|
|
| def obj(value: object) -> dict[str, object]: |
| if not isinstance(value, dict): |
| raise ValueError("expected JSON object") |
| return cast(dict[str, object], value) |
|
|
|
|
| def payload(model: str, salt: str, wave: int, index: int) -> dict[str, object]: |
| |
| |
| prompt = f"Internal startup {salt} wave {wave} stream {index}.\n" |
| prompt += "".join( |
| f"Paragraph {i}. The system reads a numbered record and checks each entry against its stored value. " |
| for i in range(240 if wave == 1 else 1) |
| ) |
| return {"model": model, "prompt": prompt, "temperature": 0, "n": 1, |
| "max_tokens": BUDGET, "min_tokens": BUDGET, "ignore_eos": True, |
| "stop": [], "stop_token_ids": [], "stream": False, "priority": 20} |
|
|
|
|
| def summary(raw: str) -> dict[str, object]: |
| data = obj(json.loads(raw)) |
| usage = obj(data.get("usage")) |
| count = usage.get("completion_tokens") |
| choices = data.get("choices") |
| if not isinstance(choices, list) or len(choices) != 1: |
| raise ValueError("warm-up requires exactly one completion per request") |
| choice = obj(choices[0]) |
| reason = choice.get("finish_reason") |
| stop = choice.get("stop_reason") |
| |
| result: dict[str, object] = { |
| "prompt_tokens": usage.get("prompt_tokens"), "completion_tokens": count, |
| "cached_tokens": obj(usage.get("prompt_tokens_details") or {}).get("cached_tokens"), |
| "finish_reason": reason if reason in ("length", "stop") else "unexpected", |
| "stop_reason": stop if type(stop) is int or stop is None else "string-present", |
| "reasoning_tokens": obj(usage.get("completion_tokens_details") or {}).get("reasoning_tokens"), |
| "parser": "raw-completion; no chat/reasoning/tool parser", |
| } |
| metrics = obj(data.get("metrics") or {}) |
| spec = obj(metrics.get("speculative_decoding") or {}) |
| |
| result["speculative_metrics"] = { |
| key: value for key, value in spec.items() |
| if re.fullmatch(r"[a-z_]+", key) and type(value) in (int, float) |
| and math.isfinite(cast(float, value)) |
| } or None |
| if type(count) is not int or count != BUDGET or reason != "length" or stop is not None: |
| raise ValueError("warm-up output contract failed: " + json.dumps(result, sort_keys=True)) |
| return result |
|
|
|
|
| def idle(raw: str) -> None: |
| for name in ("vllm:num_requests_running", "vllm:num_requests_waiting"): |
| values = [] |
| for line in raw.splitlines(): |
| fields = line.split() |
| if len(fields) >= 2 and fields[0].split("{")[0] == name: |
| values.append(float(fields[1])) |
| if not values or any(not math.isfinite(v) or v != 0 for v in values): |
| raise ValueError(f"warm-up refused: {name} missing, invalid or nonzero") |
|
|
|
|
| def request_timeout(end: float) -> float: |
| """Per-request socket timeout: never past the warm-up deadline.""" |
| left = end - time.monotonic() |
| if left <= 0: |
| raise ValueError(f"warm-up deadline {DEADLINE_S:.0f} s reached") |
| return min(float(TIMEOUT), left) |
|
|
|
|
| def run(model: str, read: Reader, end: float | None = None) -> None: |
| end = time.monotonic() + DEADLINE_S if end is None else end |
| salt = uuid.uuid4().hex |
| for size in SHAPES: |
| if time.monotonic() >= end: |
| raise ValueError(f"warm-up deadline {DEADLINE_S:.0f} s reached before the concurrency-{size} wave") |
| idle(read("/metrics", None)) |
| gate = Barrier(size) |
|
|
| def one(index: int) -> dict[str, object]: |
| request = payload(model, salt, size, index) |
| gate.wait(timeout=5) |
| result = summary(read("/v1/completions", request)) |
| if size == 1: |
| tokens = result["prompt_tokens"] |
| if type(tokens) is not int or tokens <= 4096: |
| raise ValueError("warm-up did not exercise full+partial 4096-token prefill") |
| cached = result["cached_tokens"] |
| if cached is not None and cached != 0: |
| raise ValueError("warm-up long prefill unexpectedly reused cached tokens") |
| return result |
|
|
| start = time.monotonic() |
| with ThreadPoolExecutor(max_workers=size) as pool: |
| results = list(pool.map(one, range(size))) |
| print(json.dumps({"warmup": "completed", "requested_concurrency": size, |
| "actual_graph_shapes": "requires engine/worker trace; not inferred", |
| "elapsed_s": round(time.monotonic() - start, 3), |
| "requests": results}, sort_keys=True), flush=True) |
| if time.monotonic() > end: |
| raise ValueError(f"warm-up finished after its {DEADLINE_S:.0f} s deadline") |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--port", type=int, default=8000) |
| parser.add_argument("--model", required=True) |
| args = parser.parse_args() |
| if not 1 <= args.port <= 65535: |
| parser.error("invalid port") |
|
|
| end = time.monotonic() + DEADLINE_S |
|
|
| def read(path: str, body: dict[str, object] | None) -> str: |
| request = Request(f"http://localhost:{args.port}{path}", |
| data=None if body is None else json.dumps(body).encode(), |
| headers={"Content-Type": "application/json"}) |
| with urlopen(request, timeout=request_timeout(end)) as response: |
| return response.read(1 << 20).decode() |
|
|
| try: |
| run(args.model, read, end) |
| except Exception as exc: |
| |
| parser.exit(1, f"warm-up FAILED ({type(exc).__name__}): {exc}\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
