"""Build and verify content-addressed manifests for deployable ONNX bundles. The exporter writes a small legacy manifest with semantic metadata, but that manifest does not prove which bytes a browser will load. This module treats the ONNX graph, its external-data references, tokenizer files, and explicitly selected runtime assets as a single bundle contract. Source checkpoints and other files left in an experiment directory are inventoried but excluded from the deployable byte total unless explicitly selected as a runtime asset. """ from __future__ import annotations import hashlib import json import platform import subprocess import sys from collections.abc import Iterable, Mapping, Sequence from pathlib import Path, PurePosixPath, PureWindowsPath from typing import Any SCHEMA_VERSION = "vons.bundle.manifest/v1" MODEL_ASSET_ROLES = frozenset({"model_graph", "model_external_data", "tokenizer", "calibration"}) DISTRIBUTION_ROLES = MODEL_ASSET_ROLES | frozenset({"runtime_asset", "license_notice", "config"}) class BundleManifestError(ValueError): """Raised when a bundle cannot be safely enumerated or verified.""" def _require_onnx() -> Any: try: import onnx # type: ignore[import-not-found] except ImportError as exc: # pragma: no cover - exercised in minimal installs raise RuntimeError("bundle manifest generation requires the optional 'onnx' package") from exc return onnx def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _json_bytes(value: Mapping[str, Any]) -> bytes: return (json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False) + "\n").encode("utf-8") def _root_path(root: Path) -> Path: resolved = root.resolve(strict=True) if not resolved.is_dir(): raise BundleManifestError(f"bundle root is not a directory: {root}") return resolved def _safe_relative(root: Path, value: str | Path, *, label: str) -> Path: """Resolve a bundle-relative path and reject traversal and symlinks. ONNX external-data locations are specification strings, not arbitrary local paths. Rejecting both POSIX and Windows absolute forms avoids accepting a Windows drive path when a manifest is generated on another platform. """ raw = str(value).replace("\\", "/") posix = PurePosixPath(raw) windows = PureWindowsPath(raw) if posix.is_absolute() or windows.is_absolute() or windows.drive: raise BundleManifestError(f"{label} must be relative to the bundle: {value!r}") if any(part == ".." for part in posix.parts): raise BundleManifestError(f"{label} contains traversal: {value!r}") raw_candidate = root / Path(*posix.parts) if raw_candidate.is_symlink(): raise BundleManifestError(f"{label} may not be a symlink: {value!r}") candidate = raw_candidate.resolve(strict=False) try: candidate.relative_to(root) except ValueError as exc: raise BundleManifestError(f"{label} escapes the bundle root: {value!r}") from exc return candidate def _bundle_input(root: Path, value: str | Path, *, label: str) -> Path: """Accept a caller's absolute path only when it is inside the bundle.""" candidate = Path(value) if candidate.is_absolute() or PureWindowsPath(str(value)).drive: try: relative = candidate.resolve(strict=False).relative_to(root) except ValueError as exc: raise BundleManifestError(f"{label} is outside the bundle root: {value!r}") from exc return _safe_relative(root, relative, label=label) caller_relative = (Path.cwd() / candidate).resolve(strict=False) if caller_relative.exists(): try: relative = caller_relative.relative_to(root) except ValueError as exc: raise BundleManifestError(f"{label} is outside the bundle root: {value!r}") from exc return _safe_relative(root, relative, label=label) return _safe_relative(root, value, label=label) def _relative(root: Path, path: Path) -> str: try: return path.resolve(strict=True).relative_to(root).as_posix() except ValueError as exc: raise BundleManifestError(f"path is outside bundle root: {path}") from exc def _regular_files(root: Path, directory: Path) -> list[Path]: if directory.is_symlink(): raise BundleManifestError(f"bundle directory may not be a symlink: {directory}") if not directory.is_dir(): raise BundleManifestError(f"expected directory: {directory}") files: list[Path] = [] for path in sorted(directory.rglob("*")): if path.is_symlink(): raise BundleManifestError(f"bundle file may not be a symlink: {path}") if path.is_file(): _safe_relative(root, path.relative_to(root), label="bundle file") files.append(path) return files def _external_locations(model: Any, onnx: Any) -> set[str]: """Return external-data locations from initializers and nested subgraphs.""" locations: set[str] = set() def visit_tensor(tensor: Any) -> None: for item in tensor.external_data: if item.key == "location": if not item.value: raise BundleManifestError("ONNX external_data has an empty location") locations.add(item.value) def visit_graph(graph: Any) -> None: for tensor in graph.initializer: visit_tensor(tensor) for sparse in getattr(graph, "sparse_initializer", []): visit_tensor(sparse.values) visit_tensor(sparse.indices) for node in graph.node: for attribute in node.attribute: if attribute.type == onnx.AttributeProto.GRAPH: visit_graph(attribute.g) elif attribute.type == onnx.AttributeProto.GRAPHS: for nested in attribute.graphs: visit_graph(nested) elif attribute.type == onnx.AttributeProto.TENSOR: visit_tensor(attribute.t) elif attribute.type == onnx.AttributeProto.TENSORS: for tensor in attribute.tensors: visit_tensor(tensor) elif attribute.type == onnx.AttributeProto.SPARSE_TENSOR: visit_tensor(attribute.sparse_tensor.values) visit_tensor(attribute.sparse_tensor.indices) elif attribute.type == onnx.AttributeProto.SPARSE_TENSORS: for sparse in attribute.sparse_tensors: visit_tensor(sparse.values) visit_tensor(sparse.indices) visit_graph(model.graph) return locations def _shape(value_info: Any) -> list[int | str | None] | None: tensor_type = value_info.type.tensor_type if not tensor_type.HasField("shape"): return None result: list[int | str | None] = [] for dimension in tensor_type.shape.dim: if dimension.HasField("dim_value"): result.append(int(dimension.dim_value)) elif dimension.HasField("dim_param"): result.append(dimension.dim_param) else: result.append(None) return result def _value_infos(values: Iterable[Any], onnx: Any) -> list[dict[str, Any]]: result = [] for value in values: tensor_type = value.type.tensor_type result.append( { "name": value.name, "dtype": onnx.TensorProto.DataType.Name(tensor_type.elem_type), "shape": _shape(value), } ) return result def _graph_metadata(model: Any, onnx: Any, external_locations: Sequence[str]) -> dict[str, Any]: return { "ir_version": int(model.ir_version), "opset_imports": [ {"domain": item.domain, "version": int(item.version)} for item in model.opset_import ], "inputs": _value_infos(model.graph.input, onnx), "outputs": _value_infos(model.graph.output, onnx), "external_data_locations": sorted(external_locations), "node_count": len(model.graph.node), } def _git_revision(root: Path) -> dict[str, str | None]: try: value = subprocess.run( ["git", "-C", str(root), "rev-parse", "HEAD"], capture_output=True, check=True, text=True, timeout=2, ).stdout.strip() except (OSError, subprocess.SubprocessError): value = "" return {"value": value or None, "reason": None if value else "not a git checkout or revision unavailable"} def _snapshot_hash( root: Path | None, paths: Sequence[str | Path] | None = None, ) -> dict[str, Any]: if root is None: return { "sha256": None, "reason": "source root was not supplied", "scope": None, "paths": [], } source_root = root.resolve(strict=True) if not source_root.is_dir(): raise BundleManifestError(f"source root is not a directory: {root}") digest = hashlib.sha256() ignored_directories = {".git", "__pycache__", ".pytest_cache", ".ruff_cache"} files: list[tuple[Path, Path]] = [] if paths: for supplied in paths: candidate = Path(supplied) if not candidate.is_absolute(): candidate = source_root / candidate if candidate.is_symlink(): raise BundleManifestError(f"source snapshot path may not be a symlink: {supplied}") resolved = candidate.resolve(strict=True) try: relative_path = resolved.relative_to(source_root) except ValueError as exc: raise BundleManifestError( f"source snapshot path escapes the source root: {supplied}" ) from exc if not resolved.is_file(): raise BundleManifestError(f"source snapshot path is not a file: {supplied}") if any(part in ignored_directories for part in relative_path.parts) or resolved.suffix == ".pyc": raise BundleManifestError(f"source snapshot path is ignored: {supplied}") files.append((relative_path, resolved)) if len({relative for relative, _ in files}) != len(files): raise BundleManifestError("source snapshot paths must be unique") scope = "explicit" else: for path in source_root.rglob("*"): if not path.is_file() or path.is_symlink() or path.suffix == ".pyc": continue relative_path = path.relative_to(source_root) if any(part in ignored_directories for part in relative_path.parts): continue files.append((relative_path, path)) scope = "all_eligible_files" for relative_path, path in sorted(files): relative = relative_path.as_posix().encode("utf-8") digest.update(len(relative).to_bytes(8, "big")) digest.update(relative) digest.update(path.stat().st_size.to_bytes(8, "big")) with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return { "sha256": digest.hexdigest(), "reason": None, "scope": scope, "paths": [relative.as_posix() for relative, _ in sorted(files)], } def _legacy_metadata(root: Path, output: Path) -> dict[str, Any]: candidate = root / "bundle-manifest.json" if candidate == output or not candidate.is_file(): return {} try: value = json.loads(candidate.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): return {} return value if isinstance(value, dict) else {} def _file_record(root: Path, path: Path, role: str, *, runtime_loaded: bool, model_asset: bool) -> dict[str, Any]: relative = _relative(root, path) size = path.stat().st_size return { "path": relative, "role": role, "bytes": size, "sha256": _sha256(path), "compression": {"encoding": "identity", "measured": False, "bytes": size}, "runtime_loaded": runtime_loaded, "model_asset": model_asset, } def _stable_manifest(manifest: dict[str, Any], output: Path, entries: Sequence[dict[str, Any]]) -> None: """Write the manifest until its reported own size reaches a fixed point.""" payload_bytes = sum(int(item["bytes"]) for item in entries if item["model_asset"]) distribution_bytes = sum(int(item["bytes"]) for item in entries if item["role"] in DISTRIBUTION_ROLES) manifest["summary"]["model_asset_bytes"] = payload_bytes manifest["summary"]["manifest_bytes"] = 0 manifest["summary"]["complete_download_bytes"] = distribution_bytes output.parent.mkdir(parents=True, exist_ok=True) for _ in range(8): output.write_bytes(_json_bytes(manifest)) manifest_bytes = output.stat().st_size complete = distribution_bytes + manifest_bytes if ( manifest["summary"]["manifest_bytes"] == manifest_bytes and manifest["summary"]["complete_download_bytes"] == complete ): return manifest["summary"]["manifest_bytes"] = manifest_bytes manifest["summary"]["complete_download_bytes"] = complete raise BundleManifestError("manifest size did not converge") def build_bundle_manifest( onnx_path: str | Path, tokenizer_path: str | Path, output_path: str | Path, *, checkpoint_path: str | Path | None = None, model_manifest_path: str | Path | None = None, calibration_paths: Sequence[str | Path] = (), runtime_asset_paths: Sequence[str | Path] = (), license_paths: Sequence[str | Path] = (), config_paths: Sequence[str | Path] = (), source_root: str | Path | None = None, source_paths: Sequence[str | Path] = (), ) -> Path: """Create a deployable manifest and inventory every regular file in its root.""" output = Path(output_path).resolve(strict=False) root = _root_path(output.parent) tokenizer = _bundle_input(root, tokenizer_path, label="tokenizer path") graph_path = _bundle_input(root, onnx_path, label="ONNX path") if not graph_path.is_file(): raise BundleManifestError(f"ONNX path is not a file: {onnx_path}") if output.parent != root: raise BundleManifestError("manifest output must be directly inside the bundle root") onnx = _require_onnx() try: model = onnx.load_model(str(graph_path), load_external_data=False) except Exception as exc: # onnx raises several protobuf-specific exceptions raise BundleManifestError(f"could not parse ONNX graph {graph_path}: {exc}") from exc external_locations = _external_locations(model, onnx) external_files: dict[Path, str] = {} for location in sorted(external_locations): path = _safe_relative(root, location, label="ONNX external-data location") if not path.is_file(): raise BundleManifestError(f"ONNX external-data file does not exist: {location}") external_files[path] = location paths_by_role: dict[Path, str] = {graph_path: "model_graph"} if tokenizer.is_file(): paths_by_role[tokenizer] = "tokenizer" else: tokenizer_files = _regular_files(root, tokenizer) if not tokenizer_files: raise BundleManifestError("tokenizer directory contains no regular files") for path in tokenizer_files: paths_by_role[path] = "tokenizer" for path in external_files: paths_by_role[path] = "model_external_data" explicit_groups = [ (checkpoint_path, "source_checkpoint"), (model_manifest_path, "source_manifest"), ] for values, role in ( (calibration_paths, "calibration"), (runtime_asset_paths, "runtime_asset"), (license_paths, "license_notice"), (config_paths, "config"), ): explicit_groups.extend((value, role) for value in values) for value, role in explicit_groups: if value is None: continue path = _bundle_input(root, value, label=f"{role} path") if not path.is_file(): raise BundleManifestError(f"{role} path is not a file: {value}") paths_by_role[path] = role if checkpoint_path is None: candidate = root / "model.pt" if candidate.is_file(): paths_by_role.setdefault(candidate, "source_checkpoint") if model_manifest_path is None: candidate = root / "model-manifest.json" if candidate.is_file(): paths_by_role.setdefault(candidate, "source_manifest") all_files = _regular_files(root, root) records: list[dict[str, Any]] = [] for path in all_files: if path == output: continue role = paths_by_role.get(path) if role is None: if path.suffix == ".onnx": role = "unselected_model_artifact" elif path.suffix == ".pt": role = "source_checkpoint" elif path.name.endswith(".onnx.data"): role = "unreferenced_model_artifact" elif path.name.endswith("manifest.json"): role = "source_manifest" else: role = "extra" records.append( _file_record( root, path, role, runtime_loaded=role in {"model_graph", "model_external_data", "tokenizer", "runtime_asset"}, model_asset=role in MODEL_ASSET_ROLES, ) ) records.sort(key=lambda item: item["path"]) if not any(item["role"] == "model_graph" for item in records): raise BundleManifestError("manifest inventory does not contain the selected ONNX graph") legacy = _legacy_metadata(root, output) metadata = {key: value for key, value in legacy.items() if key not in {"onnx", "tokenizer", "checkpoint"}} metadata["graph"] = _graph_metadata(model, onnx, sorted(external_locations)) metadata["checkpoint_sha256"] = next( (item["sha256"] for item in records if item["role"] == "source_checkpoint"), None ) manifest = { "schema_version": SCHEMA_VERSION, "manifest": { "path": output.name, "sha256": None, "self_hash_excluded": True, }, "bundle_root": ".", "metadata": metadata, "source": { "git_revision": _git_revision(root), "source_snapshot": _snapshot_hash( Path(source_root) if source_root is not None else None, source_paths or None, ), "builder": {"python": sys.version.split()[0], "platform": platform.platform()}, }, "files": records, "summary": { "model_asset_bytes": 0, "manifest_bytes": 0, "complete_download_bytes": 0, "model_asset_limit_bytes": 67_108_864, "within_model_asset_limit": False, "runtime_loaded_file_count": sum(item["runtime_loaded"] for item in records), }, } _stable_manifest(manifest, output, records) manifest["summary"]["within_model_asset_limit"] = ( manifest["summary"]["model_asset_bytes"] <= manifest["summary"]["model_asset_limit_bytes"] ) output.write_bytes(_json_bytes(manifest)) # Updating the boolean can alter the byte count; converge once more without # including the manifest in its own digest or model-asset total. _stable_manifest(manifest, output, records) return output def verify_bundle_manifest(manifest_path: str | Path) -> dict[str, Any]: """Verify every recorded file, graph reference, and byte summary.""" manifest_file = Path(manifest_path).resolve(strict=True) payload = json.loads(manifest_file.read_text(encoding="utf-8")) if payload.get("schema_version") != SCHEMA_VERSION: raise BundleManifestError(f"unsupported bundle manifest schema: {payload.get('schema_version')!r}") root = manifest_file.parent manifest_info = payload.get("manifest", {}) if manifest_info.get("path") != manifest_file.name: raise BundleManifestError("manifest path does not match the verification target") records = payload.get("files") if not isinstance(records, list) or not records: raise BundleManifestError("manifest has no file inventory") seen: set[str] = set() for item in records: if not isinstance(item, dict): raise BundleManifestError("manifest file entry is not an object") relative = item.get("path") if not isinstance(relative, str) or relative in seen: raise BundleManifestError(f"invalid or duplicate manifest path: {relative!r}") seen.add(relative) path = _safe_relative(root, relative, label="manifest file path") if not path.is_file(): raise BundleManifestError(f"manifest file is missing: {relative}") actual_bytes = path.stat().st_size actual_hash = _sha256(path) if actual_bytes != item.get("bytes") or actual_hash != item.get("sha256"): raise BundleManifestError(f"manifest file changed: {relative}") graph_item = next((item for item in records if item.get("role") == "model_graph"), None) if graph_item is None: raise BundleManifestError("manifest has no model_graph entry") onnx = _require_onnx() model = onnx.load_model(str(root / graph_item["path"]), load_external_data=False) external_locations = _external_locations(model, onnx) inventory_paths = {item["path"] for item in records} for location in external_locations: path = _safe_relative(root, location, label="ONNX external-data location") if _relative(root, path) not in inventory_paths: raise BundleManifestError(f"external-data file is not in inventory: {location}") summary = payload.get("summary", {}) model_asset_bytes = sum(item["bytes"] for item in records if item.get("model_asset")) distribution_bytes = sum(item["bytes"] for item in records if item.get("role") in DISTRIBUTION_ROLES) if summary.get("model_asset_bytes") != model_asset_bytes: raise BundleManifestError("model_asset_bytes does not match the inventory") if summary.get("manifest_bytes") != manifest_file.stat().st_size: raise BundleManifestError("manifest_bytes does not match the manifest file") if summary.get("complete_download_bytes") != distribution_bytes + manifest_file.stat().st_size: raise BundleManifestError("complete_download_bytes does not match the inventory") if summary.get("within_model_asset_limit") != ( model_asset_bytes <= int(summary.get("model_asset_limit_bytes", 67_108_864)) ): raise BundleManifestError("within_model_asset_limit is stale") return { "manifest": str(manifest_file), "files": len(records), "external_data_files": len(external_locations), "model_asset_bytes": model_asset_bytes, "complete_download_bytes": distribution_bytes + manifest_file.stat().st_size, "pass": True, } def build_and_verify(*args: Any, **kwargs: Any) -> tuple[Path, dict[str, Any]]: """Build a manifest and immediately run the independent verifier.""" output = build_bundle_manifest(*args, **kwargs) return output, verify_bundle_manifest(output)