# 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)