Re-land asseet management changes

Author: metascroy

backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm

@@ -437,7 +437,13 @@ - (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL
         }
 
         // If a file already exists at `dstURL`, move it to the trash for removal.
-        move_to_directory(dstURL, self.trashDirectoryURL, self.fileManager, nil);
+        if (dstURL && [self.fileManager fileExistsAtPath:dstURL.path]) {
+            if (!move_to_directory(dstURL, self.trashDirectoryURL, self.fileManager, error)) {
+                // If we failed to move the file to the trash, we can't proceed
+                return false;
+            }
+        }
+        
         // Move the asset to assets directory.
         if (![self.fileManager moveItemAtURL:srcURL toURL:dstURL error:error]) {
             return false;

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).

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
@@ -55,10 +71,28 @@ + (nullable ETCoreMLModel *)loadModelWithContentsOfURL:(NSURL *)compiledModelURL
     if ([assetManager hasAssetWithIdentifier:identifier error:&localError]) {
         asset = [assetManager assetWithIdentifier:identifier error:&localError];
     } else {
+        ETCoreMLLogInfo("Storing asset with identifier=%@ in assetManager", identifier);
         asset = [assetManager storeAssetAtURL:compiledModelURL withIdentifier:identifier error:&localError];
     }
-    
-    ETCoreMLModel *model = (asset != nil) ? get_model_from_asset(asset, configuration, metadata, &localError) : nil;
+
+    if (asset == nil) {
+        ETCoreMLLogInfo("Failed to retrieve or store asset with identifier=%@ in assetManager", identifier);
+        ETCoreMLLogInfo("compiledModelURL=%@", compiledModelURL);
+        ETCoreMLLogInfo("error=%@", localError);
+        ETCoreMLLogInfo("Attempting to fall back by loading model from compiledModelURL without transferring to assetManager");
+        auto backingAsset = Asset::make(compiledModelURL, identifier, assetManager.fileManager, &localError);
+        if (!backingAsset) {
+            ETCoreMLLogInfo("Failed to create a backing asset with error=%@", localError);
+            return nil;
+        }
+        asset = [[ETCoreMLAsset alloc] initWithBackingAsset:backingAsset.value()];
+    }
+
+    ETCoreMLModel *model;
+    if (asset != nil) {
+        model = [self loadModelWithCompiledAsset:asset configuration:configuration metadata:metadata error:&localError];
+    }
+
     if (model) {
         return model;
     }

backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm

@@ -472,28 +472,44 @@ - (nullable ETCoreMLAsset *)compiledModelAssetWithMetadata:(const ModelMetadata&
     __block ETCoreMLAsset *compiledModelAsset = [self assetWithIdentifier:identifier];
     if (compiledModelAsset) {
         ETCoreMLLogInfo("Cache Hit: Successfully retrieved compiled model with identifier=%@ from the models cache.", identifier);
-    } else {
-        ETCoreMLLogInfo("Cache Miss: Compiled Model with identifier=%@ was not found in 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) {
-        if (compiledModelAsset) {
-            return;
-        }
-
         // The directory specified by `directoryURL` is unique and will be automatically cleaned up
         // once the enclosing block completes.
-        NSURL *compiledModelURL = [self compiledModelURLWithIdentifier:identifier
+        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 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;
 }
 
@@ -582,10 +598,9 @@ - (nullable ETCoreMLAsset *)modelAssetWithMetadata:(const ModelMetadata&)metadat
         return nil;
     }
 
-    ETCoreMLModel *model = [ETCoreMLModelLoader loadModelWithContentsOfURL:compiledModelAsset.contentURL
+    ETCoreMLModel *model = [ETCoreMLModelLoader loadModelWithCompiledAsset:compiledModelAsset
                                                              configuration:configuration
                                                                   metadata:metadata
-                                                              assetManager:self.assetManager
                                                                      error:error];
     if (!model) {
         return nil;