[CoreAI] Coordinate model acquisition

backends/apple/coreai/CMakeLists.txt

@@ -27,8 +27,9 @@ endforeach()
 
 enable_language(OBJCXX)
 
-set(_coreai_runtime_sources runtime/coreai_assets.mm runtime/coreai_storage.mm
-                            runtime/coreai_bookmarks.mm
+set(_coreai_runtime_sources
+    runtime/coreai_assets.mm runtime/coreai_storage.mm
+    runtime/coreai_bookmarks.mm runtime/coreai_load_coordinator.mm
 )
 
 # Keep ARC and C++ exception settings away from the private Swift module.
@@ -56,6 +57,10 @@ if(EXECUTORCH_BUILD_TESTS AND CMAKE_SYSTEM_NAME STREQUAL "Darwin")
     runtime/test/coreai_bookmark_test.mm
     runtime/test/coreai_source_fixture.mm
     runtime/test/coreai_bookmark_fixture.mm
+    runtime/test/coreai_fake_loader.mm
+    runtime/test/coreai_acquisition_test.mm
+    runtime/test/coreai_acquisition_fixture.mm
+    runtime/ETCoreAITensor.mm
     ${_coreai_runtime_sources}
   )
   coreai_configure_objc_target(coreai_host_test)
@@ -308,6 +313,32 @@ if(EXECUTORCH_BUILD_TESTS
     PROPERTIES TIMEOUT 600 ENVIRONMENT
                "COREAI_BRIDGE_FIXTURE_DIR=${_coreai_fixture_dir}"
   )
