benchmark.sh 3.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283
  1. #!/usr/bin/env bash
  2. # Benchmark the vLLM server: time-to-first-token and decode tok/s from a streaming
  3. # /v1/chat/completions response (usage included in the final chunk).
  4. #
  5. # Usage: ./benchmark.sh [num_predict]
  6. # BASE_URL=http://localhost:18080 # direct; use :11433 to go through the proxy
  7. set -euo pipefail
  8. BASE_URL="${BASE_URL:-http://localhost:18080}"
  9. MODEL="${MODEL:-foundation}"
  10. NUM_PREDICT="${1:-300}"
  11. 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."
  12. if ! curl -sf -m 10 "$BASE_URL/v1/models" >/dev/null; then
  13. echo "ERROR: cannot reach $BASE_URL/v1/models — is the stack up? (docker compose up -d)" >&2
  14. exit 1
  15. fi
  16. python3 - "$BASE_URL" "$MODEL" "$PROMPT" "$NUM_PREDICT" <<'PY'
  17. import json, sys, time, urllib.request
  18. base, model, prompt, n = sys.argv[1], sys.argv[2], sys.argv[3], int(sys.argv[4])
  19. def chat(stream, max_tokens):
  20. body = json.dumps({
  21. "model": model,
  22. "messages": [{"role": "user", "content": prompt}],
  23. "max_tokens": max_tokens,
  24. "stream": stream,
  25. "stream_options": {"include_usage": True} if stream else None,
  26. }).encode()
  27. req = urllib.request.Request(base + "/v1/chat/completions", data=body,
  28. headers={"Content-Type": "application/json",
  29. "Authorization": "Bearer x"})
  30. return urllib.request.urlopen(req, timeout=1800)
  31. # warmup (also loads/compiles if the server is fresh)
  32. t0 = time.monotonic()
  33. with chat(False, 16) as r:
  34. json.load(r)
  35. print(f"warmup done in {time.monotonic() - t0:.1f}s")
  36. t0 = time.monotonic()
  37. t_first = None
  38. usage = {}
  39. answer = ""
  40. with chat(True, n) as r:
  41. for raw in r:
  42. line = raw.decode().strip()
  43. if not line.startswith("data: "):
  44. continue
  45. payload = line[6:]
  46. if payload == "[DONE]":
  47. break
  48. chunk = json.loads(payload)
  49. if chunk.get("usage"):
  50. usage = chunk["usage"]
  51. for ch in chunk.get("choices", []):
  52. delta = ch.get("delta") or {}
  53. # vLLM 0.27 puts the reasoning trace in `reasoning` (older builds:
  54. # `reasoning_content`); either is the first generated token.
  55. piece = (delta.get("content") or delta.get("reasoning")
  56. or delta.get("reasoning_content") or "")
  57. if piece:
  58. if t_first is None:
  59. t_first = time.monotonic()
  60. answer += piece
  61. t_end = time.monotonic()
  62. pt = usage.get("prompt_tokens", 0)
  63. ct = usage.get("completion_tokens", 0)
  64. rt = (usage.get("completion_tokens_details") or {}).get("reasoning_tokens", 0) or 0
  65. ttft = (t_first - t0) if t_first else float("nan")
  66. decode_window = (t_end - t_first) if t_first else 0.0
  67. print(f"prompt tokens : {pt}")
  68. print(f"completion tokens : {ct} (reasoning: {rt})")
  69. print(f"TTFT : {ttft:.2f} s")
  70. print(f"wall : {t_end - t0:.2f} s")
  71. if decode_window > 0 and ct > 1:
  72. print(f"decode : {(ct - 1) / decode_window:.2f} tok/s")
  73. print(f"sample : {answer[:140].replace(chr(10), ' ')!r}")
  74. PY