| """Identical-request, warmed-cache throughput comparison at fixed concurrency.""" |
| import argparse |
| import concurrent.futures |
| import json |
| import time |
| import urllib.request |
| from pathlib import Path |
|
|
|
|
| def main(): |
| p=argparse.ArgumentParser(); p.add_argument('--input',required=True); p.add_argument('--out',required=True) |
| p.add_argument('--label',required=True); p.add_argument('--workers',type=int,default=32) |
| args=p.parse_args() |
| requests=[json.loads(l)['request'] for l in Path(args.input).read_text().splitlines()[:64]] |
| def call(payload): |
| payload=dict(payload); payload.update(temperature=.7,top_p=.8,top_k=20,presence_penalty=0,repetition_penalty=1.0) |
| request=urllib.request.Request('http://127.0.0.1:18000/v1/chat/completions', |
| data=json.dumps(payload).encode(),headers={'Content-Type':'application/json'}) |
| with urllib.request.urlopen(request,timeout=180) as response: result=json.load(response) |
| try: json.loads(result['choices'][0]['message']['content']); valid=True |
| except Exception: valid=False |
| return {'tokens':result['usage']['completion_tokens'],'json_valid':valid,'finish_reason':result['choices'][0]['finish_reason']} |
| rounds=[] |
| with concurrent.futures.ThreadPoolExecutor(args.workers) as pool: |
| for repeat in range(3): |
| start=time.time(); outputs=list(pool.map(call,requests)); seconds=time.time()-start |
| entry={'round':repeat,'warmup':repeat==0,'seconds':seconds,'tokens':sum(o['tokens'] for o in outputs), |
| 'json_valid':sum(o['json_valid'] for o in outputs),'requests':len(outputs), |
| 'truncated':sum(o['finish_reason']=='length' for o in outputs)} |
| entry['output_tps']=entry['tokens']/seconds; rounds.append(entry); print(json.dumps(entry),flush=True) |
| with urllib.request.urlopen('http://127.0.0.1:18000/metrics',timeout=10) as response: |
| metrics=response.read().decode() |
| result={'label':args.label,'workers':args.workers,'rounds':rounds, |
| 'mtp_metrics':[l for l in metrics.splitlines() if 'spec_decode' in l and not l.startswith('#')]} |
| Path(args.out).write_text(json.dumps(result,indent=2)) |
|
|
|
|
| if __name__=='__main__': main() |
|
|
