File size: 9,310 Bytes
2183ffa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | """Live exact-token cold/hot/reference and branch-isolation validation.
All prompts are synthetic; no benchmark-answer tuning. Latency includes API
service overhead. reference has identical padded tokens and unique cache salts.
"""
import argparse
import asyncio
import json
import math
from pathlib import Path
import secrets
import statistics
import sys
import time
import httpx
from transformers import AutoTokenizer
from shared_material import Backend, QuestionsRequest, Question
def fixture(name, rows, *, beta=False):
color = 'amber' if beta else 'blue'
pens, books = (9, 4) if beta else (7, 5)
a, b = ('07:10', '08:50') if beta else ('09:20', '08:50')
noise = [f'Archive entry {i:03d}: shipment R{i:03d} contains routine printed forms; its inspection was completed and filed.' for i in range(rows)]
first = f'ACTIVE RECORD. Parcel Kestrel is {color}. It contains {pens} pens and {books} notebooks.\n'
middle = f'ACTIVE TIMETABLE. Train A departs at {a}; train B departs at {b}. Both depart on the same day.\n'
last = ('ACTIVE POLICY. Visitors need an escort unless they are registered contractors. '
'Lin is a registered contractor and has no escort. '
'Refunds require a receipt and purchase within 30 days. Mina has a receipt and bought 12 days ago. '
'This material does not specify Kestrel\'s arrival date.\nEND OF ACTIVE RECORDS.')
state = first + '\n'.join(noise[:rows//2]) + '\n' + middle + '\n'.join(noise[rows//2:]) + '\n' + last
questions = {
'color': {'instructions': 'What is the color of parcel Kestrel?', 'options': ['blue', 'amber', 'green']},
'total': {'instructions': 'How many pens and notebooks does Kestrel contain in total?', 'options': ['12', '13', '7']},
'train': {'instructions': 'Which train departs earliest?', 'options': ['A', 'B', 'They depart at the same time']},
'entry': {'instructions': 'Under the escort policy, may Lin enter without an escort?', 'options': ['No', 'Yes', 'The policy does not say']},
'refund': {'instructions': 'Does Mina meet all the stated refund requirements?', 'options': ['No', 'Insufficient information', 'Yes']},
'arrival': {'instructions': 'On what date did parcel Kestrel arrive? Use only the material.', 'options': ['May 12', 'The arrival date is not specified', 'June 30']},
}
expected = {'color': 2 if beta else 1, 'total': 2 if beta else 1, 'train': 1 if beta else 2,
'entry': 2, 'refund': 3, 'arrival': 2}
return {'name': name, 'state': state, 'questions': questions, 'expected': expected}
def picks(result):
return {key: value['selected_index'] for key, value in result['answers'].items()}
async def main(args):
args.out.mkdir(parents=True, exist_ok=False)
tokenizer = AutoTokenizer.from_pretrained(args.model, local_files_only=True)
backend = Backend(tokenizer, url=args.backend)
api = httpx.AsyncClient(base_url=args.api, timeout=180)
cases = [fixture('short-alpha', 0), fixture('long-alpha', 80), fixture('long-beta', 80, beta=True)]
(args.out/'fixtures.json').write_text(json.dumps(cases, indent=2))
records = []
stream = (args.out/'runs.jsonl').open('w')
async def post(path, body):
start = time.perf_counter()
response = await api.post(path, json=body)
response.raise_for_status()
return response.json(), (time.perf_counter() - start)*1000
async def reference(case):
request = QuestionsRequest(questions=case['questions'], concurrency=1)
material = backend.material(case['state'])
compiled = backend.compile(material, request.questions)
start = time.perf_counter()
answers = {}
uncached = 0
for key, index, ids in compiled:
row = await backend.complete(ids, secrets.token_hex(16))
assert row['cached_tokens'] == 0, row
uncached += row['uncached_tokens']
answers.setdefault(key, []).append(row)
result = {'answers': {key: {'selected_index': 1+max(range(len(rows)), key=lambda i:rows[i]['log_odds']),
'options': rows} for key, rows in answers.items()},
'metrics': {'candidate_uncached_tokens': uncached}}
return result, (time.perf_counter()-start)*1000
def save(case, mode, repeat, result, ms, **extra):
correct = sum(picks(result)[key] == value for key, value in case['expected'].items())
row = {'case': case['name'], 'mode': mode, 'repeat': repeat, 'correct': correct,
'n_questions': len(case['questions']), 'wall_ms': ms, 'result': result, **extra}
stream.write(json.dumps(row)+'\n');stream.flush();records.append(row)
print(case['name'],mode,repeat,correct,'/',len(case['questions']),round(ms,1),flush=True)
# Shape warm-up: excluded, recorded. Do not select fixtures based on outputs.
warm = []
for case in cases:
result, ms = await post('/v1/rerank', {'state':case['state'],'questions':case['questions'],'concurrency':4})
warm.append({'case':case['name'],'wall_ms':ms,'result':result})
(args.out/'warmup.json').write_text(json.dumps(warm,indent=2))
for repeat in range(3):
for case in cases:
if repeat%2==0:
r,ms=await reference(case);save(case,'reference',repeat,r,ms)
prep,pms=await post('/v1/materials',{'state':case['state']})
assert prep['warmup']['cached_tokens']==0
material_id=prep['material_id']
body={'questions':case['questions'],'concurrency':1}
r,ms=await post(f'/v1/materials/{material_id}/questions',body)
save(case,'shared-cold-serial',repeat,r,pms+ms,preparation=prep)
assert r['metrics']['all_candidates_reused_entire_material']
# Serial and parallel hot runs use different question suffixes from
# priming but identical full inputs to reference.
for mode,c in [('shared-hot-serial',1),('shared-hot-parallel',4)]:
body['concurrency']=c;r,ms=await post(f'/v1/materials/{material_id}/questions',body)
save(case,mode,repeat,r,ms)
assert r['metrics']['all_candidates_reused_entire_material']
if repeat%2:
r,ms=await reference(case);save(case,'reference',repeat,r,ms)
deleted=await api.delete(f'/v1/materials/{material_id}');deleted.raise_for_status()
missing=await api.post(f'/v1/materials/{material_id}/questions',json=body);assert missing.status_code==404
# Interleave same-shaped, contradictory materials to exercise cache ownership.
ids=[]
for case in cases[1:]:
prep,_=await post('/v1/materials',{'state':case['state']});ids.append(prep['material_id'])
isolation=[]
for i in [0,1,0,1]:
case=cases[i+1];r,ms=await post(f'/v1/materials/{ids[i]}/questions',{'questions':case['questions'],'concurrency':4})
isolation.append({'case':case['name'],'result':r,'wall_ms':ms});assert picks(r)==case['expected']
assert r['metrics']['all_candidates_reused_entire_material']
(args.out/'isolation.json').write_text(json.dumps(isolation,indent=2))
# Backend restart/eviction simulation on this isolated local test engine.
# This route exists only because the TEST container has DEV_MODE=1.
reset=await backend.client.post(args.backend+'/reset_prefix_cache');reset.raise_for_status()
case=cases[1];r,ms=await post(f'/v1/materials/{ids[0]}/questions',{'questions':case['questions'],'concurrency':4})
assert r['metrics']['warmup']['cached_tokens']==0
assert r['metrics']['all_candidates_reused_entire_material']
assert picks(r)==case['expected']
(args.out/'eviction-recovery.json').write_text(json.dumps({'result':r,'wall_ms':ms},indent=2))
for material_id in ids:
response=await api.delete('/v1/materials/'+material_id);response.raise_for_status()
# Report model errors, including reference errors. Cache correctness is
# checked separately above; quantized inference is not numerically deterministic.
summary={}
for case in cases:
summary[case['name']]={}
for mode in ['reference','shared-cold-serial','shared-hot-serial','shared-hot-parallel']:
rs=[r for r in records if r['case']==case['name'] and r['mode']==mode]
summary[case['name']][mode]={'median_ms':statistics.median(r['wall_ms'] for r in rs),
'min_ms':min(r['wall_ms'] for r in rs),'max_ms':max(r['wall_ms'] for r in rs),
'correct':sum(r['correct'] for r in rs),'questions':sum(r['n_questions'] for r in rs),
'candidate_uncached_tokens':rs[0]['result']['metrics']['candidate_uncached_tokens']}
(args.out/'summary.json').write_text(json.dumps(summary,indent=2));print(json.dumps(summary,indent=2),flush=True)
stream.close();await api.aclose();await backend.client.aclose()
if __name__=='__main__':
parser=argparse.ArgumentParser();parser.add_argument('--model',required=True);parser.add_argument('--out',type=Path,required=True)
parser.add_argument('--api',default='http://127.0.0.1:8243');parser.add_argument('--backend',default='http://127.0.0.1:8238')
asyncio.run(main(parser.parse_args()))
|