File size: 7,374 Bytes
1b4ed07 | 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 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 | # Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
"""Trusted Hub identity and local verification of the FP32 model bundle."""
from __future__ import annotations
import json
import os
from pathlib import Path
from ax_doclayout import _hub_pin
from ax_doclayout.artifact import (
MODEL_ID,
STAGED_ONNX_SHA256,
STAGED_SAFETENSORS_SHA256,
ArtifactReceipt,
_bundle_files,
_open_regular_file,
build_artifact_receipt,
verify_artifact_bundle,
verify_artifact_manifest,
)
from ax_doclayout.contract import HUB_COMMIT_RE
from ax_doclayout.errors import (
ArtifactError,
ArtifactUnavailableError,
BackendUnavailableError,
InvalidRevisionError,
)
PINNED_HUB_RECEIPT = _hub_pin.PINNED_HUB_RECEIPT
_MAX_SAFETENSORS_HEADER = 8 * 1024 * 1024
# FP32 weights plus the integer batch counters and the one bool mask that
# PyTorch stores beside them. BF16, FP16, and FP64 are rejected.
_ALLOWED_SAFETENSORS_DTYPES = frozenset({"F32", "I64", "BOOL"})
def safetensors_dtypes(path: Path) -> tuple[str, ...]:
"""Return tensor dtypes from a safetensors header, without reading weights."""
descriptor = _open_regular_file(path)
try:
with os.fdopen(descriptor, "rb") as handle:
prefix = handle.read(8)
if len(prefix) != 8:
raise ArtifactError("staging_safetensors_header_invalid")
length = int.from_bytes(prefix, "little")
if length <= 0 or length > _MAX_SAFETENSORS_HEADER:
raise ArtifactError("staging_safetensors_header_invalid")
raw = handle.read(length)
except OSError as exc:
raise ArtifactError("staging_safetensors_unreadable") from exc
if len(raw) != length:
raise ArtifactError("staging_safetensors_header_invalid")
try:
header = json.loads(raw.decode("utf-8"))
except (UnicodeError, ValueError) as exc:
raise ArtifactError("staging_safetensors_header_invalid") from exc
if not isinstance(header, dict):
raise ArtifactError("staging_safetensors_header_invalid")
dtypes: list[str] = []
for key, value in header.items():
if not isinstance(key, str):
raise ArtifactError("staging_safetensors_header_invalid")
if key.startswith("__"):
continue
if not isinstance(value, dict) or not isinstance(value.get("dtype"), str):
raise ArtifactError("staging_safetensors_header_invalid")
dtypes.append(value["dtype"])
if not dtypes:
raise ArtifactError("staging_safetensors_empty")
return tuple(dtypes)
def _require_fp32_weights(path: Path) -> None:
dtypes = set(safetensors_dtypes(path))
if "F32" not in dtypes or not dtypes <= _ALLOWED_SAFETENSORS_DTYPES:
raise ArtifactError("staging_safetensors_not_fp32")
def verify_staged_v01_bundle(
directory: Path,
receipt_path: Path,
*,
require_staged_digests: bool = True,
) -> None:
"""Fail closed unless a nine-file bundle is FP32 and still unpublished."""
from ax_doclayout._local_pytorch import read_publishable
if read_publishable(receipt_path):
raise ArtifactError("staging_publishable_must_stay_false")
verify_artifact_manifest(directory)
observed = _bundle_files(directory)
by_name = {item.name: item.sha256 for item in observed}
if require_staged_digests and (
by_name["model.safetensors"] != STAGED_SAFETENSORS_SHA256
or by_name["model.onnx"] != STAGED_ONNX_SHA256
):
raise ArtifactError("staging_weight_digest_mismatch")
_require_fp32_weights(directory / "model.safetensors")
_require_receipt_agreement(receipt_path, by_name)
def _require_receipt_agreement(receipt_path: Path, by_name: dict[str, str]) -> None:
try:
payload = json.loads(receipt_path.read_text(encoding="utf-8"))
except (OSError, UnicodeError, ValueError) as exc:
raise ArtifactError("staging_receipt_unreadable") from exc
if not isinstance(payload, dict):
raise ArtifactError("staging_receipt_unreadable")
model_id = payload.get("model_id", MODEL_ID)
if model_id != MODEL_ID:
raise ArtifactError("artifact_model_id_invalid")
files = payload.get("artifact_files")
if files is None:
return
if not isinstance(files, list):
raise ArtifactError("staging_receipt_unreadable")
for item in files:
if not isinstance(item, dict):
raise ArtifactError("staging_receipt_unreadable")
name = item.get("name")
digest = item.get("sha256")
if name in by_name and digest != by_name[name]:
raise ArtifactError("staging_receipt_digest_mismatch")
def qualify_published_bundle(
*,
model_id: str,
revision: str,
backend: str,
receipt: ArtifactReceipt | None,
bundle: Path | None,
) -> Path | None:
"""Verify one pinned receipt against local bytes.
``None`` means no production receipt is pinned. A pinned receipt with a
missing or mismatched bundle raises. This function does not download and
does not run a backend.
"""
if receipt is None:
return None
if backend not in ("pytorch", "onnx"):
raise BackendUnavailableError(
"backend must be given explicitly as one of pytorch, onnx"
)
if model_id != MODEL_ID or model_id != receipt.model_id:
raise ArtifactError("artifact_model_id_invalid")
if revision != receipt.model_revision:
raise ArtifactError("artifact_revision_invalid")
if bundle is None:
raise ArtifactUnavailableError(
"published bundle is not available for hash verification"
)
verify_artifact_bundle(bundle, receipt)
_require_fp32_weights(bundle / "model.safetensors")
return bundle
SIGN_OFF_FIELDS = ("adr_0031", "license_pack", "model_card")
def require_publication_signoff(path: Path) -> None:
"""Refuse publication until ADR-0031 and the license pack are approved."""
try:
payload = json.loads(path.read_text(encoding="utf-8"))
except (OSError, UnicodeError, ValueError) as exc:
raise ArtifactError("publication_signoff_missing") from exc
if not isinstance(payload, dict):
raise ArtifactError("publication_signoff_missing")
for field in SIGN_OFF_FIELDS:
if payload.get(field) != "approved":
raise ArtifactError("publication_signoff_incomplete")
def publish_redownloaded_bundle(
*,
bundle: Path,
staging_receipt: Path,
signoff: Path,
revision: str,
downloaded: Path,
require_staged_digests: bool = True,
) -> ArtifactReceipt:
"""Build a receipt only after a redownload matches the staged nine files."""
require_publication_signoff(signoff)
verify_staged_v01_bundle(
bundle,
staging_receipt,
require_staged_digests=require_staged_digests,
)
if not isinstance(revision, str) or HUB_COMMIT_RE.fullmatch(revision) is None:
raise InvalidRevisionError(
"revision must be a 40-character immutable Hub commit"
)
if _bundle_files(bundle) != _bundle_files(downloaded):
raise ArtifactError("hub_redownload_digest_mismatch")
return build_artifact_receipt(downloaded, model_revision=revision)
|