199 lines
9.7 KiB
Python
199 lines
9.7 KiB
Python
"""Generate reproducible long-context retrieval fixtures and test a running vLLM.
|
|
|
|
Run inside the model container. Never prints the mounted API key.
|
|
"""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import random
|
|
import signal
|
|
import time
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
SNAPSHOT = '/root/.cache/huggingface/hub/models--nvidia--Qwen3.8-Flash-Next-NVFP4/snapshots/fc694b54fb0174e0913e6adf86691ef85a4ead47'
|
|
MODEL = 'qwen3.8-flash-next'
|
|
MAX_OUTPUT = 2048
|
|
MAX_CONTEXT = 262144
|
|
BASE = 'http://127.0.0.1:8000'
|
|
|
|
|
|
def emit(record):
|
|
print(json.dumps(record, ensure_ascii=False), flush=True)
|
|
|
|
|
|
def digest(value):
|
|
return hashlib.sha256(value.encode()).hexdigest()
|
|
|
|
|
|
def headers():
|
|
key = Path('/run/secrets/qwen_api_key').read_text().strip()
|
|
return {'Authorization': 'Bearer ' + key, 'Content-Type': 'application/json'}
|
|
|
|
|
|
def request(path, payload=None, timeout=120):
|
|
req = urllib.request.Request(BASE + path, headers=headers(),
|
|
data=json.dumps(payload, ensure_ascii=False).encode() if payload is not None else None)
|
|
with urllib.request.urlopen(req, timeout=timeout) as response:
|
|
return json.load(response)
|
|
|
|
|
|
def metrics():
|
|
req = urllib.request.Request(BASE + '/metrics', headers=headers())
|
|
with urllib.request.urlopen(req, timeout=15) as response:
|
|
lines = response.read().decode().splitlines()
|
|
names = ['vllm:prefix_cache_hits_total', 'vllm:prefix_cache_queries_total',
|
|
'vllm:num_preemptions_total']
|
|
return {name: sum(float(l.split()[-1]) for l in lines if l.split('{')[0].split(' ')[0] == name)
|
|
for name in names}
|
|
|
|
|
|
def messages(document, project):
|
|
prompt = ('以下是一份虚构项目预算档案，请按项目编号检索其中的数据。\n' + document +
|
|
'\n档案结束。\n问题：项目编号 ' + project + ' 的核定预算是多少元？只输出数字。')
|
|
return [{'role': 'user', 'content': prompt}]
|
|
|
|
|
|
def prepare(folder):
|
|
from transformers import AutoTokenizer
|
|
tok = AutoTokenizer.from_pretrained(SNAPSHOT, local_files_only=True)
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
for target in [65536, 131072, 260000]:
|
|
for sample in [1, 2]:
|
|
name = f'ctx-{target}-{sample}'
|
|
rng = random.Random(target * 10 + sample)
|
|
amounts = [str(rng.randrange(10000, 99999)) for _ in range(3)]
|
|
projects = [f'TARGET-{position}-{sample}' for position in ['FRONT', 'MIDDLE', 'END']]
|
|
filler = [f'项目编号 BG-{i:06d}；核定预算 {rng.randrange(10000,99999)} 元；类别 C{i%19:02d}。' for i in range(20000)]
|
|
|
|
def build(n):
|
|
rows = filler[:n].copy()
|
|
positions = [int(n * f) for f in [0.05, 0.5, 0.95]]
|
|
for pos, project, amount in zip(positions, projects, amounts):
|
|
rows[pos] = f'项目编号 {project}；核定预算 {amount} 元；类别 C20。'
|
|
return '文档编号：CTX262K-B-' + name + '\n' + '\n'.join(rows), positions, rows
|
|
|
|
low, high = 20, len(filler)
|
|
while low < high:
|
|
n = (low + high + 1) // 2
|
|
doc, _, _ = build(n)
|
|
count = len(tok.apply_chat_template(messages(doc, projects[1]), tokenize=True, add_generation_prompt=True, return_dict=False))
|
|
if count <= target:
|
|
low = n
|
|
else:
|
|
high = n - 1
|
|
doc, positions, rows = build(low)
|
|
cases = []
|
|
for label, index in [('front', 0), ('middle', 1), ('end', 2), ('front_repeat', 0)]:
|
|
msgs = messages(doc, projects[index])
|
|
count = request('/tokenize', {'model': MODEL, 'messages': msgs, 'add_generation_prompt': True, 'chat_template_kwargs': {'reasoning_effort': 'low'}})['count']
|
|
assert target - 128 <= count <= target + 32, (name, count, target)
|
|
assert count + MAX_OUTPUT <= MAX_CONTEXT, (name, count)
|
|
cases.append({'case': label, 'messages': msgs, 'expected': amounts[index],
|
|
'prompt_tokens_preflight': count, 'prompt_sha256': digest(msgs[0]['content']),
|
|
'max_tokens': MAX_OUTPUT})
|
|
fractions = [round(len(tok.encode('文档编号：CTX262K-B-'+name+'\n'+'\n'.join(rows[:pos]), add_special_tokens=False)) / cases[0]['prompt_tokens_preflight'], 4)
|
|
for pos in positions]
|
|
fixture = {'document': name, 'target_input_tokens': target, 'sample': sample,
|
|
'document_sha256': digest(doc), 'needle_token_fractions': fractions, 'cases': cases}
|
|
(folder / (name + '.json')).write_text(json.dumps(fixture, ensure_ascii=False))
|
|
emit({'event': 'FIXTURE_READY', **{k:v for k,v in fixture.items() if k != 'cases'},
|
|
'prompt_tokens': [c['prompt_tokens_preflight'] for c in cases]})
|
|
|
|
|
|
def deadline(signum, frame):
|
|
raise TimeoutError('Request exceeded the 1200 second wall-clock limit')
|
|
|
|
|
|
def run_case(label, fixture, case, attempt):
|
|
before = metrics()
|
|
payload = {'model': MODEL, 'messages': case['messages'], 'temperature': 0, 'seed': 42,
|
|
'reasoning_effort': 'low', 'max_tokens': MAX_OUTPUT,
|
|
'stream': True, 'stream_options': {'include_usage': True}}
|
|
req = urllib.request.Request(BASE + '/v1/chat/completions', headers=headers(),
|
|
data=json.dumps(payload, ensure_ascii=False).encode())
|
|
start = time.monotonic()
|
|
first = first_content = last = None
|
|
content = reasoning = ''
|
|
usage = {}
|
|
finish = None
|
|
signal.alarm(1200)
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=1200) 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('error'):
|
|
raise RuntimeError('Server returned a streaming error')
|
|
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.monotonic()
|
|
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
|
|
finally:
|
|
signal.alarm(0)
|
|
if finish is None or not usage:
|
|
raise RuntimeError('Incomplete response or missing usage')
|
|
after = metrics()
|
|
elapsed = time.monotonic() - start
|
|
record = {'event': 'RESULT', 'label': label, 'document': fixture['document'],
|
|
'target_input_tokens': fixture['target_input_tokens'], 'sample': fixture['sample'],
|
|
'case': case['case'], 'attempt': attempt, 'expected': case['expected'],
|
|
'content': content, 'correct': content.strip() == case['expected'] and finish == 'stop',
|
|
'finish_reason': finish, 'usage': usage, 'max_tokens': MAX_OUTPUT,
|
|
'prompt_tokens_preflight': case['prompt_tokens_preflight'], 'prompt_sha256': case['prompt_sha256'],
|
|
'output_sha256': digest(reasoning + '\0' + content),
|
|
'ttft_s': round(first-start, 4) if first else None,
|
|
'first_content_s': round(first_content-start, 4) if first_content else None,
|
|
'elapsed_s': round(elapsed, 4),
|
|
'decode_tps_approx': round((usage['completion_tokens']-1)/(last-first), 3) if last and first and last>first else None,
|
|
'metrics_delta': {k: after[k]-before[k] for k in after}}
|
|
emit(record)
|
|
assert usage['prompt_tokens'] + MAX_OUTPUT <= MAX_CONTEXT
|
|
assert usage['prompt_tokens'] == case['prompt_tokens_preflight'], 'Server tokenization differs from preflight'
|
|
return record
|
|
|
|
|
|
def run(folder, label):
|
|
signal.signal(signal.SIGALRM, deadline)
|
|
failures = []
|
|
for target in [65536, 131072, 260000]:
|
|
emit({'event': 'STAGE_START', 'label': label, 'target_input_tokens': target})
|
|
for sample in [1, 2]:
|
|
fixture = json.loads((folder / f'ctx-{target}-{sample}.json').read_text())
|
|
emit({'event': 'DOCUMENT_START', 'label': label, 'document': fixture['document'],
|
|
'document_sha256': fixture['document_sha256'], 'needle_token_fractions': fixture['needle_token_fractions']})
|
|
for case in fixture['cases']:
|
|
emit({'event': 'REQUEST_START', 'label': label, 'document': fixture['document'], 'case': case['case']})
|
|
result = run_case(label, fixture, case, 1)
|
|
if not result['correct']:
|
|
failures.append([fixture['document'], case['case']])
|
|
run_case(label, fixture, case, 2)
|
|
emit({'event': 'STAGE_COMPLETED', 'label': label, 'target_input_tokens': target})
|
|
emit({'event': 'CONTEXT_SUITE_COMPLETED', 'label': label, 'first_attempt_failures': failures})
|
|
# Quality failures are reported above, not hidden; infrastructure failures raise.
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('mode', choices=['prepare', 'run'])
|
|
parser.add_argument('--data', type=Path, default=Path('/tmp/context262k-data'))
|
|
parser.add_argument('--label', default='current')
|
|
args = parser.parse_args()
|
|
if args.mode == 'prepare':
|
|
prepare(args.data)
|
|
else:
|
|
run(args.data, args.label)
|