Add option for background session, but use default normally (#64)

Author: ZachNagengast

Sources/Hub/Downloader.swift

@@ -11,34 +11,42 @@ import Combine
 
 class Downloader: NSObject, ObservableObject {
     private(set) var destination: URL
-    
+
     enum DownloadState {
         case notStarted
         case downloading(Double)
         case completed(URL)
         case failed(Error)
     }
-    
+
     enum DownloadError: Error {
         case invalidDownloadLocation
         case unexpectedError
     }
-    
+
     private(set) lazy var downloadState: CurrentValueSubject<DownloadState, Never> = CurrentValueSubject(.notStarted)
     private var stateSubscriber: Cancellable?
-    
+
     private var urlSession: URLSession? = nil
-    
-    init(from url: URL, to destination: URL, using authToken: String? = nil) {
+
+    init(from url: URL, to destination: URL, using authToken: String? = nil, inBackground: Bool = false) {
         self.destination = destination
         super.init()
-        
-        let config = URLSessionConfiguration.background(withIdentifier: url.path)
-#if targetEnvironment(simulator)
-        urlSession = URLSession(configuration: .default, delegate: self, delegateQueue: OperationQueue())
-#else
-        urlSession = URLSession(configuration: config, delegate: self, delegateQueue: OperationQueue())
-#endif
+        let sessionIdentifier = "swift-transformers.hub.downloader"
+
+        var config = URLSessionConfiguration.default
+        if inBackground {
+            config = URLSessionConfiguration.background(withIdentifier: sessionIdentifier)
+            config.isDiscretionary = false
+            config.sessionSendsLaunchEvents = true
+        }
+
+        self.urlSession = URLSession(configuration: config, delegate: self, delegateQueue: nil)
+
+        setupDownload(from: url, with: authToken)
+    }
+
+    private func setupDownload(from url: URL, with authToken: String?) {
         downloadState.value = .downloading(0)
         urlSession?.getAllTasks { tasks in
             // If there's an existing pending background task with the same URL, let it proceed.
@@ -70,7 +78,7 @@ class Downloader: NSObject, ObservableObject {
             self.urlSession?.downloadTask(with: request).resume()
         }
     }
-    
+
     @discardableResult
     func waitUntilDone() throws -> URL {
         // It's either this, or stream the bytes ourselves (add to a buffer, save to disk, etc; boring and finicky)
@@ -83,31 +91,28 @@ class Downloader: NSObject, ObservableObject {
             }
         }
         semaphore.wait()
-        
+
         switch downloadState.value {
         case .completed(let url): return url
         case .failed(let error):  throw error
         default:                  throw DownloadError.unexpectedError
         }
     }
-    
+
     func cancel() {
         urlSession?.invalidateAndCancel()
     }
 }
 
-extension Downloader: URLSessionDelegate, URLSessionDownloadDelegate {
+extension Downloader: URLSessionDownloadDelegate {
     func urlSession(_: URLSession, downloadTask: URLSessionDownloadTask, didWriteData _: Int64, totalBytesWritten _: Int64, totalBytesExpectedToWrite _: Int64) {
         downloadState.value = .downloading(downloadTask.progress.fractionCompleted)
     }
 
     func urlSession(_: URLSession, downloadTask _: URLSessionDownloadTask, didFinishDownloadingTo location: URL) {
-        guard FileManager.default.fileExists(atPath: location.path) else {
-            downloadState.value = .failed(DownloadError.invalidDownloadLocation)
-            return
-        }
         do {
-            try FileManager.default.moveItem(at: location, to: destination)
+            // If the downloaded file already exists on the filesystem, overwrite it
+            try FileManager.default.moveDownloadedFile(from: location, to: self.destination)
             downloadState.value = .completed(destination)
         } catch {
             downloadState.value = .failed(error)
@@ -124,3 +129,12 @@ extension Downloader: URLSessionDelegate, URLSessionDownloadDelegate {
         }
     }
 }
+
+extension FileManager {
+    func moveDownloadedFile(from srcURL: URL, to dstURL: URL) throws {
+        if fileExists(atPath: dstURL.path) {
+            try removeItem(at: dstURL)
+        }
+        try moveItem(at: srcURL, to: dstURL)
+    }
+}

Sources/Hub/HubApi.swift

@@ -8,18 +8,15 @@
 import Foundation
 
 public struct HubApi {
-    public let downloadBase: URL
-    public let hfToken: String?
-    public let endpoint: String
-    
+    var downloadBase: URL
+    var hfToken: String?
+    var endpoint: String
+    var useBackgroundSession: Bool
+
     public typealias RepoType = Hub.RepoType
     public typealias Repo = Hub.Repo
 
-    public init(
-        downloadBase: URL? = nil,
-        hfToken: String? = nil,
-        endpoint: String = "https://huggingface.co"
-    ) {
+    public init(downloadBase: URL? = nil, hfToken: String? = nil, endpoint: String = "https://huggingface.co", useBackgroundSession: Bool = false) {
         self.hfToken = hfToken
         if let downloadBase {
             self.downloadBase = downloadBase
@@ -28,6 +25,7 @@ public struct HubApi {
             self.downloadBase = documents.appending(component: "huggingface")
         }
         self.endpoint = endpoint
+        self.useBackgroundSession = useBackgroundSession
     }
 
     public static let shared = HubApi()
@@ -129,7 +127,8 @@ public extension HubApi {
         let relativeFilename: String
         let hfToken: String?
         let endpoint: String?
-        
+        let backgroundSession: Bool
+
         var source: URL {
             // https://huggingface.co/coreml-projects/Llama-2-7b-chat-coreml/resolve/main/tokenizer.json?download=true
             var url = URL(string: endpoint ?? "https://huggingface.co")!
@@ -163,7 +162,7 @@ public extension HubApi {
             guard !downloaded else { return destination }
 
             try prepareDestination()
-            let downloader = Downloader(from: source, to: destination, using: hfToken)
+            let downloader = Downloader(from: source, to: destination, using: hfToken, inBackground: backgroundSession)
             let downloadSubscriber = downloader.downloadState.sink { state in
                 if case .downloading(let progress) = state {
                     progressHandler(progress)
@@ -188,7 +187,8 @@ public extension HubApi {
                 repoDestination: repoDestination,
                 relativeFilename: filename,
                 hfToken: hfToken,
-                endpoint: endpoint
+                endpoint: endpoint,
+                backgroundSession: useBackgroundSession
             )
             try await downloader.download { fractionDownloaded in
                 fileProgress.completedUnitCount = Int64(100 * fractionDownloaded)

Tests/HubTests/HubApiTests.swift

@@ -149,6 +149,29 @@ class SnapshotDownloadTests: XCTestCase {
         )
     }
 
+    /// Background sessions get rate limited by the OS, see discussion here: https://github.com/huggingface/swift-transformers/issues/61
+    /// Test only one file at a time
+    func testDownloadInBackground() async throws {
+        let hubApi = HubApi(downloadBase: downloadDestination, useBackgroundSession: true)
+        var lastProgress: Progress? = nil
+        let downloadedTo = try await hubApi.snapshot(from: repo, matching: "llama-2-7b-chat.mlpackage/Data/com.apple.CoreML/Metadata.json") { progress in
+            print("Total Progress: \(progress.fractionCompleted)")
+            print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)")
+            lastProgress = progress
+        }
+        XCTAssertEqual(lastProgress?.fractionCompleted, 1)
+        XCTAssertEqual(lastProgress?.completedUnitCount, 1)
+        XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)"))
+
+        let downloadedFilenames = getRelativeFiles(url: downloadDestination)
+        XCTAssertEqual(
+            Set(downloadedFilenames),
+            Set([
+                "llama-2-7b-chat.mlpackage/Data/com.apple.CoreML/Metadata.json",
+            ])
+        )
+    }
+
     func testCustomEndpointDownload() async throws {
         let hubApi = HubApi(downloadBase: downloadDestination, endpoint: "https://hf-mirror.com")
         var lastProgress: Progress? = nil