From 84780b831c9653c1bd66cd534da25def71df018b Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Sun, 4 Oct 2026 19:49:35 -0700 Subject: [PATCH] [CoreAI] Prepare source bundles Summary: Materialize the selected Core AI source bundle on disk under the assets root, so the SDK can specialize it. - Source preparation (`runtime/coreai_assets.{h,mm}`): read inline bundle files through `NamedDataMap` and write them unchanged under the assets root, keyed by the export-time bundle digest (no runtime payload hashing). Reuse checks the exact file set, sizes, entry types and directory structure. Incomplete or malformed existing sources fail without being replaced, since an SDK model may still be using them. Same-size content changes are not detected. - Cold publication stages into `.staging-` and publishes with `renameatx_np(RENAME_EXCL)`, with inspection, staging and publication all relative to one open root descriptor. If another loader publishes first, its bundle is validated and used. - Tests: adds the filesystem snapshot helpers used from here on. - SDK acquisition is not included here. - README: the "Asset storage" materialization, publication and completeness text. No changes outside `backends/apple/coreai`. Test Plan: ``` # Local: configure and build only cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_OSX_DEPLOYMENT_TARGET=27.0 -DEXECUTORCH_BUILD_COREAI=ON \ -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON cmake --build build --target backends/apple/coreai/all # CI (macOS 27 runner, .github/workflows/coreai.yml) ctest --test-dir build/backends/apple/coreai --output-on-failure --no-tests=error ``` CTest runs `coreai_host_test`, which adds the source suite. ExclusivePublicationValidatesWinnerAndRemovesLosingStaging covers the concurrent `RENAME_EXCL` race. --- backends/apple/coreai/CMakeLists.txt | 2 + backends/apple/coreai/README.md | 23 ++ backends/apple/coreai/runtime/coreai_assets.h | 8 + .../apple/coreai/runtime/coreai_assets.mm | 242 +++++++++++ .../runtime/test/coreai_filesystem_fixture.h | 7 + .../runtime/test/coreai_filesystem_fixture.mm | 76 ++++ .../runtime/test/coreai_source_fixture.h | 53 +++ .../runtime/test/coreai_source_fixture.mm | 94 +++++ .../coreai/runtime/test/coreai_source_test.mm | 385 ++++++++++++++++++ 9 files changed, 890 insertions(+) create mode 100644 backends/apple/coreai/runtime/test/coreai_source_fixture.h create mode 100644 backends/apple/coreai/runtime/test/coreai_source_fixture.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_source_test.mm diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index c8b05189631..baaae9a50a9 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -37,7 +37,9 @@ if(EXECUTORCH_BUILD_TESTS) runtime/test/coreai_manifest_test.mm runtime/test/coreai_storage_test.mm runtime/test/coreai_filesystem_fixture.mm + runtime/test/coreai_source_test.mm runtime/test/coreai_bookmark_test.mm + runtime/test/coreai_source_fixture.mm runtime/test/coreai_bookmark_fixture.mm ${_coreai_runtime_sources} ) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 09534810ced..a53e6eb1709 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -51,6 +51,29 @@ processes of the same user rebinding paths. Written files and changed directorie are synced with `fsync`; atomic publication also uses `F_FULLFSYNC` before its rename. +Inline bundles are read through `NamedDataMap` and materialized as unchanged +files under the assets root. Named data can be supplied externally; inline +packaging does not require all bytes to reside in one physical PTE. The exporter +computes each bundle digest from all delivered relative filenames and file bytes, +not just the SDK's bytecode-only `main.hash`. Each AOT architecture has an +independent digest; only the selected architecture is materialized. NDS key names +and raw asset bytes are unchanged. + +When source recovery is needed, existing bundles are checked for the expected +file set, sizes, entry types and directory structure without reading asset +contents or requesting NDS payloads. +Missing bundles are written to a staging directory and published with an +exclusive rename; each selected NDS payload is fetched once and checked for size +before writing. If another loader publishes first, its bundle is validated and +used. Incomplete, size-mismatched or otherwise malformed existing bundles fail +loading without replacement, since an SDK model may still be using them. + +These are completeness checks, not full content-integrity verification. Same-size +content changes are not detected, and a complete stored bundle does not cause its +NDS source to be reread. The export-time digest is a trusted artifact identifier; +the runtime does not rehash source or stored bytes. Core AI may reject invalid +contents, but SDK load failures do not automatically evict the extracted copy. + Preparing the assets root sets `NSURLIsExcludedFromBackupKey` on it, which covers everything beneath it. Ancestors are not modified. This is backup exclusion, not a control for iCloud Drive synchronization. diff --git a/backends/apple/coreai/runtime/coreai_assets.h b/backends/apple/coreai/runtime/coreai_assets.h index ebf84a3bf01..4a60dc411ee 100644 --- a/backends/apple/coreai/runtime/coreai_assets.h +++ b/backends/apple/coreai/runtime/coreai_assets.h @@ -9,6 +9,7 @@ #pragma once #import +#include #include namespace executorch::backends::coreai { @@ -33,5 +34,12 @@ runtime::Result select_assets( const Manifest& manifest, NSString* device_architecture, NSString* platform); +// Caller holds the key's disk lock and has prepared staging_root. +// Inline bundles are atomically published at staging_root/key/bundle. +runtime::Result prepare_source_bundle( + const Manifest& manifest, + const runtime::NamedDataMap* named_data, + NSString* staging_root, + NSString* key); } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_assets.mm b/backends/apple/coreai/runtime/coreai_assets.mm index d28cb8e1d55..f2b547b9217 100644 --- a/backends/apple/coreai/runtime/coreai_assets.mm +++ b/backends/apple/coreai/runtime/coreai_assets.mm @@ -7,14 +7,29 @@ */ #import "coreai_assets.h" +#include "coreai_file.h" + +#include +#include +#include +#include +#include #include +#include #include +#import "coreai_storage.h" namespace executorch::backends::coreai { namespace { using runtime::Error; using runtime::Result; +Error path_error(int error) { + return error == ENOENT || error == ENOTDIR || error == ELOOP + ? Error::InvalidExternalData + : Error::AccessFailed; +} + bool nonempty_string(id value) { return [value isKindOfClass:NSString.class] && [value length] > 0 && [value rangeOfString:@"\0"].location == NSNotFound; @@ -121,6 +136,111 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { return numbers[0] > 0; } +Error validate_directory(int fd, NSString* relative, + NSDictionary* files, + NSSet* expected_directories, + NSMutableSet* actual_files) { + auto children = storage_children(fd); + if (!children.ok()) { + return children.error(); + } + for (NSString* name in children.get()) { + @autoreleasepool { + NSString* path = [relative stringByAppendingPathComponent:name]; + struct stat info; + if (retry_eintr([&] { + return fstatat(fd, name.fileSystemRepresentation, &info, + AT_SYMLINK_NOFOLLOW); + }) != 0) { + return path_error(errno); + } + const bool directory = S_ISDIR(info.st_mode); + ET_CHECK_OR_RETURN_ERROR( + (directory && [expected_directories containsObject:path]) || + (S_ISREG(info.st_mode) && files[path] != nil), + InvalidExternalData, "Unexpected Core AI bundle entry: %s", + path.UTF8String); + if (directory) { + FileDescriptor child(retry_eintr([&] { + return openat(fd, name.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + if (child.get() < 0) { + return path_error(errno); + } + ET_CHECK_OK_OR_RETURN_ERROR(validate_directory( + child.get(), path, files, expected_directories, actual_files)); + } else { + ET_CHECK_OR_RETURN_ERROR( + static_cast(info.st_size) == + files[path].unsignedLongLongValue, + InvalidExternalData, "Core AI asset size does not match manifest"); + [actual_files addObject:path]; + } + } + } + return Error::Ok; +} + +// Checks the exact file set, entry types and sizes. Contents are not hashed. +Error validate_tree(int fd, NSDictionary* files) { + NSMutableSet* expected_directories = [NSMutableSet set]; + for (NSString* file in files) { + NSString* directory = file.stringByDeletingLastPathComponent; + while (directory.length > 0) { + [expected_directories addObject:directory]; + directory = directory.stringByDeletingLastPathComponent; + } + } + NSMutableSet* actual_files = [NSMutableSet set]; + ET_CHECK_OK_OR_RETURN_ERROR(validate_directory( + fd, @"", files, expected_directories, actual_files)); + ET_CHECK_OR_RETURN_ERROR( + [actual_files isEqualToSet:[NSSet setWithArray:files.allKeys]], + InvalidExternalData, "Core AI bundle file set does not match manifest"); + return Error::Ok; +} + +enum class EntryState { Missing, Valid, Corrupt }; + +Result inspect_entry( + int root_fd, + NSString* name, + NSDictionary* files) { + FileDescriptor entry(retry_eintr([&] { + return openat(root_fd, name.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + if (entry.get() < 0) { + if (errno == ENOENT) return EntryState::Missing; + ET_CHECK_OR_RETURN_ERROR(errno == ENOTDIR || errno == ELOOP, AccessFailed, + "Cannot inspect Core AI stored bundle"); + return EntryState::Corrupt; + } + const auto error = validate_tree(entry.get(), files); + if (error == Error::Ok) { + return EntryState::Valid; + } + if (error == Error::InvalidExternalData) { + return EntryState::Corrupt; + } + return error; +} + +struct StagingDirectory { + NSURL* url; + ~StagingDirectory() { + if (url != nil) { + NSError* error = nil; + if (![NSFileManager.defaultManager removeItemAtURL:url error:&error]) { + ET_LOG( + Error, + "Cannot remove Core AI staging directory: %s", + error.localizedDescription.UTF8String); + } + } + } +}; } // namespace Result parse_manifest(NSData* data) { @@ -293,4 +413,126 @@ bool parse_version(id value, NSOperatingSystemVersion& version) { return selected; } +Result prepare_source_bundle( + const Manifest& manifest, + const runtime::NamedDataMap* named_data, + NSString* staging_root, + NSString* key) { + ET_CHECK_OR_RETURN_ERROR( + relative_path(manifest.path), + InvalidProgram, + "Core AI bundle must be selected before preparing its source"); + ET_CHECK_OR_RETURN_ERROR( + nonempty_string(key) && key.length == 64 && + [key rangeOfCharacterFromSet: + [[NSCharacterSet characterSetWithCharactersInString: + @"0123456789abcdef"] invertedSet]] + .location == NSNotFound, + InvalidArgument, + "Core AI staging key must be a 64-character lowercase hex digest"); + auto prepared = prepare_storage_root(staging_root, false); + if (!prepared.ok()) { + return prepared.error(); + } + NSURL* root = [NSURL fileURLWithPath:prepared.get() isDirectory:YES]; + ET_CHECK_OK_OR_RETURN_ERROR(validate_files( + manifest.files, [NSSet setWithObject:manifest.path.lastPathComponent])); + NSDictionary* files = manifest.files; + NSURL* model_url = [[root URLByAppendingPathComponent:key isDirectory:YES] + URLByAppendingPathComponent:manifest.path.lastPathComponent + isDirectory:YES]; + // Inspect, stage and publish relative to one root descriptor. + FileDescriptor root_fd(retry_eintr([&] { + return open(root.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC); + })); + ET_CHECK_OR_RETURN_ERROR( + root_fd.get() >= 0, AccessFailed, "Cannot open Core AI storage root"); + auto state = inspect_entry(root_fd.get(), key, files); + if (!state.ok()) { + return state.error(); + } + if (*state == EntryState::Valid) { + return model_url; + } + // Recovery sources may still be in use by a previously acquired SDK model. + ET_CHECK_OR_RETURN_ERROR( + *state == EntryState::Missing, + InvalidExternalData, + "Cannot replace a damaged bookmark recovery source"); + ET_CHECK_OR_RETURN_ERROR( + named_data != nullptr, + InvalidProgram, + "Missing NamedDataMap for Core AI source recovery"); + NSString* staging_name = + [@".staging-" stringByAppendingString:NSUUID.UUID.UUIDString]; + ET_CHECK_OR_RETURN_ERROR( + retry_eintr([&] { + return mkdirat( + root_fd.get(), staging_name.fileSystemRepresentation, 0700); + }) == 0, + AccessFailed, + "Cannot create Core AI staging directory"); + StagingDirectory staging{[root URLByAppendingPathComponent:staging_name + isDirectory:YES]}; + FileDescriptor staging_fd(retry_eintr([&] { + return openat( + root_fd.get(), + staging_name.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + ET_CHECK_OR_RETURN_ERROR( + staging_fd.get() >= 0, + AccessFailed, + "Cannot open Core AI staging directory"); + for (NSString* file in files) { + @autoreleasepool { + NSString* data_key = + [NSString stringWithFormat:@"coreai/%@/%@", manifest.hash, file]; + auto buffer = named_data->get_data(data_key.UTF8String); + ET_CHECK_OR_RETURN_ERROR( + buffer.ok(), + InvalidExternalData, + "Missing selected Core AI named data: %s", + data_key.UTF8String); + ET_CHECK_OR_RETURN_ERROR( + buffer->size() == files[file].unsignedLongLongValue && + (buffer->size() == 0 || buffer->data() != nullptr), + InvalidExternalData, + "Invalid Core AI named data buffer"); + ET_CHECK_OK_OR_RETURN_ERROR(write_storage_file( + staging_fd.get(), file, buffer->data(), buffer->size())); + } + } + ET_CHECK_OK_OR_RETURN_ERROR(validate_tree(staging_fd.get(), files)); + ET_CHECK_OR_RETURN_ERROR( + storage_fault(StorageOperation::Rename) == 0, + AccessFailed, + "Core AI source publication interrupted"); + if (retry_eintr([&] { + return renameatx_np( + root_fd.get(), + staging_name.fileSystemRepresentation, + root_fd.get(), + key.fileSystemRepresentation, + RENAME_EXCL); + }) == 0) { + staging.url = nil; + ET_CHECK_OK_OR_RETURN_ERROR(sync_storage_directory(root_fd.get())); + return model_url; + } + ET_CHECK_OR_RETURN_ERROR( + errno == EEXIST, AccessFailed, "Cannot publish Core AI stored bundle"); + // Another loader published first; use its source only if it is complete. + auto winner = inspect_entry(root_fd.get(), key, files); + if (!winner.ok()) { + return winner.error(); + } + ET_CHECK_OR_RETURN_ERROR( + *winner == EntryState::Valid, + InvalidExternalData, + "Concurrently published Core AI source is incomplete"); + return model_url; +} + } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h index 8a09e5068ac..2ecb5740e1a 100644 --- a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h +++ b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.h @@ -31,5 +31,12 @@ struct BookmarkDirectory { ::testing::AssertionResult backup_excluded(NSURL* url); ::testing::AssertionResult set_backup_excluded(NSURL* url, bool excluded); +::testing::AssertionResult asset_tree_snapshot( + NSURL* root, + NSDictionary* __strong& result); +::testing::AssertionResult snapshot_matches( + NSURL* root, + NSDictionary* expected); +NSArray* staging_directories(NSURL* root); } // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm index 11c1a190e0d..d2c36e5e712 100644 --- a/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm +++ b/backends/apple/coreai/runtime/test/coreai_filesystem_fixture.mm @@ -7,6 +7,7 @@ */ #include "coreai_filesystem_fixture.h" +#include #include #include @@ -80,4 +81,79 @@ EXPECT_TRUE([NSFileManager.defaultManager removeItemAtURL:url error:nil]); } +::testing::AssertionResult asset_tree_snapshot(NSURL* root, NSDictionary* __strong& result) { + result = nil; + if (root == nil) return ::testing::AssertionFailure() << "Missing snapshot root"; + NSFileManager* manager = NSFileManager.defaultManager; + NSError* error = nil; + NSArray* children = [manager subpathsOfDirectoryAtPath:root.path error:&error]; + if (children == nil || error != nil) { + return ::testing::AssertionFailure() << "Cannot list " << root.path.UTF8String; + } + NSMutableDictionary* snapshot = [NSMutableDictionary dictionary]; + for (NSString* path in [@[ @"" ] arrayByAddingObjectsFromArray:children]) { + NSURL* url = path.length == 0 ? root : [root URLByAppendingPathComponent:path]; + struct stat info; + if (lstat(url.fileSystemRepresentation, &info) != 0) { + return ::testing::AssertionFailure() << "Cannot stat " << url.path.UTF8String; + } + NSMutableDictionary* item = [@{ + @"inode" : @(info.st_ino), + @"device" : @(info.st_dev), + @"mode" : @(info.st_mode), + @"size" : @(info.st_size), + @"mtime_seconds" : @(info.st_mtimespec.tv_sec), + @"mtime_nanos" : @(info.st_mtimespec.tv_nsec) + } mutableCopy]; + error = nil; + if (S_ISREG(info.st_mode)) { + NSData* bytes = [NSData dataWithContentsOfURL:url options:0 error:&error]; + if (bytes == nil || error != nil) { + return ::testing::AssertionFailure() << "Cannot read " << url.path.UTF8String; + } + item[@"bytes"] = bytes; + } else if (S_ISLNK(info.st_mode)) { + NSString* destination = [manager destinationOfSymbolicLinkAtPath:url.path error:&error]; + if (destination == nil || error != nil) { + return ::testing::AssertionFailure() << "Cannot read link " << url.path.UTF8String; + } + item[@"link"] = destination; + } + snapshot[path] = item; + } + result = snapshot; + return ::testing::AssertionSuccess(); +} + +::testing::AssertionResult snapshot_matches(NSURL* root, NSDictionary* expected) { + if (expected == nil) return ::testing::AssertionFailure() << "Missing expected snapshot"; + NSDictionary* actual = nil; + auto read = asset_tree_snapshot(root, actual); + if (!read) return read; + return [actual isEqual:expected] ? ::testing::AssertionSuccess() + : ::testing::AssertionFailure() + << "Filesystem snapshot changed: " << root.path.UTF8String; +} + +NSArray* staging_directories(NSURL* root) { + if (root == nil) { + ADD_FAILURE() << "Missing staging root"; + return nil; + } + NSError* error = nil; + NSArray* children = [NSFileManager.defaultManager contentsOfDirectoryAtURL:root + includingPropertiesForKeys:nil + options:0 + error:&error]; + if (children == nil || error != nil) { + ADD_FAILURE() << "Cannot list staging directories: " << root.path.UTF8String; + return nil; + } + NSMutableArray* staging = [NSMutableArray array]; + for (NSURL* child in children) { + if ([child.lastPathComponent hasPrefix:@".staging-"]) [staging addObject:child]; + } + return staging; +} + } // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_source_fixture.h b/backends/apple/coreai/runtime/test/coreai_source_fixture.h new file mode 100644 index 00000000000..c637c817e40 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_fixture.h @@ -0,0 +1,53 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include "coreai_filesystem_fixture.h" +#include "coreai_manifest_fixture.h" + +namespace executorch::backends::coreai::testing { + +class TestData final : public runtime::NamedDataMap { + public: + mutable std::atomic releases{0}; + mutable std::atomic requests{0}; + mutable std::atomic attempts{0}; + mutable std::atomic metadata_requests{0}; + int fail_at = 0; + std::string required_prefix; + std::map files{ + {"coreai/ab/model.aimodel/graph.bin", std::string("model data", 11)}}; + runtime::Result get_tensor_layout( + std::string_view) const override; + runtime::Result get_data( + std::string_view key) const override; + runtime::Error load_data_into(std::string_view, void*, size_t) const override; + runtime::Result get_num_keys() const override; + runtime::Result get_key(uint32_t i) const override; +}; + +void aot_data(TestData& data, NSString* arch); + +class CoreAISourceTest : public ::testing::Test { + protected: + void SetUp() override; + void check_concurrent_source(); + runtime::Result prepare_source_bundle( + const Manifest& manifest, + const runtime::NamedDataMap* data, + NSString* staging_root); + + private: + TestDirectory locks_; +}; + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_source_fixture.mm b/backends/apple/coreai/runtime/test/coreai_source_fixture.mm new file mode 100644 index 00000000000..156c57a877a --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_fixture.mm @@ -0,0 +1,94 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include "coreai_source_fixture.h" +#include +#include +#include +#include +#include +#include +#include "coreai_file.h" +#include "coreai_storage.h" + +namespace executorch::backends::coreai::testing { +using runtime::Error; +using runtime::Result; + +Result TestData::get_tensor_layout(std::string_view) const { + return Error::NotSupported; +} +Result TestData::get_data(std::string_view key) const { + if (!required_prefix.empty() && key.substr(0, required_prefix.size()) != required_prefix) { + ADD_FAILURE() << "NamedDataMap key " << key << " is outside selected prefix " + << required_prefix; + return Error::InvalidArgument; + } + if (++attempts == fail_at) return Error::AccessFailed; + auto found = files.find(std::string(key)); + if (found == files.end()) return Error::NotFound; + ++requests; + return runtime::FreeableBuffer( + found->second.data(), found->second.size(), + [](void* context, void*, size_t) { ++static_cast(context)->releases; }, + const_cast(this)); +} +Error TestData::load_data_into(std::string_view, void*, size_t) const { + return Error::NotSupported; +} +Result TestData::get_num_keys() const { + ++metadata_requests; + return static_cast(files.size()); +} +Result TestData::get_key(uint32_t i) const { + ++metadata_requests; + if (i >= files.size()) return Error::NotFound; + auto found = files.begin(); + std::advance(found, i); + return found->first.c_str(); +} + +void aot_data(TestData& data, NSString* arch) { + data.required_prefix = + [NSString stringWithFormat:@"coreai/ab/model.%@.aimodelc/", arch].UTF8String; + data.files.clear(); + data.files[data.required_prefix + "graph.bin"] = std::string("compiled\0graph", 14); + data.files[data.required_prefix + "nested/weights.bin"] = "weights"; +} + +void CoreAISourceTest::SetUp() { + ASSERT_NE(locks_.url, nil); + struct stat info; + ASSERT_EQ(lstat(locks_.url.fileSystemRepresentation, &info), 0); + ASSERT_TRUE(S_ISDIR(info.st_mode)); +} + +Result CoreAISourceTest::prepare_source_bundle(const Manifest& manifest, + const runtime::NamedDataMap* data, + NSString* staging_root) { + NSString* key = + manifest.path == nil ? nil : manifest.bundle_digests[manifest.path.lastPathComponent]; + FileDescriptor lock(-1); + if (key != nil) { + FileDescriptor parent( + open(locks_.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC)); + if (parent.get() < 0) { + ADD_FAILURE() << "Cannot open source fixture lock root: " << strerror(errno); + return Error::AccessFailed; + } + NSString* name = [key stringByAppendingString:@".lock"]; + lock.reset(open_storage_shared_file(parent.get(), name.fileSystemRepresentation)); + if (lock.get() < 0 || retry_eintr([&] { return flock(lock.get(), LOCK_EX); }) != 0) { + ADD_FAILURE() << "Cannot acquire source fixture lock: " << strerror(errno); + return Error::AccessFailed; + } + } + return executorch::backends::coreai::prepare_source_bundle(manifest, data, staging_root, key); +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_source_test.mm b/backends/apple/coreai/runtime/test/coreai_source_test.mm new file mode 100644 index 00000000000..58ea949e955 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_source_test.mm @@ -0,0 +1,385 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include "coreai_bookmarks.h" +#include "coreai_fault_scope.h" +#include "coreai_source_fixture.h" + +namespace executorch::backends::coreai::testing { +using runtime::Error; +using runtime::Result; + +namespace { +Result source_manifest(bool aot = false) { + auto parsed = parse_manifest(encode(aot ? aot_manifest_dict() : manifest_dict())); + if (!parsed.ok()) return parsed.error(); + return select_assets(parsed.get(), @"arch_b", @"macOS"); +} +} // namespace + +TEST_F(CoreAISourceTest, RejectsUnselectedAotSource) { + auto parsed = parse_manifest(encode(aot_manifest_dict())); + ASSERT_TRUE(parsed.ok()); + ASSERT_TRUE(parsed->aot_compiled); + ASSERT_EQ(parsed->path, nil); + EXPECT_FALSE(prepare_source_bundle(parsed.get(), nullptr, nil).ok()); +} + +TEST_F(CoreAISourceTest, AotMaterializesOnlySelectedArchitecture) { + auto parsed = parse_manifest(encode(aot_manifest_dict())); + ASSERT_TRUE(parsed.ok()); + TestDirectory storage; + ASSERT_NE(storage.url, nil); + for (NSString* arch in @[ @"arch_a", @"arch_b" ]) { + SCOPED_TRACE(arch.UTF8String); + TestData data; + aot_data(data, arch); + auto selected = select_assets(parsed.get(), arch, @"macOS"); + ASSERT_TRUE(selected.ok()); + auto first = prepare_source_bundle(selected.get(), &data, storage.url.path); + ASSERT_TRUE(first.ok()); + EXPECT_EQ(data.attempts, 2); + EXPECT_EQ(data.requests, 2); + NSURL* url = first.get(); + EXPECT_TRUE([url.lastPathComponent isEqualToString:selected->path.lastPathComponent]); + NSData* contents = [NSData dataWithContentsOfURL:[url URLByAppendingPathComponent:@"graph.bin"]]; + ASSERT_TRUE([contents isEqual:[NSData dataWithBytes:"compiled\0graph" length:14]]); + + NSString* other_arch = [arch isEqualToString:@"arch_a"] ? @"arch_b" : @"arch_a"; + NSString* other_key = + [NSString stringWithFormat:@"coreai/ab/model.%@.aimodelc/graph.bin", other_arch]; + data.files[other_key.UTF8String] = "must not affect identity"; + auto reused = prepare_source_bundle(selected.get(), &data, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:url]); + EXPECT_EQ(data.attempts, 2); + + auto unselected_changed = aot_manifest_dict(); + NSMutableDictionary* digests = [unselected_changed[@"bundle_digests"] mutableCopy]; + digests[[NSString stringWithFormat:@"model.%@.aimodelc", other_arch]] = fixture_identity('f'); + unselected_changed[@"bundle_digests"] = digests; + NSMutableDictionary* files = [unselected_changed[@"files"] mutableCopy]; + files[[NSString stringWithFormat:@"model.%@.aimodelc/extra", other_arch]] = @0; + unselected_changed[@"files"] = files; + auto changed_parsed = parse_manifest(encode(unselected_changed)); + ASSERT_TRUE(changed_parsed.ok()); + auto unchanged = select_assets(changed_parsed.get(), arch, @"macOS"); + ASSERT_TRUE(unchanged.ok()); + auto unchanged_source = prepare_source_bundle(unchanged.get(), &data, storage.url.path); + ASSERT_TRUE(unchanged_source.ok()); + EXPECT_TRUE([unchanged_source.get() isEqual:url]); + EXPECT_EQ(data.attempts, 2); + EXPECT_EQ(data.releases, data.requests); + } +} + +void CoreAISourceTest::check_concurrent_source() { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto selected = source_manifest(); + ASSERT_TRUE(selected.ok()); + for (int pass = 0; pass < 2; ++pass) { + SCOPED_TRACE(pass == 0 ? "cold" : "warm"); + using FixturePointer = decltype(this); + struct Worker { + FixturePointer fixture = nullptr; + const Manifest* manifest = nullptr; + TestData* data = nullptr; + NSString* root = nil; + dispatch_semaphore_t start = nullptr; + bool started = false; + Error error = Error::Internal; + std::string path; + } workers[4]; + pthread_t threads[4]; + dispatch_semaphore_t start = dispatch_semaphore_create(0); + int launch_error = 0; + size_t launched = 0; + for (; launched < 4; ++launched) { + auto& worker = workers[launched]; + worker.fixture = this; + worker.manifest = &selected.get(); + worker.data = &data; + worker.root = storage.url.path; + worker.start = start; + launch_error = pthread_create( + &threads[launched], nullptr, + [](void* context) -> void* { + auto& worker = *static_cast(context); + @autoreleasepool { + worker.started = + dispatch_semaphore_wait(worker.start, + dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC)) == 0; + if (!worker.started) return nullptr; + auto assets = + worker.fixture->prepare_source_bundle(*worker.manifest, worker.data, worker.root); + worker.error = assets.error(); + if (assets.ok()) worker.path = assets.get().path.UTF8String; + } + return nullptr; + }, + &worker); + if (launch_error != 0) break; + } + for (size_t i = 0; i < launched; ++i) dispatch_semaphore_signal(start); + int join_errors[4] = {}; + for (size_t i = 0; i < launched; ++i) join_errors[i] = pthread_join(threads[i], nullptr); + ASSERT_EQ(launch_error, 0); + for (size_t i = 0; i < 4; ++i) { + SCOPED_TRACE(i); + ASSERT_EQ(join_errors[i], 0); + ASSERT_TRUE(workers[i].started); + ASSERT_EQ(workers[i].error, Error::Ok); + EXPECT_EQ(workers[i].path, workers[0].path); + } + EXPECT_EQ(data.attempts, 1); + EXPECT_EQ(data.requests, 1); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(data.metadata_requests, 0); + EXPECT_EQ([NSFileManager.defaultManager contentsOfDirectoryAtPath:storage.url.path error:nil] + .count, + 1u); + } +} + +TEST_F(CoreAISourceTest, ConcurrentColdAndWarmSource) { + check_concurrent_source(); +} + +TEST_F(CoreAISourceTest, RejectsInvalidKeysAndMissingSourceInputs) { + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + NSString* valid_key = fixture_identity('a'); + for (NSString* key in @[ @"../escape", [valid_key substringToIndex:63] ]) { + SCOPED_TRACE(key.UTF8String); + TestDirectory invalid_storage; + ASSERT_NE(invalid_storage.url, nil); + TestData invalid_data; + EXPECT_EQ(executorch::backends::coreai::prepare_source_bundle( + manifest.get(), &invalid_data, invalid_storage.url.path, key) + .error(), + Error::InvalidArgument); + EXPECT_EQ(invalid_data.attempts, 0); + } + + TestData data; + TestDirectory storage; + ASSERT_NE(storage.url, nil); + EXPECT_FALSE(prepare_source_bundle(manifest.get(), nullptr, storage.url.path).ok()); + NSURL* absent = [storage.url URLByAppendingPathComponent:@"absent"]; + EXPECT_FALSE(prepare_source_bundle(manifest.get(), &data, absent.path).ok()); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:absent.path]); + auto recovered = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(recovered.ok()); + EXPECT_EQ( + [NSData dataWithContentsOfURL:[recovered.get() URLByAppendingPathComponent:@"graph.bin"]] + .length, + 11u); + EXPECT_EQ(data.requests, data.releases); +} + +TEST_F(CoreAISourceTest, PublishedSourceOutlivesResult) { + TestData data; + TestDirectory storage; + ASSERT_NE(storage.url, nil); + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + NSURL* first_url = nil; + { + auto assets = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(assets.ok()); + EXPECT_EQ(data.releases, data.requests); + first_url = assets.get(); + EXPECT_EQ( + [NSData dataWithContentsOfURL:[first_url URLByAppendingPathComponent:@"graph.bin"]].length, + 11u); + } + EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:first_url.path]); + auto reused = prepare_source_bundle(manifest.get(), nullptr, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:first_url]); +} + +TEST_F(CoreAISourceTest, PersistentIdentityTracksDigestNotFileOrder) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto dict = manifest_dict(); + dict[@"files"] = @{@"model.aimodel/graph.bin" : @11, @"model.aimodel/nested/weights.bin" : @7}; + dict[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity('2')}; + data.files["coreai/ab/model.aimodel/nested/weights.bin"] = "weights"; + auto manifest = parse_manifest(encode(dict)); + ASSERT_TRUE(manifest.ok()); + auto first = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(first.ok()); + NSURL* first_url = first.get(); + dict[@"files"] = @{@"model.aimodel/nested/weights.bin" : @7, @"model.aimodel/graph.bin" : @11}; + auto reordered_manifest = parse_manifest(encode(dict)); + ASSERT_TRUE(reordered_manifest.ok()); + auto reordered = prepare_source_bundle(reordered_manifest.get(), &data, storage.url.path); + ASSERT_TRUE(reordered.ok()); + EXPECT_TRUE([reordered.get() isEqual:first_url]); + data.files["coreai/ab/model.aimodel/nested/weights.bin"] = "changed"; + dict[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity('3')}; + auto changed_manifest = parse_manifest(encode(dict)); + ASSERT_TRUE(changed_manifest.ok()); + auto changed = prepare_source_bundle(changed_manifest.get(), &data, storage.url.path); + ASSERT_TRUE(changed.ok()); + EXPECT_FALSE([changed.get() isEqual:first_url]); + EXPECT_TRUE( + [[NSData dataWithContentsOfURL:[first_url URLByAppendingPathComponent:@"nested/weights.bin"]] + isEqual:[@"weights" dataUsingEncoding:NSUTF8StringEncoding]]); + EXPECT_EQ(data.releases, data.requests); +} + +TEST_F(CoreAISourceTest, PartialMaterializationCleansStagingAndCanRetry) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto dict = manifest_dict(); + dict[@"files"] = @{@"model.aimodel/graph.bin" : @11, @"model.aimodel/nested/weights.bin" : @7}; + dict[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity('2')}; + data.files["coreai/ab/model.aimodel/nested/weights.bin"] = "weights"; + auto manifest = parse_manifest(encode(dict)); + ASSERT_TRUE(manifest.ok()); + data.fail_at = 2; + EXPECT_FALSE(prepare_source_bundle(manifest.get(), &data, storage.url.path).ok()); + EXPECT_EQ( + [NSFileManager.defaultManager contentsOfDirectoryAtPath:storage.url.path error:nil].count, + 0u); + EXPECT_EQ(data.releases, data.requests); + data.fail_at = 0; + EXPECT_TRUE(prepare_source_bundle(manifest.get(), &data, storage.url.path).ok()); + EXPECT_EQ( + [NSFileManager.defaultManager contentsOfDirectoryAtPath:storage.url.path error:nil].count, + 1u); +} + +TEST_F(CoreAISourceTest, NestedEmptyAndLargeFilesMaterializeAndReuseWithoutNamedData) { + TestDirectory storage; + ASSERT_NE(storage.url, nil); + TestData data; + auto dict = manifest_dict(); + constexpr size_t large_size = 2 * 1024 * 1024 + 17; + dict[@"files"] = @{ + @"model.aimodel/graph.bin" : @11, + @"model.aimodel/nested/empty" : @0, + @"model.aimodel/nested/deeper/large.bin" : @(large_size) + }; + data.files["coreai/ab/model.aimodel/nested/empty"] = ""; + data.files["coreai/ab/model.aimodel/nested/deeper/large.bin"] = std::string(large_size, 'x'); + auto manifest = parse_manifest(encode(dict)); + ASSERT_TRUE(manifest.ok()); + auto first = prepare_source_bundle(manifest.get(), &data, storage.url.path); + ASSERT_TRUE(first.ok()); + EXPECT_EQ(data.attempts, 3); + EXPECT_EQ(data.requests, data.releases); + NSData* empty = [NSData + dataWithContentsOfURL:[first.get() URLByAppendingPathComponent:@"nested/empty"]]; + ASSERT_NE(empty, nil); + EXPECT_EQ(empty.length, 0u); + NSData* large = [NSData + dataWithContentsOfURL:[first.get() URLByAppendingPathComponent:@"nested/deeper/large.bin"]]; + ASSERT_NE(large, nil); + EXPECT_EQ(large.length, large_size); + + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, snapshot)); + const int attempts = data.attempts; + auto reused = prepare_source_bundle(manifest.get(), nullptr, storage.url.path); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:first.get()]); + EXPECT_EQ(data.attempts, attempts); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(data.metadata_requests, 0); + EXPECT_TRUE(snapshot_matches(storage.url, snapshot)); +} + +TEST_F(CoreAISourceTest, ExplicitStagingKeyReusesSourceWithoutWritesOrRepublication) { + BookmarkDirectory directory; + ASSERT_NE(directory.url, nil); + TestData data; + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + auto root = resolve_bookmark_root(directory.url.path); + ASSERT_TRUE(root.ok()); + NSString* key = fixture_identity('9'); + auto lock = lock_bookmark(root.get(), key); + ASSERT_TRUE(lock.ok()); + auto staging = prepare_bookmark_staging(*lock.get()); + ASSERT_TRUE(staging.ok()); + auto first = executorch::backends::coreai::prepare_source_bundle(manifest.get(), &data, + staging.get(), key); + ASSERT_TRUE(first.ok()); + EXPECT_EQ(data.attempts, 1); + EXPECT_TRUE(([first.get().path + isEqualToString:[root.get() + stringByAppendingPathComponent: + [NSString stringWithFormat:@"staging/%@/model.aimodel", key]]])); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(directory.url, snapshot)); + const int attempts = data.attempts; + const int requests = data.requests; + data.files.clear(); + auto reused = executorch::backends::coreai::prepare_source_bundle(manifest.get(), &data, + staging.get(), key); + ASSERT_TRUE(reused.ok()); + EXPECT_TRUE([reused.get() isEqual:first.get()]); + EXPECT_EQ(data.attempts, attempts); + EXPECT_EQ(data.requests, requests); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(data.metadata_requests, 0); + EXPECT_TRUE(snapshot_matches(directory.url, snapshot)); +} + +TEST_F(CoreAISourceTest, ExclusivePublicationValidatesWinnerAndRemovesLosingStaging) { + TestDirectory storage; + TestDirectory winner_storage; + ASSERT_NE(storage.url, nil); + ASSERT_NE(winner_storage.url, nil); + TestData data; + TestData winner_data; + auto manifest = source_manifest(); + ASSERT_TRUE(manifest.ok()); + auto winner = prepare_source_bundle(manifest.get(), &winner_data, winner_storage.url.path); + ASSERT_TRUE(winner.ok()); + NSURL* winner_entry = winner.get().URLByDeletingLastPathComponent; + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(winner_entry, snapshot)); + NSURL* root = storage.url; + NSURL* final = + [root URLByAppendingPathComponent:manifest->bundle_digests[manifest->path.lastPathComponent]]; + __block int publications = 0; + __block bool callback_ok = true; + { + StorageFaultScope fault(StorageOperation::Rename); + fault.armed = false; + fault.observe = ^(StorageOperation operation) { + if (operation != StorageOperation::Rename) return; + ++publications; + callback_ok &= staging_directories(root).count == 1; + callback_ok &= [NSFileManager.defaultManager moveItemAtURL:winner_entry toURL:final error:nil]; + }; + auto result = prepare_source_bundle(manifest.get(), &data, storage.url.path); + EXPECT_TRUE(callback_ok); + ASSERT_TRUE(result.ok()); + EXPECT_TRUE( + [result.get() isEqual:[final URLByAppendingPathComponent:manifest->path.lastPathComponent]]); + EXPECT_EQ(publications, 1); + EXPECT_EQ(fault.hits, 0); + } + EXPECT_EQ(data.attempts, manifest->files.count); + EXPECT_EQ(data.requests, data.releases); + EXPECT_EQ(staging_directories(storage.url).count, 0u); + EXPECT_TRUE(snapshot_matches(final, snapshot)); +} + +} // namespace executorch::backends::coreai::testing