Skip to content

Commit 4c3d3bd

Browse files
authored
[CoreAI] Record file sizes and bundle digests in the delegate manifest (pytorch#23316)
The Core AI runtime checks the payloads it reads against the manifest and keys its asset cache on bundle contents. This PR records both in the delegate manifest. * Asset metadata (`compiler/preprocess.py`): the new `_asset_metadata` walks the staged bundles (`model.aimodel` or `model.<arch>.aimodelc`) and returns: * `files`: each file's path relative to the asset root, mapped to its byte size. This replaces the previous inline-only list of paths. * `bundle_digests`: one SHA-256 per bundle, over a fixed prefix, the file count, and each file's path length, path, size and bytes in UTF-8 path order. * `_deliver` merges this metadata into the manifest before choosing a delivery, so every delivery reports identical metadata. * Tests (`test_preprocess.py`): * `AssetMetadataTest` checks that `files` exactly covers the delivered files and sizes, with one digest per bundle, for source and AOT. * A second test checks that changing a file's bytes changes its digest when the size stays the same. * The real-toolchain AOT sidecar test no longer greps the `.pte` for a compiled filename, which now appears in the manifest by design. Its `named_data_store_output is None` check still covers "contents not embedded". Test plan: `python -m pytest backends/apple/coreai` gives 125 passed (coreai-core 1.0.0b3, coreai-torch 0.4.3).
1 parent b908580 commit 4c3d3bd

2 files changed

Lines changed: 77 additions & 11 deletions

File tree

‎backends/apple/coreai/compiler/preprocess.py‎

Lines changed: 40 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
import contextlib
1111
import copy
12+
import hashlib
1213
import json
1314
import logging
1415
import os
@@ -262,6 +263,35 @@ def _maybe_warn_sidecar_env_ignored() -> None:
262263
)
263264

264265

