File size: 8,508 Bytes
e1acc15 | 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 | """Offline export of frozen generation state; never makes inference/network requests.
The online validator plus frozen rows establish the completions-path top20 proof.
A real image-bound decode smoke still establishes the image/capture provenance;
its chat result alone cannot certify the different completions formatter.
"""
from __future__ import annotations
import argparse
import fcntl
import hashlib
import json
import re
from collections.abc import Mapping
from pathlib import Path
from typing import cast
import numpy as np
from numpy.typing import NDArray
from . import generation_top20 as T
from .assembly_paths import AssemblyPaths
from .replay import load_seq, read_plan, tokens_sha256
SCHEMA = 'd2.generation-top20-export.v1'
def digest(path: Path) -> str:
with path.open('rb') as fh:
return hashlib.file_digest(fh, 'sha256').hexdigest()
def obj(path: Path) -> dict[str, object]:
value = json.loads(path.read_text())
if not isinstance(value, dict):
raise TypeError('expected metadata object: ' + str(path))
return cast(dict[str, object], value)
def write(path: Path, value: object) -> None:
path.write_text(json.dumps(value))
path.chmod(0o600)
def smoke(path: Path) -> None:
check = obj(path)
capture = Path(__file__).parents[1] / 'capture/decode_capture.py'
if (check.get('schema') != 'd2.decode-smoke.v1' or check.get('synthetic') is not False
or check.get('pass') is not True or check.get('option_ii') is not True
or check.get('image') != T.IMAGE or check.get('capture_sha256') != digest(capture)):
raise ValueError('generation top20 requires a successful real image-bound smoke')
def saved_arrays(state: Mapping[str, object]) -> dict[str, NDArray[np.int32] | NDArray[np.float32]]:
if state.get('top20') != 'generation':
raise ValueError('not a generation-top20 state')
ids = cast(list[int], state['gen_ids'])
rows = cast(list[T.Row], state['logprob_rows'])
T.validate_saved(ids, rows)
return {key: np.asarray([r[key] for r in rows], dtype=np.int32 if key in ('top_ids', 'tok_rank') else np.float32)
for key in ('top_ids', 'top_lp', 'tok_lp', 'tok_rank', 'resid')}
def export(state: Path, out: Path, verification: Path) -> None:
"""Require the driver's exclusive lock, copy snapshots, then export. Publish receipt last."""
if out.exists():
raise ValueError('generation export requires a NEW output directory')
smoke(verification)
# r+ opens only an existing driver lock, and never modifies source state.
with (state / 'LOCK').open('r+') as lock:
try:
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
except BlockingIOError as exc:
raise ValueError('stop generator before frozen export: driver still holds LOCK') from exc
out.mkdir(parents=True, mode=0o700)
for sub in ('seqs', 'logprobs'):
(out / sub).mkdir(mode=0o700)
entries: list[dict[str, object]] = []
receipts: dict[str, object] = {}
for path in sorted((state / 'seqs').glob('*.json')):
raw = obj(path)
if raw.get('status') == 'running':
raise ValueError('unreaped running state: restart and stop generator cleanly before export')
seq = load_seq(str(path))
if not seq:
continue
if len(seq['tokens']) - seq['gen_start'] < 3:
continue
sid = seq['seq_id']
if not isinstance(sid, str) or not re.fullmatch('[A-Za-z0-9_-]+', sid):
raise ValueError('unsafe sequence identity')
arrays = saved_arrays(raw)
# Bind the plan to immutable private copies, never a resumable live state.
frozen = out / 'seqs' / (sid + '.json')
write(frozen, raw)
L, g = len(seq['tokens']), seq['gen_start']
entry = {k: seq[k] for k in ('seq_id', 'split', 'family', 'kind', 'prompt_sha256', 'finish_reason', 'segments')}
entry.update(src=str(frozen.resolve()), L=L, gen_start=g, **{'from': max(0, g - 128)},
tokens_sha256=tokens_sha256(seq['tokens']), request_id=f'd2g.top20.{sid}.{L}')
zp = out / 'logprobs' / (sid + '.npz')
np.savez(zp, top_ids=arrays['top_ids'], top_lp=arrays['top_lp'], tok_lp=arrays['tok_lp'],
tok_rank=arrays['tok_rank'], resid=arrays['resid'])
zp.chmod(0o600)
write(out / 'logprobs' / (sid + '.json'), entry | {
'schema': T.SCHEMA, 'source': 'generation_verified', 'rank_semantics': 'top20_position_else_21',
'rows': L - g, 'state_sha256': digest(frozen)})
entries.append(entry)
receipts[sid] = {'state_sha256': digest(frozen), 'npz_sha256': digest(zp), 'rows': L - g,
'metadata_sha256': digest(out / 'logprobs' / (sid + '.json'))}
if not entries:
raise ValueError('no exportable generation sequences')
plan = out / 'plan.jsonl'
plan.write_text(''.join(json.dumps(e) + '\n' for e in entries))
plan.chmod(0o600)
write(out / 'VERIFIED.json', {
'schema': SCHEMA, 'pass': True, 'endpoint': '/v1/completions', 'image': T.IMAGE,
'logprobs_mode': 'raw_logprobs', 'rank_semantics': 'top20_position_else_21',
'validator_sha256': digest(Path(T.__file__)), 'plan_sha256': digest(plan),
'smoke': str(verification.resolve()), 'smoke_sha256': digest(verification),
'sequences': receipts, 'rows': sum(cast(int, cast(dict[str, object], r)['rows']) for r in receipts.values())})
verify(out, out / 'VERIFIED.json', full=True)
def verify(root: Path, receipt: Path, *, full: bool = False,
paths: AssemblyPaths | None = None) -> dict[str, object]:
if paths:
for path in (root, receipt, root / 'plan.jsonl'):
paths.inside(path)
check = obj(receipt)
if (check.get('schema') != SCHEMA or check.get('pass') is not True
or check.get('endpoint') != '/v1/completions' or check.get('image') != T.IMAGE
or check.get('logprobs_mode') != 'raw_logprobs'
or check.get('rank_semantics') != 'top20_position_else_21'
or check.get('validator_sha256') != digest(Path(T.__file__))
or check.get('plan_sha256') != digest(root / 'plan.jsonl')):
raise ValueError('generation top20 requires a successful real image-bound smoke and completions validation receipt')
smoke_path = paths.resolve(str(check['smoke'])) if paths else Path(str(check['smoke']))
if digest(smoke_path) != check.get('smoke_sha256'):
raise ValueError('generation image smoke changed')
smoke(smoke_path)
seqs = check.get('sequences')
entries = read_plan(str(root))
if not isinstance(seqs, dict) or not entries or set(seqs) != {e['seq_id'] for e in entries}:
raise ValueError('generation receipt sequence coverage mismatch')
if full:
for e in entries:
verify_sequence(root, e, check, paths=paths)
return check
def verify_sequence(root: Path, entry: Mapping[str, object], check: Mapping[str, object], *,
paths: AssemblyPaths | None = None) -> None:
sid = str(entry['seq_id'])
if not re.fullmatch('[A-Za-z0-9_-]+', sid):
raise ValueError('unsafe sequence identity')
if paths:
for path in (root / 'logprobs' / (sid + '.npz'), root / 'logprobs' / (sid + '.json')):
paths.inside(path)
seqs = cast(dict[str, object], check['sequences'])
r = seqs[sid]
if (not isinstance(r, dict) or r.get('rows') != cast(int, entry['L']) - cast(int, entry['gen_start'])
or r.get('state_sha256') != digest(paths.resolve(str(entry['src'])) if paths else Path(str(entry['src'])))
or r.get('npz_sha256') != digest(root / 'logprobs' / (sid + '.npz'))
or r.get('metadata_sha256') != digest(root / 'logprobs' / (sid + '.json'))):
raise ValueError('generation export changed after online-validation receipt')
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument('--state', type=Path, required=True)
ap.add_argument('--out', type=Path, required=True)
ap.add_argument('--image-smoke', type=Path, required=True)
a = ap.parse_args()
export(a.state, a.out, a.image_smoke)
if __name__ == '__main__':
main()
|