| 1234567891011121314151617181920212223242526272829303132333435363738 |
- # Aggregate and per-stream decode tok/s with N concurrent streams (streamed,
- # timed from each stream's first token). Prompts rotate across workloads.
- import json, os, statistics, sys, threading, time, urllib.request
- BASE = os.environ.get("GATE_URL", "http://127.0.0.1:8002")
- MODEL = os.environ.get("GATE_MODEL", "glm53")
- PROMPTS = ["Write a red-black tree in Python with insert, delete and rebalancing. Code only.",
- "Explain how a hash map works, in flowing prose. No code, no lists.",
- "Write a complete Python implementation of a red-black tree with insert, delete and in-order traversal.",
- "Explain in detail how a hash map handles collisions, with examples."]
- def one(prompt, out, i):
- body = json.dumps({"model": MODEL, "max_tokens": 512, "temperature": 0, "stream": True,
- "chat_template_kwargs": {"thinking": False, "enable_thinking": False},
- "stream_options": {"include_usage": True},
- "messages": [{"role": "user", "content": prompt}]}).encode()
- r = urllib.request.urlopen(urllib.request.Request(BASE + "/v1/chat/completions", data=body,
- headers={"Content-Type": "application/json"}), timeout=900)
- first = n = None
- for line in r:
- if not line.startswith(b"data: {"):
- continue
- d = json.loads(line[6:])
- if first is None and d.get("choices"):
- first = time.monotonic()
- if d.get("usage"):
- n = d["usage"]["completion_tokens"]
- out[i] = (n, first, time.monotonic())
- for N in [int(a) for a in sys.argv[1:]] or [1, 2, 4, 8]:
- out = [None] * N
- ts = [threading.Thread(target=one, args=(PROMPTS[i % 4], out, i)) for i in range(N)]
- for t in ts: t.start()
- for t in ts: t.join()
- t0 = min(o[1] for o in out); t1 = max(o[2] for o in out)
- per = [o[0] / (o[2] - o[1]) for o in out]
- print(f"streams {N:2d}: aggregate {sum(o[0] for o in out) / (t1 - t0):6.1f} tok/s, per stream median {statistics.median(per):6.1f} (min {min(per):.1f})", flush=True)
|
