inference/src/ax_doclayout/hub_release.py

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