| |
| """Task-class x concurrency decode benchmark for the AEON serve (OpenAI chat API, vLLM metrics). |
| gen tok/s = completion tokens / (wall - TTFT) per stream; aggregate = sum tokens / max wall (decode window). |
| tok/pass = delta(spec accepted)/delta(drafts) + 1 from /metrics.""" |
| import json, os, sys, time, threading, urllib.request, statistics, random |
| BASE=os.environ.get("BASE","http://localhost:8000"); MODEL=os.environ.get("MODEL","aeon"); OUT=int(os.environ.get("OUT","512")) |
| THINK=os.environ.get("THINK","0")=="1" |
| TOPICS=["a lighthouse keeper in 1890s Norway","a courier robot in a flooded Bangkok","a violinist who loses her hearing","a mountain village's last winter","two brothers repairing a boat","a night-shift nurse in Lagos","a chess prodigy in exile","a beekeeper during a drought"] |
| TASKS={ |
| "prose": lambda i: f"Write a vivid literary short story (at least 700 words) about {TOPICS[i%8]}. No headings, flowing prose only.", |
| "code": lambda i: f"Write a complete, well-documented Python module implementing an LRU cache with TTL expiry, thread safety, statistics, and a small CLI demo. Variant #{i}: name the class Cache{i}. Code only, no prose.", |
| "list": lambda i: f"List 120 distinct {['programming concepts','world cities','chemical compounds','historical events','musical instruments','mathematical theorems','plant species','cooking techniques'][i%8]}, numbered, each with a one-line description.", |
| "essay": lambda i: f"Write a structured argumentative essay (about 800 words, with an introduction, three sections with headers, and a conclusion) on: '{['Should cities ban private cars?','Is remote work good for innovation?','Should AI systems have a right to refuse tasks?','Is nuclear power essential for decarbonization?','Should voting be mandatory?','Do social networks make us lonelier?','Should space exploration be privatized?','Is a four-day work week viable?'][i%8]}'", |
| } |
| def metrics(): |
| t=urllib.request.urlopen(BASE+"/metrics",timeout=10).read().decode(); d={} |
| for l in t.splitlines(): |
| for k in ("spec_decode_num_drafts_total","spec_decode_num_accepted_tokens_total"): |
| if l.startswith("vllm:"+k+"{"): d[k]=float(l.rsplit(" ",1)[1]) |
| return d |
| def one(prompt, res, idx): |
| body={"model":MODEL,"messages":[{"role":"user","content":prompt}],"max_tokens":OUT,"stream":True, |
| "stream_options":{"include_usage":True},"chat_template_kwargs":{"enable_thinking":THINK}} |
| body.update(json.loads(os.environ.get("SAMPLING","{}"))) |
| req=urllib.request.Request(BASE+"/v1/chat/completions",data=json.dumps(body).encode(),headers={"Content-Type":"application/json"}) |
| t0=time.time(); ttft=None; ntok=0 |
| with urllib.request.urlopen(req,timeout=1800) as r: |
| for line in r: |
| line=line.decode().strip() |
| if not line.startswith("data:") or line=="data: [DONE]": continue |
| ev=json.loads(line[5:]) |
| if ev.get("usage"): ntok=ev["usage"]["completion_tokens"] |
| ch=ev.get("choices") |
| d=(ch[0].get("delta") or {}) if ch else {} |
| if ttft is None and (d.get("content") or d.get("reasoning_content") or d.get("reasoning")): ttft=time.time()-t0 |
| wall=time.time()-t0; res[idx]=(ntok,ttft or wall,wall) |
| def run(task,conc): |
| m0=metrics(); res=[None]*conc; th=[threading.Thread(target=one,args=(TASKS[task](i+conc*7),res,i)) for i in range(conc)] |
| T=time.time(); [t.start() for t in th]; [t.join() for t in th]; W=time.time()-T; m1=metrics() |
| toks=sum(r[0] for r in res); per=[r[0]/(r[2]-r[1]) for r in res if r[2]>r[1]] or [float("nan")] |
| g=lambda m,k: m.get(k,0.0) |
| dr=g(m1,"spec_decode_num_drafts_total")-g(m0,"spec_decode_num_drafts_total"); ac=g(m1,"spec_decode_num_accepted_tokens_total")-g(m0,"spec_decode_num_accepted_tokens_total") |
| tpp=(ac/dr+1) if dr else float("nan"); ttft=statistics.median(r[1] for r in res) |
| return dict(task=task,conc=conc,agg=toks/W,per=statistics.mean(per),ttft=ttft,tpp=tpp,tok=toks/conc) |
| if __name__=="__main__": |
| tasks=sys.argv[1].split(",") if len(sys.argv)>1 else list(TASKS); concs=[int(c) for c in (sys.argv[2] if len(sys.argv)>2 else "1,2,4,8").split(",")] |
| run("prose",1) |
| print(f"thinking={'on' if THINK else 'off'} out={OUT} sampling=server default (temp 0.6/top_p .95/top_k 20/rep 1.05)") |
| print(f"{'task':6} {'C':>2} {'agg t/s':>8} {'per-stream':>10} {'TTFT':>6} {'tok/pass':>8} {'tok/req':>7}") |
| for task in tasks: |
| for c in concs: |
| r=run(task,c); print(f"{r['task']:6} {r['conc']:>2} {r['agg']:8.1f} {r['per']:10.1f} {r['ttft']:6.2f} {r['tpp']:8.2f} {r['tok']:7.0f}",flush=True) |
|
|
