"""Serving performance measurement. A self-contained async load generator rather than a wrapper around `vllm bench serve`. Reasons: * It works against ANY OpenAI-compatible endpoint -- vLLM, SGLang, llama.cpp, ollama, or a submitted model served however we choose to serve it. Submissions are the point of the site; we cannot assume vLLM. * It needs nothing installed on the target box. The runner talks HTTP. * We control the metric definitions, so numbers stay comparable across engines rather than inheriting each engine's benchmarking conventions. TWO MEASUREMENT TRAPS THIS CODE AVOIDS -- both silently produce numbers that look great and mean nothing: 1. PREFIX CACHING. vLLM caches shared prompt prefixes. If every concurrent request sends the same prompt, prefill is nearly free after the first and throughput is wildly overstated. Every request here gets a unique nonce prefix so each one actually does its own prefill. 2. VARIABLE OUTPUT LENGTH. If the model decides when to stop, concurrency levels finish at different token counts and tok/s is not comparable between runs. We send `ignore_eos` so every request emits exactly `output_tokens` tokens. """ from __future__ import annotations import asyncio import json import statistics import time import uuid from dataclasses import dataclass import httpx from .results import PerfPoint, PerfResult from .specdec import read_counters # Filler vocabulary for synthetic prompts. Ordinary words rather than repeated # junk so tokenization stays close to what real traffic looks like. _FILLER = ( "the quick brown fox jumps over a lazy dog while distant thunder rolls " "across open water and the harbour lights flicker against low cloud cover " "as fishing boats return with the morning tide carrying nets and crates " ).split() def _synthetic_prompt(approx_tokens: int, nonce: str) -> str: """Build a prompt of roughly `approx_tokens` tokens, unique per request. The nonce leads so it lands in the first block and defeats prefix-cache reuse across concurrent requests. """ # ~0.75 tokens per word for ordinary English text. n_words = max(1, int(approx_tokens / 0.75)) body = " ".join(_FILLER[i % len(_FILLER)] for i in range(n_words)) return f"[{nonce}] {body}" @dataclass class RequestResult: ok: bool ttft_ms: float | None = None total_s: float | None = None output_tokens: int = 0 input_tokens: int = 0 error: str | None = None @property def tpot_ms(self) -> float | None: """Inter-token latency: decode time divided by tokens after the first.""" if self.ttft_ms is None or self.total_s is None or self.output_tokens < 2: return None decode_ms = self.total_s * 1000 - self.ttft_ms return decode_ms / (self.output_tokens - 1) async def _one_request( client: httpx.AsyncClient, base_url: str, model: str, input_tokens: int, output_tokens: int, prompt: str | None = None, pin_output: bool = True, ) -> RequestResult: nonce = uuid.uuid4().hex[:12] # A supplied prompt still gets the nonce, because prefix-cache defeat is orthogonal to # whether the text is representative — reusing one real prompt across a concurrency level # would let the engine serve most of it from cache and report an inflated number. content = f"[{nonce}] {prompt}" if prompt else _synthetic_prompt(input_tokens, nonce) payload = { "model": model, "messages": [{"role": "user", "content": content}], "max_tokens": output_tokens, "temperature": 0.0, "stream": True, "stream_options": {"include_usage": True}, } if pin_output: # vLLM/SGLang extension: force exactly max_tokens of output so runs are # comparable. Harmless (ignored) on engines that do not support it. # # Turned OFF for speculative-decoding measurement. Forcing generation past a natural # stop produces degenerate continuation, and a draft model's acceptance rate on # degenerate text is not its acceptance rate on real work — so pinning the length here # buys comparability at the cost of measuring the wrong thing. payload["ignore_eos"] = True start = time.perf_counter() ttft: float | None = None counted = 0 usage_in = 0 usage_out = 0 try: async with client.stream( "POST", f"{base_url}/chat/completions", json=payload ) as resp: if resp.status_code != 200: body = (await resp.aread()).decode()[:200] return RequestResult(ok=False, error=f"HTTP {resp.status_code}: {body}") async for line in resp.aiter_lines(): if not line.startswith("data: "): continue data = line[6:] if data == "[DONE]": break try: chunk = json.loads(data) except json.JSONDecodeError: continue if usage := chunk.get("usage"): usage_in = usage.get("prompt_tokens", 0) usage_out = usage.get("completion_tokens", 0) for choice in chunk.get("choices") or []: delta = choice.get("delta") or {} # Engines disagree on where reasoning text lands: vLLM with # --reasoning-parser emits `reasoning`, others use # `reasoning_content`. Count all of them -- reasoning tokens # cost the same decode time and the user waits for them # either way. Checking only `content` would report TTFT as # the time to the *answer*, which on a reasoning model is # thousands of tokens late. if any( isinstance(delta.get(k), str) and delta[k] for k in ("content", "reasoning", "reasoning_content") ): # The opening chunk carries role + empty content; it is # the stream opening, not a token, and is excluded by # the emptiness check above. if ttft is None: ttft = (time.perf_counter() - start) * 1000 counted += 1 except Exception as exc: # noqa: BLE001 - report, do not abort the sweep return RequestResult(ok=False, error=f"{type(exc).__name__}: {exc}") total = time.perf_counter() - start return RequestResult( ok=True, ttft_ms=ttft, total_s=total, # Prefer server-reported usage; fall back to counted chunks. output_tokens=usage_out or counted, input_tokens=usage_in, ) async def _sweep_point( base_url: str, model: str, concurrency: int, input_tokens: int, output_tokens: int, timeout_s: float, prompts: list[str] | None = None, pin_output: bool = True, ) -> PerfPoint: limits = httpx.Limits(max_connections=concurrency + 8) async with httpx.AsyncClient(timeout=timeout_s, limits=limits) as client: # Warm up so the first measured request does not absorb graph capture, # weight paging, or connection setup. await _one_request(client, base_url, model, 32, 8) # Bracket the measured window. These are lifetime counters, so the rate for THIS point # is the delta — reading the total would fold in warm-up and every earlier point. before = await read_counters(client, base_url) start = time.perf_counter() results = await asyncio.gather( *( _one_request( client, base_url, model, input_tokens, output_tokens, # Cycle rather than repeat: every concurrent request in a point gets a # different prompt, so one unusually easy or hard example cannot set the # whole level's number. prompt=prompts[i % len(prompts)] if prompts else None, pin_output=pin_output, ) for i in range(concurrency) ) ) wall = time.perf_counter() - start after = await read_counters(client, base_url) ok = [r for r in results if r.ok] failed = [r for r in results if not r.ok] total_out = sum(r.output_tokens for r in ok) ttfts = [r.ttft_ms for r in ok if r.ttft_ms is not None] tpots = [r.tpot_ms for r in ok if r.tpot_ms is not None] per_stream = [ r.output_tokens / r.total_s for r in ok if r.total_s and r.output_tokens ] def pct(values: list[float], p: float) -> float | None: if not values: return None s = sorted(values) idx = min(len(s) - 1, round(p * (len(s) - 1))) return round(s[idx], 2) # Prefill rate: input tokens processed per second, inferred from TTFT. # At concurrency > 1 this is per-stream and includes queueing, so it reads # low; treat the concurrency-1 value as the true prefill capability. prefill = None if ttfts and ok: reported = [r.input_tokens for r in ok if r.input_tokens] mean_in = statistics.mean(reported) if reported else input_tokens mean_ttft_s = statistics.mean(ttfts) / 1000 if mean_ttft_s > 0: prefill = round(mean_in / mean_ttft_s, 1) # Delta across the measured window only. `None` when the engine exposes no counters, which # is an ordinary outcome (no speculative decoding, or metrics disabled) and not an error. spec = (after - before) if (before and after) else None acceptance = spec.acceptance_rate if spec else None drafted = int(spec.draft) if spec and spec.draft > 0 else None # With a natural stop, requests no longer emit identical token counts, so the spread is # itself a caveat on comparability and has to travel with the number. out_counts = [r.output_tokens for r in ok if r.output_tokens] stdev = ( round(statistics.stdev(out_counts), 2) if not pin_output and len(out_counts) > 1 else None ) return PerfPoint( concurrency=concurrency, input_tokens=input_tokens, output_tokens=output_tokens, n_requests=concurrency, completed=len(ok), failed=len(failed), duration_s=round(wall, 3), output_tps_total=round(total_out / wall, 2) if wall else 0.0, output_tps_per_stream=round(statistics.mean(per_stream), 2) if per_stream else 0.0, ttft_p50_ms=pct(ttfts, 0.50), ttft_p95_ms=pct(ttfts, 0.95), tpot_p50_ms=pct(tpots, 0.50), prefill_tps=prefill, spec_acceptance_rate=acceptance, spec_draft_tokens=drafted, natural_stop=not pin_output, output_token_stdev=stdev, error=failed[0].error if failed else None, ) async def run_sweep( base_url: str, model: str, concurrencies: tuple[int, ...] = (1, 8, 32), input_tokens: int = 1024, output_tokens: int = 256, timeout_s: float = 600.0, engine: str = "unknown", progress: bool = True, prompts: list[str] | None = None, pin_output: bool = True, ) -> PerfResult: """Run a concurrency sweep and return a PerfResult. `prompts` and `pin_output=False` together give the mode a speculative-decoding target needs: representative text, stopping naturally. Both defaults stay as they were, so an ordinary throughput sweep is unchanged and remains directly comparable to every result already committed. """ points: list[PerfPoint] = [] for c in concurrencies: if progress: print(f" concurrency {c:>3} ...", end="", flush=True) point = await _sweep_point( base_url, model, c, input_tokens, output_tokens, timeout_s, prompts=prompts, pin_output=pin_output, ) points.append(point) if progress: status = f"{point.output_tps_total:>8.1f} tok/s total" if point.spec_acceptance_rate is not None: status += f" accept {point.spec_acceptance_rate:.0%}" status += f" | {point.output_tps_per_stream:>6.1f} /stream" if point.ttft_p50_ms is not None: status += f" | TTFT p50 {point.ttft_p50_ms:>7.0f}ms" if point.failed: status += f" | {point.failed} FAILED" print(status, flush=True) return PerfResult(engine=engine, points=points)