| |
| |
| |
| """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 |
| |
| |
| _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) |
|
|
