File size: 16,160 Bytes
1e251f3 | 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 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 | """Coordinator-only public prompt matrix. Dry plan by default; never save text."""
from __future__ import annotations
import argparse
import hashlib
import json
import random
import time
import urllib.request
import uuid
from collections.abc import Callable, Iterable
from pathlib import Path
from typing import cast
from offload.prefill_maxlen.evidence import Json, number, record
TEMPLATE: dict[str, Json] = {'thinking': True, 'reasoning_effort': 'high'}
WARMUP, WARMUP_TOKENS = 32000, 64
HEADER = ('PUBLIC MATERIALS BENCHMARK. Summarize these public documents, compare recurring themes, '
'and identify inconsistencies using only the materials. Do not invent sources.\n\n')
def post(url: str, endpoint: str, body: dict[str, Json]) -> dict[str, Json]:
req = urllib.request.Request(url + endpoint, json.dumps(body).encode(), {'Content-Type': 'application/json'})
with urllib.request.urlopen(req, timeout=900) as response:
return record(cast(Json, json.load(response)))
def corpus(root: Path) -> tuple[str, dict[str, Json]]:
docs: list[str] = []
hashes: dict[str, Json] = {}
for name in ('opinions_cl.jsonl', 'filings_recap.jsonl'):
path = root / name
hashes[name] = hashlib.sha256(path.read_bytes()).hexdigest()
for line in path.read_text().splitlines():
row = record(cast(Json, json.loads(line)))
text = row.get('text')
if not isinstance(text, str):
raise ValueError(f'{name}: public record missing text')
docs.append(text)
random.Random(42).shuffle(docs)
return HEADER + '\n\n'.join(docs), hashes
def exact_cut(materials: str, target: int, count: Callable[[str], int]) -> tuple[int, int]:
"""(characters of the shared family, appended ' x' pads) whose complete chat is exactly target tokens."""
# All lengths draw from the same family/order; only the truncation boundary differs.
lo, hi = len(HEADER), len(materials)
best = 0
if count(materials) < target:
raise ValueError(f'public corpus too short for {target}; do not repeat filler')
while lo <= hi:
mid = (lo + hi) // 2
n = count(materials[:mid])
if n == target:
return mid, 0
if n < target:
best = mid
lo = mid + 1
else:
hi = mid - 1
# Token boundaries can skip an integer under character bisection. Bound correction.
for pads in range(32):
n = count(render(materials, best, pads))
if n == target:
return best, pads
if n > target or not best:
break
raise ValueError(f'cannot construct exact {target}-token chat; no approximate 500K claim')
def render(materials: str, chars: int, pads: int) -> str:
return materials[:chars] + ' x' * pads
def exact_prompt(materials: str, target: int, count: Callable[[str], int]) -> str:
return render(materials, *exact_cut(materials, target, count))
def reused_cuts(path: Path) -> dict[int, tuple[int, int, str]]:
"""Numeric cuts from an earlier arm of the SAME window; each is re-hashed and recounted before use."""
cuts: dict[int, tuple[int, int, str]] = {}
for line in path.read_text().splitlines():
row = record(cast(Json, json.loads(line)))
if row.get('type') != 'request':
continue
cut = (int(number(row['prompt_chars'])), int(number(row['prompt_pad_x'])), str(row['prompt_sha256']))
if cuts.setdefault(int(number(row['target_ctx'])), cut) != cut:
raise ValueError('earlier arm used two different prompts for one target')
return cuts
def cache_salt(prefix: str | None, target: int, rep: int, mode: str) -> str:
"""Deterministic coordinator pairing; default retains fresh random salts."""
if prefix is None:
return uuid.uuid4().hex
if not prefix or mode not in ('main', 'partial', 'warmup') or target < 1 or rep < 0:
raise ValueError('invalid fixed cache salt identity')
return f'{prefix}:{target}:{rep}:{mode}'
def stream_summary(lines: Iterable[bytes], started: float) -> dict[str, Json]:
first: dict[str, Json] = {}
last: dict[str, Json] = {}
events: list[Json] = []
usage: dict[str, Json] | None = None
finish: Json = None
done = False
response_id: Json = None
for raw in lines:
if time.monotonic() - started > 900:
raise RuntimeError('request exceeded 900-second wall budget; coordinator cancellation boundary required')
if not raw.startswith(b'data:'):
continue
payload = raw[5:].strip()
if payload == b'[DONE]':
done = True
break
obj = record(cast(Json, json.loads(payload)))
if obj.get('error'):
raise RuntimeError('server returned stream error; inspect coordinator server logs')
response_id = obj.get('id', response_id)
if obj.get('usage'):
usage = record(obj['usage'])
pieces = False
choices = obj.get('choices', [])
if not isinstance(choices, list):
raise ValueError('choices must be a list')
for choice in choices:
row = record(choice)
delta = record(row.get('delta') or {})
for kind, value in (('reasoning', delta.get('reasoning_content') or delta.get('reasoning')),
('answer', delta.get('content'))):
if value:
elapsed = time.monotonic() - started
first.setdefault(kind, elapsed)
last[kind] = elapsed
pieces = True
finish = row.get('finish_reason') or finish
if pieces:
events.append((time.monotonic() - started) * 1000)
total = time.monotonic() - started
if not done or usage is None or finish is None:
raise RuntimeError('incomplete stream: missing DONE, usage or finish')
details = record(usage.get('completion_tokens_details') or {})
reason = details.get('reasoning_tokens')
answer = number(usage['completion_tokens']) - number(reason) if reason is not None else None
return {'response_id': response_id, 'usage': usage, 'finish': finish,
'ttft_s': min(number(v) for v in first.values()) if first else None,
'total_s': total, 'first_by_kind_s': first, 'last_by_kind_s': last,
'reasoning_tokens': reason, 'answer_and_control_tokens': answer,
'sse_times_ms': events}
def idle(url: str) -> None:
with urllib.request.urlopen(url + '/metrics', timeout=10) as response:
lines = response.read().decode().splitlines()
for name in ('vllm:num_requests_running', 'vllm:num_requests_waiting'):
values = [float(line.split()[-1]) for line in lines if line.startswith(name + '{') or line.startswith(name + ' ')]
if not values or sum(values) != 0:
raise RuntimeError(f'serve not verified idle: {name}')
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument('--execute', action='store_true')
ap.add_argument('--url', default='http://127.0.0.1:18001')
ap.add_argument('--cap', type=int, choices=(327680, 524288), required=True)
ap.add_argument('--contexts', type=int, nargs='+', default=[150000, 300000, 400000, 450000, 500000])
ap.add_argument('--reps', type=int, default=2)
ap.add_argument('--modes', nargs='+', choices=('cold', 'full', 'partial'), default=['cold', 'full', 'partial'])
ap.add_argument('--data-dir', type=Path, default=Path('/opt/dsv41/data/bench-data'))
ap.add_argument('--output', type=Path, required=True)
ap.add_argument('--trace-run', required=True)
ap.add_argument('--cache-salt-prefix', help='same fixed value across restarted paired arms; cold usage must be zero')
ap.add_argument('--deadline-utc', type=float, default=0, help='shared absolute window submission deadline')
ap.add_argument('--warmup', action='store_true', help='first one discarded exact-32000 cold request, 64 tokens, EOS ignored')
ap.add_argument('--reuse-cuts', type=Path, help="earlier arm's output in this window: numeric cuts, re-hashed and recounted")
ap.add_argument('--ledger', type=Path, help='append send/open/complete/fail events: ids, kinds, sizes and times only')
args = ap.parse_args()
if args.reps < 1 or any(n < 100001 or n > 500000 or n + 1500 > args.cap for n in args.contexts):
ap.error('choose in-scope points with 1500 output reserve; partial suffix is 100000 tokens')
if not args.execute:
print(json.dumps({'executed': False, 'cap': args.cap, 'points': args.contexts,
'reps': args.reps, 'modes': args.modes, 'output_budget': 1500, 'warmup': args.warmup,
'fixed_cache_salts': args.cache_salt_prefix is not None}))
return
if args.deadline_utc <= time.time():
ap.error('--execute requires a future --deadline-utc; stop submissions at the window boundary')
if args.url != 'http://127.0.0.1:18001':
ap.error('coordinator must use the reserved experiment tunnel at 127.0.0.1:18001')
if args.output.exists():
ap.error('refusing an existing output; use a fresh window directory')
idle(args.url)
with urllib.request.urlopen(args.url + '/v1/models', timeout=10) as response:
info = record(cast(Json, json.load(response)))
models = info['data']
if not isinstance(models, list) or len(models) != 1:
raise ValueError('expected one served target')
model = record(models[0])
if model.get('max_model_len') != args.cap:
raise ValueError('served cap differs from requested arm')
model_id = model['id']
materials, hashes = corpus(args.data_dir)
def count(text: str) -> int:
result = post(args.url, '/tokenize', {'model': model_id, 'messages': [{'role': 'user', 'content': text}],
'add_generation_prompt': True, 'chat_template_kwargs': TEMPLATE})
return int(number(result['count']))
cuts = reused_cuts(args.reuse_cuts) if args.reuse_cuts is not None else {}
def prompt_for(target: int) -> tuple[str, int, int]:
if target in cuts:
chars, pads, digest = cuts[target]
prompt = render(materials, chars, pads)
# One /tokenize instead of a bisection; any drift refuses rather than re-bisecting silently.
if hashlib.sha256(prompt.encode()).hexdigest() != digest or count(prompt) != target:
raise ValueError('reused prompt cut no longer reproduces the earlier prompt/token count')
return prompt, chars, pads
chars, pads = exact_cut(materials, target, count)
return render(materials, chars, pads), chars, pads
def event(row: dict[str, Json]) -> None:
if args.ledger is not None:
with args.ledger.open('a') as handle:
handle.write(json.dumps({**row, 'time': time.time()}) + '\n')
with args.output.open('x') as output:
output.write(json.dumps({'type': 'identity', 'corpus_sha256': hashes, 'seed': 42,
'cap': args.cap, 'trace_run': args.trace_run, 'model': model_id}) + '\n')
output.flush()
def request(built: tuple[str, int, int], target: int, salt: str, kind: str, rep: int, budget: int = 1500,
force_length: bool = False) -> None:
prompt, chars, pads = built
if time.time() >= args.deadline_utc:
raise RuntimeError('window submission deadline reached; matrix incomplete, restore production')
idle(args.url)
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': 1.0, 'top_p': 1.0, 'seed': 42, 'priority': 10, 'cache_salt': salt,
'chat_template_kwargs': TEMPLATE, 'stream': True, 'stream_options': {'include_usage': True}}
if force_length:
body['ignore_eos'] = True
req = urllib.request.Request(args.url + '/v1/chat/completions', json.dumps(body).encode(),
{'Content-Type': 'application/json'})
event({'event': 'send_intent', 'id': request_id, 'kind': kind, 'target': target})
try:
started = time.monotonic()
with urllib.request.urlopen(req, timeout=900) as response:
event({'event': 'http_opened', 'id': request_id})
summary = stream_summary(response, started)
usage = record(summary['usage'])
cached = record(usage.get('prompt_tokens_details') or {}).get('cached_tokens')
if usage.get('prompt_tokens') != target or cached is None:
raise ValueError('usage disagrees with exact occupied prompt or lacks cache count')
cached_n = int(number(cached))
if kind in ('cold', 'warmup', 'prime-full', 'prime-partial') and cached_n != 0:
raise ValueError('fresh cache salt produced a hit')
if force_length and (usage.get('completion_tokens') != budget or summary.get('finish') != 'length'):
raise ValueError('warm-up did not generate its full forced length')
if kind == 'full' and not target - 2048 <= cached_n < target:
raise ValueError('full prefix hit outside expected alignment/replay bound')
if kind == 'partial' and not target - 102048 <= cached_n <= target - 100000:
raise ValueError('partial prefix hit outside controlled 100K suffix/replay bound')
response_id = summary['response_id']
if not isinstance(response_id, str):
raise ValueError('missing response id for two-rank correlation')
trace_request = hashlib.sha256((args.trace_run + ':' + response_id).encode()).hexdigest()[:24]
output.write(json.dumps({'type': 'request', 'kind': kind, 'rep': rep, 'target_ctx': target,
'cap': args.cap, 'trace_request': trace_request,
'cache_salt_sha256': hashlib.sha256(salt.encode()).hexdigest(), 'prompt_sha256': hashlib.sha256(prompt.encode()).hexdigest(),
'prompt_chars': chars, 'prompt_pad_x': pads,
'cached_tokens': cached_n, 'uncached_prompt_tokens': target - cached_n,
'sampling': {'temperature': 1, 'top_p': 1, 'seed': 42, 'priority': 10, 'max_tokens': budget},
**summary}, allow_nan=False) + '\n')
output.flush()
except BaseException as exc:
event({'event': 'failed', 'id': request_id, 'error_class': type(exc).__name__})
raise
event({'event': 'completed', 'id': request_id})
if args.warmup:
request(prompt_for(WARMUP), WARMUP, cache_salt(args.cache_salt_prefix, WARMUP, 0, 'warmup'), 'warmup', 0,
WARMUP_TOKENS, force_length=True)
for target in args.contexts:
prompt = prompt_for(target)
partial = prompt_for(target - 100000) if 'partial' in args.modes else ('', 0, 0)
for rep in range(args.reps):
salt = cache_salt(args.cache_salt_prefix, target, rep, 'main')
if 'cold' in args.modes:
request(prompt, target, salt, 'cold', rep)
elif 'full' in args.modes:
request(prompt, target, salt, 'prime-full', rep, 64)
if 'full' in args.modes:
request(prompt, target, salt, 'full', rep)
if 'partial' in args.modes:
partial_salt = cache_salt(args.cache_salt_prefix, target, rep, 'partial')
request(partial, target - 100000, partial_salt, 'prime-partial', rep, 64)
request(prompt, target, partial_salt, 'partial', rep)
if __name__ == '__main__':
main()
|