From c4747f45eabfc2b37631ac039e26bae168b3005c Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Sun, 4 Oct 2026 19:50:13 -0700 Subject: [PATCH 1/4] [CoreAI] Coordinate model acquisition Summary: Acquire a prepared Core AI model for a selected bundle: restore it from the persisted bookmark, or materialize the source and specialize it. - `runtime/coreai_load_coordinator.{h,mm}`: `acquire_bookmark_model` restores from the persisted bookmark, or on a miss materializes the source and specializes it with the persistent default SDK cache, then publishes the new bookmark atomically. Restore errors fail without a source fallback. Sources and SDK pins are retained across completion. - Tests: an acquisition-only fake SDK loader (`runtime/test/coreai_fake_loader.{h,mm}`) lets the coordinator run in `coreai_host_test` without Swift or CoreAI. The acquisition suite calls `acquire_bookmark_model` directly and covers cold specialization, warm restores, restore errors and misses, bookmark-save failures, contention on one key versus independent keys across processes, and a crash after the SDK call. - The backend that uses this lands separately. - README: the "Asset storage" acquisition flow. No changes outside `backends/apple/coreai`. Test Plan: ``` cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_OSX_DEPLOYMENT_TARGET=27.0 -DEXECUTORCH_BUILD_COREAI=ON \ -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON cmake --build build --target backends/apple/coreai/all ctest --test-dir build/backends/apple/coreai --output-on-failure --no-tests=error ``` CTest runs `coreai_host_test` (adding the acquisition suite) and `coreai_swift_bridge_test`. --- backends/apple/coreai/CMakeLists.txt | 9 +- backends/apple/coreai/README.md | 13 + .../coreai/runtime/coreai_load_coordinator.h | 24 ++ .../coreai/runtime/coreai_load_coordinator.mm | 96 ++++++ .../runtime/test/coreai_acquisition_fixture.h | 61 ++++ .../test/coreai_acquisition_fixture.mm | 285 ++++++++++++++++++ .../runtime/test/coreai_acquisition_test.mm | 180 +++++++++++ .../coreai/runtime/test/coreai_fake_loader.h | 91 ++++++ .../coreai/runtime/test/coreai_fake_loader.mm | 243 +++++++++++++++ .../coreai/runtime/test/coreai_host_test.mm | 7 + 10 files changed, 1007 insertions(+), 2 deletions(-) create mode 100644 backends/apple/coreai/runtime/coreai_load_coordinator.h create mode 100644 backends/apple/coreai/runtime/coreai_load_coordinator.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_acquisition_fixture.h create mode 100644 backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_acquisition_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_fake_loader.h create mode 100644 backends/apple/coreai/runtime/test/coreai_fake_loader.mm diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index fc5d5f6117d..060ca8081a1 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) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 7b45b877410..6d38e011ce7 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. Every load first attempts SDK +bookmark restoration. A missing bookmark or SDK-confirmed cache miss falls back +to source materialization and persistent specialization, then binds the requested +function and ordered I/O on that same acquired model. + ### Asset storage The assets root must be an absolute path. The application chooses it, should @@ -123,6 +128,14 @@ 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 inspect staging or consult NDS; the acquired SDK +model is retained through function binding. On a missing bookmark or confirmed +SDK miss, the runtime materializes the source, specializes it and atomically +writes the returned bookmark bytes. 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; the next load retries the ordinary flow. 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. 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..85827ea0fec --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.h @@ -0,0 +1,24 @@ +/* + * 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. +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..a34e0c50ba9 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -0,0 +1,96 @@ +/* + * 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]; + }); +} + +} // 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(); + auto locked = lock_bookmark(root, key.get()); + 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) return model; + } + 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()); + if (!source.ok()) return source.error(); + NSURL* url = source.get(); + auto model = acquire(^(ETCoreAIAcquisitionCompletion completion) { + [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_CHECK_OK_OR_RETURN_ERROR( + write_bookmark(*locked.get(), [model.get() copyBookmarkData])); + // SDK pins do not prove source independence; leave staged sources in place. + return model; +} + +} // namespace executorch::backends::coreai 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..743102fac99 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.h @@ -0,0 +1,61 @@ +/* + * 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); + +// Uses the existing bounded, scoped child process owner. +void set_acquisition_test_executable(const char* path); +::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 bookmark_child_mode and 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..154fd7e2ae4 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm @@ -0,0 +1,285 @@ +/* + * 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; +} + +void set_acquisition_test_executable(const char* path) { set_bookmark_test_executable(path); } + +::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(@"--bookmark-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], "--bookmark-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::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..6c1a91eea02 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm @@ -0,0 +1,180 @@ +/* + * 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 "coreai_fault_scope.h" +#include "coreai_load_coordinator.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +namespace { + +class CoreAIAcquisitionTest : public ::testing::Test { + protected: + // Declared first so each case's loader, models and data die before it. + ScopedFakeBridgeState bridge; + BookmarkDirectory 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; + } +}; + +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]); +} + +TEST_F(CoreAIAcquisitionTest, WarmHitRestoresWithoutNamedData) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + NSData* original = saved(); + ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->restores.load(), 1); + EXPECT_TRUE([saved() isEqual:original]); +} + +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(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); + NSData* replaced = saved(); + EXPECT_NE(replaced, nil); + EXPECT_FALSE([replaced isEqual:original]); +} + +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); + fake.loader.restoreMiss = YES; + } + { + // A cold acquisition publishes its staged source before the bookmark. + StorageFaultScope fault(StorageOperation::Rename); + fault.match_index = prior ? 1 : 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)); +} +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, 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); +} + +} // 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..7732fee3507 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.h @@ -0,0 +1,91 @@ +/* + * 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, strong) NSError* evictError; +@property(nonatomic) BOOL restoreMiss; +@property(nonatomic) BOOL emptyBookmark; +@property(nonatomic, copy) NSData* specializationBookmark; +@property(nonatomic, copy) void (^onSpecialize)(NSURL*); +@property(nonatomic, copy) void (^onEvict)(NSData*); +@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..a997ba361fa --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm @@ -0,0 +1,243 @@ +/* + * 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 = self.specializationBookmark != nil + ? self.specializationBookmark + : [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, ^{ + if (self.onEvict != nil) { + self.onEvict(bookmark); + } + @synchronized(self) { + if (self.evictError == nil) { + [self.models removeObjectForKey:bookmark]; + } + completion(self.evictError); + } + }); +} +- (void)clearRetainedState { + self.onSpecialize = nil; + self.onEvict = 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..f388bc5089f 100644 --- a/backends/apple/coreai/runtime/test/coreai_host_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_host_test.mm @@ -9,11 +9,18 @@ #import #include #include +#include "coreai_acquisition_fixture.h" #include "coreai_bookmark_fixture.h" int main(int argc, char** argv) { executorch::runtime::runtime_init(); @autoreleasepool { + executorch::backends::coreai::testing::set_acquisition_test_executable( + argv[0]); + const int acquisition_result = + executorch::backends::coreai::testing::acquisition_child_mode(argc, argv); + if (acquisition_result >= 0) + return acquisition_result; executorch::backends::coreai::testing::set_bookmark_test_executable( argv[0]); const int child_result = From 0d9f80f9cddf55f2a7c0430892b4977d65628eed Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Mon, 26 Oct 2026 17:06:55 -0700 Subject: [PATCH 2/4] [CoreAI] Remove staged sources after durable bookmark publication Summary: Avoid retaining redundant staged model bundles once persistent specialization and durable bookmark publication succeed. Preserve the bookmark and return cleanup failures without automatic retries or eviction, keeping warm restoration independent of staging. Tighten acquisition error coverage and test-fixture behavior. Test Plan: All three Core AI CTest suites passed on macOS 27. The source-lifetime regression verified portable and AOT model inference after source deletion, bookmark restoration in fresh processes, and a restore miss after explicit eviction. Native tests covered publication ordering, unsafe staging entries, and cleanup failures preserving durable bookmarks. The full Core AI Python suite passed: 113 tests. # Conflicts: # backends/apple/coreai/README.md # backends/apple/coreai/runtime/test/coreai_host_test.mm # backends/apple/coreai/runtime/test/export_bridge_fixture.py # Please enter the commit message for your changes. Lines starting # with '#' will be kept; you may remove them yourself if you want to. # An empty message aborts the commit. # # interactive rebase in progress; onto aedabbb695 # Last commands done (2 commands done): # pick 664cd4abb7 # [CoreAI] Coordinate model acquisition # pick 61e2774c4d # [CoreAI] Remove staged sources after durable bookmark publication # No commands remaining. # You are currently rebasing branch 'coreai-v2/05a-load-coordinator' on 'aedabbb695'. # # Changes to be committed: # modified: backends/apple/coreai/CMakeLists.txt # modified: backends/apple/coreai/README.md # modified: backends/apple/coreai/runtime/coreai_bookmarks.h # modified: backends/apple/coreai/runtime/coreai_bookmarks.mm # modified: backends/apple/coreai/runtime/coreai_load_coordinator.mm # modified: backends/apple/coreai/runtime/coreai_storage.h # modified: backends/apple/coreai/runtime/coreai_storage.mm # new file: backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift # modified: backends/apple/coreai/runtime/test/coreai_acquisition_fixture.h # modified: backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm # modified: backends/apple/coreai/runtime/test/coreai_acquisition_test.mm # modified: backends/apple/coreai/runtime/test/coreai_fake_loader.h # modified: backends/apple/coreai/runtime/test/coreai_fake_loader.mm # modified: backends/apple/coreai/runtime/test/coreai_host_test.mm # modified: backends/apple/coreai/runtime/test/coreai_storage_test.mm # modified: backends/apple/coreai/runtime/test/export_bridge_fixture.py # new file: backends/apple/coreai/runtime/test/test_sdk_source_lifetime.py # # Changes not staged for commit: # modified: backends/mlx/third-party/mlx (modified content) # --- backends/apple/coreai/CMakeLists.txt | 22 + backends/apple/coreai/README.md | 27 +- .../apple/coreai/runtime/coreai_bookmarks.h | 4 + .../apple/coreai/runtime/coreai_bookmarks.mm | 4 + .../coreai/runtime/coreai_load_coordinator.mm | 2 +- .../apple/coreai/runtime/coreai_storage.h | 2 + .../apple/coreai/runtime/coreai_storage.mm | 111 +++++ .../runtime/test/SDKSourceLifetimeTest.swift | 113 +++++ .../runtime/test/coreai_acquisition_fixture.h | 5 +- .../test/coreai_acquisition_fixture.mm | 6 +- .../runtime/test/coreai_acquisition_test.mm | 397 +++++++++++++++++- .../coreai/runtime/test/coreai_fake_loader.h | 3 - .../coreai/runtime/test/coreai_fake_loader.mm | 15 +- .../coreai/runtime/test/coreai_host_test.mm | 4 +- .../runtime/test/coreai_storage_test.mm | 178 ++++++++ .../runtime/test/export_bridge_fixture.py | 59 ++- .../runtime/test/test_sdk_source_lifetime.py | 63 +++ 17 files changed, 969 insertions(+), 46 deletions(-) create mode 100644 backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift create mode 100644 backends/apple/coreai/runtime/test/test_sdk_source_lifetime.py diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index 060ca8081a1..af5a2e15614 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -313,6 +313,28 @@ 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 + ) + set_target_properties( + coreai_sdk_source_lifetime_test PROPERTIES Swift_LANGUAGE_VERSION 6 + ) + coreai_require_os27(coreai_sdk_source_lifetime_test) + target_compile_options( + coreai_sdk_source_lifetime_test + PRIVATE "$<$:-parse-as-library>" + ) + target_link_libraries( + coreai_sdk_source_lifetime_test PRIVATE executorch::coreai_dependencies + ) + add_test( + NAME coreai_sdk_source_lifetime_test + COMMAND + "${Python3_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/runtime/test/test_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 6d38e011ce7..ae3f17a6d87 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -131,11 +131,21 @@ serialized lifecycle record. A healthy bookmark hit does not inspect staging or consult NDS; the acquired SDK model is retained through function binding. On a missing bookmark or confirmed SDK miss, the runtime materializes the source, specializes it and atomically -writes the returned bookmark bytes. 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; the next load retries the ordinary flow. The -backend does not track historical or unrecorded SDK entries. +writes the returned bookmark bytes. After durable publication, it removes the +entire keyed staging directory while holding the per-key lock. Persistent SDK +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. If staged-source cleanup fails, loading reports the +error but retains the durable bookmark and any remaining staged files. Neither +failure triggers automatic retries, SDK eviction or source repair. A subsequent +bookmark hit can load without inspecting or retrying staging cleanup. A partial +cleanup followed by an SDK miss fails on the damaged staged entry until it is +explicitly cleared. 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. @@ -182,6 +192,8 @@ 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 after removing disposable portable and AOT source copies. 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; @@ -204,8 +216,9 @@ covers arm64 macOS and iOS device builds. x86_64 builds are blocked by Swift 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_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.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm index a34e0c50ba9..9ad967ccc0b 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.mm +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -89,7 +89,7 @@ Error sdk_error(NSError* error) { "Core AI specialization returned no model"); ET_CHECK_OK_OR_RETURN_ERROR( write_bookmark(*locked.get(), [model.get() copyBookmarkData])); - // SDK pins do not prove source independence; leave staged sources in place. + ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark_staging(*locked.get())); return model; } diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h index daf32376fab..a367c31b680 100644 --- a/backends/apple/coreai/runtime/coreai_storage.h +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -47,6 +47,8 @@ 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); enum class StorageOperation { Write, diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm index 0a42b22f859..24e865068d2 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,68 @@ Error sync_storage_directory(int fd, bool full_sync) { return Error::Ok; } +Error remove_storage_staging(int root_fd, NSString* key) { + ET_CHECK_OR_RETURN_ERROR( + [key isKindOfClass:NSString.class] && key.length == 64 && + [key rangeOfCharacterFromSet: + [[NSCharacterSet + characterSetWithCharactersInString:@"0123456789abcdef"] + invertedSet]] + .location == NSNotFound, + 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"); + struct stat info; + const int inspected = retry_eintr([&] { + return fstatat( + staging.get(), + 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( + staging.get(), + 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( + staging.get(), key.fileSystemRepresentation, AT_REMOVEDIR); + }) == 0, + AccessFailed, + "Cannot remove Core AI keyed staging directory"); + return sync_storage_directory(staging.get()); +} + 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..8691b47591b --- /dev/null +++ b/backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift @@ -0,0 +1,113 @@ +// 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 +import Foundation + +private struct Fixture: Decodable { + let function: String + let inputs: [String] + let outputs: [String] + let expected: [[Float]] +} + +private enum Failure: Error { + case invalidArguments + case missingFunction + case incorrectOutput(String) + case unexpectedSource +} + +private func checkInference(_ model: AIModel, fixture: Fixture) async throws { + guard let function = try model.loadFunction(named: fixture.function) else { + throw Failure.missingFunction + } + guard fixture.inputs.count == 2, fixture.outputs.count == fixture.expected.count else { + throw Failure.invalidArguments + } + var inputs: [String: NDArray] = [:] + for (index, values) in [[Float(5), 9], [Float(1), 2]].enumerated() { + var array = NDArray(shape: [2], scalarType: .float32) + do { + var view = array.mutableView(as: Float.self) + view.withUnsafeMutablePointer { pointer, _, _ in + pointer[0] = values[0] + pointer[1] = values[1] + } + } + inputs[fixture.inputs[index]] = array + } + var outputs = try await function.run(inputs: inputs) + guard outputs.count == fixture.outputs.count else { + throw Failure.incorrectOutput("count") + } + for (index, name) in fixture.outputs.enumerated() { + guard let value = outputs.remove(name), let array = value.ndArray, + array.shape == [2], array.scalarType == .float32, + fixture.expected[index].count == 2 + else { throw Failure.incorrectOutput(name) } + let values = array.view(as: Float.self).withUnsafePointer { pointer, _, strides in + [pointer[0], pointer[strides[0]]] + } + 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() async { + do { + let args = CommandLine.arguments + if args.count == 2, args[1] == "architecture" { + print(AIModel.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 } + if mode == "evict" { + try AIModelCache.deleteEntry(referencedBy: Data(contentsOf: bookmarkURL)) + return + } + let fixture = try JSONDecoder().decode(Fixture.self, from: Data(contentsOf: fixtureURL)) + if mode == "specialize" { + let model = try await AIModel.specialize( + contentsOf: source, options: .default, cache: .default, cachePolicy: .persistent) + try model.bookmarkData.write(to: bookmarkURL, options: .atomic) + try await checkInference(model, fixture: fixture) + try FileManager.default.removeItem(at: source) + guard !FileManager.default.fileExists(atPath: source.path) else { + throw Failure.unexpectedSource + } + try await checkInference(model, fixture: fixture) + print("Live inference passed after source removal") + } else if mode == "restore" { + guard !FileManager.default.fileExists(atPath: source.path) else { + throw Failure.unexpectedSource + } + guard let model = try AIModel(resolvingBookmark: Data(contentsOf: bookmarkURL)) else { + exit(3) + } + try await checkInference(model, fixture: fixture) + print("Bookmark 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 index 743102fac99..963adfd784a 100644 --- a/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.h +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.h @@ -41,8 +41,6 @@ ParallelAcquisitionObservation parallel_acquisitions( TestData& data, bool same_key); -// Uses the existing bounded, scoped child process owner. -void set_acquisition_test_executable(const char* path); ::testing::AssertionResult spawn_acquisition_crash( BookmarkChild& child, NSURL* root, @@ -54,8 +52,7 @@ ::testing::AssertionResult spawn_parallel_acquisition( NSDictionary* first, NSDictionary* second, bool same_key); -// Dispatch before bookmark_child_mode and InitGoogleTest. Unrelated modes -// return -1. +// 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 index 154fd7e2ae4..a0af1340b76 100644 --- a/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm @@ -152,8 +152,6 @@ ParallelAcquisitionObservation parallel_acquisitions(ScopedFakeBridgeState& brid return observation; } -void set_acquisition_test_executable(const char* path) { set_bookmark_test_executable(path); } - ::testing::AssertionResult spawn_acquisition_crash(BookmarkChild& child, NSURL* root, StorageOperation boundary, NSDictionary* manifest) { @@ -161,7 +159,7 @@ ParallelAcquisitionObservation parallel_acquisitions(ScopedFakeBridgeState& brid 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(@"--bookmark-crash", root.path, argument); + return child.spawn(@"--acquisition-crash", root.path, argument); } ::testing::AssertionResult spawn_parallel_acquisition(BookmarkChild& child, NSURL* root, @@ -217,7 +215,7 @@ ParallelAcquisitionObservation parallel_acquisitions(ScopedFakeBridgeState& brid int acquisition_child_mode(int argc, char** argv) { if (argc < 2) return -1; - const bool crash = std::strcmp(argv[1], "--bookmark-crash") == 0; + 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; diff --git a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm index 6c1a91eea02..af80260827e 100644 --- a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm @@ -8,14 +8,35 @@ #include "coreai_acquisition_fixture.h" +#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 { @@ -51,6 +72,28 @@ Error acquire(BookmarkFakeLoader* loader, const NamedDataMap* named_data) { 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 [[root.url URLByAppendingPathComponent:@"locks"] + URLByAppendingPathComponent:[key stringByAppendingString:@".lock"]]; + } + + ino_t lock_inode() { + struct stat info; + if (lstat(lock_url().fileSystemRepresentation, &info) != 0) + return 0; + return info.st_ino; + } }; TEST_F(CoreAIAcquisitionTest, ColdSpecializesStagedSourceAndPersistsBookmark) { @@ -63,16 +106,23 @@ Error acquire(BookmarkFakeLoader* loader, const NamedDataMap* named_data) { 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) { @@ -83,6 +133,7 @@ Error acquire(BookmarkFakeLoader* loader, const NamedDataMap* named_data) { 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]); } @@ -95,10 +146,345 @@ Error acquire(BookmarkFakeLoader* loader, const NamedDataMap* named_data) { 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); + 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::AccessFailed); + 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_TRUE(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::AccessFailed); + 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_TRUE(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::AccessFailed); + 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]); +} + +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, @@ -111,12 +497,15 @@ Error acquire(BookmarkFakeLoader* loader, const NamedDataMap* named_data) { ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); original = saved(); ASSERT_NE(original, nil); + EXPECT_FALSE(staging_key_exists()); fake.loader.restoreMiss = YES; } { - // A cold acquisition publishes its staged source before the bookmark. + // 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 = prior ? 1 : 2; + fault.match_index = 2; EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); EXPECT_EQ(fault.hits, 1); } @@ -126,6 +515,7 @@ Error acquire(BookmarkFakeLoader* loader, const NamedDataMap* named_data) { fake.loader.restoreMiss = NO; ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); EXPECT_EQ(fake.loader->specializations.load(), calls + (prior ? 0 : 1)); + EXPECT_EQ(staging_key_exists(), prior); } INSTANTIATE_TEST_SUITE_P(Boundaries, CoreAIAcquisitionSaveFailureTest, ::testing::Bool(), [](const ::testing::TestParamInfo& info) { @@ -175,6 +565,7 @@ void expect_parallel(bool same_key) { EXPECT_EQ(fake.loader->restores.load(), 0); EXPECT_EQ(data.attempts.load(), 0); EXPECT_NE(saved(), nil); + EXPECT_FALSE(staging_key_exists()); } } // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.h b/backends/apple/coreai/runtime/test/coreai_fake_loader.h index 7732fee3507..cf8f70a6a6b 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.h +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.h @@ -61,12 +61,9 @@ class ScopedFakeBridgeState final { NSMutableDictionary* models; @property(nonatomic, strong) NSError* restoreError; @property(nonatomic, strong) NSError* specializeError; -@property(nonatomic, strong) NSError* evictError; @property(nonatomic) BOOL restoreMiss; @property(nonatomic) BOOL emptyBookmark; -@property(nonatomic, copy) NSData* specializationBookmark; @property(nonatomic, copy) void (^onSpecialize)(NSURL*); -@property(nonatomic, copy) void (^onEvict)(NSData*); @property(nonatomic, strong) BookmarkFakeModel* lastAcquired; @property(nonatomic, copy) NSString* lastSpecializeURL; // Call only after workers/callbacks finish and caller-owned handles are diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm index a997ba361fa..54a14287d94 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm @@ -154,9 +154,7 @@ - (void)specializeModelAtURL:(NSURL*)url completion:(ETCoreAIAcquisitionCompleti } BookmarkFakeModel* model = [[BookmarkFakeModel alloc] initWithContext:self.context]; model.owner = self; - model.bookmark = self.specializationBookmark != nil - ? self.specializationBookmark - : [NSUUID.UUID.UUIDString dataUsingEncoding:NSUTF8StringEncoding]; + model.bookmark = [NSUUID.UUID.UUIDString dataUsingEncoding:NSUTF8StringEncoding]; self.models[model.bookmark] = model; self.lastAcquired = model; completion(ETCoreAIAcquisitionStatusHit, model, nil); @@ -166,20 +164,11 @@ - (void)specializeModelAtURL:(NSURL*)url completion:(ETCoreAIAcquisitionCompleti - (void)evictModelWithBookmark:(NSData*)bookmark completion:(void (^)(NSError*))completion { ++evictions; run_async(self.context, ^{ - if (self.onEvict != nil) { - self.onEvict(bookmark); - } - @synchronized(self) { - if (self.evictError == nil) { - [self.models removeObjectForKey:bookmark]; - } - completion(self.evictError); - } + completion(fake_error(ETCoreAIErrorUnsupported)); }); } - (void)clearRetainedState { self.onSpecialize = nil; - self.onEvict = nil; self.lastAcquired = nil; [self.models removeAllObjects]; } diff --git a/backends/apple/coreai/runtime/test/coreai_host_test.mm b/backends/apple/coreai/runtime/test/coreai_host_test.mm index f388bc5089f..1acfce4320e 100644 --- a/backends/apple/coreai/runtime/test/coreai_host_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_host_test.mm @@ -15,14 +15,12 @@ int main(int argc, char** argv) { executorch::runtime::runtime_init(); @autoreleasepool { - executorch::backends::coreai::testing::set_acquisition_test_executable( + 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; - executorch::backends::coreai::testing::set_bookmark_test_executable( - argv[0]); 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..bd2f0f3cdc0 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,24 @@ 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): + def __init__(self, first: float, second: float): + 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): + 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 +67,38 @@ 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: + # 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/test_sdk_source_lifetime.py b/backends/apple/coreai/runtime/test/test_sdk_source_lifetime.py new file mode 100644 index 00000000000..4ef03012e84 --- /dev/null +++ b/backends/apple/coreai/runtime/test/test_sdk_source_lifetime.py @@ -0,0 +1,63 @@ +# 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): + 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 main(probe: Path) -> None: + from export_bridge_fixture import export_weighted_fixture + + architecture = run([str(probe), "architecture"]) + with tempfile.TemporaryDirectory(prefix="coreai-lifetime-") as temporary: + root = Path(temporary) + for kind in ("portable", "compiled"): + directory = root / kind + export_weighted_fixture(directory) + source = directory / "model.aimodel" + if kind == "compiled": + compiled = directory / "model.aimodelc" + run( + [ + "xcrun", "coreai-build", "compile", str(source), + "--output", str(compiled), "--platform", "macOS", + "--architecture", architecture, "--min-deployment-version", "27.0", + ] + ) + shutil.rmtree(source) + source = compiled + bookmark = directory / "saved.bookmark" + arguments = [str(source), str(directory / "fixture.json"), 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) + finally: + 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) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("probe", type=Path) + main(parser.parse_args().probe.resolve()) From f04d8f108170c041b7530153fc61d3c4962fcbf9 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 27 Oct 2026 13:42:57 -0700 Subject: [PATCH 3/4] [CoreAI] Recover interrupted source cleanup and harden lifetime tests Rebuild damaged staged source under the key lock after a cache miss. Treat cleanup after successful restoration or durable publication as best-effort. Exercise the production bridge in lifetime probes and cover cleanup, interruption, and crash boundaries. Test Plan: CMake build and CTest passed (134 native tests, 18 Swift tests, portable and AOT lifetime probes). Full Python suite: 121 passed; updated driver tests: 7 passed. Seven recovery/crash/contention tests passed ten repetitions. Formatters, Ruff, and diff whitespace checks passed. Generic Buck owner validation failed without diagnostics. --- backends/apple/coreai/CMakeLists.txt | 12 +- backends/apple/coreai/README.md | 44 +-- backends/apple/coreai/runtime/coreai_assets.h | 7 +- .../apple/coreai/runtime/coreai_assets.mm | 15 +- .../coreai/runtime/coreai_load_coordinator.h | 1 + .../coreai/runtime/coreai_load_coordinator.mm | 111 ++++--- .../apple/coreai/runtime/coreai_storage.h | 2 + .../apple/coreai/runtime/coreai_storage.mm | 33 +- .../runtime/test/SDKSourceLifetimeTest.swift | 134 ++++++-- .../test/coreai_acquisition_fixture.mm | 285 ++++++++++++------ .../runtime/test/coreai_acquisition_test.mm | 161 ++++++++-- .../coreai/runtime/test/coreai_fake_loader.mm | 50 ++- .../coreai/runtime/test/coreai_host_test.mm | 3 +- .../runtime/test/export_bridge_fixture.py | 9 +- .../runtime/test/run_sdk_source_lifetime.py | 114 +++++++ .../test/test_run_sdk_source_lifetime.py | 97 ++++++ .../runtime/test/test_sdk_source_lifetime.py | 63 ---- 17 files changed, 823 insertions(+), 318 deletions(-) create mode 100644 backends/apple/coreai/runtime/test/run_sdk_source_lifetime.py create mode 100644 backends/apple/coreai/runtime/test/test_run_sdk_source_lifetime.py delete mode 100644 backends/apple/coreai/runtime/test/test_sdk_source_lifetime.py diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index af5a2e15614..10d63f2b9e0 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -314,15 +314,19 @@ if(EXECUTORCH_BUILD_TESTS "COREAI_BRIDGE_FIXTURE_DIR=${_coreai_fixture_dir}" ) add_executable( - coreai_sdk_source_lifetime_test runtime/test/SDKSourceLifetimeTest.swift + 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>" + 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 @@ -330,8 +334,8 @@ if(EXECUTORCH_BUILD_TESTS add_test( NAME coreai_sdk_source_lifetime_test COMMAND - "${Python3_EXECUTABLE}" - "${CMAKE_CURRENT_SOURCE_DIR}/runtime/test/test_sdk_source_lifetime.py" + "${PYTHON_EXECUTABLE}" + "${CMAKE_CURRENT_SOURCE_DIR}/runtime/test/run_sdk_source_lifetime.py" "$" ) set_tests_properties(coreai_sdk_source_lifetime_test PROPERTIES TIMEOUT 240) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index ae3f17a6d87..c645e685def 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -57,10 +57,10 @@ 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. Every load first attempts SDK -bookmark restoration. A missing bookmark or SDK-confirmed cache miss falls back -to source materialization and persistent specialization, then binds the requested -function and ordered I/O on that same acquired model. +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 @@ -93,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. @@ -128,24 +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 inspect staging or consult NDS; the acquired SDK -model is retained through function binding. On a missing bookmark or confirmed -SDK miss, the runtime materializes the source, specializes it and atomically -writes the returned bookmark bytes. After durable publication, it removes the -entire keyed staging directory while holding the per-key lock. Persistent SDK +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 per-key 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. If staged-source cleanup fails, loading reports the -error but retains the durable bookmark and any remaining staged files. Neither -failure triggers automatic retries, SDK eviction or source repair. A subsequent -bookmark hit can load without inspecting or retrying staging cleanup. A partial -cleanup followed by an SDK miss fails on the damaged staged entry until it is -explicitly cleared. The backend does not track historical or unrecorded SDK -entries. +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. @@ -193,7 +192,10 @@ points allow delegate tests to substitute a fake bridge. `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 after removing disposable portable and AOT source copies. +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; 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_load_coordinator.h b/backends/apple/coreai/runtime/coreai_load_coordinator.h index 85827ea0fec..018117a18f9 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.h +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.h @@ -13,6 +13,7 @@ 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, diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm index 9ad967ccc0b..56f4b5ade74 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.mm +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -14,83 +14,122 @@ 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); + 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"); + 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); - }); + 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"); + 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"); + Internal, + "Invalid Core AI acquisition result"); return acquired; } -Result> restore(id loader, - NSData* bookmark) { +Result> restore( + id loader, + NSData* bookmark) { return acquire(^(ETCoreAIAcquisitionCompletion completion) { [loader restoreModelFromBookmark:bookmark completion:completion]; }); } -} // namespace +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, + 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"); + 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(); + if (!key.ok()) + return key.error(); auto locked = lock_bookmark(root, key.get()); - if (!locked.ok()) return locked.error(); + if (!locked.ok()) + return locked.error(); auto bookmark = read_bookmark(*locked.get()); - if (!bookmark.ok()) return bookmark.error(); + 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) return model; + if (!model.ok()) + return model.error(); + if (model.get() != nil) { + cleanup_source(*locked.get()); + return model; + } } 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()); - if (!source.ok()) return source.error(); + 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) { [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"); + if (!model.ok()) + return model.error(); + ET_CHECK_OR_RETURN_ERROR( + model.get() != nil, Internal, "Core AI specialization returned no model"); ET_CHECK_OK_OR_RETURN_ERROR( write_bookmark(*locked.get(), [model.get() copyBookmarkData])); - ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark_staging(*locked.get())); + cleanup_source(*locked.get()); return model; } -} // namespace executorch::backends::coreai +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h index a367c31b680..9de9a3ced28 100644 --- a/backends/apple/coreai/runtime/coreai_storage.h +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -49,6 +49,8 @@ runtime::Result*> storage_children( 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 24e865068d2..2712b3a1253 100644 --- a/backends/apple/coreai/runtime/coreai_storage.mm +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -363,16 +363,19 @@ 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( - [key isKindOfClass:NSString.class] && key.length == 64 && - [key rangeOfCharacterFromSet: - [[NSCharacterSet - characterSetWithCharactersInString:@"0123456789abcdef"] - invertedSet]] - .location == NSNotFound, - InvalidArgument, - "Invalid Core AI staging key"); + 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); @@ -383,10 +386,16 @@ FileDescriptor staging(retry_eintr([&] { 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( - staging.get(), + source_root_fd, key.fileSystemRepresentation, &info, AT_SYMLINK_NOFOLLOW); @@ -403,7 +412,7 @@ FileDescriptor staging(retry_eintr([&] { "Core AI keyed staging entry is not a directory"); FileDescriptor source(retry_eintr([&] { return openat( - staging.get(), + source_root_fd, key.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); })); @@ -418,11 +427,11 @@ FileDescriptor source(retry_eintr([&] { StorageOperation::Remove, [&] { return unlinkat( - staging.get(), key.fileSystemRepresentation, AT_REMOVEDIR); + source_root_fd, key.fileSystemRepresentation, AT_REMOVEDIR); }) == 0, AccessFailed, "Cannot remove Core AI keyed staging directory"); - return sync_storage_directory(staging.get()); + return sync_storage_directory(source_root_fd); } Error write_storage_file( diff --git a/backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift b/backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift index 8691b47591b..505a8d39284 100644 --- a/backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift +++ b/backends/apple/coreai/runtime/test/SDKSourceLifetimeTest.swift @@ -4,8 +4,8 @@ // 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 import Foundation +import Synchronization private struct Fixture: Decodable { let function: String @@ -16,41 +16,99 @@ private struct Fixture: Decodable { private enum Failure: Error { case invalidArguments - case missingFunction + case invalidResult case incorrectOutput(String) case unexpectedSource } -private func checkInference(_ model: AIModel, fixture: Fixture) async throws { - guard let function = try model.loadFunction(named: fixture.function) else { - throw Failure.missingFunction +// 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 } - var inputs: [String: NDArray] = [:] - for (index, values) in [[Float(5), 9], [Float(1), 2]].enumerated() { - var array = NDArray(shape: [2], scalarType: .float32) - do { - var view = array.mutableView(as: Float.self) - view.withUnsafeMutablePointer { pointer, _, _ in - pointer[0] = values[0] - pointer[1] = values[1] + 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)) + } } } - inputs[fixture.inputs[index]] = array } - var outputs = try await function.run(inputs: inputs) guard outputs.count == fixture.outputs.count else { throw Failure.incorrectOutput("count") } - for (index, name) in fixture.outputs.enumerated() { - guard let value = outputs.remove(name), let array = value.ndArray, - array.shape == [2], array.scalarType == .float32, - fixture.expected[index].count == 2 + 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 = array.view(as: Float.self).withUnsafePointer { pointer, _, strides in - [pointer[0], pointer[strides[0]]] + 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])") @@ -60,11 +118,11 @@ private func checkInference(_ model: AIModel, fixture: Fixture) async throws { @main private struct SDKSourceLifetimeTest { - static func main() async { + static func main() { do { let args = CommandLine.arguments if args.count == 2, args[1] == "architecture" { - print(AIModel.deviceArchitectureName) + print(ETCoreAISwiftModelLoader.deviceArchitectureName) return } guard args.count == 5 else { throw Failure.invalidArguments } @@ -77,31 +135,41 @@ private struct SDKSourceLifetimeTest { directory == fixtureURL.deletingLastPathComponent(), directory.deletingLastPathComponent().lastPathComponent.hasPrefix("coreai-lifetime-") else { throw Failure.invalidArguments } + let loader = ETCoreAISwiftModelLoader() if mode == "evict" { - try AIModelCache.deleteEntry(referencedBy: Data(contentsOf: bookmarkURL)) + 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" { - let model = try await AIModel.specialize( - contentsOf: source, options: .default, cache: .default, cachePolicy: .persistent) - try model.bookmarkData.write(to: bookmarkURL, options: .atomic) - try await checkInference(model, fixture: fixture) + 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 await checkInference(model, fixture: fixture) - print("Live inference passed after source removal") + 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 } - guard let model = try AIModel(resolvingBookmark: Data(contentsOf: bookmarkURL)) else { + let bookmark = try Data(contentsOf: bookmarkURL) + guard + let model = try acquire({ loader.restoreModel(fromBookmark: bookmark, completion: $0) }) + else { exit(3) } - try await checkInference(model, fixture: fixture) - print("Bookmark restoration and inference passed without source") + try checkInference(model, fixture: fixture) + print("Bridge restoration and inference passed without source") } else { throw Failure.invalidArguments } diff --git a/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm b/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm index a0af1340b76..7446e9ffbe6 100644 --- a/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_fixture.mm @@ -26,7 +26,9 @@ using runtime::Error; using runtime::Result; -dispatch_time_t deadline() { return dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC); } +dispatch_time_t deadline() { + return dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC); +} struct AcquisitionWorker { Manifest manifest; @@ -42,8 +44,13 @@ 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); + 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; } @@ -62,57 +69,78 @@ ParallelAcquisitionObservation& observation; bool finished = false; - AcquisitionWorkers(ScopedFakeBridgeState& state, BookmarkFakeLoader* fake, - ParallelAcquisitionObservation& result) + 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); + 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; + 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; + if (observation.callbacks_drained) + loader.onSpecialize = nil; finished = true; } - ~AcquisitionWorkers() { finish(); } + ~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(); + 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 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; + if (written == 1) + return true; + if (written < 0 && (errno == EINTR || errno == EAGAIN)) + continue; return false; } return false; } -} // namespace +} // 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()] + 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 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); @@ -120,7 +148,8 @@ ParallelAcquisitionObservation parallel_acquisitions(ScopedFakeBridgeState& brid dispatch_semaphore_t proceed = group.proceed; loader.onSpecialize = ^(NSURL*) { dispatch_semaphore_signal(entered); - if (dispatch_semaphore_wait(proceed, deadline()) != 0) ++*timeouts; + if (dispatch_semaphore_wait(proceed, deadline()) != 0) + ++*timeouts; }; for (int i = 0; i < 2; ++i) { auto& worker = group.workers[i]; @@ -129,21 +158,24 @@ ParallelAcquisitionObservation parallel_acquisitions(ScopedFakeBridgeState& brid worker.root = root.path; worker.loader = loader; } - observation.creation[0] = - pthread_create(&group.threads[0], nullptr, AcquisitionWorker::run, &group.workers[0]); + 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; + 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]); + 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) { + 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; + const bool second_entered = + dispatch_semaphore_wait(entered, second_deadline) == 0; observation.second_observed = same_key ? !second_entered : second_entered; } } @@ -152,103 +184,152 @@ ParallelAcquisitionObservation parallel_acquisitions(ScopedFakeBridgeState& brid 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]; +::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]; +::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) { +[[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); + 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; + 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())}; + 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 (!send_child_byte(fd, observation)) + _exit(12); } - if (loader->evictions.load() != 0 || data.requests.load() != data.releases.load()) _exit(14); + 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. + 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 +} // namespace int acquisition_child_mode(int argc, char** argv) { - if (argc < 2) return -1; + 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; + 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; + 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; + if (![dictionary isKindOfClass:NSDictionary.class]) + return 3; StorageOperation boundary = StorageOperation::AfterSDK; if (crash) { - if (![number isKindOfClass:NSNumber.class]) return 3; + if (![number isKindOfClass:NSNumber.class]) + return 3; boundary = static_cast([number intValue]); - if (boundary != StorageOperation::AfterSDK) return 4; + if (boundary != StorageOperation::BeforeSDK && + boundary != StorageOperation::AfterSDK) + return 4; } auto parsed = parse_manifest(encode(dictionary)); - if (!parsed.ok()) return 5; + if (!parsed.ok()) + return 5; auto selected = select_assets(parsed.get(), @"arch_b", @"macOS"); - if (!selected.ok()) return 5; + 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) + 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; + if (fcntl(fd, F_GETFD) < 0) + return 6; signal(SIGPIPE, SIG_IGN); - if (!send_child_byte(fd, 'C')) return 7; + if (!send_child_byte(fd, 'C')) + return 7; if (parallel) { id second_dictionary = payload[@"second"]; id same_key = payload[@"same_key"]; @@ -256,28 +337,36 @@ int acquisition_child_mode(int argc, char** argv) { ![same_key isKindOfClass:NSNumber.class]) return 3; auto second_parsed = parse_manifest(encode(second_dictionary)); - if (!second_parsed.ok()) return 5; + 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 (!second.ok()) + return 5; + parallel_child( + root, selected.get(), second.get(), [same_key boolValue], fd); } - if (close(fd) != 0) return 7; + if (close(fd) != 0) + return 7; ScopedFakeBridgeState bridge; @autoreleasepool { - // Terminal child paths deliberately bypass assertion-bearing scope destructors. + // Terminal child paths deliberately bypass assertion-bearing scope + // destructors. BookmarkFakeLoader* loader = [[BookmarkFakeLoader alloc] init]; - if (loader == nil) _exit(8); + 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); + if (event == boundary) + _exit(73); }; - auto result = acquire_bookmark_model(selected.get(), &data, root, @"macOS", @"arch_b", loader); + auto result = acquire_bookmark_model( + selected.get(), &data, root, @"macOS", @"arch_b", loader); _exit(result.ok() ? 9 : 10); } } -} // namespace executorch::backends::coreai::testing +} // 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 index af80260827e..a514a7ce393 100644 --- a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm @@ -43,7 +43,7 @@ - (void)specializeModelAtURL:(NSURL*)url protected: // Declared first so each case's loader, models and data die before it. ScopedFakeBridgeState bridge; - BookmarkDirectory root; + TestDirectory root; TestData data; Manifest manifest; NSString* key = nil; @@ -60,9 +60,10 @@ void SetUp() override { 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(); + 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; } } @@ -130,7 +131,9 @@ ino_t lock_inode() { 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]; + 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); @@ -138,7 +141,9 @@ ino_t lock_inode() { EXPECT_TRUE([saved() isEqual:original]); } -TEST_F(CoreAIAcquisitionTest, RestoreMissRespecializesStagedSourceAndReplacesBookmark) { +TEST_F( + CoreAIAcquisitionTest, + RestoreMissRespecializesStagedSourceAndReplacesBookmark) { BookmarkLoaderScope fake; ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); NSData* original = saved(); @@ -393,7 +398,7 @@ ino_t lock_inode() { { StorageFaultScope fault(StorageOperation::Remove); fault.match_index = 1; - EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); + EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); EXPECT_EQ(fault.hits, 1); bookmark = saved(); ASSERT_NE(bookmark, nil); @@ -412,7 +417,7 @@ ino_t lock_inode() { EXPECT_EQ(fake.loader->restores.load(), restores + 1); EXPECT_EQ(data.attempts.load(), attempts); EXPECT_TRUE([saved() isEqual:bookmark]); - EXPECT_TRUE(staging_key_exists()); + EXPECT_FALSE(staging_key_exists()); } TEST_F( @@ -424,7 +429,7 @@ ino_t lock_inode() { { StorageFaultScope fault(StorageOperation::Remove); fault.match_index = 2; - EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); + EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); EXPECT_EQ(fault.hits, 1); bookmark = saved(); ASSERT_NE(bookmark, nil); @@ -442,7 +447,7 @@ ino_t lock_inode() { EXPECT_EQ(fake.loader->restores.load(), restores + 1); EXPECT_EQ(data.attempts.load(), attempts); EXPECT_TRUE([saved() isEqual:bookmark]); - EXPECT_TRUE(staging_key_exists()); + EXPECT_FALSE(staging_key_exists()); } TEST_F( @@ -459,7 +464,7 @@ ino_t lock_inode() { if (operation == StorageOperation::Remove) arming->armed = true; }; - EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); + EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); EXPECT_EQ(scope.hits, 1); bookmark = saved(); ASSERT_NE(bookmark, nil); @@ -476,6 +481,87 @@ ino_t lock_inode() { 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) { @@ -487,8 +573,9 @@ ino_t lock_inode() { EXPECT_EQ(lock_inode(), first); } -class CoreAIAcquisitionSaveFailureTest : public CoreAIAcquisitionTest, - public ::testing::WithParamInterface {}; +class CoreAIAcquisitionSaveFailureTest + : public CoreAIAcquisitionTest, + public ::testing::WithParamInterface {}; TEST_P(CoreAIAcquisitionSaveFailureTest, PreservesBookmarkAndRetries) { const bool prior = GetParam(); BookmarkLoaderScope fake; @@ -510,29 +597,36 @@ ino_t lock_inode() { EXPECT_EQ(fault.hits, 1); } EXPECT_TRUE(original == nil ? saved() == nil : [saved() isEqual:original]); - EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:fake.loader.lastSpecializeURL]); + 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_EQ(staging_key_exists(), prior); + EXPECT_FALSE(staging_key_exists()); } -INSTANTIATE_TEST_SUITE_P(Boundaries, CoreAIAcquisitionSaveFailureTest, ::testing::Bool(), - [](const ::testing::TestParamInfo& info) { - return info.param ? "PriorBookmark" : "Cold"; - }); +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')}; + 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( + 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)) + << "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"; @@ -550,15 +644,32 @@ void expect_parallel(bool same_key) { 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(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]); + 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); @@ -568,4 +679,4 @@ void expect_parallel(bool same_key) { EXPECT_FALSE(staging_key_exists()); } -} // namespace +} // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm index 54a14287d94..18dfe5208e7 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm @@ -48,16 +48,19 @@ void run_async(CoreAIFakeBridgeContext* context, void (^work)()) { }); } -::testing::AssertionResult wait_for_callbacks(CoreAIFakeBridgeContext* context, - dispatch_time_t deadline) { +::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 +dispatch_time_t cleanup_deadline() { + return dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC); +} +} // namespace @interface BookmarkFakeModel () @property(nonatomic, strong) CoreAIFakeBridgeContext* context; @@ -130,12 +133,16 @@ - (void)restoreModelFromBookmark:(NSData*)bookmark } BookmarkFakeModel* model = self.restoreMiss ? nil : self.models[bookmark]; self.lastAcquired = model; - completion(model == nil ? ETCoreAIAcquisitionStatusMiss : ETCoreAIAcquisitionStatusHit, model, - nil); + completion( + model == nil ? ETCoreAIAcquisitionStatusMiss + : ETCoreAIAcquisitionStatusHit, + model, + nil); } }); } -- (void)specializeModelAtURL:(NSURL*)url completion:(ETCoreAIAcquisitionCompletion)completion { +- (void)specializeModelAtURL:(NSURL*)url + completion:(ETCoreAIAcquisitionCompletion)completion { ++specializations; run_async(self.context, ^{ if (self.onSpecialize != nil) { @@ -149,19 +156,25 @@ - (void)specializeModelAtURL:(NSURL*)url completion:(ETCoreAIAcquisitionCompleti } if (![NSFileManager.defaultManager fileExistsAtPath:url.path]) { ++self.context->state.missing_bundles; - completion(ETCoreAIAcquisitionStatusError, nil, fake_error(ETCoreAIErrorInvalidModel)); + completion( + ETCoreAIAcquisitionStatusError, + nil, + fake_error(ETCoreAIErrorInvalidModel)); return; } - BookmarkFakeModel* model = [[BookmarkFakeModel alloc] initWithContext:self.context]; + BookmarkFakeModel* model = + [[BookmarkFakeModel alloc] initWithContext:self.context]; model.owner = self; - model.bookmark = [NSUUID.UUID.UUIDString dataUsingEncoding:NSUTF8StringEncoding]; + 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 { +- (void)evictModelWithBookmark:(NSData*)bookmark + completion:(void (^)(NSError*))completion { ++evictions; run_async(self.context, ^{ completion(fake_error(ETCoreAIErrorUnsupported)); @@ -177,7 +190,8 @@ - (void)clearRetainedState { namespace executorch::backends::coreai::testing { ScopedFakeBridgeState::ScopedFakeBridgeState() - : context_([[CoreAIFakeBridgeContext alloc] init]), previous_(current_context) { + : context_([[CoreAIFakeBridgeContext alloc] init]), + previous_(current_context) { current_context = context_; } @@ -199,16 +213,20 @@ - (void)clearRetainedState { current_context = previous_; } -FakeBridgeState& ScopedFakeBridgeState::state() { return context_->state; } +FakeBridgeState& ScopedFakeBridgeState::state() { + return context_->state; +} -::testing::AssertionResult ScopedFakeBridgeState::wait_for_callbacks(dispatch_time_t deadline) { +::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"; + EXPECT_NE(context_, nil) + << "BookmarkLoaderScope requires ScopedFakeBridgeState"; if (context_ != nil) { loader = [[BookmarkFakeLoader alloc] init]; context_->state.bookmark_loader = loader; @@ -229,4 +247,4 @@ - (void)clearRetainedState { loader = nil; } -} // namespace executorch::backends::coreai::testing +} // 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 1acfce4320e..d87eefdbbba 100644 --- a/backends/apple/coreai/runtime/test/coreai_host_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_host_test.mm @@ -18,7 +18,8 @@ int main(int argc, char** argv) { executorch::backends::coreai::testing::set_bookmark_test_executable( argv[0]); const int acquisition_result = - executorch::backends::coreai::testing::acquisition_child_mode(argc, argv); + executorch::backends::coreai::testing::acquisition_child_mode( + argc, argv); if (acquisition_result >= 0) return acquisition_result; const int child_result = diff --git a/backends/apple/coreai/runtime/test/export_bridge_fixture.py b/backends/apple/coreai/runtime/test/export_bridge_fixture.py index bd2f0f3cdc0..947c47cc86e 100644 --- a/backends/apple/coreai/runtime/test/export_bridge_fixture.py +++ b/backends/apple/coreai/runtime/test/export_bridge_fixture.py @@ -30,7 +30,10 @@ def forward(self, left, right): 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 @@ -39,6 +42,7 @@ def __init__(self, first: float, second: float): ) 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, @@ -74,12 +78,15 @@ 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, + 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( 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 diff --git a/backends/apple/coreai/runtime/test/test_sdk_source_lifetime.py b/backends/apple/coreai/runtime/test/test_sdk_source_lifetime.py deleted file mode 100644 index 4ef03012e84..00000000000 --- a/backends/apple/coreai/runtime/test/test_sdk_source_lifetime.py +++ /dev/null @@ -1,63 +0,0 @@ -# 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): - 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 main(probe: Path) -> None: - from export_bridge_fixture import export_weighted_fixture - - architecture = run([str(probe), "architecture"]) - with tempfile.TemporaryDirectory(prefix="coreai-lifetime-") as temporary: - root = Path(temporary) - for kind in ("portable", "compiled"): - directory = root / kind - export_weighted_fixture(directory) - source = directory / "model.aimodel" - if kind == "compiled": - compiled = directory / "model.aimodelc" - run( - [ - "xcrun", "coreai-build", "compile", str(source), - "--output", str(compiled), "--platform", "macOS", - "--architecture", architecture, "--min-deployment-version", "27.0", - ] - ) - shutil.rmtree(source) - source = compiled - bookmark = directory / "saved.bookmark" - arguments = [str(source), str(directory / "fixture.json"), 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) - finally: - 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) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("probe", type=Path) - main(parser.parse_args().probe.resolve()) From 19ef724683ea8cdfe7b6db9be47fe7ce9deaa305 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 27 Oct 2026 17:19:27 -0700 Subject: [PATCH 4/4] [CoreAI] Log acquisition and durable bookmark milestones Add key-correlated Info logs for lookup, restore hits and misses, specialization start and success, and durable bookmark publication. Keep SDK-internal reuse distinct from observable bookmark hits and cover successful and failed acquisition paths. Test Plan: All three acquisition-layer CTest targets passed with Info logging enabled, including six new lifecycle-log tests. No dependency on the later cache-clear API. --- backends/apple/coreai/README.md | 17 ++- .../coreai/runtime/coreai_load_coordinator.mm | 26 +++- .../runtime/test/coreai_acquisition_test.mm | 120 +++++++++++++++++- 3 files changed, 159 insertions(+), 4 deletions(-) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index c645e685def..567a9bed233 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -133,7 +133,7 @@ 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 per-key lock. +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. @@ -213,6 +213,21 @@ 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` diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm index 56f4b5ade74..d239ccb0f64 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.mm +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -92,7 +92,9 @@ void cleanup_source(const BookmarkLock& lock) { auto key = bookmark_key(selected, platform, architecture); if (!key.ok()) return key.error(); - auto locked = lock_bookmark(root, key.get()); + 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()); @@ -103,9 +105,20 @@ void cleanup_source(const BookmarkLock& lock) { 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()) @@ -120,14 +133,25 @@ void cleanup_source(const BookmarkLock& lock) { 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; } diff --git a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm index a514a7ce393..b70ae5380d6 100644 --- a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm @@ -8,8 +8,11 @@ #include "coreai_acquisition_fixture.h" +#include #include #include +#include +#include #include #include "coreai_fault_scope.h" #include "coreai_filesystem_fixture.h" @@ -85,8 +88,7 @@ BOOL staging_key_exists() { } NSURL* lock_url() { - return [[root.url URLByAppendingPathComponent:@"locks"] - URLByAppendingPathComponent:[key stringByAppendingString:@".lock"]]; + return test_bookmark_lock_url(root.url, key); } ino_t lock_inode() { @@ -679,4 +681,118 @@ void expect_parallel(bool same_key) { 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