| """One bounded UP2 public-corpus probe; private output contains hashes/numbers only.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import math |
| import re |
| import time |
| import urllib.request |
| import uuid |
| from collections.abc import Callable |
| from pathlib import Path |
| from typing import Protocol, cast |
|
|
| from offload.prefill_maxlen.evidence import Json, number, record |
| from offload.prefill_trace.window_requests import corpus, exact_prompt, idle, stream_summary |
| from offload.launch.windows.up2_checks import CAP, require, screen_record |
|
|
| URL = 'http://127.0.0.1:18001' |
|
|
|
|
| class NeedleBuilder(Protocol): |
| def build(self, ctx: int, n: int, seed: int) -> tuple[str, dict[str, str]]: ... |
|
|
|
|
| def needle(target: int, salt: str, output: Path, ledger: Path, deadline: float, data_dir: Path, *, screen: bool = False) -> None: |
| def event(row: dict[str, Json]) -> None: |
| write(ledger,screen_record(row,screen)) |
| from offload.kvpack.needle import Builder, score |
| make = cast(Callable[[str | None, str | None, Path], NeedleBuilder], Builder) |
| grade = cast(Callable[[str, dict[str,str], set[str]], dict[str,Json]], score) |
| idle(URL) |
| with urllib.request.urlopen(URL+'/v1/models',timeout=10) as response: |
| models = cast(list[Json],record(cast(Json,json.load(response)))['data']) |
| model = record(models[0]) |
| require(model.get('max_model_len') == CAP, 'needle cap mismatch') |
| prompt,codes = make(URL,str(model['id']),data_dir).build(target,6,42+target) |
| request_id = uuid.uuid4().hex |
| body: dict[str,Json] = {'model':model['id'],'request_id':request_id, |
| 'messages':[{'role':'user','content':prompt}], 'max_tokens':300, 'temperature':0,'top_p':1, |
| 'seed':42,'priority':10,'cache_salt':salt,'chat_template_kwargs':{'thinking':False}} |
| require(time.time() < deadline,f'needle submission deadline {deadline:.6f} reached') |
| event({'event':'send_intent','id':request_id,'target':target,'kind':'needle','time':time.time()}) |
| try: |
| req = urllib.request.Request(URL+'/v1/chat/completions',json.dumps(body).encode(),{'Content-Type':'application/json'}) |
| with urllib.request.urlopen(req,timeout=900) as response: |
| event({'event':'http_opened','id':request_id,'time':time.time()}) |
| result = record(cast(Json,json.load(response))) |
| choice = record(cast(list[Json],result['choices'])[0]) |
| text = str(record(choice['message']).get('content') or '') |
| foreign = set(re.findall(r'\b\d{6}\b',text))-set(codes.values()) |
| graded = grade(text,codes,foreign) |
| require(graded.get('recall') == 6 and not graded.get('foreign_codes'),'needle6/6/foreign-code failure') |
| usage = record(result['usage']) |
| require(number(usage['prompt_tokens'])+300 <= CAP and |
| record(usage['prompt_tokens_details']).get('cached_tokens') == 0,'needle length/cache failure') |
| require(choice.get('finish_reason') == 'stop','needle did not complete normally') |
| with output.open('x') as handle: |
| json.dump(screen_record({'context':target,'recall':6,'of':6,'foreign_codes':0,'usage':usage},screen),handle) |
| handle.write('\n') |
| event({'event':'completed','id':request_id,'time':time.time()}) |
| except Exception as exc: |
| event({'event':'failed','id':request_id,'error_class':type(exc).__name__,'time':time.time()}) |
| raise |
|
|
|
|
| def write(path: Path, row: dict[str, Json]) -> None: |
| with path.open('a') as output: |
| output.write(json.dumps(row, allow_nan=False) + '\n') |
| output.flush() |
|
|
|
|
| def execute(target: int, budget: int, salt: str, output: Path, ledger: Path, |
| deadline: float, data_dir: Path, *, numeric: bool = False, force_length: bool = False, |
| rep: int = 0, kind: str = 'diagnostic', screen: bool = False) -> dict[str, Json]: |
| def event(row: dict[str, Json]) -> None: |
| write(ledger,screen_record(row,screen)) |
| require(0 < target and target + budget <= CAP, 'prompt/output exceeds low cap') |
| require(not output.exists(), 'refusing existing request result') |
| template: dict[str, Json] = {'thinking':False} if numeric else {'thinking':True, 'reasoning_effort':'high'} |
| def post(endpoint: str, body: dict[str, Json]) -> dict[str, Json]: |
| require(time.time() < deadline, f'submission deadline {deadline:.6f} reached') |
| req = urllib.request.Request(URL+endpoint, json.dumps(body).encode(), {'Content-Type':'application/json'}) |
| with urllib.request.urlopen(req, timeout=min(900, max(1, deadline-time.time()))) as response: |
| return record(cast(Json, json.load(response))) |
| idle(URL) |
| with urllib.request.urlopen(URL+'/v1/models', timeout=10) as response: |
| model = record(cast(list[Json], record(cast(Json,json.load(response)))['data'])[0]) |
| require(model.get('max_model_len') == CAP and model.get('root') == '/srv/dsv41/legal-v2a', 'served target/cap mismatch') |
| materials, hashes = corpus(data_dir) |
| def count(text: str) -> int: |
| return int(number(post('/tokenize', {'model':model['id'], 'messages':[{'role':'user','content':text}], |
| 'add_generation_prompt':True, 'chat_template_kwargs':template})['count'])) |
| prompt = exact_prompt(materials, target, count) |
| request_id = uuid.uuid4().hex |
| body: dict[str, Json] = {'model':model['id'], 'request_id':request_id, |
| 'messages':[{'role':'user','content':prompt}], 'max_tokens':budget, |
| 'temperature':0 if numeric else 1, 'top_p':1, 'seed':42, 'priority':10, |
| 'cache_salt':salt, 'chat_template_kwargs':template, 'ignore_eos':force_length, |
| 'stream':not numeric} |
| if numeric: |
| body.update({'logprobs':True, 'top_logprobs':20}) |
| else: |
| body['stream_options'] = {'include_usage':True} |
| require(time.time() < deadline, f'submission deadline {deadline:.6f} reached') |
| req = urllib.request.Request(URL+'/v1/chat/completions', json.dumps(body).encode(), {'Content-Type':'application/json'}) |
| started = time.monotonic() |
| event( {'event':'send_intent', 'id':request_id, 'target':target, 'kind':kind, 'time':time.time()}) |
| try: |
| with urllib.request.urlopen(req, timeout=900) as response: |
| event( {'event':'http_opened', 'id':request_id, 'time':time.time()}) |
| if numeric: |
| result = record(cast(Json,json.load(response))) |
| choice = record(cast(list[Json],result['choices'])[0]) |
| tokens = cast(list[Json],record(choice['logprobs'])['content']) |
| require(bool(tokens), 'no numerical logprobs') |
| top = cast(list[Json],record(tokens[0])['top_logprobs']) |
| probabilities: dict[str, Json] = {} |
| |
| for item in top: |
| row = record(item) |
| key = hashlib.sha256(str(row['token']).encode()).hexdigest() |
| probabilities[key] = math.exp(number(row['logprob'])) |
| msg = record(choice['message']) |
| summary: dict[str, Json] = {'usage':result['usage'], 'finish':choice['finish_reason'], |
| 'probabilities':probabilities, 'text_sha256':hashlib.sha256(str(msg.get('content') or '').encode()).hexdigest(), |
| 'total_s':time.monotonic()-started} |
| else: |
| summary = stream_summary(response, started) |
| usage = record(summary['usage']) |
| cached = record(usage.get('prompt_tokens_details'))['cached_tokens'] |
| require(usage.get('prompt_tokens') == target and cached == 0, 'exact-size/cold usage failure') |
| require(number(usage['completion_tokens']) > 0 and summary.get('finish') in ('stop','length'), 'unfinished probe') |
| if force_length: |
| require(usage.get('completion_tokens') == budget and summary.get('finish') == 'length', 'near-cap/warmup length finish failure') |
| row = {'window_id':str(output.parent.resolve()), 'type':'request', 'kind':kind, 'rep':rep, 'target_ctx':target, 'cap':CAP, |
| 'prompt_sha256':hashlib.sha256(prompt.encode()).hexdigest(), |
| 'cache_salt_sha256':hashlib.sha256(salt.encode()).hexdigest(), 'corpus_sha256':hashes, |
| 'cached_tokens':cached, 'uncached_prompt_tokens':target, |
| 'sampling':{'temperature':0 if numeric else 1, 'top_p':1, 'seed':42, 'priority':10, 'max_tokens':budget}, |
| **summary} |
| row = screen_record(row,screen) |
| with output.open('x') as handle: |
| json.dump(row, handle, allow_nan=False) |
| handle.write('\n') |
| event( {'event':'completed', 'id':request_id, 'time':time.time()}) |
| return row |
| except Exception as exc: |
| event( {'event':'failed', 'id':request_id, 'error_class':type(exc).__name__, 'time':time.time()}) |
| raise |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(description=__doc__) |
| ap.add_argument('--target', type=int, required=True) |
| ap.add_argument('--budget', type=int, required=True) |
| ap.add_argument('--salt', required=True) |
| ap.add_argument('--output', type=Path, required=True) |
| ap.add_argument('--ledger', type=Path, required=True) |
| ap.add_argument('--deadline', type=float, required=True) |
| ap.add_argument('--data-dir', type=Path, required=True) |
| ap.add_argument('--screen', action='store_true') |
| ap.add_argument('--numeric', action='store_true') |
| ap.add_argument('--needle', action='store_true') |
| ap.add_argument('--force-length', action='store_true') |
| ap.add_argument('--rep', type=int, default=0) |
| ap.add_argument('--kind', default='diagnostic') |
| args = ap.parse_args() |
| if args.needle: |
| needle(args.target,args.salt,args.output,args.ledger,args.deadline,args.data_dir,screen=args.screen) |
| else: |
| execute(args.target, args.budget, args.salt, args.output, args.ledger, args.deadline, args.data_dir, |
| numeric=args.numeric, force_length=args.force_length, rep=args.rep, kind=args.kind, screen=args.screen) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
