Skip to content

[Core ML] Improve asset management #13560

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Aug 23, 2025
Merged

[Core ML] Improve asset management #13560

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
6c5bb97
Improve asset management
cymbalrush Aug 20, 2025
e1187a1
Address comemnts
cymbalrush Aug 21, 2025
de47c33
Merge branch 'main' into fix_assets
metascroy Aug 21, 2025
0bb893e
Merge branch 'main' into fix_assets
metascroy Aug 22, 2025
0541ad9
lint
metascroy Aug 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand All @@ -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;

Expand Down
104 changes: 66 additions & 38 deletions backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -254,6 +254,29 @@ BOOL is_asset_alive(NSMapTable<NSString *, ETCoreMLAsset *> *assets_in_use_map,

return assets;
}

NSURL * _Nullable move_to_directory(NSURL *url,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks similar to the pattern "[self moveURL:dstURL toUniqueURLInDirectory:self.trashDirectoryURL error:nil];"

Should we consolidate?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated to use move_to_directory

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]) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to check if url is nil when using moveItemAtURL?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added checks.

return nil;
}

return dstURL;
}

} //namespace

@interface ETCoreMLAssetManager () <NSFileManagerDelegate> {
Expand Down Expand Up @@ -299,12 +322,17 @@ - (nullable instancetype)initWithDatabase:(const std::shared_ptr<Database>&)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);
Expand All @@ -315,6 +343,7 @@ - (nullable instancetype)initWithDatabase:(const std::shared_ptr<Database>&)data
_assetsStore = std::move(assetsStore);
_assetsMetaStore = std::move(assetsMetaStore);
_assetsDirectoryURL = managedAssetsDirectoryURL;
_stagingDirectoryURL = managedStagingDirectoryURL;
_trashDirectoryURL = managedTrashDirectoryURL;
_estimatedSizeInBytes = sizeInBytes.value();
_maxAssetsSizeInBytes = maxAssetsSizeInBytes;
Expand Down Expand Up @@ -346,15 +375,15 @@ - (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];
if (![self.fileManager moveItemAtURL:url toURL:dstURL error:error]) {
return nil;
- (void)withTemporaryDirectory:(void (^)(NSURL *directoryURL))block {
NSURL *dstURL = [self.stagingDirectoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString];
block(dstURL);
if (![self.fileManager fileExistsAtPath:dstURL.path]) {
return;
}

return dstURL;

move_to_directory(dstURL, self.trashDirectoryURL, self.fileManager, nil);
[self cleanupTrashDirectory];
}

- (void)cleanupAssetIfNeeded:(ETCoreMLAsset *)asset {
Expand Down Expand Up @@ -407,9 +436,8 @@ - (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.
move_to_directory(dstURL, self.trashDirectoryURL, self.fileManager, nil);
// Move the asset to assets directory.
if (![self.fileManager moveItemAtURL:srcURL toURL:dstURL error:error]) {
return false;
Expand All @@ -433,16 +461,25 @@ - (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL
}

- (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.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What happens if something is moving to the trashDirectory when it's being deleted by another thread?

For example, withTemporaryDirectory moves staging files to the trashDirectory after compilation, but then another model triggers cleanupTrashDirectory on another thread.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That should be fine trashQueue is serial and we get the list of directories to delete at the start of the method

  NSDirectoryEnumerator *enumerator = [fileManager enumeratorAtURL:self.trashDirectoryURL
                                          includingPropertiesForKeys:@[]
                                                             options:NSDirectoryEnumerationSkipsSubdirectoryDescendants
                                                        errorHandler:errorHandler];

[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];
});
}

Expand Down Expand Up @@ -548,7 +585,7 @@ - (BOOL)_removeAssetWithIdentifier:(NSString *)identifier

NSURL *assetURL = ::get_asset_url(assetValue);
if ([self.fileManager fileExistsAtPath:assetURL.path] &&
![self moveURL:assetURL toUniqueURLInDirectory:self.trashDirectoryURL error:error]) {
!move_to_directory(assetURL, self.trashDirectoryURL, self.fileManager, error)) {
return false;
}

Expand Down Expand Up @@ -649,13 +686,7 @@ - (NSUInteger)_compact:(NSUInteger)sizeInBytes error:(NSError * __autoreleasing
identifier);
}
}

// Trigger cleanup.
__weak __typeof(self) weakSelf = self;
dispatch_async(self.trashQueue, ^{
[weakSelf removeFilesInTrashDirectory];
});


return _estimatedSizeInBytes;
}

Expand All @@ -664,7 +695,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;
}

Expand Down Expand Up @@ -708,7 +742,7 @@ - (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 (!move_to_directory(self.assetsDirectoryURL, self.trashDirectoryURL, self.fileManager, error)) {
return false;
}

Expand All @@ -724,13 +758,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<BOOL>(status);
}

Expand Down
19 changes: 5 additions & 14 deletions backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,21 +62,12 @@ + (nullable ETCoreMLModel *)loadModelWithContentsOfURL:(NSURL *)compiledModelURL
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);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why are we removing this?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If store fails then it's an error there is no point recovering from it.

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
Loading
Loading