| """Native-arm identity: source provenance, exact arm policy, K6 head and runtime sources. |
| |
| Pure Python, shared by the worker and the Spark identity tool. The worker verifies |
| the official receipt's whole files before loading. The existing checkpoint identity |
| still verifies every retained checkpoint tensor (embedding, norms, Markov, K6 head). |
| Neither path trusts an advertised digest in place of checking its own bytes. |
| """ |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import re |
| from dataclasses import asdict |
| from pathlib import Path |
| from typing import Any |
|
|
| from .arms import ARMS, batch_invariant_from_env |
|
|
| from ..worker.arm_selection import SERVING_ARMS, selected_arm |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
| def sha256(path: Path) -> str: |
| h = hashlib.sha256() |
| with path.open("rb") as f: |
| for b in iter(lambda: f.read(16 << 20), b""): |
| h.update(b) |
| return h.hexdigest() |
|
|
|
|
| def digest(recipe: dict[str, Any]) -> str: |
| return hashlib.sha256(json.dumps(recipe, sort_keys=True, separators=(",", ":")).encode()).hexdigest() |
|
|
|
|
| def official_hashes(manifest: dict[str, Any]) -> dict[str, str]: |
| pairs = re.findall(r"^([0-9a-f]{64}) ([^\n]+)$", manifest["receipt"], re.M) |
| hashes = {name: h for h, name in pairs} |
| required = {"config.json", "model.safetensors.index.json"} | {t["file"] for t in manifest["tensors"]} |
| if not required <= hashes.keys(): |
| raise ValueError(f"official receipt lacks {sorted(required - hashes.keys())}") |
| return {n: hashes[n] for n in sorted(required)} |
|
|
|
|
| def runtime_hashes(root: Path = ROOT) -> dict[str, str]: |
| """Conservative coverage, including shared kernels; excludes tools/tests/private data.""" |
| paths = [] |
| for subdir in ("worker", "native", "kernels"): |
| for p in (root / subdir).rglob("*"): |
| rel = p.relative_to(root) |
| if any(x in rel.parts for x in ("tests", "tools", "__pycache__", "results")): |
| continue |
| if p.is_file() and p.suffix in (".py", ".cu", ".cuh", ".cpp", ".h", ".c", ".json"): |
| paths.append(p) |
| if not paths: |
| raise ValueError(f"no runtime sources under {root}") |
| return {str(p.relative_to(root)): sha256(p) for p in sorted(paths)} |
|
|
|
|
| def compute(arm: str, checkpoint: dict[str, Any], *, native_root: str | None = None, |
| source_root: Path = ROOT) -> dict[str, Any]: |
| """native_root present = startup verification; absent = offline matching identity.""" |
| selected_arm(arm) |
| if arm == "k2-tp2": |
| return checkpoint |
| cfg = ARMS[arm] |
| cfg.check() |
| manifest_path = source_root / "native/NATIVE_MANIFEST.json" |
| manifest = json.loads(manifest_path.read_text()) |
| hashes = official_hashes(manifest) |
| if native_root is not None: |
| root = Path(native_root).resolve(strict=True) |
| for n, want in hashes.items(): |
| p = root / n |
| if p.resolve().parent != root or not p.is_file(): |
| raise ValueError(f"{arm}: missing or redirected official shard/config {p}") |
| got = sha256(p) |
| if got != want: |
| raise ValueError(f"{arm}: official {n} sha256 {got} != receipt {want}") |
| recipe = { |
| "schema": "dspark.native-identity.v1", |
| "arm": asdict(cfg), |
| "draft_head": "target-shared EXL3 K6 (not official BF16 head)", |
| "official_repo": manifest["official_repo"], |
| "official_revision": manifest["official_revision"], |
| "official_files": hashes, |
| "official_tensor_manifest_sha256": sha256(manifest_path), |
| |
| |
| "retained_checkpoint_files": {n: h.lower() for n, h in checkpoint["files"].items()}, |
| "runtime_sources": runtime_hashes(source_root), |
| } |
| level = batch_invariant_from_env() |
| if level: |
| if not cfg.kx2_dispatch: |
| raise ValueError(f"{arm}: batch-invariant dispatch is implemented for the kx2 dispatch only") |
| recipe["batch_invariant_dispatch"] = ("nb1-draft-v1: B1 split-K / FP8 / expert variant at every B<=4" + |
| ("; insert projections in 8-row sub-batches" if level == 2 else "")) |
| return {"drafter_identity": digest(recipe), "markov_identity": checkpoint["markov_identity"], |
| "arm": arm, "draft_head": recipe["draft_head"], "recipe": recipe} |
|
|
