diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h b/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h index 11d957044e9..a9e06efa90d 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h @@ -99,6 +99,17 @@ NS_ASSUME_NONNULL_BEGIN - (NSUInteger)compact:(NSUInteger)sizeInBytes error:(NSError* __autoreleasing*)error; +/// Executes a block with a unique temporary directory. +/// +/// A new temporary subdirectory URL is created inside the receiver’s designated +/// base directory. The directory is passed to the block, which can use it to +/// perform temporary file operations. After the block finishes executing, +/// the directory and its contents are removed. +/// +/// @param block A block to execute. The block receives a unique URL. +- (void)withTemporaryDirectory:(void (^)(NSURL* directoryURL))block; + + /// Purges the assets storage. The assets are moved to the trash directory and are asynchronously /// deleted. /// @@ -117,6 +128,12 @@ NS_ASSUME_NONNULL_BEGIN /// contents are deleted asynchronously. @property (copy, readonly, nonatomic) NSURL* trashDirectoryURL; + +/// The staging directory URL, used to hold assets that are being prepared or processed +/// before they are moved into their final location. The contents of this directory +/// are temporary and may be cleared when no longer needed. +@property (copy, readonly, nonatomic) NSURL* stagingDirectoryURL; + /// The file manager. @property (strong, readonly, nonatomic) NSFileManager* fileManager; diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm index 256026e1f09..d98a1d7331b 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm @@ -170,46 +170,74 @@ bool set_total_assets_size(size_t total_size, return true; } -bool exclude_item_from_backup(NSURL *url, NSError * __autoreleasing *error) { - return [url setResourceValue:@(YES) forKey:NSURLIsExcludedFromBackupKey error:error]; -} -NSURL * _Nullable create_directory_if_needed(NSURL *url, - NSString *name, +NSURL * _Nullable create_directory_if_needed(NSURL *dirURL, NSFileManager *fm, - NSError * __autoreleasing *error) { - NSURL *directory_url = [url URLByAppendingPathComponent:name]; - if (![fm fileExistsAtPath:directory_url.path] && - ![fm createDirectoryAtURL:directory_url withIntermediateDirectories:NO attributes:@{} error:error]) { - return nil; - } + NSError **error) { + NSCParameterAssert(dirURL); + NSCParameterAssert(dirURL.isFileURL); + NSCParameterAssert(fm); - ::exclude_item_from_backup(directory_url, nil); - - return directory_url; + NSString *dirPath = dirURL.path; + + // Fast path: is it already a directory? + BOOL isDir = NO; + if (dirPath && [fm fileExistsAtPath:dirPath isDirectory:&isDir] && isDir) { + return dirURL; + } + + // Try to create the directory and its parents. + NSDictionary *attrs = @{ NSFileProtectionKey : NSFileProtectionCompleteUntilFirstUserAuthentication }; + if (![fm createDirectoryAtURL:dirURL + withIntermediateDirectories:YES + attributes:attrs + error:error]) { + // Lost a race and creation failed because something already exists, check if it's a directory. + isDir = NO; + if (dirPath && [fm fileExistsAtPath:dirPath isDirectory:&isDir] && isDir) { + if (error) { *error = nil; } + } else { + return nil; + } + } + + // Best effort: exclude from backup (ignore failure) + (void)[dirURL setResourceValue:@YES forKey:NSURLIsExcludedFromBackupKey error:nil]; + + return dirURL; } -bool is_directory_empty(NSURL *url, NSFileManager *fm, NSError * __autoreleasing *error) { - BOOL is_directory = NO; - if (![fm fileExistsAtPath:url.path isDirectory:&is_directory] && !is_directory) { +bool is_missing_or_empty_directory(NSURL *dirURL, NSFileManager *fm, NSError * __autoreleasing *error) { + NSString *dirPath = dirURL.path; + BOOL isDir = NO; + BOOL doesFileExist = dirPath && [fm fileExistsAtPath:dirPath isDirectory:&isDir]; + if (!doesFileExist) { return true; } + if (!isDir) { + return false; + } - __block NSError *local_error = nil; - BOOL (^errorHandler)(NSURL *url, NSError *error) = ^BOOL(NSURL *url, NSError *enumeration_error) { - local_error = enumeration_error; - return NO; - }; - - NSDirectoryEnumerator *enumerator = [fm enumeratorAtURL:url + __block NSError *localError = nil; + NSDirectoryEnumerator *enumerator = [fm enumeratorAtURL:dirURL includingPropertiesForKeys:@[] options:NSDirectoryEnumerationProducesRelativePathURLs - errorHandler:errorHandler]; - if (local_error && error) { - *error = local_error; + errorHandler:^BOOL(NSURL *u, NSError *e){ localError = e; return NO; }]; + + // If enumerator failed to create, do not say the directory is empty + if (!enumerator) { + return false; + } + + id nextObject = [enumerator nextObject]; + + // Do not treat enumeration errors as empty directory + if (localError) { + if (error) { *error = localError; } + return false; } - return [enumerator nextObject] == nil; + return nextObject == nil; } NSURL * _Nullable get_asset_url(const Asset& asset) { @@ -254,6 +282,7 @@ BOOL is_asset_alive(NSMapTable *assets_in_use_map, return assets; } + } //namespace @interface ETCoreMLAssetManager () { @@ -295,18 +324,33 @@ - (nullable instancetype)initWithDatabase:(const std::shared_ptr&)data } NSFileManager *fileManager = [[NSFileManager alloc] init]; - NSURL *managedAssetsDirectoryURL = ::create_directory_if_needed(assetsDirectoryURL, @"models", fileManager, error); + + NSDictionary *attrs = @{ NSFileProtectionKey : NSFileProtectionCompleteUntilFirstUserAuthentication }; + + NSURL *managedAssetsDirectoryURL = [assetsDirectoryURL URLByAppendingPathComponent:@"models"]; + managedAssetsDirectoryURL = ::create_directory_if_needed(managedAssetsDirectoryURL, fileManager, error); if (!managedAssetsDirectoryURL) { return nil; } - - NSURL *managedTrashDirectoryURL = ::create_directory_if_needed(trashDirectoryURL, @"models", fileManager, error); + (void)[fileManager setAttributes:attrs ofItemAtPath:managedAssetsDirectoryURL.path error:nil]; // best-effort + + + NSURL *managedTrashDirectoryURL = [trashDirectoryURL URLByAppendingPathComponent:@"models"]; + managedTrashDirectoryURL = ::create_directory_if_needed(managedTrashDirectoryURL, fileManager, error); if (!managedTrashDirectoryURL) { return nil; } - + (void)[fileManager setAttributes:attrs ofItemAtPath:managedTrashDirectoryURL.path error:nil]; // best-effort + + NSURL *managedStagingDirectoryURL = [assetsDirectoryURL URLByAppendingPathComponent:@"staging"]; + managedStagingDirectoryURL = ::create_directory_if_needed(managedStagingDirectoryURL, fileManager, error); + if (!managedStagingDirectoryURL) { + return nil; + } + (void)[fileManager setAttributes:attrs ofItemAtPath:managedStagingDirectoryURL.path error:nil]; // best-effort + // If directory is empty then purge the stores - if (::is_directory_empty(managedAssetsDirectoryURL, fileManager, nil)) { + if (::is_missing_or_empty_directory(managedAssetsDirectoryURL, fileManager, nil)) { assetsMetaStore.impl()->purge(ec); assetsStore.impl()->purge(ec); } @@ -315,10 +359,10 @@ - (nullable instancetype)initWithDatabase:(const std::shared_ptr&)data _assetsStore = std::move(assetsStore); _assetsMetaStore = std::move(assetsMetaStore); _assetsDirectoryURL = managedAssetsDirectoryURL; + _stagingDirectoryURL = managedStagingDirectoryURL; _trashDirectoryURL = managedTrashDirectoryURL; _estimatedSizeInBytes = sizeInBytes.value(); _maxAssetsSizeInBytes = maxAssetsSizeInBytes; - _fileManager = fileManager; _trashQueue = dispatch_queue_create("com.executorchcoreml.assetmanager.trash", DISPATCH_QUEUE_SERIAL_WITH_AUTORELEASE_POOL); _syncQueue = dispatch_queue_create("com.executorchcoreml.assetmanager.sync", DISPATCH_QUEUE_SERIAL_WITH_AUTORELEASE_POOL); @@ -333,7 +377,35 @@ - (nullable instancetype)initWithDatabaseURL:(NSURL *)databaseURL assetsDirectoryURL:(NSURL *)assetsDirectoryURL trashDirectoryURL:(NSURL *)trashDirectoryURL maxAssetsSizeInBytes:(NSInteger)maxAssetsSizeInBytes - error:(NSError * __autoreleasing *)error { + error:(NSError * __autoreleasing *)error { + + NSURL *databaseDirectoryURL = [databaseURL URLByDeletingLastPathComponent]; + NSFileManager *fm = [[NSFileManager alloc] init]; + if (!::create_directory_if_needed(databaseDirectoryURL, fm, error)) { + return nil; + } + + // Ensure correct file protection + NSMutableArray *maybeDBPaths = [NSMutableArray array]; + NSString *databaseDirectoryPath = databaseDirectoryURL.path; + if (databaseDirectoryPath) { [maybeDBPaths addObject:databaseDirectoryPath]; } + + // Ensure correct file protection on existing database files, if any + // New database files should inherit the protection from the parent directory + NSString *databasePath = databaseURL.path; + if (databasePath) { + [maybeDBPaths addObject:databasePath]; + [maybeDBPaths addObject:[databasePath stringByAppendingString:@"-wal"]]; + [maybeDBPaths addObject:[databasePath stringByAppendingString:@"-shm"]]; + [maybeDBPaths addObject:[databasePath stringByAppendingString:@"-journal"]]; + } + NSDictionary *attrs = @{ NSFileProtectionKey : NSFileProtectionCompleteUntilFirstUserAuthentication }; + for (NSString *p in maybeDBPaths) { + if ([fm fileExistsAtPath:p]) { + (void)[fm setAttributes:attrs ofItemAtPath:p error:nil]; // best-effort + } + } + auto database = make_database(databaseURL, kBusyTimeIntervalInMS, error); if (!database) { return nil; @@ -346,14 +418,30 @@ - (nullable instancetype)initWithDatabaseURL:(NSURL *)databaseURL error:error]; } -- (nullable NSURL *)moveURL:(NSURL *)url - toUniqueURLInDirectory:(NSURL *)directoryURL - error:(NSError * __autoreleasing *)error { - NSURL *dstURL = [directoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; +- (void)withTemporaryDirectory:(void (^)(NSURL *directoryURL))block { + NSURL *dstURL = [self.stagingDirectoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; + block(dstURL); + if (![self.fileManager fileExistsAtPath:dstURL.path]) { + return; + } + [self moveItemAtURLToTrash:dstURL error:nil]; + [self cleanupTrashDirectory]; +} + +- (NSURL * _Nullable) moveItemAtURLToTrash:(NSURL *)url + error:(NSError * __autoreleasing *)error { + ::create_directory_if_needed(self.trashDirectoryURL, self.fileManager, error); + NSURL *dstURL = [self.trashDirectoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; + + if (!url) { + ETCoreMLLogErrorAndSetNSError(error, ETCoreMLErrorInternalError, "Move operation failed: source URL is nil."); + return nil; + } + if (![self.fileManager moveItemAtURL:url toURL:dstURL error:error]) { return nil; } - + return dstURL; } @@ -365,7 +453,7 @@ - (void)cleanupAssetIfNeeded:(ETCoreMLAsset *)asset { NSString *identifier = asset.identifier; dispatch_async(self.syncQueue, ^{ NSError *cleanupError = nil; - if (![self _removeAssetWithIdentifier:asset.identifier error:&cleanupError]) { + if (![self _removeAssetWithIdentifier:asset.identifier alreadyInsideTransaction:NO error:&cleanupError]) { ETCoreMLLogError(cleanupError, "Failed to remove asset with identifier = %@", identifier); @@ -378,6 +466,7 @@ - (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL error:(NSError * __autoreleasing *)error { dispatch_assert_queue(self.syncQueue); NSString *extension = srcURL.lastPathComponent.pathExtension; + ::create_directory_if_needed(self.assetsDirectoryURL, self.fileManager, error); NSURL *dstURL = [self.assetsDirectoryURL URLByAppendingPathComponent:[NSString stringWithFormat:@"%@.%@", identifier, extension]]; auto asset = Asset::make(srcURL, identifier, self.fileManager, error); if (!asset) { @@ -391,7 +480,7 @@ - (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL bool status = _assetsStore.impl()->transaction([self, &assetValue, assetSizeInBytes, srcURL, dstURL, &ec, error]() { const std::string& assetIdentifier = assetValue.identifier; // If an asset exists with the same identifier then remove it. - if (![self _removeAssetWithIdentifier:@(assetIdentifier.c_str()) error:error]) { + if (![self _removeAssetWithIdentifier:@(assetIdentifier.c_str()) alreadyInsideTransaction:YES error:error]) { return false; } @@ -407,9 +496,15 @@ - (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL return false; } - // If an asset exists move it - [self moveURL:dstURL toUniqueURLInDirectory:self.trashDirectoryURL error:nil]; - + // If a file already exists at `dstURL`, move it to the trash for removal. + if ([self.fileManager fileExistsAtPath:dstURL.path]) { + if (![self moveItemAtURLToTrash:dstURL error:error]) { + // Log error and return false + ETCoreMLLogErrorAndSetNSError(error, ETCoreMLErrorInternalError, "moveItemAtURLToTrash failed"); + return false; + } + } + // Move the asset to assets directory. if (![self.fileManager moveItemAtURL:srcURL toURL:dstURL error:error]) { return false; @@ -427,22 +522,32 @@ - (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL [self.assetsInUseMap setObject:result forKey:identifier]; } else { [self cleanupAssetIfNeeded:result]; + return nil; } return result; } - (void)triggerCompaction { - if (self.estimatedSizeInBytes < self.maxAssetsSizeInBytes) { - return; + if (self.estimatedSizeInBytes >= self.maxAssetsSizeInBytes) { + __weak __typeof(self) weakSelf = self; + dispatch_async(self.syncQueue, ^{ + NSError *localError = nil; + if (![weakSelf _compact:self.maxAssetsSizeInBytes error:&localError]) { + ETCoreMLLogError(localError, "Failed to compact asset store."); + } + }); } - + + // Always clean the trash directory to ensure a minimal footprint. + // The `trashQueue` is serialized, so only one cleanup will run at a time. + [self cleanupTrashDirectory]; +} + +- (void)cleanupTrashDirectory { __weak __typeof(self) weakSelf = self; - dispatch_async(self.syncQueue, ^{ - NSError *localError = nil; - if (![weakSelf _compact:self.maxAssetsSizeInBytes error:&localError]) { - ETCoreMLLogError(localError, "Failed to compact asset store."); - } + dispatch_async(self.trashQueue, ^{ + [weakSelf removeFilesInTrashDirectory]; }); } @@ -513,6 +618,7 @@ - (BOOL)hasAssetWithIdentifier:(NSString *)identifier } - (BOOL)_removeAssetWithIdentifier:(NSString *)identifier + alreadyInsideTransaction:(BOOL)alreadyInsideTransaction error:(NSError * __autoreleasing *)error { dispatch_assert_queue(self.syncQueue); // Asset is alive we can't delete it. @@ -536,8 +642,9 @@ - (BOOL)_removeAssetWithIdentifier:(NSString *)identifier const auto& assetValue = asset.value(); size_t assetSizeInBytes = std::min(_estimatedSizeInBytes, static_cast(assetValue.total_size_in_bytes())); - // Update the stores inside a transaction, if anything fails it will automatically rollback to the previous state. - bool status = _assetsStore.impl()->transaction([self, &assetValue, assetSizeInBytes, &ec, error]() { + + + auto transaction = [self, &assetValue, assetSizeInBytes, &ec, error]() { if (!self->_assetsStore.impl()->remove(assetValue.identifier, ec)) { return false; } @@ -548,12 +655,20 @@ - (BOOL)_removeAssetWithIdentifier:(NSString *)identifier NSURL *assetURL = ::get_asset_url(assetValue); if ([self.fileManager fileExistsAtPath:assetURL.path] && - ![self moveURL:assetURL toUniqueURLInDirectory:self.trashDirectoryURL error:error]) { + ![self moveItemAtURLToTrash:assetURL error:error]) { return false; } return true; - }, Database::TransactionBehavior::Immediate, ec); + }; + + // Update the stores inside a transaction, if anything fails it will automatically rollback to the previous state. + bool status = false; + if (alreadyInsideTransaction) { + status = transaction(); + } else { + status = _assetsStore.impl()->transaction(transaction, Database::TransactionBehavior::Immediate, ec); + } // Update the estimated size if the transaction succeeded. _estimatedSizeInBytes -= status ? assetSizeInBytes : 0; @@ -565,7 +680,7 @@ - (BOOL)removeAssetWithIdentifier:(NSString *)identifier error:(NSError * __autoreleasing *)error { __block BOOL result = NO; dispatch_sync(self.syncQueue, ^{ - result = [self _removeAssetWithIdentifier:identifier error:error]; + result = [self _removeAssetWithIdentifier:identifier alreadyInsideTransaction:NO error:error]; }); return result; @@ -643,19 +758,13 @@ - (NSUInteger)_compact:(NSUInteger)sizeInBytes error:(NSError * __autoreleasing for (const auto& asset : assets) { NSError *cleanupError = nil; NSString *identifier = @(asset.identifier.c_str()); - if (![self _removeAssetWithIdentifier:identifier error:&cleanupError] && cleanupError) { + if (![self _removeAssetWithIdentifier:identifier alreadyInsideTransaction:NO error:&cleanupError] && cleanupError) { ETCoreMLLogError(cleanupError, "Failed to remove asset with identifier = %@.", identifier); } } - - // Trigger cleanup. - __weak __typeof(self) weakSelf = self; - dispatch_async(self.trashQueue, ^{ - [weakSelf removeFilesInTrashDirectory]; - }); - + return _estimatedSizeInBytes; } @@ -664,7 +773,10 @@ - (NSUInteger)compact:(NSUInteger)sizeInBytes error:(NSError * __autoreleasing * dispatch_sync(self.syncQueue, ^{ result = [self _compact:sizeInBytes error:error]; }); - + + // Always clean the trash directory to ensure a minimal footprint. + // The `trashQueue` is serialized, so only one cleanup will run at a time. + [self cleanupTrashDirectory]; return result; } @@ -708,14 +820,14 @@ - (BOOL)_purge:(NSError * __autoreleasing *)error { } // Move the the whole assets directory to the temp directory. - if (![self moveURL:self.assetsDirectoryURL toUniqueURLInDirectory:self.trashDirectoryURL error:error]) { + if (![self moveItemAtURLToTrash:self.assetsDirectoryURL error:error]) { return false; } self->_estimatedSizeInBytes = 0; NSError *localError = nil; // Create the assets directory, if we fail here it's okay. - if (![self.fileManager createDirectoryAtURL:self.assetsDirectoryURL withIntermediateDirectories:NO attributes:@{} error:&localError]) { + if (![self.fileManager createDirectoryAtURL:self.assetsDirectoryURL withIntermediateDirectories:YES attributes:@{} error:&localError]) { ETCoreMLLogError(localError, "Failed to create assets directory."); } @@ -724,13 +836,7 @@ - (BOOL)_purge:(NSError * __autoreleasing *)error { ::set_error_from_error_code(ec, error); // Trigger cleanup - if (status) { - __weak __typeof(self) weakSelf = self; - dispatch_async(self.trashQueue, ^{ - [weakSelf removeFilesInTrashDirectory]; - }); - } - + [self cleanupTrashDirectory]; return static_cast(status); } diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.h b/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.h index 05e96ad59f5..1819710cfda 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.h +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.h @@ -9,6 +9,7 @@ @class ETCoreMLModel; @class ETCoreMLAssetManager; +@class ETCoreMLAsset; namespace executorchcoreml { struct ModelMetadata; @@ -23,6 +24,12 @@ __attribute__((objc_subclassing_restricted)) - (instancetype)init NS_UNAVAILABLE; + ++ (nullable ETCoreMLModel*)loadModelWithCompiledAsset:(ETCoreMLAsset*)compiledAsset + configuration:(MLModelConfiguration*)configuration + metadata:(const executorchcoreml::ModelMetadata&)metadata + error:(NSError* __autoreleasing*)error; + /// Synchronously loads a model given the location of its on-disk representation and configuration. /// /// @param compiledModelURL The location of the model's on-disk representation (.mlmodelc directory). diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm index 05aa910d954..731b8506f31 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm @@ -44,6 +44,22 @@ @implementation ETCoreMLModelLoader ++ (nullable ETCoreMLModel *)loadModelWithCompiledAsset:(ETCoreMLAsset *)compiledAsset + configuration:(MLModelConfiguration *)configuration + metadata:(const executorchcoreml::ModelMetadata&)metadata + error:(NSError * __autoreleasing *)error { + NSError *localError = nil; + ETCoreMLModel *model = (compiledAsset != nil) ? get_model_from_asset(compiledAsset, configuration, metadata, &localError) : nil; + if (model) { + return model; + } + if (error) { + *error = localError; + } + return nil; +} + + + (nullable ETCoreMLModel *)loadModelWithContentsOfURL:(NSURL *)compiledModelURL configuration:(MLModelConfiguration *)configuration metadata:(const executorchcoreml::ModelMetadata&)metadata @@ -58,25 +74,22 @@ + (nullable ETCoreMLModel *)loadModelWithContentsOfURL:(NSURL *)compiledModelURL asset = [assetManager storeAssetAtURL:compiledModelURL withIdentifier:identifier error:&localError]; } - ETCoreMLModel *model = (asset != nil) ? get_model_from_asset(asset, configuration, metadata, &localError) : nil; + ETCoreMLModel *model; + if (asset != nil) { + model = [self loadModelWithCompiledAsset:asset configuration:configuration metadata:metadata error:&localError]; + } else { + model = nil; + } + if (model) { return model; } - - if (localError) { - ETCoreMLLogError(localError, - "Failed to load model from compiled asset with identifier = %@", - identifier); - } - - // If store failed then we will load the model from compiledURL. - auto backingAsset = Asset::make(compiledModelURL, identifier, assetManager.fileManager, error); - if (!backingAsset) { - return nil; + + if (error) { + *error = localError; } - - asset = [[ETCoreMLAsset alloc] initWithBackingAsset:backingAsset.value()]; - return ::get_model_from_asset(asset, configuration, metadata, error); + + return nil; } @end diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm index 2347936fd34..d59890ee00f 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm @@ -345,6 +345,10 @@ void add_compute_unit(std::string& identifier, MLComputeUnits compute_units) { return [ETCoreMLModelDebugInfo modelDebugInfoFromData:file_data error:error]; } +NSString *raw_model_identifier(NSString *identifier) { + return [NSString stringWithFormat:@"raw_%@", identifier]; +} + #endif } //namespace @@ -408,7 +412,7 @@ - (nullable ETCoreMLAsset *)assetWithIdentifier:(NSString *)identifier { return modelAsset; } - NSError *localError = nil; + __block NSError *localError = nil; modelAsset = [self.assetManager assetWithIdentifier:identifier error:&localError]; if (localError) { ETCoreMLLogError(localError, @@ -420,8 +424,9 @@ - (nullable ETCoreMLAsset *)assetWithIdentifier:(NSString *)identifier { } - (nullable NSURL *)compiledModelURLWithIdentifier:(NSString *)identifier + modelURL:(nullable NSURL *)modelURL inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS - assetManager:(ETCoreMLAssetManager *)assetManager + dstURL:(NSURL *)dstURL error:(NSError * __autoreleasing *)error { auto modelAssetType = get_model_asset_type(inMemoryFS); if (!modelAssetType) { @@ -430,80 +435,148 @@ - (nullable NSURL *)compiledModelURLWithIdentifier:(NSString *)identifier "AOT blob is missing model file."); return nil; } - - NSURL *dstURL = [self.assetManager.trashDirectoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; - NSURL *modelURL = ::write_model_files(dstURL, self.fileManager, identifier, modelAssetType.value(), inMemoryFS, error); + + // If modelURL is not provided, write model files to the destination directory (dstURL) + // and obtain a URL pointing to them. Otherwise, use the provided modelURL. + modelURL = (modelURL == nil) ? ::write_model_files(dstURL, self.fileManager, identifier, modelAssetType.value(), inMemoryFS, error) : modelURL; + if (!modelURL) { + // Failed to generate or locate model files, return nil. + return nil; + } + + // Handle based on the type of the model asset. switch (modelAssetType.value()) { case ModelAssetType::CompiledModel: { // Model is already compiled. ETCoreMLLogInfo("The model in the pte file is pre-compiled. Skipping compilation."); return modelURL; } - + case ModelAssetType::Model: { // Compile the model. ETCoreMLLogInfo("The model in the pte file is not pre-compiled. Compiling with a 5 min timeout."); NSURL *compiledModelURL = [ETCoreMLModelCompiler compileModelAtURL:modelURL maxWaitTimeInSeconds:(5 * 60) error:error]; - + // Return the URL of the compiled model or nil if compilation fails. return compiledModelURL; } } } -#if ET_EVENT_TRACER_ENABLED -- (nullable id)modelExecutorWithMetadata:(const ModelMetadata&)metadata - inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS - configuration:(MLModelConfiguration *)configuration - error:(NSError * __autoreleasing *)error { +- (nullable ETCoreMLAsset *)compiledModelAssetWithMetadata:(const ModelMetadata&)metadata + modelURL:(nullable NSURL *)modelURL + inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS + error:(NSError * __autoreleasing *)error { NSString *identifier = @(metadata.identifier.c_str()); - // Otherwise try to retrieve the compiled asset. - ETCoreMLAsset *compiledModelAsset = [self assetWithIdentifier:identifier]; + __block ETCoreMLAsset *compiledModelAsset = [self assetWithIdentifier:identifier]; if (compiledModelAsset) { + ETCoreMLLogInfo("Cache Hit: Successfully retrieved compiled model with identifier=%@ from the models cache.", identifier); + return compiledModelAsset; + } + + ETCoreMLLogInfo("Cache Miss: Compiled Model with identifier=%@ was not found in the models cache.", identifier); + __block NSURL *compiledModelURL; + [self.assetManager withTemporaryDirectory:^(NSURL * _Nonnull directoryURL) { + // The directory specified by `directoryURL` is unique and will be automatically cleaned up + // once the enclosing block completes. + compiledModelURL = [self compiledModelURLWithIdentifier:identifier + modelURL:modelURL + inMemoryFS:inMemoryFS + dstURL:directoryURL + error:error]; + if (compiledModelURL) { + // Move the compiled model to the asset manager to transfer ownership. + ETCoreMLLogInfo("Successfully got compiled model with identifier=%@. Transferring ownership to assetManager.", identifier); + compiledModelAsset = [self.assetManager storeAssetAtURL:compiledModelURL withIdentifier:identifier error:error]; + } + }]; + + if (!compiledModelAsset) { + ETCoreMLLogInfo("Failed to transfer ownership of asset with identifier=%@ to assetManager", identifier); + if (compiledModelURL && [self.fileManager fileExistsAtPath:compiledModelURL.path]) { + // Log what error was since we now attempt backup path, and previous error is overwritten + if (error && *error) { + ETCoreMLLogInfo("error=%@", (*error).localizedDescription); + *error = nil; + } + ETCoreMLLogInfo("Attempting to fall back by loading model without transferring ownership"); + auto backingAsset = Asset::make(compiledModelURL, identifier, self.assetManager.fileManager, error); + if (backingAsset) { + compiledModelAsset = [[ETCoreMLAsset alloc] initWithBackingAsset:backingAsset.value()]; + } + } + } + + // compiledModelAsset can still be nil if our backup path failed + + return compiledModelAsset; +} + +#if ET_EVENT_TRACER_ENABLED +- (nullable ETCoreMLAsset *)modelAssetWithMetadata:(const ModelMetadata&)metadata + inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS + error:(NSError * __autoreleasing *)error { + NSString *identifier = @(metadata.identifier.c_str()); + NSString *rawIdentifier = raw_model_identifier(identifier); + __block ETCoreMLAsset *modelAsset = [self assetWithIdentifier:rawIdentifier]; + if (modelAsset) { ETCoreMLLogInfo("Cache Hit: Successfully retrieved model with identifier=%@ from the models cache.", identifier); } else { ETCoreMLLogInfo("Cache Miss: Model with identifier=%@ was not found in the models cache.", identifier); } - - // Create a unique directory for writing model files. - NSURL *dstURL = [self.assetManager.trashDirectoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; - auto modelAssetType = get_model_asset_type(inMemoryFS); - ETCoreMLAsset *modelAsset = nil; - // Write the model files. - if (modelAssetType == ModelAssetType::Model) { - NSURL *modelURL = ::write_model_files(dstURL, self.fileManager, identifier, modelAssetType.value(), inMemoryFS, error); + + [self.assetManager withTemporaryDirectory:^(NSURL * _Nonnull directoryURL) { + if (modelAsset) { + return; + } + + auto modelAssetType = get_model_asset_type(inMemoryFS); + if (modelAssetType != ModelAssetType::Model) { + return; + } + + // The directory specified by `directoryURL` is unique and will be automatically cleaned up + // once the enclosing block completes. + NSURL *modelURL = ::write_model_files(directoryURL, + self.fileManager, + identifier, + modelAssetType.value(), + inMemoryFS, + error); if (modelURL) { - modelAsset = make_asset(modelURL, - identifier, - self.fileManager, - error); + // Move the model to the asset manager to transfer ownership. + modelAsset = [self.assetManager storeAssetAtURL:modelURL withIdentifier:rawIdentifier error:error]; } + }]; + + return modelAsset; +} + +- (nullable id)modelExecutorWithMetadata:(const ModelMetadata&)metadata + inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS + configuration:(MLModelConfiguration *)configuration + error:(NSError * __autoreleasing *)error { + NSError *localError = nil; + ETCoreMLAsset *modelAsset = [self modelAssetWithMetadata:metadata inMemoryFS:inMemoryFS error:&localError]; + if (localError) { + if (error) { + *error = localError; + } + + return nil; } - - if (!compiledModelAsset) { - // Compile the model. - NSURL *compiledModelURL = [self compiledModelURLWithIdentifier:identifier - inMemoryFS:inMemoryFS - assetManager:self.assetManager - error:error]; - compiledModelAsset = make_asset(compiledModelURL, - identifier, - self.fileManager, - error); - } - + + ETCoreMLAsset *compiledModelAsset = [self compiledModelAssetWithMetadata:metadata + modelURL:modelAsset.contentURL + inMemoryFS:inMemoryFS + error:error]; if (!compiledModelAsset) { return nil; } - - NSError *localError = nil; - ETCoreMLModelDebugInfo *debug_info = get_model_debug_info(inMemoryFS, &localError); - if (localError) { - ETCoreMLLogError(localError, "Failed to parse debug info file"); - } - + ETCoreMLModelDebugInfo *debug_info = get_model_debug_info(inMemoryFS, error); + // The analyzer requires both the raw (uncompiled) asset and the compiled model asset to perform analysis. return [[ETCoreMLModelAnalyzer alloc] initWithCompiledModelAsset:compiledModelAsset modelAsset:modelAsset modelDebugInfo:debug_info @@ -512,41 +585,32 @@ - (nullable NSURL *)compiledModelURLWithIdentifier:(NSString *)identifier assetManager:self.assetManager error:error]; } - #else - (nullable id)modelExecutorWithMetadata:(const ModelMetadata&)metadata inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS configuration:(MLModelConfiguration *)configuration error:(NSError * __autoreleasing *)error { - NSString *identifier = @(metadata.identifier.c_str()); - // Otherwise try to retrieve the compiled asset. - ETCoreMLAsset *asset = [self assetWithIdentifier:identifier]; - ETCoreMLModel *model = asset ? get_model_from_asset(asset, configuration, metadata, error) : nil; - if (model) { - ETCoreMLLogInfo("Cache Hit: Successfully retrieved model with identifier=%@ from the models cache.", identifier); - return [[ETCoreMLDefaultModelExecutor alloc] initWithModel:model]; + ETCoreMLAsset *compiledModelAsset = [self compiledModelAssetWithMetadata:metadata + modelURL:nil + inMemoryFS:inMemoryFS + error:error]; + if (!compiledModelAsset) { + return nil; } - - ETCoreMLLogInfo("Cache Miss: Model with identifier=%@ was not found in the models cache.", identifier); - // Compile the model. - NSURL *compiledModelURL = [self compiledModelURLWithIdentifier:identifier - inMemoryFS:inMemoryFS - assetManager:self.assetManager - error:error]; - if (!compiledModelURL) { + + ETCoreMLModel *model = [ETCoreMLModelLoader loadModelWithCompiledAsset:compiledModelAsset + configuration:configuration + metadata:metadata + error:error]; + if (!model) { return nil; } - - model = [ETCoreMLModelLoader loadModelWithContentsOfURL:compiledModelURL - configuration:configuration - metadata:metadata - assetManager:self.assetManager - error:error]; - + return [[ETCoreMLDefaultModelExecutor alloc] initWithModel:model]; } #endif + - (nullable id)_modelExecutorWithAOTData:(NSData *)data configuration:(MLModelConfiguration *)configuration error:(NSError * __autoreleasing *)error { @@ -731,6 +795,7 @@ - (BOOL)executeModelWithHandle:(ModelHandle *)handle args.count); return result; } + NSError *localError = nil; @autoreleasepool { NSArray *inputs = [args subarrayWithRange:NSMakeRange(0, model.orderedInputNames.count)]; @@ -750,11 +815,11 @@ - (BOOL)executeModelWithHandle:(ModelHandle *)handle result = YES; } } - if (!result) { - if (error) { - *error = localError; - } + + if (localError && error) { + *error = localError; } + return result; } diff --git a/backends/apple/coreml/runtime/delegate/backend_delegate.mm b/backends/apple/coreml/runtime/delegate/backend_delegate.mm index 2cb274f0a89..680c5c63143 100644 --- a/backends/apple/coreml/runtime/delegate/backend_delegate.mm +++ b/backends/apple/coreml/runtime/delegate/backend_delegate.mm @@ -45,40 +45,15 @@ MLComputeUnits get_compute_units(const Buffer& buffer) { return configuration; } -NSURL * _Nullable create_directory_if_needed(NSURL *url, - NSFileManager *fileManager, - NSError * __autoreleasing *error) { - if (![fileManager fileExistsAtPath:url.path] && - ![fileManager createDirectoryAtURL:url withIntermediateDirectories:YES attributes:@{} error:error]) { - return nil; - } - - return url; -} - ETCoreMLAssetManager * _Nullable create_asset_manager(NSString *assets_directory_path, NSString *trash_directory_path, NSString *database_directory_path, NSString *database_name, NSInteger max_assets_size_in_bytes, NSError * __autoreleasing *error) { - NSFileManager *fm = [[NSFileManager alloc] init]; - NSURL *assets_directory_url = [NSURL fileURLWithPath:assets_directory_path]; - if (!create_directory_if_needed(assets_directory_url, fm, error)) { - return nil; - } - NSURL *trash_directory_url = [NSURL fileURLWithPath:trash_directory_path]; - if (!create_directory_if_needed(trash_directory_url, fm, error)) { - return nil; - } - NSURL *database_directory_url = [NSURL fileURLWithPath:database_directory_path]; - if (!create_directory_if_needed(database_directory_url, fm, error)) { - return nil; - } - NSURL *database_url = [database_directory_url URLByAppendingPathComponent:database_name]; ETCoreMLAssetManager *manager = [[ETCoreMLAssetManager alloc] initWithDatabaseURL:database_url assetsDirectoryURL:assets_directory_url