266+
def _asset_metadata(staging: Path) -> Dict[str, Any]:
267+
"""Per-file sizes and a SHA-256 digest per bundle under ``staging``.
268+
269+
Each bundle digest covers a fixed prefix, the file count, then per file (in
270+
UTF-8 byte order) the path length, path, size and contents.
271+
"""
272+
sizes: Dict[str, int] = {}
273+
digests: Dict[str, str] = {}
274+
for bundle in sorted(p for p in staging.iterdir() if p.is_dir()):
275+
files = sorted(
276+
(
277+
(path.relative_to(staging).as_posix(), path)
278+
for path in bundle.rglob("*")
279+
if path.is_file()
280+
),
281+
key=lambda item: item[0].encode("utf-8"),
282+
)
283+
digest = hashlib.sha256(b"ExecuTorch.CoreAI.raw-assets.sha256.v1\0")
284+
digest.update(len(files).to_bytes(8, "big"))
285+
for rel, path in files:
286+
payload = path.read_bytes()
287+
name = rel.encode("utf-8")
288+
digest.update(len(name).to_bytes(8, "big") + name)
289+
digest.update(len(payload).to_bytes(8, "big") + payload)
290+
sizes[rel] = len(payload)
291+
digests[bundle.name] = digest.hexdigest()
292+
return {"files": sizes, "bundle_digests": digests}
293+
294+
265295
def _deliver(
266296
staging: Path,
267297
model_hash: str,
@@ -281,7 +311,12 @@ def _deliver(
281311
cannot leave a half-populated ``<hash>/`` behind for
282312
:func:`_reject_existing_asset_dir` to trip over on the next build.
283313
"""
284-
manifest = {"packaging": packaging.value, "hash": model_hash, **manifest_extra}
314+
manifest = {
315+
"packaging": packaging.value,
316+
"hash": model_hash,
317+
**manifest_extra,
318+
**_asset_metadata(staging),
319+
}
285320
if sidecar_dir is None:
286321
return _embed_dir_inline(staging, model_hash, manifest)
287322

@@ -352,23 +387,19 @@ def _embed_dir_inline(
352387
bundles, never a bundle itself, so the bundle name survives into the keys
353388
and the asset reconstructs to the layout the sidecar routes write on disk.
354389
355-
Keys are ``coreai/{hash}/{relpath}``; the manifest lists the relpaths plus
356-
any ``manifest_extra`` (packaging, archs, ...).
390+
Keys are ``coreai/{hash}/{relpath}``; ``manifest_extra`` becomes the
391+
manifest (packaging, files, archs, ...).
357392
"""
358393
store = NamedDataStore()
359-
files: List[str] = []
360394
for path in sorted(root_dir.rglob("*")):
361395
if path.is_file():
362-
rel = path.relative_to(root_dir).as_posix()
363-
files.append(rel)
364396
store.add_named_data(
365-
_nds_key(model_hash, rel),
397+
_nds_key(model_hash, path.relative_to(root_dir).as_posix()),
366398
path.read_bytes(),
367399
alignment=_ASSET_ALIGNMENT,
368400
)
369-
manifest = {"files": files, **manifest_extra}
370401
return PreprocessResult(
371-
processed_bytes=json.dumps(manifest).encode("utf-8"),
402+
processed_bytes=json.dumps(manifest_extra).encode("utf-8"),
372403
data_store_output=store.get_named_data_store_output(),
373404
)
374405

‎backends/apple/coreai/test/test_preprocess.py‎

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626

2727
from executorch.backends.apple.coreai.compiler.preprocess import (
2828
_aot_compile_options,
29+
_asset_metadata,
2930
AOTCompileConfig,
3031
COMPILE_SPEC_KEYS,
3132
coreai_sidecar_dir,
@@ -253,6 +254,42 @@ def forward(self, *args):
253254
self.assertEqual(manifest["output_names"], [f"output_{i}" for i in range(12)])
254255

255256

257+
class AssetMetadataTest(unittest.TestCase):
258+
@mock.patch(_MOCK_BUILD, side_effect=_fake_run_coreai_build)
259+
def test_sizes_and_digests_cover_every_delivered_file(self, _build):
260+
aot = _aot_spec({"platform": "iOS", "architectures": ["h15g", "h16"]})
261+
for specs, bundles in (
262+
([], {"model.aimodel"}),
263+
([aot], {"model.h15g.aimodelc", "model.h16.aimodelc"}),
264+
):
265+
with self.subTest(aot=bool(specs)), tempfile.TemporaryDirectory() as d:
266+
with coreai_sidecar_dir(d):
267+
manifest = json.loads(
268+
CoreAIBackend.preprocess(
269+
_edge_program(), specs + [_SIDECAR_SPEC]
270+
).processed_bytes
271+
)
272+
root = Path(d) / manifest["hash"]
273+
on_disk = {
274+
p.relative_to(root).as_posix(): p.stat().st_size
275+
for p in root.rglob("*")
276+
if p.is_file()
277+
}
278+
self.assertEqual(manifest["files"], on_disk)
279+
self.assertEqual(set(manifest["bundle_digests"]), bundles)
280+
281+
def test_digest_changes_with_file_contents(self):
282+
with tempfile.TemporaryDirectory() as d:
283+
weights = Path(d) / "model.aimodel" / "weights.bin"
284+
weights.parent.mkdir()
285+
weights.write_bytes(b"\x00\x01")
286+
before = _asset_metadata(Path(d))
287+
weights.write_bytes(b"\x01\x00")
288+
after = _asset_metadata(Path(d))
289+
self.assertEqual(before["files"], after["files"])
290+
self.assertNotEqual(before["bundle_digests"], after["bundle_digests"])
291+
292+
256293
def _lower_with_break(partitioner):
257294
"""Lower a model whose middle op is tagged, forcing two delegates."""
258295
from executorch.backends.apple.coreai import get_default_passes
@@ -542,13 +579,11 @@ def test_aot_sidecar_writes_aimodelc_and_keeps_pte_small(self):
542579
)
543580
)
544581
lbm, manifest = _lowered_manifest(lowered)
545-
pte = bytes(lowered.to_executorch().buffer)
546582

547583
self.assertEqual(manifest["packaging"], "aot_compiled_sidecar")
548584
self.assertGreaterEqual(len(manifest["archs"]), 1)
549585
# Sidecar: compiled bundle *contents* must not be embedded in the .pte.
550586
self.assertIsNone(lbm.named_data_store_output)
551-
self.assertNotIn(b"mpsExecutable", pte)
552587
for rel in manifest["archs"].values():
553588
bundle = os.path.join(sidecar, rel)
554589
self.assertTrue(os.path.isdir(bundle), bundle)

0 commit comments

Comments
 (0)