Refactor modelInfo into localModelInfo and remoteModelInfo (#7163)

* Restores ML Pods after M77.

* Fix Package.swift

* Refactor model info to be in-memory.

* Add improved model download functionality w/ single set of download handlers per task.

* Rename model download to better match URLSessionTask.

* Rename model download to better match URLSessionTask.

* Update unit test to read/write to user defaults.

* Configure downloader with Firebase app.

* Minor updates to model downloader and tests.

* Singleton for model downloader instance, for default and custom app.

* Download types WIP - local model.

* Download types WIP.

* Manually add app lifecycle handling + TODOs for Firebase Components.

* Replace options with app name as a property in ModelDownloader.

* Download types WIP.

* Refactor model info.

* Add convenience inits + other minor fixes.

* Explicit enum for model info result.

Author: manjanac

FirebaseMLModelDownloader/Sources/CustomModel.swift

@@ -24,4 +24,21 @@ public struct CustomModel: Hashable {
   public let path: String
   /// Hash for the model, used for model verification.
   public let hash: String
+
+  init(name: String, size: Int, path: String, hash: String) {
+    self.name = name
+    self.size = size
+    self.path = path
+    self.hash = hash
+  }
+
+  /// Convenience init to create model from local model info.
+  init(localModelInfo: LocalModelInfo) {
+    self.init(
+      name: localModelInfo.name,
+      size: localModelInfo.size,
+      path: localModelInfo.path,
+      hash: localModelInfo.modelHash
+    )
+  }
 }

FirebaseMLModelDownloader/Sources/ModelInfo.swift renamed to FirebaseMLModelDownloader/Sources/LocalModelInfo.swift

@@ -15,12 +15,12 @@
 import Foundation
 import FirebaseCore
 
-/// Model info object with details about pending or downloaded model.
-struct ModelInfo {
+/// Model info object with details about downloaded and locally available model.
+// TODO: Can this be backed by user defaults property wrappers?
+class LocalModelInfo {
   /// Model name.
   let name: String
 
-  // TODO: revisit UserDefaultsBacked
   /// Download URL for the model file, as returned by server.
   let downloadURL: URL
 
@@ -31,25 +31,36 @@ struct ModelInfo {
   let size: Int
 
   /// Local path of the model.
-  var path: String?
+  let path: String
 
-  /// Initialize model info and create user default keys.
-  init(name: String, downloadURL: URL, modelHash: String, size: Int) {
+  /// Get user defaults key prefix.
+  private static func getUserDefaultsKeyPrefix(appName: String, modelName: String) -> String {
+    let bundleID = Bundle.main.bundleIdentifier ?? ""
+    return "\(bundleID).\(appName).\(modelName)"
+  }
+
+  init(name: String, downloadURL: URL, modelHash: String, size: Int, path: String) {
     self.name = name
     self.downloadURL = downloadURL
     self.modelHash = modelHash
     self.size = size
+    self.path = path
   }
 
-  /// Get user defaults key prefix.
-  private static func getUserDefaultsKeyPrefix(appName: String, modelName: String) -> String {
-    let bundleID = Bundle.main.bundleIdentifier ?? ""
-    return "\(bundleID).\(appName).\(modelName)"
+  /// Convenience init to create local model info from remotely downloaded model info and a local model path.
+  convenience init(from remoteModelInfo: RemoteModelInfo, path: String) {
+    self.init(
+      name: remoteModelInfo.name,
+      downloadURL: remoteModelInfo.downloadURL,
+      modelHash: remoteModelInfo.modelHash,
+      size: remoteModelInfo.size,
+      path: path
+    )
   }
 
-  // TODO: Move reading and writing to user defaults to a new file.
-  init?(fromDefaults defaults: UserDefaults, modelName: String, appName: String) {
-    let defaultsPrefix = ModelInfo.getUserDefaultsKeyPrefix(appName: appName, modelName: modelName)
+  /// Convenience init to create local model info from stored info in user defaults.
+  convenience init?(fromDefaults defaults: UserDefaults, name: String, appName: String) {
+    let defaultsPrefix = LocalModelInfo.getUserDefaultsKeyPrefix(appName: appName, modelName: name)
     guard let downloadURL = defaults
       .value(forKey: "\(defaultsPrefix).model-download-url") as? String,
       let url = URL(string: downloadURL),
@@ -58,22 +69,25 @@ struct ModelInfo {
       let path = defaults.value(forKey: "\(defaultsPrefix).model-path") as? String else {
       return nil
     }
-    name = modelName
-    self.downloadURL = url
-    self.modelHash = modelHash
-    self.size = size
-    self.path = path
+    self.init(name: name, downloadURL: url, modelHash: modelHash, size: size, path: path)
   }
 
-  func writeToDefaults(_ defaults: UserDefaults, appName: String) throws {
-    guard let modelPath = path else {
-      throw DownloadedModelError
-        .fileIOError(description: "Could not save model info to user defaults.")
-    }
-    let defaultsPrefix = ModelInfo.getUserDefaultsKeyPrefix(appName: appName, modelName: name)
+  /// Write local model info to user defaults.
+  func writeToDefaults(_ defaults: UserDefaults, appName: String) {
+    let defaultsPrefix = LocalModelInfo.getUserDefaultsKeyPrefix(appName: appName, modelName: name)
     defaults.setValue(downloadURL.absoluteString, forKey: "\(defaultsPrefix).model-download-url")
     defaults.setValue(modelHash, forKey: "\(defaultsPrefix).model-hash")
     defaults.setValue(size, forKey: "\(defaultsPrefix).model-size")
-    defaults.setValue(modelPath, forKey: "\(defaultsPrefix).model-path")
+    defaults.setValue(path, forKey: "\(defaultsPrefix).model-path")
+  }
+}
+
+/// Named user defaults for FirebaseML.
+extension UserDefaults {
+  static var firebaseMLDefaults: UserDefaults {
+    let suiteName = "com.google.firebase.ml"
+    // TODO: reconsider force unwrapping
+    let defaults = UserDefaults(suiteName: suiteName)!
+    return defaults
   }
 }

FirebaseMLModelDownloader/Sources/ModelDownloadTask.swift

@@ -15,6 +15,7 @@
 import Foundation
 import FirebaseCore
 
+/// Possible states of model downloading.
 enum DownloadStatus {
   case notStarted
   case inProgress
@@ -37,22 +38,29 @@ class DownloadHandlers {
 
 /// Manager to handle model downloading device and storing downloaded model info to persistent storage.
 class ModelDownloadTask: NSObject {
+  /// Name of the app associated with this instance of ModelDownloadTask.
   private let appName: String
-  private(set) var modelInfo: ModelInfo
+  /// Model info downloaded from server.
+  private(set) var remoteModelInfo: RemoteModelInfo
+  /// User defaults to which local model info should ultimately be written.
+  private let defaults: UserDefaults
+  /// Task to handle model file download.
   private var downloadTask: URLSessionDownloadTask?
+  /// Progress and completion handlers associated with this model download task.
   private let downloadHandlers: DownloadHandlers
-
+  /// Keeps track of download associated with this model download task.
   private(set) var downloadStatus: DownloadStatus = .notStarted
-
+  /// URLSession to handle model downloads.
   private lazy var downloadSession = URLSession(configuration: .ephemeral,
                                                 delegate: self,
                                                 delegateQueue: nil)
 
-  init(modelInfo: ModelInfo, appName: String,
+  init(remoteModelInfo: RemoteModelInfo, appName: String, defaults: UserDefaults,
        progressHandler: DownloadHandlers.ProgressHandler? = nil,
        completion: @escaping DownloadHandlers.Completion) {
-    self.modelInfo = modelInfo
+    self.remoteModelInfo = remoteModelInfo
     self.appName = appName
+    self.defaults = defaults
     downloadHandlers = DownloadHandlers(
       progressHandler: progressHandler,
       completion: completion
@@ -62,7 +70,7 @@ class ModelDownloadTask: NSObject {
   /// Asynchronously download model file to device.
   func resumeModelDownload() {
     guard downloadStatus == .notStarted else { return }
-    let downloadTask = downloadSession.downloadTask(with: modelInfo.downloadURL)
+    let downloadTask = downloadSession.downloadTask(with: remoteModelInfo.downloadURL)
     downloadTask.resume()
     downloadStatus = .inProgress
     self.downloadTask = downloadTask
@@ -71,6 +79,11 @@ class ModelDownloadTask: NSObject {
 
 /// Extension to handle delegate methods.
 extension ModelDownloadTask: URLSessionDownloadDelegate {
+  /// Name for model file stored on device.
+  var downloadedModelFileName: String {
+    return "fbml_model__\(appName)__\(remoteModelInfo.name)"
+  }
+
   func urlSession(_ session: URLSession,
                   downloadTask: URLSessionDownloadTask,
                   didFinishDownloadingTo location: URL) {
@@ -80,32 +93,27 @@ extension ModelDownloadTask: URLSessionDownloadDelegate {
       .appendingPathComponent(downloadedModelFileName)
     do {
       try ModelFileManager.moveFile(at: location, to: savedURL)
+    } catch let error as DownloadError {
+      DispatchQueue.main.async {
+        self.downloadHandlers
+          .completion(.failure(error))
+      }
     } catch {
-      downloadHandlers
-        .completion(.failure(.internalError(description: error.localizedDescription)))
-      return
+      DispatchQueue.main.async {
+        self.downloadHandlers
+          .completion(.failure(.internalError(description: error.localizedDescription)))
+      }
     }
 
-    /// Set path to local model.
-    modelInfo.path = savedURL.absoluteString
+    /// Generate local model info.
+    let localModelInfo = LocalModelInfo(from: remoteModelInfo, path: savedURL.absoluteString)
     /// Write model to user defaults.
-    do {
-      try modelInfo.writeToDefaults(.firebaseMLDefaults, appName: appName)
-    } catch {
-      downloadHandlers
-        .completion(.failure(.internalError(description: error.localizedDescription)))
-    }
+    localModelInfo.writeToDefaults(defaults, appName: appName)
     /// Build model from model info.
-    guard let model = buildModel() else {
-      downloadHandlers
-        .completion(
-          .failure(
-            .internalError(description: "Could not create model due to incomplete model info.")
-          )
-        )
-      return
+    let model = CustomModel(localModelInfo: localModelInfo)
+    DispatchQueue.main.async {
+      self.downloadHandlers.completion(.success(model))
     }
-    downloadHandlers.completion(.success(model))
   }
 
   func urlSession(_ session: URLSession,
@@ -114,50 +122,11 @@ extension ModelDownloadTask: URLSessionDownloadDelegate {
                   totalBytesWritten: Int64,
                   totalBytesExpectedToWrite: Int64) {
     assert(downloadTask == self.downloadTask)
+    /// Check if progress handler is unspecified.
     guard let progressHandler = downloadHandlers.progressHandler else { return }
     let calculatedProgress = Float(totalBytesWritten) / Float(totalBytesExpectedToWrite)
-    progressHandler(calculatedProgress)
-  }
-}
-
-/// Extension to handle post-download operations.
-extension ModelDownloadTask {
-  var downloadedModelFileName: String {
-    return "fbml_model__\(appName)__\(modelInfo.name)"
-  }
-
-  /// Build custom model object from model info.
-  // TODO: Consider moving this to CustomModel as a convenience init
-  func buildModel() -> CustomModel? {
-    /// Build custom model only if the model file is already on device.
-    guard let path = modelInfo.path else { return nil }
-    let model = CustomModel(
-      name: modelInfo.name,
-      size: modelInfo.size,
-      path: path,
-      hash: modelInfo.modelHash
-    )
-    return model
-  }
-
-  /// Get the local path to model on device.
-  func getLocalModelPath(model: CustomModel) -> URL? {
-    let fileURL: URL = ModelFileManager.modelsDirectory
-      .appendingPathComponent(downloadedModelFileName)
-    if ModelFileManager.isFileReachable(at: fileURL) {
-      return fileURL
-    } else {
-      return nil
+    DispatchQueue.main.async {
+      progressHandler(calculatedProgress)
     }
   }
 }
-
-/// Named user defaults for FirebaseML.
-extension UserDefaults {
-  static var firebaseMLDefaults: UserDefaults {
-    let suiteName = "com.google.firebase.ml"
-    // TODO: reconsider force unwrapping
-    let defaults = UserDefaults(suiteName: suiteName)!
-    return defaults
-  }
-}

FirebaseMLModelDownloader/Sources/ModelFileManager.swift

@@ -44,14 +44,14 @@ enum ModelFileManager {
       do {
         try FileManager.default.removeItem(at: destinationURL)
       } catch {
-        throw DownloadedModelError
-          .fileIOError(description: "Could not replace existing model file.")
+        throw DownloadError
+          .internalError(description: "Could not replace existing model file.")
       }
     }
     do {
       try FileManager.default.moveItem(at: sourceURL, to: destinationURL)
     } catch {
-      throw DownloadedModelError.fileIOError(description: "Unable to save model file.")
+      throw DownloadError.internalError(description: "Unable to save model file.")
     }
   }
 }