diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index fc5d5f6117d..10d63f2b9e0 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -27,8 +27,9 @@ endforeach() enable_language(OBJCXX) -set(_coreai_runtime_sources runtime/coreai_assets.mm runtime/coreai_storage.mm - runtime/coreai_bookmarks.mm +set(_coreai_runtime_sources + runtime/coreai_assets.mm runtime/coreai_storage.mm + runtime/coreai_bookmarks.mm runtime/coreai_load_coordinator.mm ) # Keep ARC and C++ exception settings away from the private Swift module. @@ -56,6 +57,10 @@ if(EXECUTORCH_BUILD_TESTS AND CMAKE_SYSTEM_NAME STREQUAL "Darwin") runtime/test/coreai_bookmark_test.mm runtime/test/coreai_source_fixture.mm runtime/test/coreai_bookmark_fixture.mm + runtime/test/coreai_fake_loader.mm + runtime/test/coreai_acquisition_test.mm + runtime/test/coreai_acquisition_fixture.mm + runtime/ETCoreAITensor.mm ${_coreai_runtime_sources} ) coreai_configure_objc_target(coreai_host_test) @@ -308,6 +313,32 @@ if(EXECUTORCH_BUILD_TESTS PROPERTIES TIMEOUT 600 ENVIRONMENT "COREAI_BRIDGE_FIXTURE_DIR=${_coreai_fixture_dir}" ) + add_executable( + coreai_sdk_source_lifetime_test + runtime/test/SDKSourceLifetimeTest.swift runtime/ETCoreAIModel.swift + runtime/ETCoreAITensor.mm + ) + set_target_properties( + coreai_sdk_source_lifetime_test PROPERTIES Swift_LANGUAGE_VERSION 6 + ) + coreai_configure_objc_target(coreai_sdk_source_lifetime_test) + coreai_require_os27(coreai_sdk_source_lifetime_test) + target_compile_options( + coreai_sdk_source_lifetime_test + PRIVATE + "$<$:-parse-as-library;-import-objc-header;${CMAKE_CURRENT_SOURCE_DIR}/runtime/ETCoreAIBridge.h>" + ) + target_link_libraries( + coreai_sdk_source_lifetime_test PRIVATE executorch::coreai_dependencies + ) + add_test( + NAME coreai_sdk_source_lifetime_test + COMMAND + "${PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/runtime/test/run_sdk_source_lifetime.py" + "$" + ) + set_tests_properties(coreai_sdk_source_lifetime_test PROPERTIES TIMEOUT 240) endif() install( diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 7b45b877410..567a9bed233 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -57,6 +57,11 @@ architectures=["h17p"])` requests that compiler architecture; omitting the list lets the compiler emit its supported architectures for the target platform. These are Core AI architecture names, not CPU names such as `arm64`. +Both formats can require device specialization. The acquisition coordinator +first attempts SDK bookmark restoration. A missing bookmark or SDK-confirmed +cache miss falls back to source materialization and persistent specialization. +It returns the acquired model; function binding is the caller's responsibility. + ### Asset storage The assets root must be an absolute path. The application chooses it, should @@ -88,10 +93,11 @@ source root, including after a previous publication's final flush failed. 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. -Only interrupted system calls are retried; other failures are returned to the -caller without automatic retries or repair. Temporary bundles are removed on +used. The acquisition coordinator can remove an incomplete or size-mismatched +regular-file tree under the key lock and rebuild it from available named data. +Cleanup never follows symlinks and rejects unexpected entry types. Standalone +source preparation rejects damaged trees unless replacement is explicitly enabled. +Named-data and SDK failures are not retried. Temporary bundles are removed on ordinary failures before publication; a caller can explicitly retry after addressing the failure. @@ -123,6 +129,22 @@ a raw bookmark in `bookmarks`, a bundle directory in `staging`, and one of 256 shared coordination files in `locks`. There is no PTE-wide bookmark list or serialized lifecycle record. +A healthy bookmark hit does not consult NDS and attempts to clean leftover +staging. On a missing bookmark or confirmed SDK miss, the coordinator prepares +the source, specializes it and atomically writes the returned bookmark bytes. +After durable publication, it attempts to remove the entire keyed staging +directory while holding the key's stripe lock. +caching supports inference and bookmark restoration without that extracted copy. +Keep the packaged named data available: a later SDK cache miss requires source +materialization again. There is no secondary source-URL cache lookup. + +SDK errors, unreadable bookmarks and publication failures fail loading without +automatic cache deletion. If publication fails after specialization, the SDK +entry and source remain. Cleanup after a bookmark hit or durable publication is +best-effort: failures are logged without discarding the usable model or bookmark. +A later load can finish cleanup or rebuild a partially removed source after an +SDK miss. No cleanup failure triggers SDK eviction or retries within that load. +The backend does not track historical or unrecorded SDK entries. Raw bookmarks have an 8 MiB backend allocation limit, not an asserted SDK format maximum, and are published atomically because losing one orphans an SDK entry. @@ -169,6 +191,11 @@ points allow delegate tests to substitute a fake bridge. `coreai_swift_bridge_test` runs the bridge's XCTests under `xctest` with `EXECUTORCH_BUILD_TESTS=ON`. It links CoreAI, so it runs only on macOS 27. The test build exports static- and dynamic-shape models for real session tests. +The SDK source-lifetime regression checks live inference and fresh-process +bookmark restoration through the bridge after removing disposable portable and +AOT source copies. The AOT case is skipped when the SDK cannot identify the +device architecture. A killed probe can leave an SDK entry if it dies before +saving its bookmark; cleanup never purges unrelated default-cache entries. CMake's `PYTHON_EXECUTABLE` must use the Python environment with the pinned Core AI requirements above. Backend Python changes regenerate the dedicated fixture directory, replacing previous contents. CTest supplies its location to XCTest; @@ -186,13 +213,29 @@ Ninja builds use one architecture per build directory. Current runtime support covers arm64 macOS and iOS device builds. x86_64 builds are blocked by Swift `Float16` availability, and the tested iOS simulator SDKs do not contain Core AI. +### Lifecycle logging + +Build with `EXECUTORCH_ENABLE_LOGGING=ON` and `EXECUTORCH_LOG_LEVEL=Info` to +observe cache lookup, bookmark restoration hits, missing-bookmark versus +SDK-confirmed misses, specialization start/success, and durable bookmark saves. Each message carries the cache key so events from +multiple partitions can be correlated. These messages do not include bookmark +contents, model data, or per-inference logs. + +A cache hit means that the SDK successfully restored a model from its bookmark. +Specialization start means a request to the SDK's persistent default cache; the +SDK does not report whether that request compiled afresh or reused an internal +artifact. The durable-save message confirms local bookmark publication, +including its final directory flush, not independent verification of SDK-private +artifact storage. Failed operations do not emit their success messages. + ## Host Tests On macOS, `EXECUTORCH_BUILD_COREAI=ON` with `EXECUTORCH_BUILD_TESTS=ON` registers the standalone `coreai_host_test` executable with CTest. It tests manifest parsing, AOT asset selection, and filesystem storage using Foundation, -without Swift or CoreAI SDK calls. The host test target is omitted for Apple -cross-builds. +without Swift or CoreAI SDK calls. Acquisition contention tests run two threads +inside a bounded child process; bookmark lock tests separately cover contention +between processes. The host test target is omitted for Apple cross-builds. CI builds and runs these tests on macOS 27. To use the same configuration from the repository root: diff --git a/backends/apple/coreai/runtime/coreai_assets.h b/backends/apple/coreai/runtime/coreai_assets.h index e8e216305b8..302fcfd729e 100644 --- a/backends/apple/coreai/runtime/coreai_assets.h +++ b/backends/apple/coreai/runtime/coreai_assets.h @@ -38,11 +38,14 @@ runtime::Result select_assets( // The key is 64 lowercase hex characters identifying the selected export digest // and bundle name, optionally including platform or SDK context. // Bundles are published at source_root/key/bundle. Reuse validates structure, -// not contents, and fully flushes source_root before returning. +// not contents, and fully flushes source_root before returning. When +// replace_corrupt is set, damaged regular-file trees may be removed and +// rebuilt. runtime::Result prepare_source_bundle( const Manifest& manifest, const runtime::NamedDataMap* named_data, NSString* source_root, - NSString* key); + NSString* key, + bool replace_corrupt = false); } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_assets.mm b/backends/apple/coreai/runtime/coreai_assets.mm index 784f91c71c3..1a66c48792b 100644 --- a/backends/apple/coreai/runtime/coreai_assets.mm +++ b/backends/apple/coreai/runtime/coreai_assets.mm @@ -445,7 +445,8 @@ FileDescriptor entry(retry_eintr([&] { const Manifest& manifest, const runtime::NamedDataMap* named_data, NSString* source_root, - NSString* key) { + NSString* key, + bool replace_corrupt) { ET_CHECK_OR_RETURN_ERROR( relative_path(manifest.path), InvalidProgram, @@ -485,11 +486,13 @@ FileDescriptor root_fd(retry_eintr([&] { ET_CHECK_OK_OR_RETURN_ERROR(sync_storage_directory(root_fd.get(), true)); 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 Core AI stored bundle"); + if (*state == EntryState::Corrupt) { + ET_CHECK_OR_RETURN_ERROR( + replace_corrupt && named_data != nullptr, + InvalidExternalData, + "Cannot replace a damaged Core AI stored bundle"); + ET_CHECK_OK_OR_RETURN_ERROR(remove_storage_source(root_fd.get(), key)); + } ET_CHECK_OR_RETURN_ERROR( named_data != nullptr, InvalidProgram, diff --git a/backends/apple/coreai/runtime/coreai_bookmarks.h b/backends/apple/coreai/runtime/coreai_bookmarks.h index 8a0181719f4..26b841a4f8e 100644 --- a/backends/apple/coreai/runtime/coreai_bookmarks.h +++ b/backends/apple/coreai/runtime/coreai_bookmarks.h @@ -37,6 +37,7 @@ class BookmarkLock final { lock_bookmark(NSString* root, NSString* key, bool create); friend runtime::Result read_bookmark(const BookmarkLock& lock); friend runtime::Error write_bookmark(const BookmarkLock& lock, NSData* data); + friend runtime::Error remove_bookmark_staging(const BookmarkLock& lock); friend runtime::Result prepare_bookmark_staging( const BookmarkLock& lock); BookmarkLock( @@ -65,6 +66,9 @@ lock_bookmark(NSString* root, NSString* key, bool create = true); runtime::Result read_bookmark(const BookmarkLock& lock); // A staging-only lock has no bookmark directory; writes return InvalidArgument. runtime::Error write_bookmark(const BookmarkLock& lock, NSData* data); +// Requires successful persistent specialization and durable bookmark +// publication. +runtime::Error remove_bookmark_staging(const BookmarkLock& lock); // Prepared only for cold loads; source preparation appends key/bundle. runtime::Result prepare_bookmark_staging(const BookmarkLock& lock); diff --git a/backends/apple/coreai/runtime/coreai_bookmarks.mm b/backends/apple/coreai/runtime/coreai_bookmarks.mm index 1213237b9df..03521ad9e91 100644 --- a/backends/apple/coreai/runtime/coreai_bookmarks.mm +++ b/backends/apple/coreai/runtime/coreai_bookmarks.mm @@ -283,6 +283,10 @@ Error write_bookmark(const BookmarkLock& lock, NSData* data) { lock.bookmarks_.get(), bookmark_name(lock.key_), data); } +Error remove_bookmark_staging(const BookmarkLock& lock) { + return remove_storage_staging(lock.root_.get(), lock.key_); +} + Result prepare_bookmark_staging(const BookmarkLock& lock) { auto staging = open_child(lock.root_.get(), @"staging", true); if (!staging.ok()) diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.h b/backends/apple/coreai/runtime/coreai_load_coordinator.h new file mode 100644 index 00000000000..018117a18f9 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.h @@ -0,0 +1,25 @@ +/* + * 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 + +#import "ETCoreAIBridge.h" +#import "coreai_bookmarks.h" + +namespace executorch::backends::coreai { + +// The key lock outlives all SDK callbacks and raw bookmark publication. +// Cleanup after a restore hit or durable publication is best-effort. +runtime::Result> acquire_bookmark_model( + const Manifest& selected, + const runtime::NamedDataMap* named_data, + NSString* root, + NSString* platform, + NSString* architecture, + id loader); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm new file mode 100644 index 00000000000..d239ccb0f64 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -0,0 +1,159 @@ +/* + * 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. + */ +#import "coreai_load_coordinator.h" +#import "coreai_storage.h" + +namespace executorch::backends::coreai { +namespace { +using runtime::Error; +using runtime::Result; + +Error sdk_error(NSError* error) { + ET_LOG( + Error, + "Core AI SDK error (%s, %ld): %s", + error.domain.UTF8String, + static_cast(error.code), + error.description.UTF8String); + return Error::Internal; +} + +Result> acquire( + void (^operation)(ETCoreAIAcquisitionCompletion)) { + ET_CHECK_OR_RETURN_ERROR( + storage_fault(StorageOperation::BeforeSDK) == 0, + AccessFailed, + "Core AI acquisition interrupted"); + dispatch_semaphore_t ready = dispatch_semaphore_create(0); + __block id model = nil; + __block NSError* error = nil; + __block ETCoreAIAcquisitionStatus status = ETCoreAIAcquisitionStatusError; + operation( + ^(ETCoreAIAcquisitionStatus result, + id value, + NSError* failure) { + status = result; + model = value; + error = failure; + dispatch_semaphore_signal(ready); + }); + // Do not abandon asynchronous source access on cancellation or timeout. + dispatch_semaphore_wait(ready, DISPATCH_TIME_FOREVER); + id acquired = model; + // A loader may retain the completion block after invoking it. + model = nil; + if (error != nil) + return sdk_error(error); + ET_CHECK_OR_RETURN_ERROR( + storage_fault(StorageOperation::AfterSDK) == 0, + AccessFailed, + "Core AI acquisition interrupted"); + ET_CHECK_OR_RETURN_ERROR( + (status == ETCoreAIAcquisitionStatusHit && acquired != nil) || + (status == ETCoreAIAcquisitionStatusMiss && acquired == nil), + Internal, + "Invalid Core AI acquisition result"); + return acquired; +} + +Result> restore( + id loader, + NSData* bookmark) { + return acquire(^(ETCoreAIAcquisitionCompletion completion) { + [loader restoreModelFromBookmark:bookmark completion:completion]; + }); +} + +void cleanup_source(const BookmarkLock& lock) { + const auto error = remove_bookmark_staging(lock); + if (error != Error::Ok) { + ET_LOG( + Error, + "Cannot clean Core AI staging: error=%u", + static_cast(error)); + } +} +} // namespace + +Result> acquire_bookmark_model( + const Manifest& selected, + const runtime::NamedDataMap* named_data, + NSString* root, + NSString* platform, + NSString* architecture, + id loader) { + ET_CHECK_OR_RETURN_ERROR( + loader != nil, InvalidArgument, "Missing Core AI model loader"); + auto key = bookmark_key(selected, platform, architecture); + if (!key.ok()) + return key.error(); + NSString* cache_key = key.get(); + ET_LOG(Info, "Core AI cache lookup: key=%s", cache_key.UTF8String); + auto locked = lock_bookmark(root, cache_key); + if (!locked.ok()) + return locked.error(); + auto bookmark = read_bookmark(*locked.get()); + if (!bookmark.ok()) + return bookmark.error(); + if (bookmark.get() != nil) { + auto model = restore(loader, bookmark.get()); + if (!model.ok()) + return model.error(); + if (model.get() != nil) { + ET_LOG( + Info, + "Core AI cache hit: restored model from bookmark, key=%s", + key.get().UTF8String); + cleanup_source(*locked.get()); + return model; + } + ET_LOG( + Info, + "Core AI cache miss: SDK could not restore bookmark, key=%s", + key.get().UTF8String); + } else { + ET_LOG( + Info, "Core AI cache miss: no bookmark, key=%s", key.get().UTF8String); + } + auto staging = prepare_bookmark_staging(*locked.get()); + if (!staging.ok()) + return staging.error(); + auto source = prepare_source_bundle( + selected, + named_data, + staging.get(), + key.get(), + /*replace_corrupt=*/true); + if (!source.ok()) + return source.error(); + NSURL* url = source.get(); + auto model = acquire(^(ETCoreAIAcquisitionCompletion completion) { + ET_LOG( + Info, + "Core AI specialization started: kind=%s, key=%s", + selected.aot_compiled ? "aot" : "source", + cache_key.UTF8String); + [loader specializeModelAtURL:url completion:completion]; + }); + if (!model.ok()) + return model.error(); + ET_CHECK_OR_RETURN_ERROR( + model.get() != nil, Internal, "Core AI specialization returned no model"); + ET_LOG( + Info, "Core AI specialization succeeded: key=%s", key.get().UTF8String); + ET_CHECK_OK_OR_RETURN_ERROR( + write_bookmark(*locked.get(), [model.get() copyBookmarkData])); + ET_LOG( + Info, + "Core AI cache bookmark saved durably: key=%s", + key.get().UTF8String); + cleanup_source(*locked.get()); + return model; +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h index daf32376fab..9de9a3ced28 100644 --- a/backends/apple/coreai/runtime/coreai_storage.h +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -47,6 +47,10 @@ runtime::Error publish_storage_data(int root_fd, NSString* name, NSData* data); runtime::Result*> storage_children( int fd, bool skip_invalid_names = false); +// The caller holds the key's lock. Missing staging is already removed. +runtime::Error remove_storage_staging(int root_fd, NSString* key); +// Removes a keyed child of source_root_fd under the same lock requirement. +runtime::Error remove_storage_source(int source_root_fd, NSString* key); enum class StorageOperation { Write, diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm index 0a42b22f859..2712b3a1253 100644 --- a/backends/apple/coreai/runtime/coreai_storage.mm +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -197,6 +197,55 @@ FileDescriptor file(retry_eintr([&] { } return result; } + +// Preflight the whole tree before unlinking anything; never follow source +// links. +Error visit_source_tree(int fd, bool remove) { + auto names = storage_children(fd); + if (!names.ok()) + return names.error(); + for (NSString* name in names.get()) { + struct stat info; + ET_CHECK_OR_RETURN_ERROR( + retry_eintr([&] { + return fstatat( + fd, name.fileSystemRepresentation, &info, AT_SYMLINK_NOFOLLOW); + }) == 0, + AccessFailed, + "Cannot inspect Core AI source entry"); + ET_CHECK_OR_RETURN_ERROR( + S_ISDIR(info.st_mode) || S_ISREG(info.st_mode), + InvalidExternalData, + "Unexpected Core AI source entry type"); + if (S_ISDIR(info.st_mode)) { + FileDescriptor child(retry_eintr([&] { + return openat( + fd, + name.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + ET_CHECK_OR_RETURN_ERROR( + child.get() >= 0, + AccessFailed, + "Cannot open Core AI source directory"); + ET_CHECK_OK_OR_RETURN_ERROR(visit_source_tree(child.get(), remove)); + } + if (remove) { + ET_CHECK_OR_RETURN_ERROR( + storage_call( + StorageOperation::Remove, + [&] { + return unlinkat( + fd, + name.fileSystemRepresentation, + S_ISDIR(info.st_mode) ? AT_REMOVEDIR : 0); + }) == 0, + AccessFailed, + "Cannot remove Core AI source entry"); + } + } + return remove ? sync_storage_directory(fd) : Error::Ok; +} } // namespace Result @@ -314,6 +363,77 @@ Error sync_storage_directory(int fd, bool full_sync) { return Error::Ok; } +namespace { +bool storage_key(NSString* key) { + return [key isKindOfClass:NSString.class] && key.length == 64 && + [key rangeOfCharacterFromSet:[[NSCharacterSet + characterSetWithCharactersInString: + @"0123456789abcdef"] invertedSet]] + .location == NSNotFound; +} +} // namespace + +Error remove_storage_staging(int root_fd, NSString* key) { + ET_CHECK_OR_RETURN_ERROR( + storage_key(key), InvalidArgument, "Invalid Core AI staging key"); + FileDescriptor staging(retry_eintr([&] { + return openat( + root_fd, "staging", O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + if (staging.get() < 0 && errno == ENOENT) + return Error::Ok; + ET_CHECK_OR_RETURN_ERROR( + staging.get() >= 0, + AccessFailed, + "Cannot open Core AI staging directory"); + return remove_storage_source(staging.get(), key); +} + +Error remove_storage_source(int source_root_fd, NSString* key) { + ET_CHECK_OR_RETURN_ERROR( + storage_key(key), InvalidArgument, "Invalid Core AI staging key"); + struct stat info; + const int inspected = retry_eintr([&] { + return fstatat( + source_root_fd, + key.fileSystemRepresentation, + &info, + AT_SYMLINK_NOFOLLOW); + }); + if (inspected < 0 && errno == ENOENT) + return Error::Ok; + ET_CHECK_OR_RETURN_ERROR( + inspected == 0, + AccessFailed, + "Cannot inspect Core AI keyed staging directory"); + ET_CHECK_OR_RETURN_ERROR( + S_ISDIR(info.st_mode), + InvalidExternalData, + "Core AI keyed staging entry is not a directory"); + FileDescriptor source(retry_eintr([&] { + return openat( + source_root_fd, + key.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + ET_CHECK_OR_RETURN_ERROR( + source.get() >= 0, + AccessFailed, + "Cannot open Core AI keyed staging directory"); + ET_CHECK_OK_OR_RETURN_ERROR(visit_source_tree(source.get(), false)); + ET_CHECK_OK_OR_RETURN_ERROR(visit_source_tree(source.get(), true)); + ET_CHECK_OR_RETURN_ERROR( + storage_call( + StorageOperation::Remove, + [&] { + return unlinkat( + source_root_fd, key.fileSystemRepresentation, AT_REMOVEDIR); + }) == 0, + AccessFailed, + "Cannot remove Core AI keyed staging directory"); + return sync_storage_directory(source_root_fd); +} + Error write_storage_file( int root_fd, NSString* relative_path, diff --git a/backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift b/backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift new file mode 100644 index 00000000000..505a8d39284 --- /dev/null +++ b/backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift @@ -0,0 +1,181 @@ +// 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. + +import Foundation +import Synchronization + +private struct Fixture: Decodable { + let function: String + let inputs: [String] + let outputs: [String] + let expected: [[Float]] +} + +private enum Failure: Error { + case invalidArguments + case invalidResult + case incorrectOutput(String) + case unexpectedSource +} + +// Bridge results are immutable and consumed only after completion. +private struct BridgeResult: @unchecked Sendable { + let value: Result +} + +// The executable waits on an ordinary thread, not Swift's cooperative pool. +private func waitForResult( + _ operation: (@escaping @Sendable (Result) -> Void) -> Void +) throws -> Value { + let result = Mutex?>(nil) + let ready = DispatchSemaphore(value: 0) + operation { value in + result.withLock { $0 = BridgeResult(value: value) } + ready.signal() + } + ready.wait() + let value = result.withLock { value in + let snapshot = value + value = nil + return snapshot + } + guard let value else { throw Failure.invalidResult } + return try value.value.get() +} + +private func acquire( + _ operation: (@escaping ETCoreAIAcquisitionCompletion) -> Void +) throws -> (any ETCoreAIPreparedModel)? { + try waitForResult { completed in + operation { status, model, error in + if let error { + completed(.failure(error)) + } else if (status == .hit && model != nil) || (status == .miss && model == nil) { + completed(.success(model)) + } else { + completed(.failure(Failure.invalidResult)) + } + } + } +} + +private func checkInference(_ model: any ETCoreAIPreparedModel, fixture: Fixture) throws { + guard fixture.inputs.count == 2, fixture.outputs.count == fixture.expected.count else { + throw Failure.invalidArguments + } + let session: any ETCoreAISession = try waitForResult { completed in + model.loadFunctionNamed( + fixture.function, inputNames: fixture.inputs, outputNames: fixture.outputs + ) { session, error in + if let error { + completed(.failure(error)) + } else if let session { + completed(.success(session)) + } else { + completed(.failure(Failure.invalidResult)) + } + } + } + let storage = [[Float(5), 9], [Float(1), 2]].map { values in + values.withUnsafeBytes { NSData(bytes: $0.baseAddress, length: $0.count) } + } + let inputs = storage.map { + ETCoreAIInputTensor( + bytes: $0.bytes, byteCount: UInt($0.length), shape: [2], scalarType: .float32) + } + let outputs: [ETCoreAITensor] = try withExtendedLifetime(storage) { + try waitForResult { completed in + session.executeInputs(inputs) { outputs, error in + if let error { + completed(.failure(error)) + } else if let outputs { + completed(.success(outputs)) + } else { + completed(.failure(Failure.invalidResult)) + } + } + } + } + guard outputs.count == fixture.outputs.count else { + throw Failure.incorrectOutput("count") + } + for (index, output) in outputs.enumerated() { + let name = fixture.outputs[index] + guard output.shape == [2], output.scalarType == .float32, + output.data.count == 8, fixture.expected[index].count == 2 + else { throw Failure.incorrectOutput(name) } + let values = output.data.withUnsafeBytes { + [$0.loadUnaligned(as: Float.self), $0.loadUnaligned(fromByteOffset: 4, as: Float.self)] + } + guard zip(values, fixture.expected[index]).allSatisfy({ abs($0 - $1) < 0.001 }) else { + throw Failure.incorrectOutput("\(name): \(values), expected \(fixture.expected[index])") + } + } +} + +@main +private struct SDKSourceLifetimeTest { + static func main() { + do { + let args = CommandLine.arguments + if args.count == 2, args[1] == "architecture" { + print(ETCoreAISwiftModelLoader.deviceArchitectureName) + return + } + guard args.count == 5 else { throw Failure.invalidArguments } + let mode = args[1] + let source = URL(fileURLWithPath: args[2]).standardizedFileURL + let fixtureURL = URL(fileURLWithPath: args[3]).standardizedFileURL + let bookmarkURL = URL(fileURLWithPath: args[4]) + let directory = source.deletingLastPathComponent() + guard ["model.aimodel", "model.aimodelc"].contains(source.lastPathComponent), + directory == fixtureURL.deletingLastPathComponent(), + directory.deletingLastPathComponent().lastPathComponent.hasPrefix("coreai-lifetime-") + else { throw Failure.invalidArguments } + let loader = ETCoreAISwiftModelLoader() + if mode == "evict" { + let bookmark = try Data(contentsOf: bookmarkURL) + let _: Void = try waitForResult { completed in + loader.evictModel(withBookmark: bookmark) { error in + if let error { completed(.failure(error)) } else { completed(.success(())) } + } + } + return + } + let fixture = try JSONDecoder().decode(Fixture.self, from: Data(contentsOf: fixtureURL)) + if mode == "specialize" { + guard let model = try acquire({ loader.specializeModel(at: source, completion: $0) }) else { + throw Failure.invalidResult + } + try model.copyBookmarkData().write(to: bookmarkURL, options: .atomic) + try checkInference(model, fixture: fixture) + try FileManager.default.removeItem(at: source) + guard !FileManager.default.fileExists(atPath: source.path) else { + throw Failure.unexpectedSource + } + try checkInference(model, fixture: fixture) + print("Bridge inference passed after source removal") + } else if mode == "restore" { + guard !FileManager.default.fileExists(atPath: source.path) else { + throw Failure.unexpectedSource + } + let bookmark = try Data(contentsOf: bookmarkURL) + guard + let model = try acquire({ loader.restoreModel(fromBookmark: bookmark, completion: $0) }) + else { + exit(3) + } + try checkInference(model, fixture: fixture) + print("Bridge restoration and inference passed without source") + } else { + throw Failure.invalidArguments + } + } catch { + print("SDK source-lifetime test failed: \(error)") + exit(1) + } + } +} diff --git a/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.h b/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.h new file mode 100644 index 00000000000..963adfd784a --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.h @@ -0,0 +1,58 @@ +/* + * 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 "coreai_bookmark_fixture.h" +#include "coreai_fake_loader.h" +#include "coreai_source_fixture.h" +#include "coreai_storage.h" + +namespace executorch::backends::coreai::testing { + +runtime::Result acquisition_staged_url( + NSURL* root, + const Manifest& manifest); + +struct ParallelAcquisitionObservation { + int creation[2] = {-1, -1}; + int joined[2] = {-1, -1}; + bool completed[2] = {false, false}; + runtime::Error errors[2] = { + runtime::Error::Internal, + runtime::Error::Internal}; + bool models[2] = {false, false}; + bool first_entered = false; + bool second_observed = false; + int gate_timeouts = 0; + bool callbacks_drained = false; +}; +ParallelAcquisitionObservation parallel_acquisitions( + ScopedFakeBridgeState& bridge, + BookmarkFakeLoader* loader, + NSURL* root, + const Manifest& first, + const Manifest& second, + TestData& data, + bool same_key); + +::testing::AssertionResult spawn_acquisition_crash( + BookmarkChild& child, + NSURL* root, + StorageOperation boundary, + NSDictionary* manifest); +::testing::AssertionResult spawn_parallel_acquisition( + BookmarkChild& child, + NSURL* root, + NSDictionary* first, + NSDictionary* second, + bool same_key); +// Dispatch before InitGoogleTest. Unrelated modes return -1. +int acquisition_child_mode(int argc, char** argv); + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm b/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm new file mode 100644 index 00000000000..7446e9ffbe6 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm @@ -0,0 +1,372 @@ +/* + * 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_acquisition_fixture.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "coreai_fault_scope.h" +#include "coreai_load_coordinator.h" + +namespace executorch::backends::coreai::testing { +namespace { +using runtime::Error; +using runtime::Result; + +dispatch_time_t deadline() { + return dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC); +} + +struct AcquisitionWorker { + Manifest manifest; + TestData* data = nullptr; + NSString* root = nil; + BookmarkFakeLoader* loader = nil; + Error error = Error::Internal; + bool model = false; + dispatch_semaphore_t started = dispatch_semaphore_create(0); + dispatch_semaphore_t done = dispatch_semaphore_create(0); + + static void* run(void* opaque) { + auto* worker = static_cast(opaque); + dispatch_semaphore_signal(worker->started); + @autoreleasepool { + auto acquired = acquire_bookmark_model( + worker->manifest, + worker->data, + worker->root, + @"macOS", + @"arch_b", + worker->loader); + worker->error = acquired.ok() ? Error::Ok : acquired.error(); + worker->model = acquired.ok() && acquired.get() != nil; + } + dispatch_semaphore_signal(worker->done); + return nullptr; + } +}; + +struct AcquisitionWorkers { + ScopedFakeBridgeState& bridge; + BookmarkFakeLoader* loader; + dispatch_semaphore_t proceed = dispatch_semaphore_create(0); + pthread_t threads[2]{}; + bool started[2] = {false, false}; + AcquisitionWorker workers[2]; + ParallelAcquisitionObservation& observation; + bool finished = false; + + AcquisitionWorkers( + ScopedFakeBridgeState& state, + BookmarkFakeLoader* fake, + ParallelAcquisitionObservation& result) + : bridge(state), loader(fake), observation(result) {} + + void finish() { + if (finished) + return; + for (int i = 0; i < 2; ++i) + dispatch_semaphore_signal(proceed); + for (int i = 0; i < 2; ++i) { + if (!started[i]) + continue; + observation.completed[i] = + dispatch_semaphore_wait(workers[i].done, deadline()) == 0; + observation.joined[i] = pthread_join(threads[i], nullptr); + observation.errors[i] = workers[i].error; + observation.models[i] = workers[i].model; + } + observation.callbacks_drained = bridge.wait_for_callbacks(deadline()); + if (observation.callbacks_drained) + loader.onSpecialize = nil; + finished = true; + } + ~AcquisitionWorkers() { + finish(); + } +}; + +bool send_child_byte(int fd, char value) { + const auto end = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (std::chrono::steady_clock::now() < end) { + const auto remaining = + std::chrono::duration_cast( + end - std::chrono::steady_clock::now()) + .count(); + pollfd entry{fd, POLLOUT, 0}; + const int ready = + poll(&entry, 1, static_cast(remaining > 0 ? remaining : 0)); + if (ready < 0 && errno == EINTR) + continue; + if (ready <= 0 || !(entry.revents & POLLOUT)) + return false; + const ssize_t written = write(fd, &value, 1); + if (written == 1) + return true; + if (written < 0 && (errno == EINTR || errno == EAGAIN)) + continue; + return false; + } + return false; +} +} // namespace + +Result acquisition_staged_url(NSURL* root, const Manifest& manifest) { + auto key = test_bookmark_key(manifest); + if (!key.ok()) + return key.error(); + return [[[root URLByAppendingPathComponent:@"staging"] + URLByAppendingPathComponent:key.get()] + URLByAppendingPathComponent:manifest.path.lastPathComponent]; +} + +ParallelAcquisitionObservation parallel_acquisitions( + ScopedFakeBridgeState& bridge, + BookmarkFakeLoader* loader, + NSURL* root, + const Manifest& first, + const Manifest& second, + TestData& data, + bool same_key) { + ParallelAcquisitionObservation observation; + auto timeouts = std::make_shared>(0); + dispatch_semaphore_t entered = dispatch_semaphore_create(0); + AcquisitionWorkers group(bridge, loader, observation); + dispatch_semaphore_t proceed = group.proceed; + loader.onSpecialize = ^(NSURL*) { + dispatch_semaphore_signal(entered); + if (dispatch_semaphore_wait(proceed, deadline()) != 0) + ++*timeouts; + }; + for (int i = 0; i < 2; ++i) { + auto& worker = group.workers[i]; + worker.manifest = i == 0 ? first : second; + worker.data = &data; + worker.root = root.path; + worker.loader = loader; + } + observation.creation[0] = pthread_create( + &group.threads[0], nullptr, AcquisitionWorker::run, &group.workers[0]); + group.started[0] = observation.creation[0] == 0; + if (group.started[0]) { + observation.first_entered = + dispatch_semaphore_wait(entered, deadline()) == 0; + } + if (observation.first_entered) { + observation.creation[1] = pthread_create( + &group.threads[1], nullptr, AcquisitionWorker::run, &group.workers[1]); + group.started[1] = observation.creation[1] == 0; + if (group.started[1] && + dispatch_semaphore_wait(group.workers[1].started, deadline()) == 0) { + const auto second_deadline = same_key + ? dispatch_time(DISPATCH_TIME_NOW, 50 * NSEC_PER_MSEC) + : deadline(); + const bool second_entered = + dispatch_semaphore_wait(entered, second_deadline) == 0; + observation.second_observed = same_key ? !second_entered : second_entered; + } + } + group.finish(); + observation.gate_timeouts = timeouts->load(); + return observation; +} + +::testing::AssertionResult spawn_acquisition_crash( + BookmarkChild& child, + NSURL* root, + StorageOperation boundary, + NSDictionary* manifest) { + // Serialize the shared fixture in the parent; its builders contain GTest + // checks. + NSData* payload = encode( + @{@"boundary" : @(static_cast(boundary)), + @"manifest" : manifest}); + if (payload == nil) + return ::testing::AssertionFailure() << "Cannot encode crash child input"; + NSString* argument = [[NSString alloc] initWithData:payload + encoding:NSUTF8StringEncoding]; + return child.spawn(@"--acquisition-crash", root.path, argument); +} + +::testing::AssertionResult spawn_parallel_acquisition( + BookmarkChild& child, + NSURL* root, + NSDictionary* first, + NSDictionary* second, + bool same_key) { + NSData* payload = encode( + @{@"manifest" : first, + @"second" : second, + @"same_key" : @(same_key)}); + if (payload == nil) + return ::testing::AssertionFailure() + << "Cannot encode parallel child input"; + NSString* argument = [[NSString alloc] initWithData:payload + encoding:NSUTF8StringEncoding]; + return child.spawn(@"--acquisition-parallel", root.path, argument); +} + +namespace { +[[noreturn]] void parallel_child( + NSString* root, + const Manifest& first, + const Manifest& second, + bool same_key, + int fd) { + ScopedFakeBridgeState bridge; + @autoreleasepool { + BookmarkFakeLoader* loader = [[BookmarkFakeLoader alloc] init]; + if (loader == nil) + _exit(8); + bridge.state().bookmark_loader = loader; + TestData data; + auto result = parallel_acquisitions( + bridge, + loader, + [NSURL fileURLWithPath:root], + first, + second, + data, + same_key); + if (!result.callbacks_drained) + _exit(11); + // These bytes are observations, not child assertions. The parent checks + // each one. + const bool harness_ok = result.creation[0] == 0 && + result.creation[1] == 0 && result.first_entered && + result.joined[0] == 0 && result.joined[1] == 0 && result.completed[0] && + result.completed[1] && result.gate_timeouts == 0; + const bool both_ok = result.errors[0] == Error::Ok && result.models[0] && + result.errors[1] == Error::Ok && result.models[1]; + const char observations[] = { + static_cast(harness_ok), + static_cast(result.second_observed), + static_cast(both_ok), + static_cast(loader->specializations.load()), + static_cast(data.attempts.load()), + static_cast(loader->restores.load())}; + for (char observation : observations) { + if (!send_child_byte(fd, observation)) + _exit(12); + } + if (loader->evictions.load() != 0 || + data.requests.load() != data.releases.load()) + _exit(14); + [loader clearRetainedState]; + bridge.state().bookmark_loader = nil; + loader = nil; + } + auto& state = bridge.state(); + const bool none_live = state.prepared_models.load() == 0 && + state.loaders.load() == 0 && state.missing_bundles.load() == 0; + if (!send_child_byte(fd, static_cast(none_live))) + _exit(12); + // No GTest-bearing destructors run before InitGoogleTest, even on a failed + // child. + _exit(close(fd) == 0 ? 0 : 13); +} +} // namespace + +int acquisition_child_mode(int argc, char** argv) { + if (argc < 2) + return -1; + const bool crash = std::strcmp(argv[1], "--acquisition-crash") == 0; + const bool parallel = std::strcmp(argv[1], "--acquisition-parallel") == 0; + if (!crash && !parallel) + return -1; + if (argc != 5) + return 1; + NSString* root = [NSString stringWithUTF8String:argv[2]]; + NSString* argument = [NSString stringWithUTF8String:argv[3]]; + if (root == nil || argument == nil) + return 2; + id payload = [NSJSONSerialization + JSONObjectWithData:[argument dataUsingEncoding:NSUTF8StringEncoding] + options:0 + error:nil]; + if (![payload isKindOfClass:NSDictionary.class]) + return 3; + id number = payload[@"boundary"]; + id dictionary = payload[@"manifest"]; + if (![dictionary isKindOfClass:NSDictionary.class]) + return 3; + StorageOperation boundary = StorageOperation::AfterSDK; + if (crash) { + if (![number isKindOfClass:NSNumber.class]) + return 3; + boundary = static_cast([number intValue]); + if (boundary != StorageOperation::BeforeSDK && + boundary != StorageOperation::AfterSDK) + return 4; + } + auto parsed = parse_manifest(encode(dictionary)); + if (!parsed.ok()) + return 5; + auto selected = select_assets(parsed.get(), @"arch_b", @"macOS"); + if (!selected.ok()) + return 5; + char* end = nullptr; + errno = 0; + const long descriptor = std::strtol(argv[4], &end, 10); + if (errno != 0 || end == argv[4] || *end != '\0' || descriptor < 0 || + descriptor > INT_MAX) + return 6; + const int fd = static_cast(descriptor); + if (fcntl(fd, F_GETFD) < 0) + return 6; + signal(SIGPIPE, SIG_IGN); + if (!send_child_byte(fd, 'C')) + return 7; + if (parallel) { + id second_dictionary = payload[@"second"]; + id same_key = payload[@"same_key"]; + if (![second_dictionary isKindOfClass:NSDictionary.class] || + ![same_key isKindOfClass:NSNumber.class]) + return 3; + auto second_parsed = parse_manifest(encode(second_dictionary)); + if (!second_parsed.ok()) + return 5; + auto second = select_assets(second_parsed.get(), @"arch_b", @"macOS"); + if (!second.ok()) + return 5; + parallel_child( + root, selected.get(), second.get(), [same_key boolValue], fd); + } + if (close(fd) != 0) + return 7; + + ScopedFakeBridgeState bridge; + @autoreleasepool { + // Terminal child paths deliberately bypass assertion-bearing scope + // destructors. + BookmarkFakeLoader* loader = [[BookmarkFakeLoader alloc] init]; + if (loader == nil) + _exit(8); + bridge.state().bookmark_loader = loader; + TestData data; + StorageFaultScope fault(boundary); + fault.armed = false; + fault.observe = ^(StorageOperation event) { + if (event == boundary) + _exit(73); + }; + auto result = acquire_bookmark_model( + selected.get(), &data, root, @"macOS", @"arch_b", loader); + _exit(result.ok() ? 9 : 10); + } +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm new file mode 100644 index 00000000000..b70ae5380d6 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm @@ -0,0 +1,798 @@ +/* + * 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_acquisition_fixture.h" + +#include +#include +#include +#include +#include +#include +#include "coreai_fault_scope.h" +#include "coreai_filesystem_fixture.h" +#include "coreai_load_coordinator.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +@interface InvalidAcquisitionLoader : BookmarkFakeLoader +@property(nonatomic) ETCoreAIAcquisitionStatus restoreStatus; +@property(nonatomic, strong) id restoreModel; +@end + +@implementation InvalidAcquisitionLoader +- (void)restoreModelFromBookmark:(NSData*)bookmark + completion:(ETCoreAIAcquisitionCompletion)completion { + ++restores; + completion(self.restoreStatus, self.restoreModel, nil); +} +- (void)specializeModelAtURL:(NSURL*)url + completion:(ETCoreAIAcquisitionCompletion)completion { + ++specializations; + completion(ETCoreAIAcquisitionStatusMiss, nil, nil); +} +@end + +namespace { + +class CoreAIAcquisitionTest : public ::testing::Test { + protected: + // Declared first so each case's loader, models and data die before it. + ScopedFakeBridgeState bridge; + TestDirectory root; + TestData data; + Manifest manifest; + NSString* key = nil; + + void SetUp() override { + ASSERT_NE(root.url, nil); + auto selected = bookmark_manifest(); + ASSERT_TRUE(selected.ok()); + manifest = selected.get(); + auto identity = test_bookmark_key(manifest); + ASSERT_TRUE(identity.ok()); + key = identity.get(); + } + + Error acquire(BookmarkFakeLoader* loader, const NamedDataMap* named_data) { + @autoreleasepool { + auto model = acquire_bookmark_model( + manifest, named_data, root.url.path, @"macOS", @"arch_b", loader); + if (!model.ok()) + return model.error(); + return model.get() != nil ? Error::Ok : Error::Internal; + } + } + + NSData* saved() { + auto bookmark = saved_bookmark(root.url, key); + EXPECT_TRUE(bookmark.ok()); + return bookmark.ok() ? bookmark.get() : nil; + } + + NSURL* staging_key_url() { + return [[[root.url URLByAppendingPathComponent:@"staging"] + URLByAppendingPathComponent:key] copy]; + } + + BOOL staging_key_exists() { + return + [NSFileManager.defaultManager fileExistsAtPath:staging_key_url().path]; + } + + NSURL* lock_url() { + return test_bookmark_lock_url(root.url, key); + } + + ino_t lock_inode() { + struct stat info; + if (lstat(lock_url().fileSystemRepresentation, &info) != 0) + return 0; + return info.st_ino; + } +}; + +TEST_F(CoreAIAcquisitionTest, ColdSpecializesStagedSourceAndPersistsBookmark) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(data.attempts.load(), 1); + auto staged = acquisition_staged_url(root.url, manifest); + ASSERT_TRUE(staged.ok()); + EXPECT_TRUE([fake.loader.lastSpecializeURL isEqual:staged.get().path]); + EXPECT_TRUE([saved() isEqual:fake.loader.lastAcquired.bookmark]); + EXPECT_FALSE(staging_key_exists()); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:staged.get().path]); +} + +TEST_F(CoreAIAcquisitionTest, WarmHitRestoresWithoutNamedData) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + NSData* original = saved(); + EXPECT_FALSE(staging_key_exists()); + const int attempts = data.attempts.load(); + ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->restores.load(), 1); + EXPECT_EQ(data.attempts.load(), attempts); + EXPECT_TRUE([saved() isEqual:original]); + EXPECT_FALSE(staging_key_exists()); +} + +TEST_F(CoreAIAcquisitionTest, RestoreErrorFailsClosedWithoutSourceFallback) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + NSData* original = saved(); + const int attempts = data.attempts; + fake.loader.restoreError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EIO + userInfo:nil]; + EXPECT_EQ(acquire(fake.loader, &data), Error::Internal); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(data.attempts.load(), attempts); + EXPECT_TRUE([saved() isEqual:original]); +} + +TEST_F( + CoreAIAcquisitionTest, + RestoreMissRespecializesStagedSourceAndReplacesBookmark) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + NSData* original = saved(); + const int attempts = data.attempts; + fake.loader.restoreMiss = YES; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 2); + EXPECT_EQ(data.attempts.load(), attempts + 1); + NSData* replaced = saved(); + EXPECT_NE(replaced, nil); + EXPECT_FALSE([replaced isEqual:original]); + EXPECT_FALSE(staging_key_exists()); +} + +TEST_F( + CoreAIAcquisitionTest, + SpecializationErrorFailsClosedPreservingBookmark) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + NSData* original = saved(); + ASSERT_NE(original, nil); + const int attempts = data.attempts.load(); + const int specializations = fake.loader->specializations.load(); + fake.loader.restoreMiss = YES; + fake.loader.specializeError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EIO + userInfo:nil]; + EXPECT_EQ(acquire(fake.loader, &data), Error::Internal); + EXPECT_EQ(fake.loader->specializations.load(), specializations + 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(data.attempts.load(), attempts + 1); + EXPECT_TRUE([saved() isEqual:original]); + EXPECT_TRUE([NSFileManager.defaultManager + fileExistsAtPath:fake.loader.lastSpecializeURL]); +} + +TEST_F(CoreAIAcquisitionTest, EmptyBookmarkFailsClosedPreservingPrior) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + NSData* original = saved(); + ASSERT_NE(original, nil); + const int specializations = fake.loader->specializations.load(); + fake.loader.restoreMiss = YES; + fake.loader.emptyBookmark = YES; + EXPECT_EQ(acquire(fake.loader, &data), Error::InvalidExternalData); + EXPECT_EQ(fake.loader->specializations.load(), specializations + 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_TRUE([saved() isEqual:original]); + EXPECT_TRUE([NSFileManager.defaultManager + fileExistsAtPath:fake.loader.lastSpecializeURL]); +} + +TEST_F(CoreAIAcquisitionTest, NilLoaderRejectsWithoutFilesystemEffects) { + NSDictionary* before = nil; + ASSERT_TRUE(asset_tree_snapshot(root.url, before)); + Error error = Error::Ok; + @autoreleasepool { + auto result = acquire_bookmark_model( + manifest, &data, root.url.path, @"macOS", @"arch_b", nil); + ASSERT_FALSE(result.ok()); + error = result.error(); + } + EXPECT_EQ(error, Error::InvalidArgument); + EXPECT_TRUE(snapshot_matches(root.url, before)); + EXPECT_EQ(saved(), nil); + EXPECT_FALSE(staging_key_exists()); +} + +TEST_F(CoreAIAcquisitionTest, RestoreHitWithoutModelFailsClosed) { + NSData* original = nil; + { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + original = saved(); + ASSERT_NE(original, nil); + } + InvalidAcquisitionLoader* loader = [[InvalidAcquisitionLoader alloc] init]; + ASSERT_NE(loader, nil); + loader.restoreStatus = ETCoreAIAcquisitionStatusHit; + Error error = Error::Ok; + @autoreleasepool { + auto result = acquire_bookmark_model( + manifest, &data, root.url.path, @"macOS", @"arch_b", loader); + ASSERT_FALSE(result.ok()); + error = result.error(); + } + EXPECT_EQ(error, Error::Internal); + loader = nil; + EXPECT_TRUE([saved() isEqual:original]); +} + +TEST_F(CoreAIAcquisitionTest, RestoreMissWithModelFailsClosed) { + NSData* original = nil; + { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + original = saved(); + ASSERT_NE(original, nil); + } + InvalidAcquisitionLoader* loader = [[InvalidAcquisitionLoader alloc] init]; + ASSERT_NE(loader, nil); + BookmarkFakeModel* model = [[BookmarkFakeModel alloc] init]; + ASSERT_NE(model, nil); + loader.restoreStatus = ETCoreAIAcquisitionStatusMiss; + loader.restoreModel = model; + Error error = Error::Ok; + @autoreleasepool { + auto result = acquire_bookmark_model( + manifest, &data, root.url.path, @"macOS", @"arch_b", loader); + ASSERT_FALSE(result.ok()); + error = result.error(); + } + EXPECT_EQ(error, Error::Internal); + loader.restoreModel = nil; + loader = nil; + EXPECT_TRUE([saved() isEqual:original]); +} + +TEST_F(CoreAIAcquisitionTest, RestoreErrorWithoutErrorFailsClosed) { + NSData* original = nil; + { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + original = saved(); + ASSERT_NE(original, nil); + } + InvalidAcquisitionLoader* loader = [[InvalidAcquisitionLoader alloc] init]; + ASSERT_NE(loader, nil); + loader.restoreStatus = ETCoreAIAcquisitionStatusError; + Error error = Error::Ok; + @autoreleasepool { + auto result = acquire_bookmark_model( + manifest, &data, root.url.path, @"macOS", @"arch_b", loader); + ASSERT_FALSE(result.ok()); + error = result.error(); + } + EXPECT_EQ(error, Error::Internal); + loader = nil; + EXPECT_TRUE([saved() isEqual:original]); +} + +TEST_F(CoreAIAcquisitionTest, SpecializeMissFailsClosed) { + InvalidAcquisitionLoader* loader = [[InvalidAcquisitionLoader alloc] init]; + ASSERT_NE(loader, nil); + Error error = Error::Ok; + @autoreleasepool { + auto result = acquire_bookmark_model( + manifest, &data, root.url.path, @"macOS", @"arch_b", loader); + ASSERT_FALSE(result.ok()); + error = result.error(); + } + EXPECT_EQ(error, Error::Internal); + EXPECT_EQ(loader->specializations.load(), 1); + EXPECT_EQ(loader->restores.load(), 0); + loader = nil; + EXPECT_EQ(saved(), nil); +} + +TEST_F( + CoreAIAcquisitionTest, + ColdSuccessRemovesWholeStagingAfterBookmarkFlush) { + BookmarkLoaderScope fake; + __block std::vector ops; + Error error = Error::Internal; + { + StorageFaultScope fault(StorageOperation::Remove); + fault.armed = false; + fault.observe = ^(StorageOperation operation) { + ops.push_back(operation); + }; + error = acquire(fake.loader, &data); + } + ASSERT_EQ(error, Error::Ok); + EXPECT_TRUE([saved() isEqual:fake.loader.lastAcquired.bookmark]); + EXPECT_FALSE(staging_key_exists()); + auto staged = acquisition_staged_url(root.url, manifest); + ASSERT_TRUE(staged.ok()); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:staged.get().path]); + size_t last_flush = 0; + size_t removes = 0; + size_t first_remove = ops.size(); + for (size_t i = 0; i < ops.size(); ++i) { + if (ops[i] == StorageOperation::FullSync) + last_flush = i; + if (ops[i] == StorageOperation::Remove) { + ++removes; + if (first_remove == ops.size()) + first_remove = i; + } + } + EXPECT_EQ(removes, 3u); + EXPECT_GT(first_remove, last_flush); +} + +TEST_F(CoreAIAcquisitionTest, SpecializeErrorPerformsNoCleanup) { + BookmarkLoaderScope fake; + __block std::vector ops; + StorageFaultScope fault(StorageOperation::Remove); + fault.armed = false; + fault.observe = ^(StorageOperation operation) { + ops.push_back(operation); + }; + fake.loader.specializeError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EIO + userInfo:nil]; + EXPECT_EQ(acquire(fake.loader, &data), Error::Internal); + size_t removes = 0; + for (auto operation : ops) { + if (operation == StorageOperation::Remove) + ++removes; + } + EXPECT_EQ(removes, 0u); + EXPECT_EQ(saved(), nil); + EXPECT_EQ(fake.loader->evictions.load(), 0); +} + +TEST_F(CoreAIAcquisitionTest, BookmarkFinalFlushFailurePerformsNoCleanup) { + BookmarkLoaderScope fake; + __block std::vector ops; + Error error = Error::Internal; + { + // Cold source contributes two FullSyncs and the bookmark two more, so + // index four is the post-rename bookmark flush. + StorageFaultScope fault(StorageOperation::FullSync); + fault.match_index = 4; + fault.observe = ^(StorageOperation operation) { + ops.push_back(operation); + }; + error = acquire(fake.loader, &data); + } + EXPECT_EQ(error, Error::AccessFailed); + EXPECT_TRUE([saved() isEqual:fake.loader.lastAcquired.bookmark]); + size_t removes = 0; + for (auto operation : ops) { + if (operation == StorageOperation::Remove) + ++removes; + } + EXPECT_EQ(removes, 0u); + auto staged = acquisition_staged_url(root.url, manifest); + ASSERT_TRUE(staged.ok()); + EXPECT_TRUE( + [NSFileManager.defaultManager fileExistsAtPath:staged.get().path]); +} + +TEST_F( + CoreAIAcquisitionTest, + CleanupFirstRemoveFailurePreservesDurableBookmark) { + BookmarkLoaderScope fake; + NSData* bookmark = nil; + int specializations = 0; + { + StorageFaultScope fault(StorageOperation::Remove); + fault.match_index = 1; + EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fault.hits, 1); + bookmark = saved(); + ASSERT_NE(bookmark, nil); + EXPECT_TRUE([bookmark isEqual:fake.loader.lastAcquired.bookmark]); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(fake.loader->evictions.load(), 0); + specializations = fake.loader->specializations.load(); + } + EXPECT_EQ(fake.loader->specializations.load(), specializations); + const int attempts = data.attempts.load(); + const int restores = fake.loader->restores.load(); + EXPECT_TRUE(staging_key_exists()); + ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), specializations); + EXPECT_EQ(fake.loader->restores.load(), restores + 1); + EXPECT_EQ(data.attempts.load(), attempts); + EXPECT_TRUE([saved() isEqual:bookmark]); + EXPECT_FALSE(staging_key_exists()); +} + +TEST_F( + CoreAIAcquisitionTest, + CleanupPartialRemoveFailurePreservesDurableBookmark) { + BookmarkLoaderScope fake; + NSData* bookmark = nil; + int specializations = 0; + { + StorageFaultScope fault(StorageOperation::Remove); + fault.match_index = 2; + EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fault.hits, 1); + bookmark = saved(); + ASSERT_NE(bookmark, nil); + EXPECT_TRUE([bookmark isEqual:fake.loader.lastAcquired.bookmark]); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + specializations = fake.loader->specializations.load(); + } + EXPECT_EQ(fake.loader->specializations.load(), specializations); + const int attempts = data.attempts.load(); + const int restores = fake.loader->restores.load(); + EXPECT_TRUE(staging_key_exists()); + ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), specializations); + EXPECT_EQ(fake.loader->restores.load(), restores + 1); + EXPECT_EQ(data.attempts.load(), attempts); + EXPECT_TRUE([saved() isEqual:bookmark]); + EXPECT_FALSE(staging_key_exists()); +} + +TEST_F( + CoreAIAcquisitionTest, + CleanupDirectorySyncFailurePreservesDurableBookmark) { + BookmarkLoaderScope fake; + NSData* bookmark = nil; + int specializations = 0; + { + StorageFaultScope scope(StorageOperation::SyncDirectory); + scope.armed = false; + StorageFaultScope* arming = &scope; + scope.observe = ^(StorageOperation operation) { + if (operation == StorageOperation::Remove) + arming->armed = true; + }; + EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(scope.hits, 1); + bookmark = saved(); + ASSERT_NE(bookmark, nil); + EXPECT_TRUE([bookmark isEqual:fake.loader.lastAcquired.bookmark]); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + specializations = fake.loader->specializations.load(); + } + const int attempts = data.attempts.load(); + const int restores = fake.loader->restores.load(); + EXPECT_TRUE(staging_key_exists()); + ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), specializations); + EXPECT_EQ(fake.loader->restores.load(), restores + 1); + EXPECT_EQ(data.attempts.load(), attempts); + EXPECT_TRUE([saved() isEqual:bookmark]); + EXPECT_FALSE(staging_key_exists()); +} + +TEST_F(CoreAIAcquisitionTest, WarmHitCleanupFailureStillReturnsModel) { + BookmarkLoaderScope fake; + { + StorageFaultScope fault(StorageOperation::Remove); + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + } + NSData* original = saved(); + StorageFaultScope fault(StorageOperation::Remove); + EXPECT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fault.hits, 1); + EXPECT_EQ(fake.loader->restores.load(), 1); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(data.attempts.load(), 1); + EXPECT_TRUE([saved() isEqual:original]); +} + +TEST_F(CoreAIAcquisitionTest, PartialCleanupThenMissRebuildsSource) { + BookmarkLoaderScope fake; + { + StorageFaultScope fault(StorageOperation::Remove); + fault.match_index = 2; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fault.hits, 1); + } + ASSERT_TRUE(staging_key_exists()); + NSData* original = saved(); + fake.loader.restoreMiss = YES; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fake.loader->restores.load(), 1); + EXPECT_EQ(fake.loader->specializations.load(), 2); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(data.attempts.load(), 2); + EXPECT_FALSE([saved() isEqual:original]); + EXPECT_FALSE(staging_key_exists()); +} + +TEST_F( + CoreAIAcquisitionTest, + DamagedSourceCleanupFailureDoesNotReadOrSpecialize) { + BookmarkLoaderScope fake; + { + StorageFaultScope fault(StorageOperation::Remove); + fault.match_index = 2; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + } + NSData* original = saved(); + fake.loader.restoreMiss = YES; + { + StorageFaultScope fault(StorageOperation::Remove); + EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_EQ(data.attempts.load(), 1); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_TRUE([saved() isEqual:original]); + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_FALSE(staging_key_exists()); +} + +TEST_F(CoreAIAcquisitionTest, MissingNamedDataIsNotRetried) { + BookmarkLoaderScope fake; + data.files.clear(); + EXPECT_EQ(acquire(fake.loader, &data), Error::InvalidExternalData); + EXPECT_EQ(data.attempts.load(), 1); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(saved(), nil); +} + +TEST_F(CoreAIAcquisitionTest, BeforeSdkFaultNeverInvokesLoader) { + BookmarkLoaderScope fake; + StorageFaultScope fault(StorageOperation::BeforeSDK); + EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(fault.hits, 1); + EXPECT_EQ(saved(), nil); +} + +TEST_F(CoreAIAcquisitionTest, KeyLockIdentitySurvivesSuccessfulAcquisition) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + ino_t first = lock_inode(); + ASSERT_NE(first, 0u); + ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(lock_inode(), first); +} + +class CoreAIAcquisitionSaveFailureTest + : public CoreAIAcquisitionTest, + public ::testing::WithParamInterface {}; +TEST_P(CoreAIAcquisitionSaveFailureTest, PreservesBookmarkAndRetries) { + const bool prior = GetParam(); + BookmarkLoaderScope fake; + NSData* original = nil; + if (prior) { + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + original = saved(); + ASSERT_NE(original, nil); + EXPECT_FALSE(staging_key_exists()); + fake.loader.restoreMiss = YES; + } + { + // Cold and rematerialized sources publish before the bookmark, so the + // bookmark rename is the second rename once the first success has + // cleaned the source. + StorageFaultScope fault(StorageOperation::Rename); + fault.match_index = 2; + EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_TRUE(original == nil ? saved() == nil : [saved() isEqual:original]); + EXPECT_TRUE([NSFileManager.defaultManager + fileExistsAtPath:fake.loader.lastSpecializeURL]); + const int calls = fake.loader->specializations; + fake.loader.restoreMiss = NO; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), calls + (prior ? 0 : 1)); + EXPECT_FALSE(staging_key_exists()); +} +INSTANTIATE_TEST_SUITE_P( + Boundaries, + CoreAIAcquisitionSaveFailureTest, + ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return info.param ? "PriorBookmark" : "Cold"; + }); + +class CoreAIAcquisitionProcessTest : public CoreAIAcquisitionTest { + protected: + void expect_parallel(bool same_key) { + NSMutableDictionary* first = manifest_dict(); + NSMutableDictionary* second = [first mutableCopy]; + if (!same_key) + second[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity('2')}; + BookmarkChild child; + ASSERT_TRUE( + spawn_parallel_acquisition(child, root.url, first, second, same_key)); + ASSERT_TRUE(child.receive('C')); + ASSERT_TRUE(child.receive(1)) << "harness completed without timeouts"; + ASSERT_TRUE(child.receive(1)) + << "second acquisition blocked or entered as expected"; + ASSERT_TRUE(child.receive(1)) << "both acquisitions succeeded"; + ASSERT_TRUE(child.receive(same_key ? 1 : 2)) << "specializations"; + ASSERT_TRUE(child.receive(same_key ? 1 : 2)) << "named-data attempts"; + ASSERT_TRUE(child.receive(same_key ? 1 : 0)) << "restores"; + ASSERT_TRUE(child.receive(1)) << "no live fake SDK objects"; + ASSERT_TRUE(child.expect_exit(0)); + } +}; + +TEST_F(CoreAIAcquisitionProcessTest, ContendedSameKeySerializesSpecialization) { + expect_parallel(true); +} + +TEST_F(CoreAIAcquisitionProcessTest, IndependentKeysSpecializeInParallel) { + expect_parallel(false); +} + +TEST_F(CoreAIAcquisitionProcessTest, CrashBeforeSdkRetainsStagingAndRetries) { + BookmarkChild child; + ASSERT_TRUE(spawn_acquisition_crash( + child, root.url, StorageOperation::BeforeSDK, manifest_dict())); + ASSERT_TRUE(child.receive('C')); + ASSERT_TRUE(child.expect_exit(73)); + EXPECT_EQ(saved(), nil); + EXPECT_TRUE(staging_key_exists()); + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(data.attempts.load(), 0); + EXPECT_FALSE(staging_key_exists()); +} + +TEST_F(CoreAIAcquisitionProcessTest, CrashAfterSdkRetainsStagingAndRetries) { + BookmarkChild child; + ASSERT_TRUE(spawn_acquisition_crash( + child, root.url, StorageOperation::AfterSDK, manifest_dict())); + ASSERT_TRUE(child.receive('C')); + ASSERT_TRUE(child.expect_exit(73)); + EXPECT_EQ(saved(), nil); + auto staged = acquisition_staged_url(root.url, manifest); + ASSERT_TRUE(staged.ok()); + EXPECT_TRUE( + [NSFileManager.defaultManager fileExistsAtPath:staged.get().path]); + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(data.attempts.load(), 0); + EXPECT_NE(saved(), nil); + EXPECT_FALSE(staging_key_exists()); +} + +class CoreAILifecycleLogTest : public CoreAIAcquisitionTest { + protected: + void SetUp() override { + CoreAIAcquisitionTest::SetUp(); + if (!ET_LOG_ENABLED || LogLevel::Info < LogLevel::ET_MIN_LOG_LEVEL) { + GTEST_SKIP() << "Info logging is compiled out"; + } + } + + template + std::string capture(Operation operation, Error expected = Error::Ok) { + ::testing::internal::CaptureStderr(); + const auto error = operation(); + auto logs = ::testing::internal::GetCapturedStderr(); + EXPECT_EQ(error, expected); + return logs; + } + + void expect_order( + const std::string& logs, + std::initializer_list events) { + size_t position = 0; + for (const char* event : events) { + const auto found = logs.find(event, position); + ASSERT_NE(found, std::string::npos) << event << "\n" << logs; + position = found + 1; + } + EXPECT_NE(logs.find(key.UTF8String), std::string::npos); + } +}; + +TEST_F(CoreAILifecycleLogTest, ColdAcquisitionLogsValidatedMilestonesInOrder) { + BookmarkLoaderScope fake; + const auto logs = capture([&] { return acquire(fake.loader, &data); }); + expect_order( + logs, + { + "Core AI cache lookup:", + "Core AI cache miss: no bookmark", + "Core AI specialization started:", + "Core AI specialization succeeded:", + "Core AI cache bookmark saved durably:", + }); + EXPECT_EQ(logs.find("Core AI cache hit:"), std::string::npos); + EXPECT_TRUE([saved() isEqual:fake.loader.lastAcquired.bookmark]); +} + +TEST_F(CoreAILifecycleLogTest, WarmHitDoesNotLogSpecializationOrSave) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + const auto logs = capture([&] { return acquire(fake.loader, nullptr); }); + expect_order( + logs, + { + "Core AI cache lookup:", + "Core AI cache hit: restored model from bookmark", + }); + EXPECT_EQ(logs.find("Core AI cache miss:"), std::string::npos); + EXPECT_EQ(logs.find("Core AI specialization"), std::string::npos); + EXPECT_EQ(logs.find("Core AI cache bookmark saved"), std::string::npos); +} + +TEST_F(CoreAILifecycleLogTest, ConfirmedMissLogsRespecialization) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + fake.loader.restoreMiss = YES; + const auto logs = capture([&] { return acquire(fake.loader, &data); }); + expect_order( + logs, + { + "Core AI cache miss: SDK could not restore bookmark", + "Core AI specialization started:", + "Core AI specialization succeeded:", + "Core AI cache bookmark saved durably:", + }); + EXPECT_EQ(logs.find("Core AI cache hit:"), std::string::npos); +} + +TEST_F(CoreAILifecycleLogTest, RestoreErrorIsNotReportedAsHitOrMiss) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + fake.loader.restoreError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EIO + userInfo:nil]; + const auto logs = + capture([&] { return acquire(fake.loader, &data); }, Error::Internal); + EXPECT_EQ(logs.find("Core AI cache hit:"), std::string::npos); + EXPECT_EQ(logs.find("Core AI cache miss:"), std::string::npos); + EXPECT_EQ(logs.find("Core AI specialization"), std::string::npos); + EXPECT_EQ(logs.find("Core AI cache bookmark saved"), std::string::npos); +} + +TEST_F(CoreAILifecycleLogTest, SpecializationErrorDoesNotLogSuccessOrSave) { + BookmarkLoaderScope fake; + fake.loader.specializeError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EIO + userInfo:nil]; + const auto logs = + capture([&] { return acquire(fake.loader, &data); }, Error::Internal); + EXPECT_NE(logs.find("Core AI specialization started:"), std::string::npos); + EXPECT_EQ(logs.find("Core AI specialization succeeded:"), std::string::npos); + EXPECT_EQ(logs.find("Core AI cache bookmark saved"), std::string::npos); +} + +TEST_F(CoreAILifecycleLogTest, BeforeSdkFaultDoesNotLogSpecializationStart) { + BookmarkLoaderScope fake; + StorageFaultScope fault(StorageOperation::BeforeSDK); + const auto logs = + capture([&] { return acquire(fake.loader, &data); }, Error::AccessFailed); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(logs.find("Core AI specialization"), std::string::npos); + EXPECT_EQ(logs.find("Core AI cache bookmark saved"), std::string::npos); +} + +} // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.h b/backends/apple/coreai/runtime/test/coreai_fake_loader.h new file mode 100644 index 00000000000..cf8f70a6a6b --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.h @@ -0,0 +1,88 @@ +/* + * 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 +#import "ETCoreAIBridge.h" + +@class CoreAIFakeBridgeContext; +@class BookmarkFakeLoader; + +namespace executorch::backends::coreai::testing { + +// Configure on the test thread before starting workers; join before changing +// it. +struct FakeBridgeState { + id bookmark_loader = nil; + std::atomic prepared_models{0}; + std::atomic loaders{0}; + // Worker-side contract failures are checked by the scope on the test thread. + std::atomic missing_bundles{0}; +}; + +// Process-visible, not thread-local. Scopes and configuration changes are +// serial. +class ScopedFakeBridgeState final { + public: + ScopedFakeBridgeState(); + ~ScopedFakeBridgeState(); + FakeBridgeState& state(); + // Join caller threads before draining callbacks. + ::testing::AssertionResult wait_for_callbacks(dispatch_time_t deadline); + ScopedFakeBridgeState(const ScopedFakeBridgeState&) = delete; + ScopedFakeBridgeState& operator=(const ScopedFakeBridgeState&) = delete; + + private: + CoreAIFakeBridgeContext* context_; + CoreAIFakeBridgeContext* previous_; +}; + +} // namespace executorch::backends::coreai::testing + +@interface BookmarkFakeModel : NSObject +@property(nonatomic, weak) BookmarkFakeLoader* owner; +@property(nonatomic, copy) NSData* bookmark; +@end + +@interface BookmarkFakeLoader : NSObject { + @public + std::atomic restores; + std::atomic specializations; + std::atomic evictions; +} +@property(nonatomic, strong) + NSMutableDictionary* models; +@property(nonatomic, strong) NSError* restoreError; +@property(nonatomic, strong) NSError* specializeError; +@property(nonatomic) BOOL restoreMiss; +@property(nonatomic) BOOL emptyBookmark; +@property(nonatomic, copy) void (^onSpecialize)(NSURL*); +@property(nonatomic, strong) BookmarkFakeModel* lastAcquired; +@property(nonatomic, copy) NSString* lastSpecializeURL; +// Call only after workers/callbacks finish and caller-owned handles are +// released. +- (void)clearRetainedState; +@end + +namespace executorch::backends::coreai::testing { + +struct BookmarkLoaderScope final { + BookmarkFakeLoader* loader = nil; + BookmarkLoaderScope(); + ~BookmarkLoaderScope(); + BookmarkLoaderScope(const BookmarkLoaderScope&) = delete; + BookmarkLoaderScope& operator=(const BookmarkLoaderScope&) = delete; + + private: + CoreAIFakeBridgeContext* context_; + id previous_; +}; + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm new file mode 100644 index 00000000000..18dfe5208e7 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm @@ -0,0 +1,250 @@ +/* + * 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. + */ + +#import "coreai_fake_loader.h" + +using executorch::backends::coreai::testing::FakeBridgeState; + +// ARC keeps failed asynchronous work from outliving its state storage. +@interface CoreAIFakeBridgeContext : NSObject { + @public + FakeBridgeState state; +} +@property(nonatomic, strong) dispatch_group_t callbacks; +@end + +@implementation CoreAIFakeBridgeContext +- (instancetype)init { + if ((self = [super init])) { + self.callbacks = dispatch_group_create(); + } + return self; +} +@end + +namespace { +CoreAIFakeBridgeContext* current_context = nil; + +NSError* fake_error(ETCoreAIErrorCode code) { + return [NSError errorWithDomain:ETCoreAIErrorDomain code:code userInfo:nil]; +} + +// Release callback captures before publishing that all fake work has drained. +void run_async(CoreAIFakeBridgeContext* context, void (^work)()) { + dispatch_group_t group = context.callbacks; + __block void (^task)() = [work copy]; + dispatch_group_enter(group); + dispatch_async(dispatch_get_global_queue(QOS_CLASS_DEFAULT, 0), ^{ + @autoreleasepool { + task(); + task = nil; + } + dispatch_group_leave(group); + }); +} + +::testing::AssertionResult wait_for_callbacks( + CoreAIFakeBridgeContext* context, + dispatch_time_t deadline) { + if (dispatch_group_wait(context.callbacks, deadline) != 0) { + return ::testing::AssertionFailure() << "Fake SDK callbacks did not finish"; + } + return ::testing::AssertionSuccess(); +} + +dispatch_time_t cleanup_deadline() { + return dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC); +} +} // namespace + +@interface BookmarkFakeModel () +@property(nonatomic, strong) CoreAIFakeBridgeContext* context; +- (instancetype)initWithContext:(CoreAIFakeBridgeContext*)context; +@end + +@interface BookmarkFakeLoader () +@property(nonatomic, strong) CoreAIFakeBridgeContext* context; +@end + +@implementation BookmarkFakeModel +- (instancetype)init { + return [self initWithContext:current_context]; +} +- (instancetype)initWithContext:(CoreAIFakeBridgeContext*)context { + if (context == nil) { + return nil; + } + if ((self = [super init])) { + self.context = context; + ++context->state.prepared_models; + } + return self; +} +- (void)dealloc { + if (_context != nil) { + --_context->state.prepared_models; + } +} +- (NSData*)copyBookmarkData { + return self.owner.emptyBookmark ? [NSData data] : self.bookmark; +} +- (void)loadFunctionNamed:(NSString*)functionName + inputNames:(NSArray*)inputNames + outputNames:(NSArray*)outputNames + completion:(void (^)(id, NSError*))completion { + // Acquisition tests never bind functions. + completion(nil, fake_error(ETCoreAIErrorUnsupported)); +} +@end + +@implementation BookmarkFakeLoader +- (instancetype)init { + if (current_context == nil) { + return nil; + } + if ((self = [super init])) { + self.context = current_context; + self.models = [NSMutableDictionary dictionary]; + restores = 0; + specializations = 0; + evictions = 0; + ++_context->state.loaders; + } + return self; +} +- (void)dealloc { + if (_context != nil) { + --_context->state.loaders; + } +} +- (void)restoreModelFromBookmark:(NSData*)bookmark + completion:(ETCoreAIAcquisitionCompletion)completion { + ++restores; + run_async(self.context, ^{ + @synchronized(self) { + if (self.restoreError != nil) { + completion(ETCoreAIAcquisitionStatusError, nil, self.restoreError); + return; + } + BookmarkFakeModel* model = self.restoreMiss ? nil : self.models[bookmark]; + self.lastAcquired = model; + completion( + model == nil ? ETCoreAIAcquisitionStatusMiss + : ETCoreAIAcquisitionStatusHit, + model, + nil); + } + }); +} +- (void)specializeModelAtURL:(NSURL*)url + completion:(ETCoreAIAcquisitionCompletion)completion { + ++specializations; + run_async(self.context, ^{ + if (self.onSpecialize != nil) { + self.onSpecialize(url); + } + @synchronized(self) { + self.lastSpecializeURL = url.path; + if (self.specializeError != nil) { + completion(ETCoreAIAcquisitionStatusError, nil, self.specializeError); + return; + } + if (![NSFileManager.defaultManager fileExistsAtPath:url.path]) { + ++self.context->state.missing_bundles; + completion( + ETCoreAIAcquisitionStatusError, + nil, + fake_error(ETCoreAIErrorInvalidModel)); + return; + } + BookmarkFakeModel* model = + [[BookmarkFakeModel alloc] initWithContext:self.context]; + model.owner = self; + model.bookmark = + [NSUUID.UUID.UUIDString dataUsingEncoding:NSUTF8StringEncoding]; + self.models[model.bookmark] = model; + self.lastAcquired = model; + completion(ETCoreAIAcquisitionStatusHit, model, nil); + } + }); +} +- (void)evictModelWithBookmark:(NSData*)bookmark + completion:(void (^)(NSError*))completion { + ++evictions; + run_async(self.context, ^{ + completion(fake_error(ETCoreAIErrorUnsupported)); + }); +} +- (void)clearRetainedState { + self.onSpecialize = nil; + self.lastAcquired = nil; + [self.models removeAllObjects]; +} +@end + +namespace executorch::backends::coreai::testing { + +ScopedFakeBridgeState::ScopedFakeBridgeState() + : context_([[CoreAIFakeBridgeContext alloc] init]), + previous_(current_context) { + current_context = context_; +} + +ScopedFakeBridgeState::~ScopedFakeBridgeState() { + auto drained = wait_for_callbacks(cleanup_deadline()); + EXPECT_TRUE(drained); + if (drained) { + id loader = context_->state.bookmark_loader; + if ([loader isKindOfClass:BookmarkFakeLoader.class]) { + [(BookmarkFakeLoader*)loader clearRetainedState]; + } + context_->state.bookmark_loader = nil; + loader = nil; + EXPECT_EQ(context_->state.prepared_models.load(), 0); + EXPECT_EQ(context_->state.loaders.load(), 0); + EXPECT_EQ(context_->state.missing_bundles.load(), 0); + } + EXPECT_EQ(current_context, context_); + current_context = previous_; +} + +FakeBridgeState& ScopedFakeBridgeState::state() { + return context_->state; +} + +::testing::AssertionResult ScopedFakeBridgeState::wait_for_callbacks( + dispatch_time_t deadline) { + return ::wait_for_callbacks(context_, deadline); +} + +BookmarkLoaderScope::BookmarkLoaderScope() + : context_(current_context), + previous_(context_ == nil ? nil : context_->state.bookmark_loader) { + EXPECT_NE(context_, nil) + << "BookmarkLoaderScope requires ScopedFakeBridgeState"; + if (context_ != nil) { + loader = [[BookmarkFakeLoader alloc] init]; + context_->state.bookmark_loader = loader; + } +} + +BookmarkLoaderScope::~BookmarkLoaderScope() { + if (context_ == nil) { + return; + } + auto drained = ::wait_for_callbacks(context_, cleanup_deadline()); + EXPECT_TRUE(drained); + if (drained) { + [loader clearRetainedState]; + } + EXPECT_EQ(context_->state.bookmark_loader, loader); + context_->state.bookmark_loader = previous_; + loader = nil; +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_host_test.mm b/backends/apple/coreai/runtime/test/coreai_host_test.mm index 679321c248b..d87eefdbbba 100644 --- a/backends/apple/coreai/runtime/test/coreai_host_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_host_test.mm @@ -9,6 +9,7 @@ #import #include #include +#include "coreai_acquisition_fixture.h" #include "coreai_bookmark_fixture.h" int main(int argc, char** argv) { @@ -16,6 +17,11 @@ int main(int argc, char** argv) { @autoreleasepool { executorch::backends::coreai::testing::set_bookmark_test_executable( argv[0]); + const int acquisition_result = + executorch::backends::coreai::testing::acquisition_child_mode( + argc, argv); + if (acquisition_result >= 0) + return acquisition_result; const int child_result = executorch::backends::coreai::testing::bookmark_child_mode(argc, argv); if (child_result >= 0) diff --git a/backends/apple/coreai/runtime/test/coreai_storage_test.mm b/backends/apple/coreai/runtime/test/coreai_storage_test.mm index dd38d9386ae..1afa11b8f8b 100644 --- a/backends/apple/coreai/runtime/test/coreai_storage_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_storage_test.mm @@ -622,6 +622,184 @@ FileDescriptor root(open( EXPECT_TRUE(no_temporary_files(root.get())); } +namespace { +class CoreAIStagingRemovalTest : public ::testing::Test { + protected: + TestDirectory storage; + FileDescriptor root{-1}; + NSString* key = [@"a" stringByPaddingToLength:64 + withString:@"a" + startingAtIndex:0]; + + void SetUp() override { + ASSERT_NE(storage.url, nil); + root.reset(open( + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + ASSERT_GE(root.get(), 0); + } + + NSString* entry() { + return [@"staging" stringByAppendingPathComponent:key]; + } + + Error populate() { + return write_storage_file( + root.get(), + [entry() + stringByAppendingPathComponent:@"model.aimodel/nested/weights"], + "weights", + 7); + } +}; + +TEST_F(CoreAIStagingRemovalTest, RejectsInvalidKeysWithoutTouchingStorage) { + ASSERT_EQ(populate(), Error::Ok); + NSDictionary* before = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, before)); + for (NSString* invalid in + @[ @"", @"../escape", @"/absolute", @"short", key.uppercaseString ]) { + EXPECT_EQ( + remove_storage_staging(root.get(), invalid), Error::InvalidArgument); + } + EXPECT_EQ(remove_storage_staging(root.get(), nil), Error::InvalidArgument); + EXPECT_TRUE(snapshot_matches(storage.url, before)); +} + +TEST_F(CoreAIStagingRemovalTest, MissingStagingIsAlreadyRemoved) { + StorageFaultScope fault(StorageOperation::Remove); + EXPECT_EQ(remove_storage_staging(root.get(), key), Error::Ok); + ASSERT_EQ(mkdirat(root.get(), "staging", 0700), 0); + EXPECT_EQ(remove_storage_staging(root.get(), key), Error::Ok); + EXPECT_EQ(fault.matches, 0); + EXPECT_EQ(fault.hits, 0); +} + +TEST_F(CoreAIStagingRemovalTest, RemovesWholeKeyWithoutTouchingOtherEntries) { + ASSERT_EQ(populate(), Error::Ok); + ASSERT_EQ( + mkdirat( + root.get(), + [entry() stringByAppendingPathComponent:@"empty"] + .fileSystemRepresentation, + 0700), + 0); + NSString* other = [@"b" stringByPaddingToLength:64 + withString:@"b" + startingAtIndex:0]; + NSString* other_path = [[@"staging" stringByAppendingPathComponent:other] + stringByAppendingPathComponent:@"keep"]; + ASSERT_EQ(write_storage_file(root.get(), other_path, "other", 5), Error::Ok); + ASSERT_EQ( + write_storage_file(root.get(), @"unrelated", "outside", 7), Error::Ok); + EXPECT_EQ(remove_storage_staging(root.get(), key), Error::Ok); + struct stat info; + EXPECT_EQ( + fstatat( + root.get(), + entry().fileSystemRepresentation, + &info, + AT_SYMLINK_NOFOLLOW), + -1); + EXPECT_EQ(errno, ENOENT); + EXPECT_EQ(fstatat(root.get(), "staging", &info, AT_SYMLINK_NOFOLLOW), 0); + EXPECT_EQ( + fstatat( + root.get(), + other_path.fileSystemRepresentation, + &info, + AT_SYMLINK_NOFOLLOW), + 0); + EXPECT_EQ(info.st_size, 5); + EXPECT_EQ(fstatat(root.get(), "unrelated", &info, AT_SYMLINK_NOFOLLOW), 0); + EXPECT_EQ(info.st_size, 7); + EXPECT_EQ(remove_storage_staging(root.get(), key), Error::Ok); +} + +TEST_F(CoreAIStagingRemovalTest, RejectsLinkedStagingParent) { + TestDirectory outside; + ASSERT_NE(outside.url, nil); + ASSERT_EQ( + symlinkat(outside.url.fileSystemRepresentation, root.get(), "staging"), + 0); + NSDictionary* before = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, before)); + EXPECT_EQ(remove_storage_staging(root.get(), key), Error::AccessFailed); + EXPECT_TRUE(snapshot_matches(storage.url, before)); +} + +class CoreAIUnsafeStagingTest : public CoreAIStagingRemovalTest, + public ::testing::WithParamInterface {}; + +TEST_P(CoreAIUnsafeStagingTest, PreflightsEntireTreeBeforeUnlinking) { + TestDirectory outside; + ASSERT_NE(outside.url, nil); + NSURL* sentinel = [outside.url URLByAppendingPathComponent:@"sentinel"]; + ASSERT_TRUE([[@"untouched" dataUsingEncoding:NSUTF8StringEncoding] + writeToURL:sentinel + atomically:YES]); + NSString* invalid = entry(); + if (GetParam() < 2) { + ASSERT_EQ(mkdirat(root.get(), "staging", 0700), 0); + } else { + ASSERT_EQ(populate(), Error::Ok); + invalid = [entry() stringByAppendingPathComponent:@"z-invalid"]; + } + if (GetParam() == 0 || GetParam() == 2) { + ASSERT_EQ( + symlinkat( + outside.url.fileSystemRepresentation, + root.get(), + invalid.fileSystemRepresentation), + 0); + } else if (GetParam() == 1) { + ASSERT_EQ(write_storage_file(root.get(), invalid, "file", 4), Error::Ok); + } else { + ASSERT_EQ(mkfifoat(root.get(), invalid.fileSystemRepresentation, 0600), 0); + } + NSDictionary* before = nil; + NSDictionary* outside_before = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, before)); + ASSERT_TRUE(asset_tree_snapshot(outside.url, outside_before)); + StorageFaultScope fault(StorageOperation::Remove); + EXPECT_EQ( + remove_storage_staging(root.get(), key), Error::InvalidExternalData); + EXPECT_EQ(fault.matches, 0); + EXPECT_TRUE(snapshot_matches(storage.url, before)); + EXPECT_TRUE(snapshot_matches(outside.url, outside_before)); +} + +INSTANTIATE_TEST_SUITE_P( + Entries, + CoreAIUnsafeStagingTest, + ::testing::Values(0, 1, 2, 3), + ([](const ::testing::TestParamInfo& info) { + const char* names[] = { + "LinkedKey", "RegularKey", "LinkedChild", "FifoChild"}; + return names[info.param]; + })); + +TEST_F(CoreAIStagingRemovalTest, RetriesOnlyInterruptedRemovalAndSync) { + for (StorageOperation operation : + {StorageOperation::Remove, StorageOperation::SyncDirectory}) { + ASSERT_EQ(populate(), Error::Ok); + StorageFaultScope fault(operation); + fault.error = EINTR; + EXPECT_EQ(remove_storage_staging(root.get(), key), Error::Ok); + EXPECT_EQ(fault.hits, 1); + struct stat info; + EXPECT_EQ( + fstatat( + root.get(), + entry().fileSystemRepresentation, + &info, + AT_SYMLINK_NOFOLLOW), + -1); + EXPECT_EQ(errno, ENOENT); + } +} +} // namespace + TEST(CoreAIStorageTest, RetriesInterruptedPublicationOperations) { TestDirectory storage; ASSERT_NE(storage.url, nil); diff --git a/backends/apple/coreai/runtime/test/export_bridge_fixture.py b/backends/apple/coreai/runtime/test/export_bridge_fixture.py index 13709de4716..947c47cc86e 100644 --- a/backends/apple/coreai/runtime/test/export_bridge_fixture.py +++ b/backends/apple/coreai/runtime/test/export_bridge_fixture.py @@ -7,6 +7,7 @@ import argparse import json import shutil +import uuid from pathlib import Path import torch @@ -28,14 +29,28 @@ def forward(self, left, right): return left + right, left - right -def export_fixture(output_dir: Path, *, dynamic: bool = False) -> None: - """Export into a dedicated fixture directory, replacing its contents.""" - size = torch.export.Dim("size", min=1, max=8) - exported = torch.export.export( - Arithmetic().eval(), - (torch.ones(2), torch.ones(2)), - dynamic_shapes=({0: size}, {0: size}) if dynamic else None, - ) +class WeightedArithmetic(torch.nn.Module): + """External weights expose any dependency on removed source files.""" + + def __init__(self, first: float, second: float): + """Build a weighted model with distinguishable output offsets.""" + super().__init__() + self.first = first + self.second = second + self.network = torch.nn.Sequential( + torch.nn.Linear(2, 4096), torch.nn.ReLU(), torch.nn.Linear(4096, 2) + ) + + def forward(self, left, right): + """Apply the shared network to two distinct input combinations.""" + return ( + self.network(left + right) + self.first, + self.network(left - right) + self.second, + ) + + +def _export_model(model, inputs, output_dir: Path, *, dynamic_shapes=None) -> dict: + exported = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes) lowered = to_edge_transform_and_lower( exported, transform_passes=get_default_passes(), @@ -56,6 +71,41 @@ def export_fixture(output_dir: Path, *, dynamic: bool = False) -> None: destination.parent.mkdir(parents=True, exist_ok=True) destination.write_bytes(store.buffers[entry.buffer_index]) (output_dir / "manifest.json").write_bytes(manifest_bytes) + return manifest + + +def export_fixture(output_dir: Path, *, dynamic: bool = False) -> None: + """Export into a dedicated fixture directory, replacing its contents.""" + size = torch.export.Dim("size", min=1, max=8) + _export_model( + Arithmetic().eval(), + (torch.ones(2), torch.ones(2)), + output_dir, + dynamic_shapes=({0: size}, {0: size}) if dynamic else None, + ) + + +def export_weighted_fixture(output_dir: Path) -> None: + """Export isolated SDK cache contents and their expected output values.""" + # Unique model contents keep cleanup from evicting another test's SDK entry. + seed = uuid.uuid4().int + model = WeightedArithmetic( + (seed & 0xFFFFFF) / 2**24, ((seed >> 24) & 0xFFFFFF) / 2**24 + ).eval() + inputs = (torch.tensor([5.0, 9.0]), torch.tensor([1.0, 2.0])) + manifest = _export_model(model, inputs, output_dir) + with torch.no_grad(): + expected = [tensor.tolist() for tensor in model(*inputs)] + (output_dir / "fixture.json").write_text( + json.dumps( + { + "function": manifest["function"], + "inputs": manifest["input_names"], + "outputs": manifest["output_names"], + "expected": expected, + } + ) + ) if __name__ == "__main__": diff --git a/backends/apple/coreai/runtime/test/run_sdk_source_lifetime.py b/backends/apple/coreai/runtime/test/run_sdk_source_lifetime.py new file mode 100644 index 00000000000..dc6c044b3aa --- /dev/null +++ b/backends/apple/coreai/runtime/test/run_sdk_source_lifetime.py @@ -0,0 +1,114 @@ +# 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. + +import argparse +import shutil +import subprocess +import tempfile +from pathlib import Path + + +def run(command, expected_status=0): + """Run a probe command, asserting its exit status.""" + result = subprocess.run(command, capture_output=True, text=True, timeout=60) + if result.returncode != expected_status: + raise RuntimeError( + f"{command} returned {result.returncode}, expected {expected_status}\n" + f"{result.stdout}\n{result.stderr}" + ) + return result.stdout.strip() + + +def has_device_architecture(architecture: str) -> bool: + """Report whether the SDK identified a device architecture.""" + normalized = architecture.strip().lower() + return bool(normalized and normalized != "unknown") + + +def _export_weighted_fixture(directory: Path) -> None: + from export_bridge_fixture import export_weighted_fixture + + export_weighted_fixture(directory) + + +def compile_source(source: Path, compiled: Path, architecture: str) -> None: + """Compile a portable source to an AOT model.""" + run( + [ + "xcrun", + "coreai-build", + "compile", + str(source), + "--output", + str(compiled), + "--platform", + "macOS", + "--architecture", + architecture, + "--min-deployment-version", + "27.0", + ] + ) + shutil.rmtree(source) + + +def run_variant( + probe: Path, kind: str, source: Path, fixture: Path, bookmark: Path +) -> None: + """Exercise specialize, restore, and evict for one source variant.""" + arguments = [str(source), str(fixture), str(bookmark)] + try: + print(run([str(probe), "specialize", *arguments]), flush=True) + if source.exists() or not bookmark.is_file(): + raise AssertionError("Expected a saved bookmark and no original source") + for _ in range(2): + print(run([str(probe), "restore", *arguments]), flush=True) + except BaseException: + try: + if bookmark.is_file(): + run([str(probe), "evict", *arguments]) + except Exception as cleanup: + print(f"cleanup evict failed: {cleanup}", flush=True) + raise + if bookmark.is_file(): + run([str(probe), "evict", *arguments]) + run([str(probe), "restore", *arguments], expected_status=3) + print(f"{kind}: persistent cache lifetime and eviction control passed", flush=True) + + +def main(probe: Path) -> None: + """Run portable and compiled lifetime probes against a test binary.""" + architecture = run([str(probe), "architecture"]) + with tempfile.TemporaryDirectory(prefix="coreai-lifetime-") as temporary: + root = Path(temporary) + for kind in ("portable", "compiled"): + # AOT compile rejects unknown architectures. + if kind == "compiled" and not has_device_architecture(architecture): + print( + f"compiled: skipped (unsupported architecture {architecture!r})", + flush=True, + ) + continue + directory = root / kind + _export_weighted_fixture(directory) + source = directory / "model.aimodel" + if kind == "compiled": + compiled = directory / "model.aimodelc" + compile_source(source, compiled, architecture) + source = compiled + run_variant( + probe, + kind, + source, + directory / "fixture.json", + directory / "saved.bookmark", + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("probe", type=Path) + main(parser.parse_args().probe.resolve()) diff --git a/backends/apple/coreai/runtime/test/test_run_sdk_source_lifetime.py b/backends/apple/coreai/runtime/test/test_run_sdk_source_lifetime.py new file mode 100644 index 00000000000..548ffe1f9c6 --- /dev/null +++ b/backends/apple/coreai/runtime/test/test_run_sdk_source_lifetime.py @@ -0,0 +1,97 @@ +# 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. + +from pathlib import Path +from unittest.mock import patch + +import pytest + +from executorch.backends.apple.coreai.runtime.test import ( + run_sdk_source_lifetime as driver, +) + + +@pytest.mark.parametrize("architecture", ["unknown", "Unknown", "", " "]) +def test_compiled_variant_skipped_without_supported_architecture( + tmp_path, capsys, architecture +): + """AOT compile rejects unknown architectures, so only portable runs.""" + calls = [] + with ( + patch.object(driver, "run", return_value=architecture), + patch.object( + driver, + "_export_weighted_fixture", + side_effect=lambda d: d.mkdir(parents=True), + ), + patch.object( + driver, "compile_source", side_effect=lambda *args: calls.append("compile") + ), + patch.object( + driver, + "run_variant", + side_effect=lambda *args: calls.append(args[1]), + ), + patch.object(driver.tempfile, "TemporaryDirectory") as temporary, + ): + temporary.return_value.__enter__.return_value = str(tmp_path) + driver.main(Path("probe")) + + assert calls == ["portable"] + assert "compiled: skipped" in capsys.readouterr().out + + +def test_compiled_variant_runs_with_known_architecture(tmp_path): + """Physical known-arch Macs still exercise the AOT path.""" + calls = [] + with ( + patch.object(driver, "run", return_value="h16c"), + patch.object( + driver, + "_export_weighted_fixture", + side_effect=lambda d: d.mkdir(parents=True), + ), + patch.object( + driver, "compile_source", side_effect=lambda *args: calls.append("compile") + ), + patch.object( + driver, + "run_variant", + side_effect=lambda *args: calls.append(args[1]), + ), + patch.object(driver.tempfile, "TemporaryDirectory") as temporary, + ): + temporary.return_value.__enter__.return_value = str(tmp_path) + driver.main(Path("probe")) + + assert calls == ["portable", "compile", "compiled"] + + +@pytest.mark.parametrize( + "failure", [RuntimeError("probe failed"), KeyboardInterrupt("interrupted")] +) +def test_cleanup_failure_preserves_original_error(tmp_path, capsys, failure): + """A failing evict must not mask the probe failure it cleans up after.""" + + def fake_run(command, expected_status=0): + if "specialize" in command: + raise failure + raise RuntimeError("evict failed") + + bookmark = tmp_path / "saved.bookmark" + bookmark.write_bytes(b"bookmark") + with patch.object(driver, "run", side_effect=fake_run): + with pytest.raises(type(failure)) as raised: + driver.run_variant( + Path("probe"), + "portable", + tmp_path / "model.aimodel", + tmp_path / "fixture.json", + bookmark, + ) + + assert raised.value is failure + assert "cleanup evict failed: evict failed" in capsys.readouterr().out