"""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)