+  add_executable(
+    coreai_sdk_source_lifetime_test
+    runtime/test/SDKSourceLifetimeTest.swift runtime/ETCoreAIModel.swift
+    runtime/ETCoreAITensor.mm
+  )
+  set_target_properties(
+    coreai_sdk_source_lifetime_test PROPERTIES Swift_LANGUAGE_VERSION 6
+  )
+  coreai_configure_objc_target(coreai_sdk_source_lifetime_test)
+  coreai_require_os27(coreai_sdk_source_lifetime_test)
+  target_compile_options(
+    coreai_sdk_source_lifetime_test
+    PRIVATE
+      "$<$<COMPILE_LANGUAGE:Swift>:-parse-as-library;-import-objc-header;${CMAKE_CURRENT_SOURCE_DIR}/runtime/ETCoreAIBridge.h>"
+  )
+  target_link_libraries(
+    coreai_sdk_source_lifetime_test PRIVATE executorch::coreai_dependencies
+  )
+  add_test(
+    NAME coreai_sdk_source_lifetime_test
+    COMMAND
+      "${PYTHON_EXECUTABLE}"
+      "${CMAKE_CURRENT_SOURCE_DIR}/runtime/test/run_sdk_source_lifetime.py"
+      "$<TARGET_FILE:coreai_sdk_source_lifetime_test>"
+  )
+  set_tests_properties(coreai_sdk_source_lifetime_test PROPERTIES TIMEOUT 240)
 endif()
 
 install(

backends/apple/coreai/README.md

@@ -57,6 +57,11 @@ architectures=["h17p"])` requests that compiler architecture; omitting the list
 lets the compiler emit its supported architectures for the target platform.
 These are Core AI architecture names, not CPU names such as `arm64`.
 
+Both formats can require device specialization. The acquisition coordinator
+first attempts SDK bookmark restoration. A missing bookmark or SDK-confirmed
+cache miss falls back to source materialization and persistent specialization.
+It returns the acquired model; function binding is the caller's responsibility.
+
 ### Asset storage
 
 The assets root must be an absolute path. The application chooses it, should
@@ -88,10 +93,11 @@ source root, including after a previous publication's final flush failed.
 Missing bundles are written to a staging directory and published with an
 exclusive rename; each selected NDS payload is fetched once and checked for size
 before writing. If another loader publishes first, its bundle is validated and
-used. Incomplete, size-mismatched or otherwise malformed existing bundles fail
-loading without replacement, since an SDK model may still be using them.
-Only interrupted system calls are retried; other failures are returned to the
-caller without automatic retries or repair. Temporary bundles are removed on
+used. The acquisition coordinator can remove an incomplete or size-mismatched
+regular-file tree under the key lock and rebuild it from available named data.
+Cleanup never follows symlinks and rejects unexpected entry types. Standalone
+source preparation rejects damaged trees unless replacement is explicitly enabled.
+Named-data and SDK failures are not retried. Temporary bundles are removed on
 ordinary failures before publication; a caller can explicitly retry after
 addressing the failure.
 
@@ -123,6 +129,22 @@ a raw bookmark in `bookmarks`, a bundle directory in `staging`, and one of
 256 shared coordination files in `locks`. There is no PTE-wide bookmark list or
 serialized lifecycle record.
 
+A healthy bookmark hit does not consult NDS and attempts to clean leftover
+staging. On a missing bookmark or confirmed SDK miss, the coordinator prepares
+the source, specializes it and atomically writes the returned bookmark bytes.
+After durable publication, it attempts to remove the entire keyed staging
+directory while holding the key's stripe lock.
+caching supports inference and bookmark restoration without that extracted copy.
+Keep the packaged named data available: a later SDK cache miss requires source
+materialization again. There is no secondary source-URL cache lookup.
+
+SDK errors, unreadable bookmarks and publication failures fail loading without
+automatic cache deletion. If publication fails after specialization, the SDK
+entry and source remain. Cleanup after a bookmark hit or durable publication is
+best-effort: failures are logged without discarding the usable model or bookmark.
+A later load can finish cleanup or rebuild a partially removed source after an
+SDK miss. No cleanup failure triggers SDK eviction or retries within that load.
+The backend does not track historical or unrecorded SDK entries.
 
 Raw bookmarks have an 8 MiB backend allocation limit, not an asserted SDK format
 maximum, and are published atomically because losing one orphans an SDK entry.
@@ -169,6 +191,11 @@ points allow delegate tests to substitute a fake bridge.
 `coreai_swift_bridge_test` runs the bridge's XCTests under `xctest` with
 `EXECUTORCH_BUILD_TESTS=ON`. It links CoreAI, so it runs only on macOS 27.
 The test build exports static- and dynamic-shape models for real session tests.
+The SDK source-lifetime regression checks live inference and fresh-process
+bookmark restoration through the bridge after removing disposable portable and
+AOT source copies. The AOT case is skipped when the SDK cannot identify the
+device architecture. A killed probe can leave an SDK entry if it dies before
+saving its bookmark; cleanup never purges unrelated default-cache entries.
 CMake's `PYTHON_EXECUTABLE` must use the Python environment with the pinned Core
 AI requirements above. Backend Python changes regenerate the dedicated fixture
 directory, replacing previous contents. CTest supplies its location to XCTest;
@@ -186,13 +213,29 @@ Ninja builds use one architecture per build directory. Current runtime support
 covers arm64 macOS and iOS device builds. x86_64 builds are blocked by Swift
 `Float16` availability, and the tested iOS simulator SDKs do not contain Core AI.
 
+### Lifecycle logging
+
+Build with `EXECUTORCH_ENABLE_LOGGING=ON` and `EXECUTORCH_LOG_LEVEL=Info` to
+observe cache lookup, bookmark restoration hits, missing-bookmark versus
+SDK-confirmed misses, specialization start/success, and durable bookmark saves. Each message carries the cache key so events from
+multiple partitions can be correlated. These messages do not include bookmark
+contents, model data, or per-inference logs.
+
+A cache hit means that the SDK successfully restored a model from its bookmark.
+Specialization start means a request to the SDK's persistent default cache; the
+SDK does not report whether that request compiled afresh or reused an internal
+artifact. The durable-save message confirms local bookmark publication,
+including its final directory flush, not independent verification of SDK-private
+artifact storage. Failed operations do not emit their success messages.
+
 ## Host Tests
 
 On macOS, `EXECUTORCH_BUILD_COREAI=ON` with `EXECUTORCH_BUILD_TESTS=ON`
 registers the standalone `coreai_host_test` executable with CTest. It tests
 manifest parsing, AOT asset selection, and filesystem storage using Foundation,
-without Swift or CoreAI SDK calls. The host test target is omitted for Apple
-cross-builds.
+without Swift or CoreAI SDK calls. Acquisition contention tests run two threads
+inside a bounded child process; bookmark lock tests separately cover contention
+between processes. The host test target is omitted for Apple cross-builds.
 
 CI builds and runs these tests on macOS 27. To use the same configuration from
 the repository root:

backends/apple/coreai/runtime/coreai_assets.h

@@ -38,11 +38,14 @@ runtime::Result<Manifest> 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<NSURL*> 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

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,

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<NSData*> 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<NSString*> prepare_bookmark_staging(
       const BookmarkLock& lock);
   BookmarkLock(
@@ -65,6 +66,9 @@ lock_bookmark(NSString* root, NSString* key, bool create = true);
 runtime::Result<NSData*> 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<NSString*> prepare_bookmark_staging(const BookmarkLock& lock);
 

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<NSString*> prepare_bookmark_staging(const BookmarkLock& lock) {
   auto staging = open_child(lock.root_.get(), @"staging", true);
   if (!staging.ok())

backends/apple/coreai/runtime/coreai_load_coordinator.h

@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#pragma once
+
+#import "ETCoreAIBridge.h"
+#import "coreai_bookmarks.h"
+
+namespace executorch::backends::coreai {
+
+// The key lock outlives all SDK callbacks and raw bookmark publication.
+// Cleanup after a restore hit or durable publication is best-effort.
+runtime::Result<id<ETCoreAIPreparedModel>> acquire_bookmark_model(
+    const Manifest& selected,
+    const runtime::NamedDataMap* named_data,
+    NSString* root,
+    NSString* platform,
+    NSString* architecture,
+    id<ETCoreAIModelLoading> loader);
+
+} // namespace executorch::backends::coreai