mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-10 03:41:36 -10:00
197 lines
11 KiB
Python
197 lines
11 KiB
Python
"""Compare untouched upstream with cache-disabled engines; dry-run by default.
|
|
|
|
Private engines run sequentially in an exclusive GPU window. This is a correctness
|
|
gate, not a throughput benchmark: state fingerprinting adds diagnostic overhead.
|
|
"""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
import re
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path[:0] = [str(ROOT), str(ROOT / 'tools')]
|
|
from conversation_cache_parity import load_tokenizer, require
|
|
from conversation_cache_soak import answer_text
|
|
from serve.frontend import ChatTemplate
|
|
from serve.server import StrataEngine, child_env
|
|
|
|
NAMES = ('short', 'A', 'A-live', 'A-checkpoint', 'B', 'A-return', 'B-return', 'A-return-again')
|
|
# (not 'version': the gate also compares a new release with the previous one, where only the version differs)
|
|
SETTINGS = ('context', 'kv', 'kv_resident', 'expert_slots', 'spec',
|
|
'mtp_max', 'lookup', 'cvec', 'pcie_frac', 'spec_min_p', 'pool_workers')
|
|
# Untouched upstream does not emit the indexer `dead` and `pooled_full` fields
|
|
# added by this PR; `pooled` is the completed-row extent in both engines.
|
|
# Draft scratch and unused page padding are not comparable authoritative state.
|
|
STATE_KEYS = ('L', 'gdn', 'ple', 'tail', 'pooled', 'kv', 'ple_prev')
|
|
|
|
|
|
def state_hashes(text, *, candidate=False):
|
|
hashes = []
|
|
for line in text.splitlines():
|
|
if 'STATE_HASH L=' in line:
|
|
fields = dict(re.findall(r'(\w+)=([0-9a-f,-]+)', line))
|
|
require(set(STATE_KEYS) <= fields.keys(), 'missing upstream-compatible state fields')
|
|
if candidate:
|
|
# Evidence that the candidate is a build with this change: it
|
|
# also fingerprints the spare row (compared by the cache-on gates).
|
|
require('pooled_full' in fields, 'candidate lacks the pooled_full fingerprint')
|
|
hashes.append({key: fields[key] for key in STATE_KEYS})
|
|
return hashes
|
|
|
|
|
|
def common_args(cfg, spec):
|
|
# Upstream has no conversation-cache options. Omit them from EVERY arm so
|
|
# the candidates exercise their disabled default with identical arguments.
|
|
args, i = [], 0
|
|
while i < len(cfg['args']):
|
|
arg = cfg['args'][i]
|
|
if arg.startswith('--conversation-cache-'):
|
|
require(i + 1 < len(cfg['args']), 'cache option lacks a value')
|
|
i += 2
|
|
else:
|
|
args.append(arg)
|
|
i += 1
|
|
return args + ['--prompt-cache', '6', '--adapt-swaps', '0', '--spec', str(max(2, spec)),
|
|
'--mtp-max-t', str(spec), '--suffix-draft', '0', '--spec-min-p', '0']
|
|
|
|
|
|
def verify(results):
|
|
arms = results['arms']
|
|
require(len(arms) >= 2 and arms[0]['name'] == 'upstream', 'missing comparison arms')
|
|
baseline = arms[0]
|
|
requests = results['requests']
|
|
require([r['name'] for r in requests] == list(NAMES), 'incomplete request sequence')
|
|
require(all(r['ids'] and r['max_new'] > 0 for r in requests), 'empty request')
|
|
digest = hashlib.sha256(json.dumps(requests, sort_keys=True).encode()).hexdigest()
|
|
require(digest == results['request_digest'], 'request evidence changed')
|
|
for arm in arms:
|
|
require(arm['args'] == baseline['args'], 'engine arguments differ')
|
|
require(arm['request_digest'] == results['request_digest'], 'input requests differ')
|
|
require([r['name'] for r in arm['records']] == list(NAMES), 'incomplete output sequence')
|
|
require(not arm['cache_events'], 'disabled cache performed snapshot operations')
|
|
for key in SETTINGS:
|
|
require(key in arm['info'] and arm['info'][key] == baseline['info'][key],
|
|
f'{arm["name"]}: resolved setting differs: {key}')
|
|
if arm['name'] != 'upstream':
|
|
require(arm['info'].get('conversation_cache_mib') == 0, 'RAM cache not disabled')
|
|
require(arm['info'].get('conversation_disk_mib', 0) == 0, 'disk cache not disabled')
|
|
for request, reference, record in zip(requests, baseline['records'], arm['records']):
|
|
require(record['ids'] and record['finish'] in ('stop', 'length'), 'incomplete generation')
|
|
require(set(STATE_KEYS) <= record['state'].keys(), 'incomplete state fingerprint')
|
|
require(record['ids'] == reference['ids'], f'{arm["name"]}/{record["name"]}: output differs')
|
|
require(record['state'] == reference['state'], f'{arm["name"]}/{record["name"]}: state differs')
|
|
require(record['reused'] == reference['reused'], 'prefix reuse differs')
|
|
require(record['prompt_tokens'] == len(request['ids']), 'wrong prompt length')
|
|
if request.get('expected'):
|
|
require(answer_text(record['text']) == request['expected'].casefold(), 'known answer is wrong')
|
|
records = {r['name']: r for r in arm['records']}
|
|
require(records['A-live']['reused'] > 0 and records['A-checkpoint']['reused'] > 0,
|
|
'continuation/checkpoint path not exercised')
|
|
require(records['A-return']['reused'] == records['A-return-again']['reused'] == 0,
|
|
'conversation switching did not force cold fallback')
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__)
|
|
ap.add_argument('--config', type=Path, required=True)
|
|
ap.add_argument('--upstream', type=Path, required=True)
|
|
ap.add_argument('--candidate', type=Path, action='append', required=True)
|
|
ap.add_argument('--output', type=Path, required=True, help='new directory; existing paths refused')
|
|
ap.add_argument('--paragraphs', type=int, default=128)
|
|
ap.add_argument('--spec', type=int, choices=(1, 4), default=1)
|
|
ap.add_argument('--run', action='store_true')
|
|
a = ap.parse_args()
|
|
if a.paragraphs < 1:
|
|
ap.error('paragraphs must be positive')
|
|
if not a.run:
|
|
print('Dry run: upstream and disabled candidates; identical inputs/settings, known answers and state parity.')
|
|
return
|
|
cfg = json.loads(a.config.read_text(encoding='utf-8'))
|
|
tok_path = Path(cfg['tokenizer'])
|
|
tok = load_tokenizer(tok_path)
|
|
tpl_path = tok_path / 'chat_template.jinja'
|
|
tpl = ChatTemplate(tpl_path if tpl_path.exists() else ROOT / 'serve/chat_template.jinja')
|
|
def encode(text):
|
|
return tok.encode(text, parse_special=True)
|
|
def prompt(text):
|
|
return encode(tpl.render([{'role': 'user', 'content': text}], enable_thinking=False))
|
|
text = 'Conversation A. Remember the exact code AZURE-314159.\n'
|
|
text += '\n'.join(f'Record {i}: blue square, green triangle, red circle.' for i in range(a.paragraphs))
|
|
A = prompt(text + '\nThe exact code to remember is AZURE-314159. Reply OK.')
|
|
B = prompt('Unrelated conversation B. Remember BRONZE-271828. Reply OK.')
|
|
suffix = encode('<|im_end|>\n<|im_start|>user\nWhat exact code did I ask you to remember? '
|
|
'Reply with only that code.<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n')
|
|
args = common_args(cfg, a.spec)
|
|
env = child_env(cfg)
|
|
env.pop('STRATA_SNAPSHOT_VERIFY', None)
|
|
env['STRATA_STATE_HASH'] = '1'
|
|
env['STRATA_MTP_BATCH'] = '1'
|
|
a.output.mkdir(mode=0o700, parents=False, exist_ok=False)
|
|
results = {'config': cfg, 'paragraphs': a.paragraphs, 'spec': a.spec, 'arms': [], 'requests': [],
|
|
'compared_state_fields': STATE_KEYS,
|
|
'environment': {k: v for k, v in env.items() if k.startswith('STRATA_')}}
|
|
for index, exe in enumerate([a.upstream, *a.candidate]):
|
|
name = 'upstream' if index == 0 else f'candidate-{index}'
|
|
log = a.output / f'{name}.log'
|
|
print(f'START {name}: {exe}', flush=True)
|
|
started = time.monotonic()
|
|
engine = StrataEngine(str(exe.resolve()), args, cwd=cfg.get('cwd'), log=str(log), env=env)
|
|
arm = {'name': name, 'exe': str(exe.resolve()), 'args': args, 'info': dict(engine.info),
|
|
'startup_seconds': time.monotonic() - started, 'records': []}
|
|
with exe.open('rb') as source:
|
|
arm['exe_sha256'] = hashlib.file_digest(source, 'sha256').hexdigest()
|
|
def generate(request):
|
|
ids = [t for t in engine.generate(request['ids'], request['max_new'], {'temperature': 0},
|
|
threading.Event()) if t is not None]
|
|
arm['records'].append({'name': request['name'], 'ids': ids, **engine.last,
|
|
'text': tok.decode(ids)})
|
|
return ids
|
|
try:
|
|
if index == 0:
|
|
def add(name, ids, max_new, expected=None):
|
|
request = dict(name=name, ids=ids, max_new=max_new, expected=expected)
|
|
results['requests'].append(request)
|
|
return generate(request)
|
|
add('short', prompt('What is 2 + 2? Reply with only the number.'), 16, '4')
|
|
head = add('A', A, 1)
|
|
require(len(head) == 1, 'A must produce one token')
|
|
continuation = A + head + suffix
|
|
add('A-live', continuation, 32, 'AZURE-314159')
|
|
add('A-checkpoint', continuation, 32, 'AZURE-314159')
|
|
add('B', B, 1)
|
|
add('A-return', continuation, 32, 'AZURE-314159')
|
|
add('B-return', B, 1)
|
|
add('A-return-again', continuation, 32, 'AZURE-314159')
|
|
payload = json.dumps(results['requests'], sort_keys=True).encode()
|
|
results['request_digest'] = hashlib.sha256(payload).hexdigest()
|
|
(a.output / 'requests.json').write_bytes(payload)
|
|
else:
|
|
for request in results['requests']:
|
|
generate(request)
|
|
finally:
|
|
engine.close()
|
|
engine.proc.wait(timeout=30)
|
|
engine.log.close()
|
|
arm['request_digest'] = hashlib.sha256(json.dumps(results['requests'], sort_keys=True).encode()).hexdigest()
|
|
log_text = log.read_text(encoding='utf-8')
|
|
hashes = state_hashes(log_text, candidate=index != 0)
|
|
require(len(hashes) == len(arm['records']), f'{name}: missing state hashes')
|
|
for record, fingerprint in zip(arm['records'], hashes):
|
|
record['state'] = fingerprint
|
|
arm['cache_events'] = [line for line in log_text.splitlines() if any(mark in line for mark in
|
|
('conversation cache: parked', 'disk cache: spilled', 'SNAPSHOT_VERIFY'))]
|
|
results['arms'].append(arm)
|
|
(a.output / 'results.json').write_text(json.dumps(results, indent=2) + '\n')
|
|
print(f'DONE {name}', flush=True)
|
|
verify(results)
|
|
(a.output / 'passed.json').write_text(json.dumps({'passed': True, 'request_digest': results['request_digest']}) + '\n')
|
|
print('PASS: untouched upstream versus disabled engines; identical tokens, main state, reuse and known answers')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|