mlx/src/solomon_mlx/artifacts.py

"""Artifact verification and distinct MLX runtime identities."""

import hashlib
import json
from importlib.metadata import version
from pathlib import Path

SOLOMON_REVISION = "5c0a4a82ddaeca6da2e3013f7045a8196c86957d"
BASE_REVISION = "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0"
ADAPTER_SHA = "2addaf841ecc58829ad73081603b8d3e52743c53c6d558a17a1dd30e9bd2dbca"
HEADS_SHA = "126a9b5487dca937a768a4f228f2d2e7d513900d8ad0e99c4137fbaa42a1aa6f"


def sha256(path):
    h = hashlib.sha256()
    with Path(path).open("rb") as f:
        for block in iter(lambda: f.read(8 << 20), b""):
            h.update(block)
    return h.hexdigest()


def digest(value):
    return hashlib.sha256(
        json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()
    ).hexdigest()


def verify_release(root):
    root = Path(root)
    manifest = json.loads((root / "MANIFEST.json").read_text())
    checked = {}
    for row in manifest["files"]:
        rel = row["destination"]
        path = (root / rel).resolve()
        if not path.is_relative_to(root.resolve()):
            raise ValueError("Manifest path escapes source directory")
        expected = row.get("staged_sha256")
        size = row.get("staged_bytes")
        actual = sha256(path)
        if expected and actual != expected:
            raise ValueError(f"Source checksum mismatch: {rel}")
        if size is not None and path.stat().st_size != size:
            raise ValueError(f"Source size mismatch: {rel}")
        checked[rel] = actual
    return checked


def code_identity():
    root = Path(__file__).parent
    return digest({str(p.relative_to(root)): sha256(p) for p in sorted(root.rglob("*.py"))})


def runtime_identity(binding, *, chunk_size=2048, max_tokens=40960):
    import platform

    import mlx.core as mx

    value = {
        "backend": "mlx-metal",
        "contract": "solomon-mlx-v1",
        "source_contract": "solomon-v1",
        "profile": binding["profile"],
        "chunk_size": chunk_size,
        "max_tokens": max_tokens,
        "metal_device": mx.device_info(),
        "macos_version": platform.mac_ver()[0],
        "model_binding": digest(binding),
        "code_sha256": code_identity(),
        "versions": {p: version(p) for p in ("mlx", "mlx-vlm", "transformers", "numpy", "pillow")},
        "placement": "question",
        "lora_scale": 2.0,
        "answer_projection": "trained-semantic-head-float32",
        "recurrence": "mlx-float32",
        "solomon_revision": SOLOMON_REVISION,
        "base_revision": BASE_REVISION,
    }
    return {**value, "fingerprint": digest(value)}