| """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) |
| |
| 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'] |
| |
| |
| 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 |
| |
| 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)) |
| |
| |
| 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() |
| |
| |
| 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())) |
|
|
