From 57d9517ad1f1cd5a860385483690ff7e60e341cb Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 29 Sep 2026 14:55:03 -0700 Subject: [PATCH] [CoreAI] Record file sizes and bundle digests in the delegate manifest Summary: List every delivered file with its byte size, and a SHA-256 digest per bundle, in the delegate manifest for both inline and sidecar delivery, so a runtime can check the payloads it reads and key its cache on bundle contents. The digest covers a fixed prefix, the file count, and each file's path, size and bytes in UTF-8 path order. `files` changes from a list of paths (inline only) to a map of path to size. Test Plan: ``` python -m pytest backends/apple/coreai ``` 126 passed (coreai-core 1.0.0b3, coreai-torch 0.4.3). --- backends/apple/coreai/compiler/preprocess.py | 49 +++++++++++++++---- backends/apple/coreai/test/test_preprocess.py | 39 ++++++++++++++- 2 files changed, 77 insertions(+), 11 deletions(-) diff --git a/backends/apple/coreai/compiler/preprocess.py b/backends/apple/coreai/compiler/preprocess.py index 4a914a0d526..dd6902671c6 100644 --- a/backends/apple/coreai/compiler/preprocess.py +++ b/backends/apple/coreai/compiler/preprocess.py @@ -9,6 +9,7 @@ import contextlib import copy +import hashlib import json import logging import os @@ -262,6 +263,35 @@ def _maybe_warn_sidecar_env_ignored() -> None: ) +def _asset_metadata(staging: Path) -> Dict[str, Any]: + """Per-file sizes and a SHA-256 digest per bundle under ``staging``. + + Each bundle digest covers a fixed prefix, the file count, then per file (in + UTF-8 byte order) the path length, path, size and contents. + """ + sizes: Dict[str, int] = {} + digests: Dict[str, str] = {} + for bundle in sorted(p for p in staging.iterdir() if p.is_dir()): + files = sorted( + ( + (path.relative_to(staging).as_posix(), path) + for path in bundle.rglob("*") + if path.is_file() + ), + key=lambda item: item[0].encode("utf-8"), + ) + digest = hashlib.sha256(b"ExecuTorch.CoreAI.raw-assets.sha256.v1\0") + digest.update(len(files).to_bytes(8, "big")) + for rel, path in files: + payload = path.read_bytes() + name = rel.encode("utf-8") + digest.update(len(name).to_bytes(8, "big") + name) + digest.update(len(payload).to_bytes(8, "big") + payload) + sizes[rel] = len(payload) + digests[bundle.name] = digest.hexdigest() + return {"files": sizes, "bundle_digests": digests} + + def _deliver( staging: Path, model_hash: str, @@ -281,7 +311,12 @@ def _deliver( cannot leave a half-populated ``/`` behind for :func:`_reject_existing_asset_dir` to trip over on the next build. """ - manifest = {"packaging": packaging.value, "hash": model_hash, **manifest_extra} + manifest = { + "packaging": packaging.value, + "hash": model_hash, + **manifest_extra, + **_asset_metadata(staging), + } if sidecar_dir is None: return _embed_dir_inline(staging, model_hash, manifest) @@ -352,23 +387,19 @@ def _embed_dir_inline( bundles, never a bundle itself, so the bundle name survives into the keys and the asset reconstructs to the layout the sidecar routes write on disk. - Keys are ``coreai/{hash}/{relpath}``; the manifest lists the relpaths plus - any ``manifest_extra`` (packaging, archs, ...). + Keys are ``coreai/{hash}/{relpath}``; ``manifest_extra`` becomes the + manifest (packaging, files, archs, ...). """ store = NamedDataStore() - files: List[str] = [] for path in sorted(root_dir.rglob("*")): if path.is_file(): - rel = path.relative_to(root_dir).as_posix() - files.append(rel) store.add_named_data( - _nds_key(model_hash, rel), + _nds_key(model_hash, path.relative_to(root_dir).as_posix()), path.read_bytes(), alignment=_ASSET_ALIGNMENT, ) - manifest = {"files": files, **manifest_extra} return PreprocessResult( - processed_bytes=json.dumps(manifest).encode("utf-8"), + processed_bytes=json.dumps(manifest_extra).encode("utf-8"), data_store_output=store.get_named_data_store_output(), ) diff --git a/backends/apple/coreai/test/test_preprocess.py b/backends/apple/coreai/test/test_preprocess.py index 527b1299300..27674b27b64 100644 --- a/backends/apple/coreai/test/test_preprocess.py +++ b/backends/apple/coreai/test/test_preprocess.py @@ -26,6 +26,7 @@ from executorch.backends.apple.coreai.compiler.preprocess import ( _aot_compile_options, + _asset_metadata, AOTCompileConfig, COMPILE_SPEC_KEYS, coreai_sidecar_dir, @@ -253,6 +254,42 @@ def forward(self, *args): self.assertEqual(manifest["output_names"], [f"output_{i}" for i in range(12)]) +class AssetMetadataTest(unittest.TestCase): + @mock.patch(_MOCK_BUILD, side_effect=_fake_run_coreai_build) + def test_sizes_and_digests_cover_every_delivered_file(self, _build): + aot = _aot_spec({"platform": "iOS", "architectures": ["h15g", "h16"]}) + for specs, bundles in ( + ([], {"model.aimodel"}), + ([aot], {"model.h15g.aimodelc", "model.h16.aimodelc"}), + ): + with self.subTest(aot=bool(specs)), tempfile.TemporaryDirectory() as d: + with coreai_sidecar_dir(d): + manifest = json.loads( + CoreAIBackend.preprocess( + _edge_program(), specs + [_SIDECAR_SPEC] + ).processed_bytes + ) + root = Path(d) / manifest["hash"] + on_disk = { + p.relative_to(root).as_posix(): p.stat().st_size + for p in root.rglob("*") + if p.is_file() + } + self.assertEqual(manifest["files"], on_disk) + self.assertEqual(set(manifest["bundle_digests"]), bundles) + + def test_digest_changes_with_file_contents(self): + with tempfile.TemporaryDirectory() as d: + weights = Path(d) / "model.aimodel" / "weights.bin" + weights.parent.mkdir() + weights.write_bytes(b"\x00\x01") + before = _asset_metadata(Path(d)) + weights.write_bytes(b"\x01\x00") + after = _asset_metadata(Path(d)) + self.assertEqual(before["files"], after["files"]) + self.assertNotEqual(before["bundle_digests"], after["bundle_digests"]) + + def _lower_with_break(partitioner): """Lower a model whose middle op is tagged, forcing two delegates.""" from executorch.backends.apple.coreai import get_default_passes @@ -542,13 +579,11 @@ def test_aot_sidecar_writes_aimodelc_and_keeps_pte_small(self): ) ) lbm, manifest = _lowered_manifest(lowered) - pte = bytes(lowered.to_executorch().buffer) self.assertEqual(manifest["packaging"], "aot_compiled_sidecar") self.assertGreaterEqual(len(manifest["archs"]), 1) # Sidecar: compiled bundle *contents* must not be embedded in the .pte. self.assertIsNone(lbm.named_data_store_output) - self.assertNotIn(b"mpsExecutable", pte) for rel in manifest["archs"].values(): bundle = os.path.join(sidecar, rel) self.assertTrue(os.path.isdir(bundle), bundle)