From 23832018369d2f5cb09892b9c5f8a082a24369a1 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Wed, 17 Sep 2025 17:13:04 -0700 Subject: [PATCH] Back out "Improve asset management" Differential Revision: D82581443 Pull Request resolved: https://github.com/pytorch/executorch/pull/14383 (cherry picked from commit 82c1d772f74beca46dd43380126dcb34500902ff) --- .../runtime/delegate/ETCoreMLAssetManager.h | 17 -- .../runtime/delegate/ETCoreMLAssetManager.mm | 104 +++---- .../runtime/delegate/ETCoreMLModelLoader.mm | 19 +- .../runtime/delegate/ETCoreMLModelManager.mm | 264 +++++++----------- 4 files changed, 158 insertions(+), 246 deletions(-) diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h b/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h index a9e06efa90d..11d957044e9 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h @@ -99,17 +99,6 @@ 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. /// @@ -128,12 +117,6 @@ 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 53c3d1cdc69..256026e1f09 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm @@ -254,29 +254,6 @@ BOOL is_asset_alive(NSMapTable *assets_in_use_map, return assets; } - -NSURL * _Nullable move_to_directory(NSURL *url, - NSURL *directoryURL, - NSFileManager *fileManager, - NSError * __autoreleasing *error) { - if (!url) { - ETCoreMLLogErrorAndSetNSError(error, ETCoreMLErrorInternalError, "Move operation failed: source URL is nil."); - return nil; - } - - if (!directoryURL) { - ETCoreMLLogErrorAndSetNSError(error, ETCoreMLErrorInternalError, "Move operation failed: destination URL is nil."); - return nil; - } - - NSURL *dstURL = [directoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; - if (![fileManager moveItemAtURL:url toURL:dstURL error:error]) { - return nil; - } - - return dstURL; -} - } //namespace @interface ETCoreMLAssetManager () { @@ -322,17 +299,12 @@ - (nullable instancetype)initWithDatabase:(const std::shared_ptr&)data if (!managedAssetsDirectoryURL) { return nil; } - + NSURL *managedTrashDirectoryURL = ::create_directory_if_needed(trashDirectoryURL, @"models", fileManager, error); if (!managedTrashDirectoryURL) { return nil; } - - NSURL *managedStagingDirectoryURL = ::create_directory_if_needed(assetsDirectoryURL, @"staging", fileManager, error); - if (!managedStagingDirectoryURL) { - return nil; - } - + // If directory is empty then purge the stores if (::is_directory_empty(managedAssetsDirectoryURL, fileManager, nil)) { assetsMetaStore.impl()->purge(ec); @@ -343,7 +315,6 @@ - (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; @@ -375,15 +346,15 @@ - (nullable instancetype)initWithDatabaseURL:(NSURL *)databaseURL error:error]; } -- (void)withTemporaryDirectory:(void (^)(NSURL *directoryURL))block { - NSURL *dstURL = [self.stagingDirectoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; - block(dstURL); - if (![self.fileManager fileExistsAtPath:dstURL.path]) { - return; +- (nullable NSURL *)moveURL:(NSURL *)url + toUniqueURLInDirectory:(NSURL *)directoryURL + error:(NSError * __autoreleasing *)error { + NSURL *dstURL = [directoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; + if (![self.fileManager moveItemAtURL:url toURL:dstURL error:error]) { + return nil; } - - move_to_directory(dstURL, self.trashDirectoryURL, self.fileManager, nil); - [self cleanupTrashDirectory]; + + return dstURL; } - (void)cleanupAssetIfNeeded:(ETCoreMLAsset *)asset { @@ -436,8 +407,9 @@ - (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL return false; } - // If a file already exists at `dstURL`, move it to the trash for removal. - move_to_directory(dstURL, self.trashDirectoryURL, self.fileManager, nil); + // If an asset exists move it + [self moveURL:dstURL toUniqueURLInDirectory:self.trashDirectoryURL error:nil]; + // Move the asset to assets directory. if (![self.fileManager moveItemAtURL:srcURL toURL:dstURL error:error]) { return false; @@ -461,25 +433,16 @@ - (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL } - (void)triggerCompaction { - 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."); - } - }); + if (self.estimatedSizeInBytes < self.maxAssetsSizeInBytes) { + return; } - - // 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.trashQueue, ^{ - [weakSelf removeFilesInTrashDirectory]; + dispatch_async(self.syncQueue, ^{ + NSError *localError = nil; + if (![weakSelf _compact:self.maxAssetsSizeInBytes error:&localError]) { + ETCoreMLLogError(localError, "Failed to compact asset store."); + } }); } @@ -585,7 +548,7 @@ - (BOOL)_removeAssetWithIdentifier:(NSString *)identifier NSURL *assetURL = ::get_asset_url(assetValue); if ([self.fileManager fileExistsAtPath:assetURL.path] && - !move_to_directory(assetURL, self.trashDirectoryURL, self.fileManager, error)) { + ![self moveURL:assetURL toUniqueURLInDirectory:self.trashDirectoryURL error:error]) { return false; } @@ -686,7 +649,13 @@ - (NSUInteger)_compact:(NSUInteger)sizeInBytes error:(NSError * __autoreleasing identifier); } } - + + // Trigger cleanup. + __weak __typeof(self) weakSelf = self; + dispatch_async(self.trashQueue, ^{ + [weakSelf removeFilesInTrashDirectory]; + }); + return _estimatedSizeInBytes; } @@ -695,10 +664,7 @@ - (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; } @@ -742,7 +708,7 @@ - (BOOL)_purge:(NSError * __autoreleasing *)error { } // Move the the whole assets directory to the temp directory. - if (!move_to_directory(self.assetsDirectoryURL, self.trashDirectoryURL, self.fileManager, error)) { + if (![self moveURL:self.assetsDirectoryURL toUniqueURLInDirectory:self.trashDirectoryURL error:error]) { return false; } @@ -758,7 +724,13 @@ - (BOOL)_purge:(NSError * __autoreleasing *)error { ::set_error_from_error_code(ec, error); // Trigger cleanup - [self cleanupTrashDirectory]; + if (status) { + __weak __typeof(self) weakSelf = self; + dispatch_async(self.trashQueue, ^{ + [weakSelf removeFilesInTrashDirectory]; + }); + } + return static_cast(status); } diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm index 9e8ae04842e..05aa910d954 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm @@ -62,12 +62,21 @@ + (nullable ETCoreMLModel *)loadModelWithContentsOfURL:(NSURL *)compiledModelURL if (model) { return model; } - - if (error) { - *error = localError; + + if (localError) { + ETCoreMLLogError(localError, + "Failed to load model from compiled asset with identifier = %@", + identifier); } - - return nil; + + // If store failed then we will load the model from compiledURL. + auto backingAsset = Asset::make(compiledModelURL, identifier, assetManager.fileManager, error); + if (!backingAsset) { + return nil; + } + + asset = [[ETCoreMLAsset alloc] initWithBackingAsset:backingAsset.value()]; + return ::get_model_from_asset(asset, configuration, metadata, error); } @end diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm index c27b42566dc..f697d8c63ab 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm @@ -56,7 +56,7 @@ for (NSString *component in components) { result.emplace_back(component.UTF8String); } - + return result; } @@ -69,7 +69,7 @@ NSString *input_name = [enumerator nextObject]; features[input_name] = [MLFeatureValue featureValueWithMultiArray:input]; } - + return [[MLDictionaryFeatureProvider alloc] initWithDictionary:features error:error]; } @@ -84,7 +84,7 @@ BOOL is_backed_by_same_buffer(MLMultiArray *array1, MLMultiArray *array2) { } else { result = (array1.dataPointer == array2.dataPointer); } - + return result; } @@ -105,7 +105,7 @@ BOOL is_backed_by_same_buffer(MLMultiArray *array1, MLMultiArray *array2) { } options.outputBackings = output_backings; } - + return options; } @@ -113,7 +113,7 @@ void copy(MLMultiArray *src, MLMultiArray *dst) { if (::is_backed_by_same_buffer(src, dst)) { return; } - + [src copyInto:dst]; } @@ -150,10 +150,10 @@ void copy(MLMultiArray *src, executorchcoreml::MultiArray& dst) { if (bytes == dst.data()) { return; } - + MultiArray::MemoryLayout src_layout( - get_data_type(src.dataType).value(), - to_vector(src.shape), + get_data_type(src.dataType).value(), + to_vector(src.shape), to_vector(src.strides) ); MultiArray(const_cast(bytes), std::move(src_layout)).copy(dst); @@ -184,13 +184,13 @@ void set_outputs(std::vector& outputs, if (!buffer || buffer->size() == 0) { return nil; } - + NSData *file_data = [[NSData alloc] initWithBytesNoCopy:buffer->data() length:buffer->size() deallocator:^(void * _Nonnull __unused bytes, NSUInteger __unused length) { buffer.reset(); }]; - + return file_data; } @@ -199,7 +199,7 @@ void set_outputs(std::vector& outputs, if (!file_data) { return std::nullopt; } - + std::string contents; contents.assign(static_cast(file_data.bytes), file_data.length); ModelMetadata metadata; @@ -207,7 +207,7 @@ void set_outputs(std::vector& outputs, if (metadata.is_valid()) { return metadata; } - + return std::nullopt; } @@ -216,7 +216,7 @@ void set_outputs(std::vector& outputs, for (const auto& value : values) { [result addObject:@(value.c_str())]; } - + return result; } @@ -235,7 +235,7 @@ void set_outputs(std::vector& outputs, identifier); return nil; } - + std::filesystem::path model_path(dst_url.fileSystemRepresentation); std::error_code ec; std::vector file_path; @@ -244,13 +244,13 @@ void set_outputs(std::vector& outputs, file_path = canonical_path(ETCoreMLStrings.modelFileRelativePath); break; } - + case ModelAssetType::CompiledModel: { file_path = canonical_path(ETCoreMLStrings.compiledModelFileRelativePath); break; } } - + if (!inmemory_fs->write_item_to_disk(file_path, model_path, true, ec)) { ETCoreMLLogErrorAndSetNSError(error, ETCoreMLErrorModelSaveFailed, @@ -258,7 +258,7 @@ void set_outputs(std::vector& outputs, identifier); return nil; } - + switch (model_asset_type) { case ModelAssetType::Model: { return [dst_url URLByAppendingPathComponent:[NSString stringWithFormat:@"model.%@", ETCoreMLStrings.modelExtensionName]]; @@ -274,11 +274,11 @@ void set_outputs(std::vector& outputs, if (inmemory_fs->exists(canonical_path(ETCoreMLStrings.compiledModelFileRelativePath))) { return ModelAssetType::CompiledModel; } - + if (inmemory_fs->exists(canonical_path(ETCoreMLStrings.modelFileRelativePath))) { return ModelAssetType::Model; } - + return std::nullopt; } @@ -331,7 +331,7 @@ void add_compute_unit(std::string& identifier, MLComputeUnits compute_units) { if (!backingAsset) { return nil; } - + return [[ETCoreMLAsset alloc] initWithBackingAsset:std::move(backingAsset.value())]; } @@ -345,10 +345,6 @@ 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 @@ -379,7 +375,7 @@ - (instancetype)initWithAssetManager:(ETCoreMLAssetManager *)assetManager { dispatch_queue_attr_t attr = dispatch_queue_attr_make_with_qos_class(DISPATCH_QUEUE_SERIAL, QOS_CLASS_DEFAULT, -1); _prewarmQueue = dispatch_queue_create("com.executorchcoreml.modelmanager.prewarm", attr); } - + return self; } @@ -391,7 +387,7 @@ - (instancetype)initWithAssetManager:(ETCoreMLAssetManager *)assetManager { executor = self.handleToExecutorMap[key]; os_unfair_lock_unlock(&_lock); } - + return executor; } @@ -407,26 +403,25 @@ - (nullable ETCoreMLAsset *)assetWithIdentifier:(NSString *)identifier { modelAsset = self.modelIdentifierToPrewarmedAssetMap[identifier]; os_unfair_lock_unlock(&_lock); } - + if (modelAsset) { return modelAsset; } - - __block NSError *localError = nil; + + NSError *localError = nil; modelAsset = [self.assetManager assetWithIdentifier:identifier error:&localError]; if (localError) { ETCoreMLLogError(localError, "Failed to retrieve asset with identifier = %@.", identifier); } - + return modelAsset; } - (nullable NSURL *)compiledModelURLWithIdentifier:(NSString *)identifier - modelURL:(nullable NSURL *)modelURL inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS - dstURL:(NSURL *)dstURL + assetManager:(ETCoreMLAssetManager *)assetManager error:(NSError * __autoreleasing *)error { auto modelAssetType = get_model_asset_type(inMemoryFS); if (!modelAssetType) { @@ -436,131 +431,77 @@ - (nullable NSURL *)compiledModelURLWithIdentifier:(NSString *)identifier return nil; } - // 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. + NSURL *dstURL = [self.assetManager.trashDirectoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; + NSURL *modelURL = ::write_model_files(dstURL, self.fileManager, identifier, modelAssetType.value(), inMemoryFS, error); switch (modelAssetType.value()) { case ModelAssetType::CompiledModel: { - // The model is already compiled; no further action needed. - // Return the existing model URL. + // Model is already compiled. return modelURL; } case ModelAssetType::Model: { - // The model is not compiled yet. - // Compile the model at the specified URL with a maximum wait time of 5 minutes. + // Compile the model. NSURL *compiledModelURL = [ETCoreMLModelCompiler compileModelAtURL:modelURL maxWaitTimeInSeconds:(5 * 60) error:error]; - // Return the URL of the compiled model or nil if compilation fails. + return compiledModelURL; } } } -- (nullable ETCoreMLAsset *)compiledModelAssetWithMetadata:(const ModelMetadata&)metadata - modelURL:(nullable NSURL *)modelURL - inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS - error:(NSError * __autoreleasing *)error { +#if ET_EVENT_TRACER_ENABLED +- (nullable id)modelExecutorWithMetadata:(const ModelMetadata&)metadata + inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS + configuration:(MLModelConfiguration *)configuration + error:(NSError * __autoreleasing *)error { NSString *identifier = @(metadata.identifier.c_str()); - __block ETCoreMLAsset *compiledModelAsset = [self assetWithIdentifier:identifier]; + // Otherwise try to retrieve the compiled asset. + ETCoreMLAsset *compiledModelAsset = [self assetWithIdentifier:identifier]; if (compiledModelAsset) { - ETCoreMLLogInfo("Cache Hit: Successfully retrieved compiled model with identifier=%@ from the models cache.", identifier); + ETCoreMLLogInfo("Cache Hit: Successfully retrieved model with identifier=%@ from the models cache.", identifier); } else { - ETCoreMLLogInfo("Cache Miss: Compiled Model with identifier=%@ was not found in the models cache.", identifier); + ETCoreMLLogInfo("Cache Miss: Model with identifier=%@ was not found in the models cache.", identifier); } - [self.assetManager withTemporaryDirectory:^(NSURL * _Nonnull directoryURL) { - if (compiledModelAsset) { - return; + // 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); + if (modelURL) { + modelAsset = make_asset(modelURL, + identifier, + self.fileManager, + error); } + } - // The directory specified by `directoryURL` is unique and will be automatically cleaned up - // once the enclosing block completes. + if (!compiledModelAsset) { + // Compile the model. NSURL *compiledModelURL = [self compiledModelURLWithIdentifier:identifier - modelURL:modelURL inMemoryFS:inMemoryFS - dstURL:directoryURL + assetManager:self.assetManager error:error]; - if (compiledModelURL) { - // Move the compiled model to the asset manager to transfer ownership. - compiledModelAsset = [self.assetManager storeAssetAtURL:compiledModelURL withIdentifier:identifier error:error]; - } - }]; - - 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); + compiledModelAsset = make_asset(compiledModelURL, + identifier, + self.fileManager, + 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) { - // Move the model to the asset manager to transfer ownership. - modelAsset = [self.assetManager storeAssetAtURL:modelURL withIdentifier:rawIdentifier error:error]; - } - }]; - - return modelAsset; -} + if (!compiledModelAsset) { + return nil; + } -- (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]; + ETCoreMLModelDebugInfo *debug_info = get_model_debug_info(inMemoryFS, &localError); if (localError) { - if (error) { - *error = localError; - } - - return nil; + ETCoreMLLogError(localError, "Failed to parse debug info file"); } - ETCoreMLAsset *compiledModelAsset = [self compiledModelAssetWithMetadata:metadata - modelURL:modelAsset.contentURL - inMemoryFS:inMemoryFS - error:error]; - if (!compiledModelAsset) { - return nil; - } - 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 @@ -569,38 +510,46 @@ - (nullable ETCoreMLAsset *)modelAssetWithMetadata:(const ModelMetadata&)metadat assetManager:self.assetManager error:error]; } + #else - (nullable id)modelExecutorWithMetadata:(const ModelMetadata&)metadata inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS configuration:(MLModelConfiguration *)configuration error:(NSError * __autoreleasing *)error { - ETCoreMLAsset *compiledModelAsset = [self compiledModelAssetWithMetadata:metadata - modelURL:nil - inMemoryFS:inMemoryFS - error:error]; - if (!compiledModelAsset) { - return nil; + 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]; } - ETCoreMLModel *model = [ETCoreMLModelLoader loadModelWithContentsOfURL:compiledModelAsset.contentURL - configuration:configuration - metadata:metadata - assetManager:self.assetManager - error:error]; - if (!model) { + 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) { 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 { using namespace inmemoryfs; - + auto buffer = MemoryBuffer::make_unowned(const_cast(data.bytes), data.length); std::unique_ptr inMemoryFS = inmemoryfs::make_from_buffer(std::move(buffer)); if (!inMemoryFS) { @@ -609,7 +558,7 @@ - (nullable ETCoreMLAsset *)modelAssetWithMetadata:(const ModelMetadata&)metadat "Model data is corrupted."); return nil; } - + std::optional metadata = ::get_model_metadata(inMemoryFS.get()); if (!metadata) { ETCoreMLLogErrorAndSetNSError(error, @@ -617,7 +566,7 @@ - (nullable ETCoreMLAsset *)modelAssetWithMetadata:(const ModelMetadata&)metadat "Metadata is invalid or missing."); return nil; } - + auto metadataValue = metadata.value(); add_compute_unit(metadataValue.identifier, configuration.computeUnits); NSString *identifier = @(metadataValue.identifier.c_str()); @@ -631,7 +580,7 @@ - (nullable ETCoreMLAsset *)modelAssetWithMetadata:(const ModelMetadata&)metadat configuration:configuration error:error]; }); - + return executor; } @@ -643,7 +592,7 @@ - (dispatch_queue_t)queueForLoadingModelWithIdentifier:(NSString *)identifier { [self.modelIdentifierToLoadingQueueMap setObject:queue forKey:identifier]; } os_unfair_lock_unlock(&_lock); - + return queue; } @@ -661,7 +610,7 @@ - (ModelHandle *)loadModelFromAOTData:(NSData*)data } os_unfair_lock_unlock(&_lock); } - + return (__bridge ModelHandle *)executor.model; } @@ -678,15 +627,15 @@ - (BOOL)prewarmModelWithHandle:(ModelHandle *)handle - (void)prewarmRecentlyUsedAssetsWithMaxCount:(NSUInteger)maxCount { NSError *localError = nil; NSArray *assets = [self.assetManager mostRecentlyUsedAssetsWithMaxCount:maxCount error:&localError]; - + if (localError) { ETCoreMLLogError(localError, "Failed to retrieve recently used assets."); } - + if (assets.count == 0) { return; } - + for (ETCoreMLAsset *asset in assets) { __weak __typeof(self) weakSelf = self; dispatch_async(self.prewarmQueue, ^{ @@ -694,7 +643,7 @@ - (void)prewarmRecentlyUsedAssetsWithMaxCount:(NSUInteger)maxCount { if (!strongSelf) { return; } - + NSError *prewarmError = nil; if (![asset prewarmAndReturnError:&prewarmError]) { ETCoreMLLogError(prewarmError, @@ -702,7 +651,7 @@ - (void)prewarmRecentlyUsedAssetsWithMaxCount:(NSUInteger)maxCount { asset.identifier); return; } - + [strongSelf addPrewarmedAsset:asset]; }); } @@ -726,12 +675,12 @@ - (void)addPrewarmedAsset:(ETCoreMLAsset *)asset { if (!predictionOptions) { return nil; } - + id inputFeatures = ::get_feature_provider(inputs, model.orderedInputNames, error); if (!inputFeatures) { return nil; } - + NSArray *modelOutputs = [executor executeModelWithInputs:inputFeatures predictionOptions:predictionOptions loggingOptions:loggingOptions @@ -753,7 +702,7 @@ - (void)addPrewarmedAsset:(ETCoreMLAsset *)asset { if (error) { *error = localError; } - + return modelOutputs; } @@ -780,7 +729,6 @@ - (BOOL)executeModelWithHandle:(ModelHandle *)handle args.count); return result; } - NSError *localError = nil; @autoreleasepool { NSArray *inputs = [args subarrayWithRange:NSMakeRange(0, model.orderedInputNames.count)]; @@ -800,11 +748,11 @@ - (BOOL)executeModelWithHandle:(ModelHandle *)handle result = YES; } } - - if (localError && error) { - *error = localError; + if (!result) { + if (error) { + *error = localError; + } } - return result; } @@ -877,7 +825,7 @@ - (BOOL)unloadModelWithHandle:(ModelHandle *)handle { [self.handleToExecutorMap removeObjectForKey:key]; os_unfair_lock_unlock(&_lock); } - + return result; }