| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283 |
- #!/usr/bin/env bash
- # Benchmark the vLLM server: time-to-first-token and decode tok/s from a streaming
- # /v1/chat/completions response (usage included in the final chunk).
- #
- # Usage: ./benchmark.sh [num_predict]
- # BASE_URL=http://localhost:18080 # direct; use :11433 to go through the proxy
- set -euo pipefail
- BASE_URL="${BASE_URL:-http://localhost:18080}"
- MODEL="${MODEL:-foundation}"
- NUM_PREDICT="${1:-300}"
- PROMPT="Write a comprehensive explanation of how quantum computing works, including the principles of superposition, entanglement, and quantum gates. Provide examples of real-world applications and discuss the current state of quantum hardware."
- if ! curl -sf -m 10 "$BASE_URL/v1/models" >/dev/null; then
- echo "ERROR: cannot reach $BASE_URL/v1/models — is the stack up? (docker compose up -d)" >&2
- exit 1
- fi
- python3 - "$BASE_URL" "$MODEL" "$PROMPT" "$NUM_PREDICT" <<'PY'
- import json, sys, time, urllib.request
- base, model, prompt, n = sys.argv[1], sys.argv[2], sys.argv[3], int(sys.argv[4])
- def chat(stream, max_tokens):
- body = json.dumps({
- "model": model,
- "messages": [{"role": "user", "content": prompt}],
- "max_tokens": max_tokens,
- "stream": stream,
- "stream_options": {"include_usage": True} if stream else None,
- }).encode()
- req = urllib.request.Request(base + "/v1/chat/completions", data=body,
- headers={"Content-Type": "application/json",
- "Authorization": "Bearer x"})
- return urllib.request.urlopen(req, timeout=1800)
- # warmup (also loads/compiles if the server is fresh)
- t0 = time.monotonic()
- with chat(False, 16) as r:
- json.load(r)
- print(f"warmup done in {time.monotonic() - t0:.1f}s")
- t0 = time.monotonic()
- t_first = None
- usage = {}
- answer = ""
- with chat(True, n) as r:
- for raw in r:
- line = raw.decode().strip()
- if not line.startswith("data: "):
- continue
- payload = line[6:]
- if payload == "[DONE]":
- break
- chunk = json.loads(payload)
- if chunk.get("usage"):
- usage = chunk["usage"]
- for ch in chunk.get("choices", []):
- delta = ch.get("delta") or {}
- # vLLM 0.27 puts the reasoning trace in `reasoning` (older builds:
- # `reasoning_content`); either is the first generated token.
- piece = (delta.get("content") or delta.get("reasoning")
- or delta.get("reasoning_content") or "")
- if piece:
- if t_first is None:
- t_first = time.monotonic()
- answer += piece
- t_end = time.monotonic()
- pt = usage.get("prompt_tokens", 0)
- ct = usage.get("completion_tokens", 0)
- rt = (usage.get("completion_tokens_details") or {}).get("reasoning_tokens", 0) or 0
- ttft = (t_first - t0) if t_first else float("nan")
- decode_window = (t_end - t_first) if t_first else 0.0
- print(f"prompt tokens : {pt}")
- print(f"completion tokens : {ct} (reasoning: {rt})")
- print(f"TTFT : {ttft:.2f} s")
- print(f"wall : {t_end - t0:.2f} s")
- if decode_window > 0 and ct > 1:
- print(f"decode : {(ct - 1) / decode_window:.2f} tok/s")
- print(f"sample : {answer[:140].replace(chr(10), ' ')!r}")
- PY
|