Document long-context limits and host OOM at near-262K input

2026-09-18 00:03:19 +08:00
parent 48a1c9d7c4
commit f1a8964072
16 changed files with 1054 additions and 0 deletions
@@ -0,0 +1,198 @@
+"""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)