| """Amendment-4 fixed pilot sample. CPU only; never reads model outcomes. |
| |
| The population is complete five-token rounds on the same fixed K2 trajectory |
| used by G2. Reading one token vector per sequence avoids materializing taps for |
| all non-sampled rounds. Keys and their hash stay in the private pilot plan. |
| """ |
| from __future__ import annotations |
|
|
| import fcntl |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import random |
| import tempfile |
| from typing import TypedDict |
| from collections.abc import Iterable |
|
|
| from safetensors import safe_open |
|
|
| from .data import D2Blocks, sha256 |
| from .evaluate import match_length |
| from .replay import g1_drafts |
|
|
| |
| SPLITS = ("acap", "heldout_legal") |
| DEV_POPULATION = ("dev", "heldout_legal") |
|
|
|
|
| def population_splits(keys: Iterable[str]) -> tuple[str, ...]: |
| if set(keys) == set(DEV_POPULATION): |
| return DEV_POPULATION |
| if set(keys) == set(SPLITS): |
| return SPLITS |
| raise ValueError("pilot population requires dev + heldout_legal (or historical analytical acap)") |
|
|
| SAMPLE_SEED = 20261024 |
| SAMPLE_LIMIT = 1000 |
|
|
|
|
| class RoundKey(TypedDict): |
| request: str |
| P: int |
|
|
|
|
| class Sample(TypedDict): |
| schema: str |
| seed: int |
| limit: int |
| manifest_sha256: str |
| k2_sha256: dict[str, str] |
| population_sha256: dict[str, str] |
| population_rounds: dict[str, int] |
| keys: dict[str, list[RoundKey]] |
| sha256: str |
|
|
|
|
| def hash_json(value: object) -> str: |
| return hashlib.sha256(json.dumps(value, sort_keys=True, allow_nan=False, separators=(",", ":")).encode()).hexdigest() |
|
|
|
|
| def key_hash(keys: list[RoundKey]) -> str: |
| """Matches the scorer's ordered (request, position) population digest.""" |
| return hash_json([(k["request"], k["P"]) for k in keys]) |
|
|
|
|
| def validate_sample(sample: Sample, manifest_hash: str) -> None: |
| if (sample.get("schema") != "dsv41.d2.pilot-sample.v1" or sample.get("seed") != SAMPLE_SEED |
| or sample.get("limit") != SAMPLE_LIMIT or sample.get("manifest_sha256") != manifest_hash): |
| raise ValueError("pilot sample schema/seed/limit/manifest mismatch") |
| splits = population_splits(sample["keys"]) |
| for mapping in (sample["keys"], sample["k2_sha256"], sample["population_sha256"], sample["population_rounds"]): |
| if set(mapping) != set(splits): |
| raise ValueError("pilot sample requires both held-out splits") |
| for split in splits: |
| keys = sample["keys"][split] |
| count = sample["population_rounds"][split] |
| if type(count) is not int or count <= 0 or len(keys) != min(SAMPLE_LIMIT, count): |
| raise ValueError("pilot sample has the wrong complete-round count") |
| if any(set(k) != {"request", "P"} or not isinstance(k["request"], str) or not k["request"] |
| or type(k["P"]) is not int or k["P"] < 0 for k in keys): |
| raise ValueError("invalid pilot sample key") |
| if len({(k["request"], k["P"]) for k in keys}) != len(keys): |
| raise ValueError("duplicate pilot sample key") |
| if sample.get("sha256") != hash_json(sample["keys"]): |
| raise ValueError("pilot sample key-list hash mismatch") |
|
|
|
|
| def complete_rounds(data: D2Blocks, k2: Path) -> list[tuple[int, RoundKey]]: |
| drafts = g1_drafts(k2, data) |
| result: list[tuple[int, RoundKey]] = [] |
| for seq_index, (path, seq, request) in enumerate(data.sequences): |
| token_offset, length, gen = seq[:3] |
| start = data.ends[seq_index-1] if seq_index else 0 |
| with safe_open(path, framework="pt", device="cpu") as f: |
| tokens = f.get_slice("tokens")[token_offset:token_offset+length].tolist() |
| pos = gen-1 |
| while pos+6 < length: |
| truth = tokens[pos+2:pos+7] |
| accepted = match_length(drafts[request][pos], truth) |
| if accepted is None: |
| raise ValueError("indeterminate complete K2 round") |
| result.append((start+pos-(gen-1), {"request": request, "P": pos})) |
| pos += accepted+1 |
| return result |
|
|
|
|
| def draw_sample(manifest_hash: str, populations: dict[str, list[RoundKey]], k2_hashes: dict[str, str]) -> Sample: |
| splits = population_splits(populations) |
| keys: dict[str, list[RoundKey]] = {} |
| for split in splits: |
| population = populations[split] |
| |
| |
| selected = sorted(random.Random(SAMPLE_SEED).sample(range(len(population)), min(SAMPLE_LIMIT, len(population)))) |
| keys[split] = [population[i] for i in selected] |
| sample: Sample = {"schema": "dsv41.d2.pilot-sample.v1", "seed": SAMPLE_SEED, "limit": SAMPLE_LIMIT, |
| "manifest_sha256": manifest_hash, "k2_sha256": k2_hashes, |
| "population_sha256": {s: key_hash(populations[s]) for s in splits}, |
| "population_rounds": {s: len(populations[s]) for s in splits}, |
| "keys": keys, "sha256": hash_json(keys)} |
| validate_sample(sample, manifest_hash) |
| return sample |
|
|
|
|
| def sample_plan(plan: Path, manifest: Path, k2: dict[str, Path]) -> None: |
| from .pilot import manifest_is_synthetic, read_spec |
| lock_fd = os.open(str(plan)+".sample.lock", os.O_CREAT | os.O_RDWR, 0o600) |
| with os.fdopen(lock_fd, "a") as lock: |
| fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) |
| before = plan.read_bytes() |
| if "sample" in json.loads(before): |
| raise FileExistsError("pilot sample already exists; refusing to redraw") |
| spec = read_spec(plan) |
| manifest_hash = sha256(manifest) |
| if manifest_hash != spec["manifest_sha256"]: |
| raise ValueError("pilot sample manifest differs from prepared plan") |
| if manifest_is_synthetic(manifest) != spec["synthetic"]: |
| raise ValueError("pilot sample synthetic label differs from prepared plan") |
| splits = population_splits(k2) |
| if not spec["synthetic"] and splits != DEV_POPULATION: |
| raise ValueError("real pilot requires amendment-7 dev population") |
| if set(k2) != set(splits): |
| raise ValueError("pilot sampling requires both K2 sources") |
| hashes = {s: sha256(k2[s]) for s in splits} |
| populations = {s: [key for _, key in complete_rounds(D2Blocks(manifest, s), k2[s])] for s in splits} |
| spec["sample"] = draw_sample(manifest_hash, populations, hashes) |
| if hashes != {s: sha256(k2[s]) for s in splits} or sha256(manifest) != manifest_hash or plan.read_bytes() != before: |
| raise ValueError("pilot sampling inputs changed during the draw") |
| |
| |
| fd, name = tempfile.mkstemp(prefix=plan.name+".sample-", dir=plan.parent) |
| temporary = Path(name) |
| try: |
| with os.fdopen(fd, "w") as f: |
| json.dump(spec, f, indent=2, allow_nan=False) |
| f.write("\n") |
| f.flush() |
| os.fsync(f.fileno()) |
| os.replace(temporary, plan) |
| directory = os.open(plan.parent, os.O_RDONLY | os.O_DIRECTORY) |
| try: |
| os.fsync(directory) |
| finally: |
| os.close(directory) |
| finally: |
| temporary.unlink(missing_ok=True) |
|
|
