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