129 lines
6.5 KiB
Python
129 lines
6.5 KiB
Python
"""Synthetic streaming benchmark; run inside serving container, writes JSONL to stdout.
|
|
|
|
No external datasets or user prompts. Usage: python3 benchmark.py --label eager
|
|
TTFT includes reasoning tokens; first_content_s separately measures visible answer.
|
|
Decode tok/s is an approximation from total completion tokens and streaming time.
|
|
"""
|
|
import argparse
|
|
import concurrent.futures
|
|
import hashlib
|
|
import json
|
|
import random
|
|
import re
|
|
import threading
|
|
import time
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--label', required=True)
|
|
parser.add_argument('--long-max-tokens', type=int, default=1024)
|
|
parser.add_argument('--phase', choices=['all', 'short', 'long'], default='all')
|
|
args = parser.parse_args()
|
|
BASE = 'http://127.0.0.1:8000'
|
|
KEY = Path('/run/secrets/qwen_api_key').read_text().strip()
|
|
HEADERS = {'Authorization': 'Bearer ' + KEY, 'Content-Type': 'application/json'}
|
|
|
|
def emit(record):
|
|
print(json.dumps({'label': args.label, **record}, ensure_ascii=False), flush=True)
|
|
|
|
def metrics():
|
|
request = urllib.request.Request(BASE + '/metrics', headers=HEADERS)
|
|
with urllib.request.urlopen(request, timeout=10) as response:
|
|
text = response.read().decode()
|
|
values = {}
|
|
for line in text.splitlines():
|
|
if line.startswith('#') or 'prefix_cache_' not in line:
|
|
continue
|
|
name = line.split('{')[0].split()[0]
|
|
if name.endswith('_total'):
|
|
values[name] = values.get(name, 0) + float(line.split()[-1])
|
|
return values
|
|
|
|
def run(name, prompt, expected=None, barrier=None, max_tokens=512):
|
|
payload = {'model': 'qwen3.8-flash-next', 'messages': [{'role': 'user', 'content': prompt}],
|
|
'temperature': 0, 'seed': 42, 'max_tokens': max_tokens,
|
|
'reasoning_effort': 'low', 'stream': True,
|
|
'stream_options': {'include_usage': True}}
|
|
request = urllib.request.Request(BASE + '/v1/chat/completions',
|
|
data=json.dumps(payload).encode(), headers=HEADERS)
|
|
if barrier:
|
|
barrier.wait()
|
|
start = time.perf_counter()
|
|
first = first_content = last = None
|
|
content, reasoning = '', ''
|
|
usage, finish = {}, None
|
|
with urllib.request.urlopen(request, timeout=900) as response:
|
|
for raw in response:
|
|
line = raw.decode().strip()
|
|
if not line.startswith('data: ') or line == 'data: [DONE]':
|
|
continue
|
|
item = json.loads(line[6:])
|
|
if item.get('usage'):
|
|
usage = item['usage']
|
|
for choice in item.get('choices', []):
|
|
delta = choice.get('delta', {})
|
|
answer = delta.get('content') or ''
|
|
thought = delta.get('reasoning') or delta.get('reasoning_content') or ''
|
|
if answer or thought:
|
|
now = time.perf_counter()
|
|
first = first if first is not None else now
|
|
last = now
|
|
if answer:
|
|
first_content = first_content if first_content is not None else now
|
|
content += answer
|
|
reasoning += thought
|
|
finish = choice.get('finish_reason') or finish
|
|
elapsed = time.perf_counter() - start
|
|
count = usage.get('completion_tokens', 0)
|
|
record = {'test': name, 'max_tokens': max_tokens, 'elapsed_s': round(elapsed, 4),
|
|
'ttft_s': round(first-start, 4) if first else None,
|
|
'first_content_s': round(first_content-start, 4) if first_content else None,
|
|
'decode_tps_approx': round((count-1)/(last-first), 3) if last and first and last>first else None,
|
|
'usage': usage, 'finish_reason': finish, 'content': content,
|
|
'output_sha256': hashlib.sha256((reasoning+'\0'+content).encode()).hexdigest(),
|
|
'prompt_sha256': hashlib.sha256(prompt.encode()).hexdigest(),
|
|
'correct': bool(content.strip()) and (expected is None or expected in content)}
|
|
emit(record)
|
|
if not record['correct'] or finish != 'stop':
|
|
emit({'event':'ANSWER_CHECK_FAILED','test':name})
|
|
return record
|
|
|
|
emit({'event': 'start', 'phase': args.phase})
|
|
if args.phase in ('all', 'short'):
|
|
run('warmup', '计算17乘19，只输出结果。', '323')
|
|
prompt = '用中文写一段约150字的说明，解释数据库索引为什么能加快查询，以及它对写入有什么影响。'
|
|
for i in range(3):
|
|
run('short_'+str(i), prompt)
|
|
run('code', '写一个Python函数is_even(n)，判断整数是否为偶数，只输出代码。', '%')
|
|
for concurrency in [2, 4]:
|
|
barrier = threading.Barrier(concurrency)
|
|
start = time.perf_counter()
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool:
|
|
futures = [pool.submit(run, f'concurrent_{concurrency}_{i}',
|
|
f'场景编号{i}。'+prompt, None, barrier) for i in range(concurrency)]
|
|
results = [f.result() for f in futures]
|
|
seconds = time.perf_counter()-start
|
|
emit({'test': 'concurrency_summary', 'concurrency': concurrency,
|
|
'elapsed_s': round(seconds, 3),
|
|
'aggregate_completion_tps': round(sum(r['usage']['completion_tokens'] for r in results)/seconds, 3)})
|
|
if args.phase in ('all', 'long'):
|
|
from transformers import AutoTokenizer
|
|
snapshot = '/root/.cache/huggingface/hub/models--nvidia--Qwen3.8-Flash-Next-NVFP4/snapshots/fc694b54fb0174e0913e6adf86691ef85a4ead47'
|
|
tokenizer = AutoTokenizer.from_pretrained(snapshot, local_files_only=True)
|
|
rng = random.Random(42)
|
|
filler = '\n'.join(f'Record {i}: item={rng.randrange(100000,999999)}; status=archived; category={i%17}.' for i in range(6000))
|
|
ids = tokenizer.encode(filler, add_special_tokens=False)
|
|
for size in [8192, 32768]:
|
|
text = tokenizer.decode(ids[:size-150])
|
|
pos = len(text)//2
|
|
text = text[:pos]+'\n唯一有效的校验口令是：青松739251。\n'+text[pos:]
|
|
prompt = '从下面的档案里找出唯一有效的校验口令。只输出口令，不要解释。\n<档案>\n'+text+'\n</档案>\n请输出校验口令。'
|
|
before = metrics()
|
|
a = run(f'long_{size}_cold', prompt, '青松739251', max_tokens=args.long_max_tokens)
|
|
b = run(f'long_{size}_repeat', prompt, '青松739251', max_tokens=args.long_max_tokens)
|
|
after = metrics()
|
|
emit({'test':f'prefix_{size}', 'same_output':a['output_sha256']==b['output_sha256'],
|
|
'metrics_delta':{key:after[key]-before.get(key,0) for key in after}})
|
|
emit({'event':'BENCHMARK_COMPLETED'})
|