From 2b9b8733a27c7a3913a7952ede97ff1cbbcf43a2 Mon Sep 17 00:00:00 2001 From: Anthony DePasquale Date: Thu, 6 Mar 2025 09:50:00 +0100 Subject: [PATCH 1/7] Add resumable downloads across app sessions --- Sources/Hub/Downloader.swift | 175 +++++++++++++++++++++------ Sources/Hub/Hub.swift | 6 +- Sources/Hub/HubApi.swift | 46 +++++-- Tests/HubTests/DownloaderTests.swift | 24 ++++ 4 files changed, 202 insertions(+), 49 deletions(-) diff --git a/Sources/Hub/Downloader.swift b/Sources/Hub/Downloader.swift index f52c596..521ad15 100644 --- a/Sources/Hub/Downloader.swift +++ b/Sources/Hub/Downloader.swift @@ -11,6 +11,7 @@ import Foundation class Downloader: NSObject, ObservableObject { private(set) var destination: URL + private(set) var sourceURL: URL private let chunkSize = 10 * 1024 * 1024 // 10MB @@ -24,25 +25,69 @@ class Downloader: NSObject, ObservableObject { enum DownloadError: Error { case invalidDownloadLocation case unexpectedError + case tempFileNotFound } private(set) lazy var downloadState: CurrentValueSubject = CurrentValueSubject(.notStarted) private var stateSubscriber: Cancellable? + + private(set) var tempFilePath: URL? + private(set) var expectedSize: Int? + private(set) var downloadedSize: Int = 0 private var urlSession: URLSession? = nil + + /// Creates the incomplete file path for a given destination URL + /// This is similar to the Hugging Face Hub approach of using .incomplete files + static func incompletePath(for destination: URL) -> URL { + destination.appendingPathExtension("incomplete") + } + + /// Check if an incomplete file exists for the destination and returns its size + /// - Parameter destination: The destination URL for the download + /// - Returns: Size of the incomplete file if it exists, otherwise 0 + static func checkForIncompleteFile(at destination: URL) -> Int { + let incompletePath = Self.incompletePath(for: destination) + + if FileManager.default.fileExists(atPath: incompletePath.path) { + if let attributes = try? FileManager.default.attributesOfItem(atPath: incompletePath.path), + let fileSize = attributes[.size] as? Int + { + print("[Downloader] Found existing incomplete file for \(destination.lastPathComponent): \(fileSize) bytes") + return fileSize + } + } + + return 0 + } init( from url: URL, to destination: URL, using authToken: String? = nil, inBackground: Bool = false, - resumeSize: Int = 0, + resumeSize: Int = 0, // Can be specified manually, but will also check for incomplete files headers: [String: String]? = nil, expectedSize: Int? = nil, timeout: TimeInterval = 10, numRetries: Int = 5 ) { self.destination = destination + sourceURL = url + self.expectedSize = expectedSize + + // Create incomplete file path based on destination + tempFilePath = Downloader.incompletePath(for: destination) + + // If resume size wasn't specified, check for an existing incomplete file + let actualResumeSize: Int = if resumeSize > 0 { + resumeSize + } else { + Downloader.checkForIncompleteFile(at: destination) + } + + downloadedSize = actualResumeSize + super.init() let sessionIdentifier = "swift-transformers.hub.downloader" @@ -55,7 +100,7 @@ class Downloader: NSObject, ObservableObject { urlSession = URLSession(configuration: config, delegate: self, delegateQueue: nil) - setupDownload(from: url, with: authToken, resumeSize: resumeSize, headers: headers, expectedSize: expectedSize, timeout: timeout, numRetries: numRetries) + setUpDownload(from: url, with: authToken, resumeSize: actualResumeSize, headers: headers, expectedSize: expectedSize, timeout: timeout, numRetries: numRetries) } /// Sets up and initiates a file download operation @@ -68,7 +113,7 @@ class Downloader: NSObject, ObservableObject { /// - expectedSize: Expected file size in bytes for validation /// - timeout: Time interval before the request times out /// - numRetries: Number of retry attempts for failed downloads - private func setupDownload( + private func setUpDownload( from url: URL, with authToken: String?, resumeSize: Int, @@ -77,61 +122,101 @@ class Downloader: NSObject, ObservableObject { timeout: TimeInterval, numRetries: Int ) { - downloadState.value = .downloading(0) + print("[Downloader] Setting up download for \(url.lastPathComponent)") + print("[Downloader] Destination: \(destination.path)") + print("[Downloader] Incomplete file: \(tempFilePath?.path ?? "none")") + urlSession?.getAllTasks { tasks in // If there's an existing pending background task with the same URL, let it proceed. if let existing = tasks.filter({ $0.originalRequest?.url == url }).first { switch existing.state { case .running: - // print("Already downloading \(url)") + print("[Downloader] Task already running for \(url.lastPathComponent)") return case .suspended: - // print("Resuming suspended download task for \(url)") + print("[Downloader] Resuming suspended download task for \(url.lastPathComponent)") existing.resume() return - case .canceling: - // print("Starting new download task for \(url), previous was canceling") - break - case .completed: - // print("Starting new download task for \(url), previous is complete but the file is no longer present (I think it's cached)") - break + case .canceling, .completed: + existing.cancel() @unknown default: - // print("Unknown state for running task; cancelling and creating a new one") existing.cancel() } } - var request = URLRequest(url: url) - - // Use headers from argument else create an empty header dictionary - var requestHeaders = headers ?? [:] - - // Populate header auth and range fields - if let authToken { - requestHeaders["Authorization"] = "Bearer \(authToken)" - } - if resumeSize > 0 { - requestHeaders["Range"] = "bytes=\(resumeSize)-" - } - request.timeoutInterval = timeout - request.allHTTPHeaderFields = requestHeaders - Task { do { - // Create a temp file to write - let tempURL = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) - FileManager.default.createFile(atPath: tempURL.path, contents: nil) - let tempFile = try FileHandle(forWritingTo: tempURL) + // Check if incomplete file exists and get its size + var existingSize = 0 + guard let incompleteFilePath = self.tempFilePath else { + throw DownloadError.unexpectedError + } + + let fileManager = FileManager.default + if fileManager.fileExists(atPath: incompleteFilePath.path) { + let attributes = try fileManager.attributesOfItem(atPath: incompleteFilePath.path) + existingSize = attributes[.size] as? Int ?? 0 + print("[Downloader] Found incomplete file with \(existingSize) bytes") + self.downloadedSize = existingSize + } else { + // Create parent directory if needed + try fileManager.createDirectory(at: incompleteFilePath.deletingLastPathComponent(), withIntermediateDirectories: true) + + // Create empty incomplete file + fileManager.createFile(atPath: incompleteFilePath.path, contents: nil) + print("[Downloader] Created new incomplete file at \(incompleteFilePath.path)") + } + + // Set up the request with appropriate headers + var request = URLRequest(url: url) + var requestHeaders = headers ?? [:] + + if let authToken { + requestHeaders["Authorization"] = "Bearer \(authToken)" + } + + // Set Range header if we're resuming + if existingSize > 0 { + requestHeaders["Range"] = "bytes=\(existingSize)-" + + // Calculate and show initial progress + if let expectedSize, expectedSize > 0 { + let initialProgress = Double(existingSize) / Double(expectedSize) + self.downloadState.value = .downloading(initialProgress) + print("[Downloader] Resuming from \(existingSize)/\(expectedSize) bytes (\(Int(initialProgress * 100))%)") + } else { + self.downloadState.value = .downloading(0) + print("[Downloader] Resuming download from byte \(existingSize)") + } + } else { + self.downloadState.value = .downloading(0) + print("[Downloader] Starting new download") + } + + request.timeoutInterval = timeout + request.allHTTPHeaderFields = requestHeaders + + // Open the incomplete file for writing + let tempFile = try FileHandle(forWritingTo: incompleteFilePath) + + // If resuming, seek to end of file + if existingSize > 0 { + try tempFile.seekToEnd() + } defer { tempFile.closeFile() } - try await self.httpGet(request: request, tempFile: tempFile, resumeSize: resumeSize, numRetries: numRetries, expectedSize: expectedSize) + try await self.httpGet(request: request, tempFile: tempFile, resumeSize: self.downloadedSize, numRetries: numRetries, expectedSize: expectedSize) // Clean up and move the completed download to its final destination tempFile.closeFile() - try FileManager.default.moveDownloadedFile(from: tempURL, to: self.destination) + print("[Downloader] Download completed with total size \(self.downloadedSize) bytes") + print("[Downloader] Moving incomplete file to destination: \(self.destination.path)") + try fileManager.moveDownloadedFile(from: incompleteFilePath, to: self.destination) + print("[Downloader] Download successfully completed") self.downloadState.value = .completed(self.destination) } catch { + print("[Downloader] Error: \(error)") self.downloadState.value = .failed(error) } } @@ -164,20 +249,31 @@ class Downloader: NSObject, ObservableObject { var newRequest = request if resumeSize > 0 { newRequest.setValue("bytes=\(resumeSize)-", forHTTPHeaderField: "Range") + print("[Downloader] Adding Range header: bytes=\(resumeSize)-") } // Start the download and get the byte stream let (asyncBytes, response) = try await session.bytes(for: newRequest) - guard let response = response as? HTTPURLResponse else { + guard let httpResponse = response as? HTTPURLResponse else { + print("[Downloader] Error: Non-HTTP response received") throw DownloadError.unexpectedError } + + print("[Downloader] Received HTTP \(httpResponse.statusCode) response") + if let contentRange = httpResponse.value(forHTTPHeaderField: "Content-Range") { + print("[Downloader] Content-Range: \(contentRange)") + } + if let contentLength = httpResponse.value(forHTTPHeaderField: "Content-Length") { + print("[Downloader] Content-Length: \(contentLength)") + } - guard (200..<300).contains(response.statusCode) else { + guard (200..<300).contains(httpResponse.statusCode) else { + print("[Downloader] Error: HTTP status code \(httpResponse.statusCode)") throw DownloadError.unexpectedError } - var downloadedSize = resumeSize + downloadedSize = resumeSize // Create a buffer to collect bytes before writing to disk var buffer = Data(capacity: chunkSize) @@ -218,7 +314,7 @@ class Downloader: NSObject, ObservableObject { try await httpGet( request: request, tempFile: tempFile, - resumeSize: downloadedSize, + resumeSize: self.downloadedSize, numRetries: newNumRetries - 1, expectedSize: expectedSize ) @@ -227,7 +323,10 @@ class Downloader: NSObject, ObservableObject { // Verify the downloaded file size matches the expected size let actualSize = try tempFile.seekToEnd() if let expectedSize, expectedSize != actualSize { + print("[Downloader] Error: Size mismatch - expected \(expectedSize) bytes but got \(actualSize) bytes") throw DownloadError.unexpectedError + } else { + print("[Downloader] Final verification passed, size: \(actualSize) bytes") } } diff --git a/Sources/Hub/Hub.swift b/Sources/Hub/Hub.swift index 1c2cd22..f81febf 100644 --- a/Sources/Hub/Hub.swift +++ b/Sources/Hub/Hub.swift @@ -51,13 +51,13 @@ public extension Hub { } } - enum RepoType: String { + enum RepoType: String, Codable { case models case datasets case spaces } - - struct Repo { + + struct Repo: Codable { public let id: String public let type: RepoType diff --git a/Sources/Hub/HubApi.swift b/Sources/Hub/HubApi.swift index 8102dc6..a38647a 100644 --- a/Sources/Hub/HubApi.swift +++ b/Sources/Hub/HubApi.swift @@ -425,19 +425,49 @@ public extension HubApi { try prepareDestination() try prepareMetadataDestination() - let downloader = Downloader(from: source, to: destination, using: hfToken, inBackground: backgroundSession, expectedSize: remoteSize) + // Check for an existing incomplete file + let incompleteFile = Downloader.incompletePath(for: destination) + var resumeSize = 0 + + if FileManager.default.fileExists(atPath: incompleteFile.path) { + if let fileAttributes = try? FileManager.default.attributesOfItem(atPath: incompleteFile.path) { + resumeSize = (fileAttributes[FileAttributeKey.size] as? Int) ?? 0 + print("[HubApi] Found existing incomplete file for \(destination.lastPathComponent): \(resumeSize) bytes at \(incompleteFile.path)") + } + } else { + print("[HubApi] No existing incomplete file found for \(destination.lastPathComponent)") + } + + let downloader = Downloader( + from: source, + to: destination, + using: hfToken, + inBackground: backgroundSession, + resumeSize: resumeSize, + expectedSize: remoteSize + ) + let downloadSubscriber = downloader.downloadState.sink { state in - if case let .downloading(progress) = state { + switch state { + case let .downloading(progress): progressHandler(progress) + case .completed, .failed, .notStarted: + break } } - _ = try withExtendedLifetime(downloadSubscriber) { - try downloader.waitUntilDone() + do { + _ = try withExtendedLifetime(downloadSubscriber) { + try downloader.waitUntilDone() + } + + try HubApi.shared.writeDownloadMetadata(commitHash: remoteCommitHash, etag: remoteEtag, metadataPath: metadataDestination) + + return destination + } catch { + // If download fails, leave the incomplete file in place for future resume + print("[HubApi] Download failed but incomplete file is preserved for future resume: \(error.localizedDescription)") + throw error } - - try HubApi.shared.writeDownloadMetadata(commitHash: remoteCommitHash, etag: remoteEtag, metadataPath: metadataDestination) - - return destination } } diff --git a/Tests/HubTests/DownloaderTests.swift b/Tests/HubTests/DownloaderTests.swift index 2c6ed02..1ef7c0f 100644 --- a/Tests/HubTests/DownloaderTests.swift +++ b/Tests/HubTests/DownloaderTests.swift @@ -168,4 +168,28 @@ final class DownloaderTests: XCTestCase { throw error } } + + func testAutomaticIncompleteFileDetection() throws { + let url = URL(string: "https://huggingface.co/coreml-projects/sam-2-studio/resolve/main/SAM%202%20Studio%201.1.zip")! + let destination = tempDir.appendingPathComponent("SAM%202%20Studio%201.1.zip") + + // Create a sample incomplete file with test content + let incompletePath = Downloader.incompletePath(for: destination) + try FileManager.default.createDirectory(at: incompletePath.deletingLastPathComponent(), withIntermediateDirectories: true) + let testContent = Data(repeating: 65, count: 1024) // 1KB of data + FileManager.default.createFile(atPath: incompletePath.path, contents: testContent) + + // Create a downloader for the same destination + // It should automatically detect and use the incomplete file + let downloader = Downloader( + from: url, + to: destination + ) + + // Verify the downloader found and is using the incomplete file + XCTAssertEqual(downloader.downloadedSize, 1024, "Should have detected the incomplete file and set resumeSize") + + // Clean up + try? FileManager.default.removeItem(at: incompletePath) + } } From ba49d2f271a2bf6bfc3b9b488c82dee963683e7f Mon Sep 17 00:00:00 2001 From: Anthony DePasquale Date: Tue, 25 Mar 2025 21:38:09 +0100 Subject: [PATCH 2/7] Remove debug print statements --- Sources/Hub/Downloader.swift | 32 -------------------------------- Sources/Hub/HubApi.swift | 4 ---- 2 files changed, 36 deletions(-) diff --git a/Sources/Hub/Downloader.swift b/Sources/Hub/Downloader.swift index 521ad15..dd6a921 100644 --- a/Sources/Hub/Downloader.swift +++ b/Sources/Hub/Downloader.swift @@ -53,7 +53,6 @@ class Downloader: NSObject, ObservableObject { if let attributes = try? FileManager.default.attributesOfItem(atPath: incompletePath.path), let fileSize = attributes[.size] as? Int { - print("[Downloader] Found existing incomplete file for \(destination.lastPathComponent): \(fileSize) bytes") return fileSize } } @@ -122,19 +121,13 @@ class Downloader: NSObject, ObservableObject { timeout: TimeInterval, numRetries: Int ) { - print("[Downloader] Setting up download for \(url.lastPathComponent)") - print("[Downloader] Destination: \(destination.path)") - print("[Downloader] Incomplete file: \(tempFilePath?.path ?? "none")") - urlSession?.getAllTasks { tasks in // If there's an existing pending background task with the same URL, let it proceed. if let existing = tasks.filter({ $0.originalRequest?.url == url }).first { switch existing.state { case .running: - print("[Downloader] Task already running for \(url.lastPathComponent)") return case .suspended: - print("[Downloader] Resuming suspended download task for \(url.lastPathComponent)") existing.resume() return case .canceling, .completed: @@ -156,7 +149,6 @@ class Downloader: NSObject, ObservableObject { if fileManager.fileExists(atPath: incompleteFilePath.path) { let attributes = try fileManager.attributesOfItem(atPath: incompleteFilePath.path) existingSize = attributes[.size] as? Int ?? 0 - print("[Downloader] Found incomplete file with \(existingSize) bytes") self.downloadedSize = existingSize } else { // Create parent directory if needed @@ -164,7 +156,6 @@ class Downloader: NSObject, ObservableObject { // Create empty incomplete file fileManager.createFile(atPath: incompleteFilePath.path, contents: nil) - print("[Downloader] Created new incomplete file at \(incompleteFilePath.path)") } // Set up the request with appropriate headers @@ -183,14 +174,11 @@ class Downloader: NSObject, ObservableObject { if let expectedSize, expectedSize > 0 { let initialProgress = Double(existingSize) / Double(expectedSize) self.downloadState.value = .downloading(initialProgress) - print("[Downloader] Resuming from \(existingSize)/\(expectedSize) bytes (\(Int(initialProgress * 100))%)") } else { self.downloadState.value = .downloading(0) - print("[Downloader] Resuming download from byte \(existingSize)") } } else { self.downloadState.value = .downloading(0) - print("[Downloader] Starting new download") } request.timeoutInterval = timeout @@ -209,14 +197,9 @@ class Downloader: NSObject, ObservableObject { // Clean up and move the completed download to its final destination tempFile.closeFile() - print("[Downloader] Download completed with total size \(self.downloadedSize) bytes") - print("[Downloader] Moving incomplete file to destination: \(self.destination.path)") try fileManager.moveDownloadedFile(from: incompleteFilePath, to: self.destination) - - print("[Downloader] Download successfully completed") self.downloadState.value = .completed(self.destination) } catch { - print("[Downloader] Error: \(error)") self.downloadState.value = .failed(error) } } @@ -249,27 +232,15 @@ class Downloader: NSObject, ObservableObject { var newRequest = request if resumeSize > 0 { newRequest.setValue("bytes=\(resumeSize)-", forHTTPHeaderField: "Range") - print("[Downloader] Adding Range header: bytes=\(resumeSize)-") } // Start the download and get the byte stream let (asyncBytes, response) = try await session.bytes(for: newRequest) guard let httpResponse = response as? HTTPURLResponse else { - print("[Downloader] Error: Non-HTTP response received") throw DownloadError.unexpectedError } - - print("[Downloader] Received HTTP \(httpResponse.statusCode) response") - if let contentRange = httpResponse.value(forHTTPHeaderField: "Content-Range") { - print("[Downloader] Content-Range: \(contentRange)") - } - if let contentLength = httpResponse.value(forHTTPHeaderField: "Content-Length") { - print("[Downloader] Content-Length: \(contentLength)") - } - guard (200..<300).contains(httpResponse.statusCode) else { - print("[Downloader] Error: HTTP status code \(httpResponse.statusCode)") throw DownloadError.unexpectedError } @@ -323,10 +294,7 @@ class Downloader: NSObject, ObservableObject { // Verify the downloaded file size matches the expected size let actualSize = try tempFile.seekToEnd() if let expectedSize, expectedSize != actualSize { - print("[Downloader] Error: Size mismatch - expected \(expectedSize) bytes but got \(actualSize) bytes") throw DownloadError.unexpectedError - } else { - print("[Downloader] Final verification passed, size: \(actualSize) bytes") } } diff --git a/Sources/Hub/HubApi.swift b/Sources/Hub/HubApi.swift index a38647a..d3a1e3d 100644 --- a/Sources/Hub/HubApi.swift +++ b/Sources/Hub/HubApi.swift @@ -432,10 +432,7 @@ public extension HubApi { if FileManager.default.fileExists(atPath: incompleteFile.path) { if let fileAttributes = try? FileManager.default.attributesOfItem(atPath: incompleteFile.path) { resumeSize = (fileAttributes[FileAttributeKey.size] as? Int) ?? 0 - print("[HubApi] Found existing incomplete file for \(destination.lastPathComponent): \(resumeSize) bytes at \(incompleteFile.path)") } - } else { - print("[HubApi] No existing incomplete file found for \(destination.lastPathComponent)") } let downloader = Downloader( @@ -465,7 +462,6 @@ public extension HubApi { return destination } catch { // If download fails, leave the incomplete file in place for future resume - print("[HubApi] Download failed but incomplete file is preserved for future resume: \(error.localizedDescription)") throw error } } From dfaed4f219a4707fa638ac3b9b9a6726ee22e3a5 Mon Sep 17 00:00:00 2001 From: Arda Atahan Ibis Date: Mon, 14 Apr 2025 12:22:42 -0700 Subject: [PATCH 3/7] cleanup and add new test cases --- Sources/Hub/Downloader.swift | 63 +++++------------------ Sources/Hub/HubApi.swift | 26 +++++----- Tests/HubTests/DownloaderTests.swift | 49 +++++++++--------- Tests/HubTests/HubApiTests.swift | 74 ++++++++++++++++++++++++++++ 4 files changed, 122 insertions(+), 90 deletions(-) diff --git a/Sources/Hub/Downloader.swift b/Sources/Hub/Downloader.swift index dd6a921..1f27c4b 100644 --- a/Sources/Hub/Downloader.swift +++ b/Sources/Hub/Downloader.swift @@ -11,7 +11,6 @@ import Foundation class Downloader: NSObject, ObservableObject { private(set) var destination: URL - private(set) var sourceURL: URL private let chunkSize = 10 * 1024 * 1024 // 10MB @@ -31,28 +30,18 @@ class Downloader: NSObject, ObservableObject { private(set) lazy var downloadState: CurrentValueSubject = CurrentValueSubject(.notStarted) private var stateSubscriber: Cancellable? - private(set) var tempFilePath: URL? + private(set) var tempFilePath: URL private(set) var expectedSize: Int? private(set) var downloadedSize: Int = 0 private var urlSession: URLSession? = nil - /// Creates the incomplete file path for a given destination URL - /// This is similar to the Hugging Face Hub approach of using .incomplete files - static func incompletePath(for destination: URL) -> URL { - destination.appendingPathExtension("incomplete") - } - /// Check if an incomplete file exists for the destination and returns its size /// - Parameter destination: The destination URL for the download /// - Returns: Size of the incomplete file if it exists, otherwise 0 - static func checkForIncompleteFile(at destination: URL) -> Int { - let incompletePath = Self.incompletePath(for: destination) - + static func incompleteFileSize(at incompletePath: URL) -> Int { if FileManager.default.fileExists(atPath: incompletePath.path) { - if let attributes = try? FileManager.default.attributesOfItem(atPath: incompletePath.path), - let fileSize = attributes[.size] as? Int - { + if let attributes = try? FileManager.default.attributesOfItem(atPath: incompletePath.path), let fileSize = attributes[.size] as? Int { return fileSize } } @@ -63,29 +52,22 @@ class Downloader: NSObject, ObservableObject { init( from url: URL, to destination: URL, + incompleteDestination: URL, using authToken: String? = nil, inBackground: Bool = false, - resumeSize: Int = 0, // Can be specified manually, but will also check for incomplete files headers: [String: String]? = nil, expectedSize: Int? = nil, timeout: TimeInterval = 10, numRetries: Int = 5 ) { self.destination = destination - sourceURL = url self.expectedSize = expectedSize // Create incomplete file path based on destination - tempFilePath = Downloader.incompletePath(for: destination) + self.tempFilePath = incompleteDestination // If resume size wasn't specified, check for an existing incomplete file - let actualResumeSize: Int = if resumeSize > 0 { - resumeSize - } else { - Downloader.checkForIncompleteFile(at: destination) - } - - downloadedSize = actualResumeSize + let resumeSize = Self.incompleteFileSize(at: incompleteDestination) super.init() let sessionIdentifier = "swift-transformers.hub.downloader" @@ -99,7 +81,7 @@ class Downloader: NSObject, ObservableObject { urlSession = URLSession(configuration: config, delegate: self, delegateQueue: nil) - setUpDownload(from: url, with: authToken, resumeSize: actualResumeSize, headers: headers, expectedSize: expectedSize, timeout: timeout, numRetries: numRetries) + setUpDownload(from: url, with: authToken, resumeSize: resumeSize, headers: headers, expectedSize: expectedSize, timeout: timeout, numRetries: numRetries) } /// Sets up and initiates a file download operation @@ -139,25 +121,6 @@ class Downloader: NSObject, ObservableObject { Task { do { - // Check if incomplete file exists and get its size - var existingSize = 0 - guard let incompleteFilePath = self.tempFilePath else { - throw DownloadError.unexpectedError - } - - let fileManager = FileManager.default - if fileManager.fileExists(atPath: incompleteFilePath.path) { - let attributes = try fileManager.attributesOfItem(atPath: incompleteFilePath.path) - existingSize = attributes[.size] as? Int ?? 0 - self.downloadedSize = existingSize - } else { - // Create parent directory if needed - try fileManager.createDirectory(at: incompleteFilePath.deletingLastPathComponent(), withIntermediateDirectories: true) - - // Create empty incomplete file - fileManager.createFile(atPath: incompleteFilePath.path, contents: nil) - } - // Set up the request with appropriate headers var request = URLRequest(url: url) var requestHeaders = headers ?? [:] @@ -167,12 +130,12 @@ class Downloader: NSObject, ObservableObject { } // Set Range header if we're resuming - if existingSize > 0 { - requestHeaders["Range"] = "bytes=\(existingSize)-" + if resumeSize > 0 { + requestHeaders["Range"] = "bytes=\(resumeSize)-" // Calculate and show initial progress if let expectedSize, expectedSize > 0 { - let initialProgress = Double(existingSize) / Double(expectedSize) + let initialProgress = Double(resumeSize) / Double(expectedSize) self.downloadState.value = .downloading(initialProgress) } else { self.downloadState.value = .downloading(0) @@ -185,10 +148,10 @@ class Downloader: NSObject, ObservableObject { request.allHTTPHeaderFields = requestHeaders // Open the incomplete file for writing - let tempFile = try FileHandle(forWritingTo: incompleteFilePath) + let tempFile = try FileHandle(forWritingTo: self.tempFilePath) // If resuming, seek to end of file - if existingSize > 0 { + if resumeSize > 0 { try tempFile.seekToEnd() } @@ -197,7 +160,7 @@ class Downloader: NSObject, ObservableObject { // Clean up and move the completed download to its final destination tempFile.closeFile() - try fileManager.moveDownloadedFile(from: incompleteFilePath, to: self.destination) + try FileManager.default.moveDownloadedFile(from: self.tempFilePath, to: self.destination) self.downloadState.value = .completed(self.destination) } catch { self.downloadState.value = .failed(error) diff --git a/Sources/Hub/HubApi.swift b/Sources/Hub/HubApi.swift index d3a1e3d..1430229 100644 --- a/Sources/Hub/HubApi.swift +++ b/Sources/Hub/HubApi.swift @@ -361,6 +361,10 @@ public extension HubApi { repoMetadataDestination.appending(path: relativeFilename + ".metadata") } + var incompleteDestination: URL { + repoMetadataDestination.appending(path: relativeFilename + ".incomplete") + } + var downloaded: Bool { FileManager.default.fileExists(atPath: destination.path) } @@ -370,9 +374,13 @@ public extension HubApi { try FileManager.default.createDirectory(at: directoryURL, withIntermediateDirectories: true, attributes: nil) } - func prepareMetadataDestination() throws { - let directoryURL = metadataDestination.deletingLastPathComponent() + // We're using incomplete destination to prepare cache destination because incomplete files include lfs + non-lfs files (vs only lfs for metadata files) + func prepareCacheDestination() throws { + let directoryURL = incompleteDestination.deletingLastPathComponent() try FileManager.default.createDirectory(at: directoryURL, withIntermediateDirectories: true, attributes: nil) + if !FileManager.default.fileExists(atPath: incompleteDestination.path) { + try "".write(to: incompleteDestination, atomically: true, encoding: .utf8) + } } /// Note we go from Combine in Downloader to callback-based progress reporting @@ -423,24 +431,14 @@ public extension HubApi { // Otherwise, let's download the file! try prepareDestination() - try prepareMetadataDestination() + try prepareCacheDestination() - // Check for an existing incomplete file - let incompleteFile = Downloader.incompletePath(for: destination) - var resumeSize = 0 - - if FileManager.default.fileExists(atPath: incompleteFile.path) { - if let fileAttributes = try? FileManager.default.attributesOfItem(atPath: incompleteFile.path) { - resumeSize = (fileAttributes[FileAttributeKey.size] as? Int) ?? 0 - } - } - let downloader = Downloader( from: source, to: destination, + incompleteDestination: incompleteDestination, using: hfToken, inBackground: backgroundSession, - resumeSize: resumeSize, expectedSize: remoteSize ) diff --git a/Tests/HubTests/DownloaderTests.swift b/Tests/HubTests/DownloaderTests.swift index 1ef7c0f..3e7df3e 100644 --- a/Tests/HubTests/DownloaderTests.swift +++ b/Tests/HubTests/DownloaderTests.swift @@ -59,9 +59,16 @@ final class DownloaderTests: XCTestCase { """ + let cacheDir = tempDir.appendingPathComponent("cache") + try? FileManager.default.createDirectory(at: cacheDir, withIntermediateDirectories: true) + + let incompleteDestination = cacheDir.appendingPathComponent("config.json.incomplete") + FileManager.default.createFile(atPath: incompleteDestination.path, contents: nil, attributes: nil) + let downloader = Downloader( from: url, - to: destination + to: destination, + incompleteDestination: incompleteDestination ) // Store subscriber outside the continuation to maintain its lifecycle @@ -95,10 +102,17 @@ final class DownloaderTests: XCTestCase { let url = URL(string: "https://huggingface.co/coreml-projects/Llama-2-7b-chat-coreml/resolve/main/config.json")! let destination = tempDir.appendingPathComponent("config.json") + let cacheDir = tempDir.appendingPathComponent("cache") + try? FileManager.default.createDirectory(at: cacheDir, withIntermediateDirectories: true) + + let incompleteDestination = cacheDir.appendingPathComponent("config.json.incomplete") + FileManager.default.createFile(atPath: incompleteDestination.path, contents: nil, attributes: nil) + // Create downloader with incorrect expected size let downloader = Downloader( from: url, to: destination, + incompleteDestination: incompleteDestination, expectedSize: 999999 // Incorrect size ) @@ -120,10 +134,17 @@ final class DownloaderTests: XCTestCase { // Create parent directory if it doesn't exist try FileManager.default.createDirectory(at: destination.deletingLastPathComponent(), withIntermediateDirectories: true) - + + let cacheDir = tempDir.appendingPathComponent("cache") + try? FileManager.default.createDirectory(at: cacheDir, withIntermediateDirectories: true) + + let incompleteDestination = cacheDir.appendingPathComponent("config.json.incomplete") + FileManager.default.createFile(atPath: incompleteDestination.path, contents: nil, attributes: nil) + let downloader = Downloader( from: url, to: destination, + incompleteDestination: incompleteDestination, expectedSize: 73194001 // Correct size for verification ) @@ -168,28 +189,4 @@ final class DownloaderTests: XCTestCase { throw error } } - - func testAutomaticIncompleteFileDetection() throws { - let url = URL(string: "https://huggingface.co/coreml-projects/sam-2-studio/resolve/main/SAM%202%20Studio%201.1.zip")! - let destination = tempDir.appendingPathComponent("SAM%202%20Studio%201.1.zip") - - // Create a sample incomplete file with test content - let incompletePath = Downloader.incompletePath(for: destination) - try FileManager.default.createDirectory(at: incompletePath.deletingLastPathComponent(), withIntermediateDirectories: true) - let testContent = Data(repeating: 65, count: 1024) // 1KB of data - FileManager.default.createFile(atPath: incompletePath.path, contents: testContent) - - // Create a downloader for the same destination - // It should automatically detect and use the incomplete file - let downloader = Downloader( - from: url, - to: destination - ) - - // Verify the downloader found and is using the incomplete file - XCTAssertEqual(downloader.downloadedSize, 1024, "Should have detected the incomplete file and set resumeSize") - - // Clean up - try? FileManager.default.removeItem(at: incompletePath) - } } diff --git a/Tests/HubTests/HubApiTests.swift b/Tests/HubTests/HubApiTests.swift index 451816f..cf2c253 100644 --- a/Tests/HubTests/HubApiTests.swift +++ b/Tests/HubTests/HubApiTests.swift @@ -968,4 +968,78 @@ class SnapshotDownloadTests: XCTestCase { XCTFail("Unexpected error: \(error)") } } + + func testResumeDownloadFromEmptyIncomplete() async throws { + let hubApi = HubApi(downloadBase: downloadDestination) + var lastProgress: Progress? = nil + var downloadedTo = FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent("Library/Caches/huggingface-tests/models/coreml-projects/Llama-2-7b-chat-coreml") + + let metadataDestination = downloadedTo.appending(component: ".cache/huggingface/download") + + try FileManager.default.createDirectory(at: metadataDestination, withIntermediateDirectories: true, attributes: nil) + try "".write(to: metadataDestination.appendingPathComponent("config.json.incomplete"), atomically: true, encoding: .utf8) + downloadedTo = try await hubApi.snapshot(from: repo, matching: "config.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 fileContents = try String(contentsOfFile: downloadedTo.appendingPathComponent("config.json").path) + + let expected = """ + { + "architectures": [ + "LlamaForCausalLM" + ], + "bos_token_id": 1, + "eos_token_id": 2, + "model_type": "llama", + "pad_token_id": 0, + "vocab_size": 32000 + } + + """ + XCTAssertTrue(fileContents.contains(expected)) + } + + func testResumeDownloadFromNonEmptyIncomplete() async throws { + let hubApi = HubApi(downloadBase: downloadDestination) + var lastProgress: Progress? = nil + var downloadedTo = FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent("Library/Caches/huggingface-tests/models/coreml-projects/Llama-2-7b-chat-coreml") + + let metadataDestination = downloadedTo.appending(component: ".cache/huggingface/download") + + try FileManager.default.createDirectory(at: metadataDestination, withIntermediateDirectories: true, attributes: nil) + try "X".write(to: metadataDestination.appendingPathComponent("config.json.incomplete"), atomically: true, encoding: .utf8) + downloadedTo = try await hubApi.snapshot(from: repo, matching: "config.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 fileContents = try String(contentsOfFile: downloadedTo.appendingPathComponent("config.json").path) + + let expected = """ + X + "architectures": [ + "LlamaForCausalLM" + ], + "bos_token_id": 1, + "eos_token_id": 2, + "model_type": "llama", + "pad_token_id": 0, + "vocab_size": 32000 + } + + """ + XCTAssertTrue(fileContents.contains(expected)) + } } From fbe20cc564890b0f89e5ffff1607e6cba7ba1db0 Mon Sep 17 00:00:00 2001 From: Arda Atahan Ibis Date: Wed, 23 Apr 2025 12:59:37 -0700 Subject: [PATCH 4/7] add download cancellation handling and fix download progress update --- Sources/Hub/Downloader.swift | 5 +++- Sources/Hub/HubApi.swift | 40 ++++++++++++++++++-------------- Tests/HubTests/HubApiTests.swift | 21 +++++++++++++++++ 3 files changed, 48 insertions(+), 18 deletions(-) diff --git a/Sources/Hub/Downloader.swift b/Sources/Hub/Downloader.swift index 1f27c4b..26708b0 100644 --- a/Sources/Hub/Downloader.swift +++ b/Sources/Hub/Downloader.swift @@ -129,6 +129,8 @@ class Downloader: NSObject, ObservableObject { requestHeaders["Authorization"] = "Bearer \(authToken)" } + self.downloadedSize = resumeSize + // Set Range header if we're resuming if resumeSize > 0 { requestHeaders["Range"] = "bytes=\(resumeSize)-" @@ -207,7 +209,7 @@ class Downloader: NSObject, ObservableObject { throw DownloadError.unexpectedError } - downloadedSize = resumeSize +// downloadedSize = resumeSize // Create a buffer to collect bytes before writing to disk var buffer = Data(capacity: chunkSize) @@ -283,6 +285,7 @@ class Downloader: NSObject, ObservableObject { func cancel() { urlSession?.invalidateAndCancel() + downloadState.value = .failed(URLError(.cancelled)) } } diff --git a/Sources/Hub/HubApi.swift b/Sources/Hub/HubApi.swift index 1430229..d5f5d2a 100644 --- a/Sources/Hub/HubApi.swift +++ b/Sources/Hub/HubApi.swift @@ -442,25 +442,31 @@ public extension HubApi { expectedSize: remoteSize ) - let downloadSubscriber = downloader.downloadState.sink { state in - switch state { - case let .downloading(progress): - progressHandler(progress) - case .completed, .failed, .notStarted: - break + return try await withTaskCancellationHandler { + let downloadSubscriber = downloader.downloadState.sink { state in + switch state { + case let .downloading(progress): + progressHandler(progress) + case .completed, .failed, .notStarted: + break + } } - } - do { - _ = try withExtendedLifetime(downloadSubscriber) { - try downloader.waitUntilDone() + do { + _ = try withExtendedLifetime(downloadSubscriber) { + do { try Task.checkCancellation() } catch { print("Task cancelled") } + try downloader.waitUntilDone() + } + + try HubApi.shared.writeDownloadMetadata(commitHash: remoteCommitHash, etag: remoteEtag, metadataPath: metadataDestination) + + return destination + } catch { + // If download fails, leave the incomplete file in place for future resume + throw error } - - try HubApi.shared.writeDownloadMetadata(commitHash: remoteCommitHash, etag: remoteEtag, metadataPath: metadataDestination) - - return destination - } catch { - // If download fails, leave the incomplete file in place for future resume - throw error + } onCancel: { + print("download canceled") + downloader.cancel() } } } diff --git a/Tests/HubTests/HubApiTests.swift b/Tests/HubTests/HubApiTests.swift index cf2c253..f2263bb 100644 --- a/Tests/HubTests/HubApiTests.swift +++ b/Tests/HubTests/HubApiTests.swift @@ -1042,4 +1042,25 @@ class SnapshotDownloadTests: XCTestCase { """ XCTAssertTrue(fileContents.contains(expected)) } + + func testRealDownloadInterruptionAndResumption() async throws { + // Use the DepthPro model weights file + let targetFile = "DepthProNormalizedInverseDepth.mlpackage/Data/com.apple.CoreML/weights/weight.bin" + let repo = "coreml-projects/DepthPro-coreml-normalized-inverse-depth" + let hubApi = HubApi(downloadBase: downloadDestination) + + // Create a task for the download + let downloadTask = Task { + try await hubApi.snapshot(from: repo, matching: targetFile) { progress in + print("Progress reached 1 \(progress.fractionCompleted * 100)%") + } + } + + try await Task.sleep(nanoseconds: 15_000_000_000) + downloadTask.cancel() + + try await hubApi.snapshot(from: repo, matching: targetFile) { progress in + print("Progress reached 2 \(progress.fractionCompleted * 100)%") + } + } } From e5fe3b8ee983a08a425ec1b43350d393199107a7 Mon Sep 17 00:00:00 2001 From: Arda Atahan Ibis Date: Thu, 24 Apr 2025 11:04:04 -0700 Subject: [PATCH 5/7] create target files right before moving and add etag to incomplete file names for versioning --- Sources/Hub/Downloader.swift | 18 +++-- Sources/Hub/HubApi.swift | 17 +--- Tests/HubTests/DownloaderTests.swift | 17 +++- Tests/HubTests/HubApiTests.swift | 117 ++++++++++++++------------- 4 files changed, 89 insertions(+), 80 deletions(-) diff --git a/Sources/Hub/Downloader.swift b/Sources/Hub/Downloader.swift index 26708b0..6b893cd 100644 --- a/Sources/Hub/Downloader.swift +++ b/Sources/Hub/Downloader.swift @@ -34,7 +34,7 @@ class Downloader: NSObject, ObservableObject { private(set) var expectedSize: Int? private(set) var downloadedSize: Int = 0 - private var urlSession: URLSession? = nil + internal var session: URLSession? = nil /// Check if an incomplete file exists for the destination and returns its size /// - Parameter destination: The destination URL for the download @@ -79,7 +79,7 @@ class Downloader: NSObject, ObservableObject { config.sessionSendsLaunchEvents = true } - urlSession = URLSession(configuration: config, delegate: self, delegateQueue: nil) + session = URLSession(configuration: config, delegate: self, delegateQueue: nil) setUpDownload(from: url, with: authToken, resumeSize: resumeSize, headers: headers, expectedSize: expectedSize, timeout: timeout, numRetries: numRetries) } @@ -103,7 +103,7 @@ class Downloader: NSObject, ObservableObject { timeout: TimeInterval, numRetries: Int ) { - urlSession?.getAllTasks { tasks in + session?.getAllTasks { tasks in // If there's an existing pending background task with the same URL, let it proceed. if let existing = tasks.filter({ $0.originalRequest?.url == url }).first { switch existing.state { @@ -189,7 +189,7 @@ class Downloader: NSObject, ObservableObject { numRetries: Int, expectedSize: Int? ) async throws { - guard let session = urlSession else { + guard let session = session else { throw DownloadError.unexpectedError } @@ -209,8 +209,6 @@ class Downloader: NSObject, ObservableObject { throw DownloadError.unexpectedError } -// downloadedSize = resumeSize - // Create a buffer to collect bytes before writing to disk var buffer = Data(capacity: chunkSize) @@ -245,7 +243,7 @@ class Downloader: NSObject, ObservableObject { try await Task.sleep(nanoseconds: 1_000_000_000) let config = URLSessionConfiguration.default - self.urlSession = URLSession(configuration: config, delegate: self, delegateQueue: nil) + self.session = URLSession(configuration: config, delegate: self, delegateQueue: nil) try await httpGet( request: request, @@ -284,7 +282,7 @@ class Downloader: NSObject, ObservableObject { } func cancel() { - urlSession?.invalidateAndCancel() + session?.invalidateAndCancel() downloadState.value = .failed(URLError(.cancelled)) } } @@ -320,6 +318,10 @@ extension FileManager { if fileExists(atPath: dstURL.path) { try removeItem(at: dstURL) } + + let directoryURL = dstURL.deletingLastPathComponent() + try createDirectory(at: directoryURL, withIntermediateDirectories: true, attributes: nil) + try moveItem(at: srcURL, to: dstURL) } } diff --git a/Sources/Hub/HubApi.swift b/Sources/Hub/HubApi.swift index d5f5d2a..accfd81 100644 --- a/Sources/Hub/HubApi.swift +++ b/Sources/Hub/HubApi.swift @@ -361,21 +361,12 @@ public extension HubApi { repoMetadataDestination.appending(path: relativeFilename + ".metadata") } - var incompleteDestination: URL { - repoMetadataDestination.appending(path: relativeFilename + ".incomplete") - } - var downloaded: Bool { FileManager.default.fileExists(atPath: destination.path) } - func prepareDestination() throws { - let directoryURL = destination.deletingLastPathComponent() - try FileManager.default.createDirectory(at: directoryURL, withIntermediateDirectories: true, attributes: nil) - } - // We're using incomplete destination to prepare cache destination because incomplete files include lfs + non-lfs files (vs only lfs for metadata files) - func prepareCacheDestination() throws { + func prepareCacheDestination(_ incompleteDestination: URL) throws { let directoryURL = incompleteDestination.deletingLastPathComponent() try FileManager.default.createDirectory(at: directoryURL, withIntermediateDirectories: true, attributes: nil) if !FileManager.default.fileExists(atPath: incompleteDestination.path) { @@ -430,8 +421,8 @@ public extension HubApi { } // Otherwise, let's download the file! - try prepareDestination() - try prepareCacheDestination() + let incompleteDestination = repoMetadataDestination.appending(path: relativeFilename + ".\(remoteEtag).incomplete") + try prepareCacheDestination(incompleteDestination) let downloader = Downloader( from: source, @@ -453,7 +444,6 @@ public extension HubApi { } do { _ = try withExtendedLifetime(downloadSubscriber) { - do { try Task.checkCancellation() } catch { print("Task cancelled") } try downloader.waitUntilDone() } @@ -465,7 +455,6 @@ public extension HubApi { throw error } } onCancel: { - print("download canceled") downloader.cancel() } } diff --git a/Tests/HubTests/DownloaderTests.swift b/Tests/HubTests/DownloaderTests.swift index 3e7df3e..89ff611 100644 --- a/Tests/HubTests/DownloaderTests.swift +++ b/Tests/HubTests/DownloaderTests.swift @@ -24,6 +24,12 @@ enum DownloadError: LocalizedError { } } +fileprivate extension Downloader { + func interruptDownload() { + self.session?.invalidateAndCancel() + } +} + final class DownloaderTests: XCTestCase { var tempDir: URL! @@ -44,6 +50,7 @@ final class DownloaderTests: XCTestCase { func testSuccessfulDownload() async throws { // Create a test file let url = URL(string: "https://huggingface.co/coreml-projects/Llama-2-7b-chat-coreml/resolve/main/config.json")! + let etag = try await Hub.getFileMetadata(fileURL: url).etag! let destination = tempDir.appendingPathComponent("config.json") let fileContent = """ { @@ -62,7 +69,7 @@ final class DownloaderTests: XCTestCase { let cacheDir = tempDir.appendingPathComponent("cache") try? FileManager.default.createDirectory(at: cacheDir, withIntermediateDirectories: true) - let incompleteDestination = cacheDir.appendingPathComponent("config.json.incomplete") + let incompleteDestination = cacheDir.appendingPathComponent("config.json.\(etag).incomplete") FileManager.default.createFile(atPath: incompleteDestination.path, contents: nil, attributes: nil) let downloader = Downloader( @@ -100,12 +107,13 @@ final class DownloaderTests: XCTestCase { /// This test attempts to download with incorrect expected file, verifies the download fails, ensures no partial file is left behind func testDownloadFailsWithIncorrectSize() async throws { let url = URL(string: "https://huggingface.co/coreml-projects/Llama-2-7b-chat-coreml/resolve/main/config.json")! + let etag = try await Hub.getFileMetadata(fileURL: url).etag! let destination = tempDir.appendingPathComponent("config.json") let cacheDir = tempDir.appendingPathComponent("cache") try? FileManager.default.createDirectory(at: cacheDir, withIntermediateDirectories: true) - let incompleteDestination = cacheDir.appendingPathComponent("config.json.incomplete") + let incompleteDestination = cacheDir.appendingPathComponent("config.json.\(etag).incomplete") FileManager.default.createFile(atPath: incompleteDestination.path, contents: nil, attributes: nil) // Create downloader with incorrect expected size @@ -129,6 +137,7 @@ final class DownloaderTests: XCTestCase { /// verifies the download can resume and complete successfully, checks the final file exists and has content func testSuccessfulInterruptedDownload() async throws { let url = URL(string: "https://huggingface.co/coreml-projects/sam-2-studio/resolve/main/SAM%202%20Studio%201.1.zip")! + let etag = try await Hub.getFileMetadata(fileURL: url).etag! let destination = tempDir.appendingPathComponent("SAM%202%20Studio%201.1.zip") // Create parent directory if it doesn't exist @@ -138,7 +147,7 @@ final class DownloaderTests: XCTestCase { let cacheDir = tempDir.appendingPathComponent("cache") try? FileManager.default.createDirectory(at: cacheDir, withIntermediateDirectories: true) - let incompleteDestination = cacheDir.appendingPathComponent("config.json.incomplete") + let incompleteDestination = cacheDir.appendingPathComponent("config.json.\(etag).incomplete") FileManager.default.createFile(atPath: incompleteDestination.path, contents: nil, attributes: nil) let downloader = Downloader( @@ -163,7 +172,7 @@ final class DownloaderTests: XCTestCase { if threshold != 1.0, progress >= threshold { // Move to next threshold and interrupt threshold = threshold == 0.5 ? 0.75 : 1.0 - downloader.cancel() + downloader.interruptDownload() } case .completed: continuation.resume() diff --git a/Tests/HubTests/HubApiTests.swift b/Tests/HubTests/HubApiTests.swift index f2263bb..84182cc 100644 --- a/Tests/HubTests/HubApiTests.swift +++ b/Tests/HubTests/HubApiTests.swift @@ -173,9 +173,9 @@ class SnapshotDownloadTests: XCTestCase { let base = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first! return base.appending(component: "huggingface-tests") }() - + override func setUp() { } - + override func tearDown() { do { try FileManager.default.removeItem(at: downloadDestination) @@ -183,11 +183,11 @@ class SnapshotDownloadTests: XCTestCase { print("Can't remove test download destination \(downloadDestination), error: \(error)") } } - + func getRelativeFiles(url: URL, repo: String) -> [String] { var filenames: [String] = [] let prefix = downloadDestination.appending(path: "models/\(repo)").path.appending("/") - + if let enumerator = FileManager.default.enumerator(at: url, includingPropertiesForKeys: [.isRegularFileKey], options: [.skipsHiddenFiles], errorHandler: nil) { for case let fileURL as URL in enumerator { do { @@ -202,7 +202,7 @@ class SnapshotDownloadTests: XCTestCase { } return filenames } - + func testDownload() async throws { let hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil @@ -226,7 +226,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 6) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)")) - + XCTAssertEqual( Set(downloadedFilenames), Set([ @@ -237,7 +237,7 @@ 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 { @@ -251,7 +251,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 1) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)")) - + let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: repo) XCTAssertEqual( Set(downloadedFilenames), @@ -260,7 +260,7 @@ class SnapshotDownloadTests: XCTestCase { ]) ) } - + func testCustomEndpointDownload() async throws { let hubApi = HubApi(downloadBase: downloadDestination, endpoint: "https://hf-mirror.com") var lastProgress: Progress? = nil @@ -272,7 +272,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 6) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)")) - + let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: repo) XCTAssertEqual( Set(downloadedFilenames), @@ -326,7 +326,7 @@ class SnapshotDownloadTests: XCTestCase { func testDownloadFileMetadataExists() async throws { let hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil - + let downloadedTo = try await hubApi.snapshot(from: repo, matching: "*.json") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -336,7 +336,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 6) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)")) - + let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: repo) XCTAssertEqual( Set(downloadedFilenames), @@ -375,7 +375,7 @@ class SnapshotDownloadTests: XCTestCase { attributes = try FileManager.default.attributesOfItem(atPath: configPath.path) let secondDownloadTimestamp = attributes[.modificationDate] as! Date - + // File will not be downloaded again thus last modified date will remain unchanged XCTAssertTrue(originalTimestamp == secondDownloadTimestamp) } @@ -383,7 +383,7 @@ class SnapshotDownloadTests: XCTestCase { func testDownloadFileMetadataSame() async throws { let hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil - + let downloadedTo = try await hubApi.snapshot(from: repo, matching: "tokenizer.json") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -393,7 +393,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 1) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)")) - + let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: repo) XCTAssertEqual(Set(downloadedFilenames), Set(["tokenizer.json"])) @@ -430,7 +430,7 @@ class SnapshotDownloadTests: XCTestCase { func testDownloadFileMetadataCorrupted() async throws { let hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil - + let downloadedTo = try await hubApi.snapshot(from: repo, matching: "*.json") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -440,7 +440,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 6) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)")) - + let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: repo) XCTAssertEqual( Set(downloadedFilenames), @@ -483,7 +483,7 @@ class SnapshotDownloadTests: XCTestCase { attributes = try FileManager.default.attributesOfItem(atPath: configPath.path) let secondDownloadTimestamp = attributes[.modificationDate] as! Date - + // File will be downloaded again thus last modified date will change XCTAssertTrue(originalTimestamp != secondDownloadTimestamp) @@ -499,7 +499,7 @@ class SnapshotDownloadTests: XCTestCase { attributes = try FileManager.default.attributesOfItem(atPath: configPath.path) let thirdDownloadTimestamp = attributes[.modificationDate] as! Date - + // File will be downloaded again thus last modified date will change XCTAssertTrue(originalTimestamp != thirdDownloadTimestamp) } @@ -507,7 +507,7 @@ class SnapshotDownloadTests: XCTestCase { func testDownloadLargeFileMetadataCorrupted() async throws { let hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil - + let downloadedTo = try await hubApi.snapshot(from: repo, matching: "*.mlmodel") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -517,7 +517,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 1) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)")) - + let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: repo) XCTAssertEqual( Set(downloadedFilenames), @@ -552,7 +552,7 @@ class SnapshotDownloadTests: XCTestCase { attributes = try FileManager.default.attributesOfItem(atPath: modelPath.path) let thirdDownloadTimestamp = attributes[.modificationDate] as! Date - + // File will not be downloaded again because this is an LFS file. // While downloading LFS files, we first check if local file ETag is the same as remote ETag. // If that's the case we just update the metadata and keep the local file. @@ -577,7 +577,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 1) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)")) - + let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: repo) XCTAssertEqual(Set(downloadedFilenames), Set(["llama-2-7b-chat.mlpackage/Data/com.apple.CoreML/model.mlmodel"])) @@ -590,7 +590,7 @@ class SnapshotDownloadTests: XCTestCase { let metadataFile = metadataDestination.appendingPathComponent("llama-2-7b-chat.mlpackage/Data/com.apple.CoreML/model.mlmodel.metadata") let metadataString = try String(contentsOfFile: metadataFile.path) - + let expected = "eaf97358a37d03fd48e5a87d15aff2e8423c1afb\nfc329090bfbb2570382c9af997cffd5f4b78b39b8aeca62076db69534e020107" XCTAssertTrue(metadataString.contains(expected)) } @@ -607,7 +607,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 1) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(lfsRepo)")) - + let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: lfsRepo) XCTAssertEqual(Set(downloadedFilenames), Set(["x.bin"])) @@ -637,7 +637,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 1) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(lfsRepo)")) - + let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: lfsRepo) XCTAssertEqual(Set(downloadedFilenames), Set(["x.bin"])) @@ -665,7 +665,7 @@ class SnapshotDownloadTests: XCTestCase { func testLFSFileNoMetadata() async throws { let hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil - + let downloadedTo = try await hubApi.snapshot(from: lfsRepo, matching: "x.bin") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -675,7 +675,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 1) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(lfsRepo)")) - + let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: lfsRepo) XCTAssertEqual(Set(downloadedFilenames), Set(["x.bin"])) @@ -702,7 +702,7 @@ class SnapshotDownloadTests: XCTestCase { attributes = try FileManager.default.attributesOfItem(atPath: filePath.path) let secondDownloadTimestamp = attributes[.modificationDate] as! Date - + // File will not be downloaded again thus last modified date will remain unchanged XCTAssertTrue(originalTimestamp == secondDownloadTimestamp) XCTAssertTrue(FileManager.default.fileExists(atPath: metadataDestination.path)) @@ -716,7 +716,7 @@ class SnapshotDownloadTests: XCTestCase { func testLFSFileCorruptedMetadata() async throws { let hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil - + let downloadedTo = try await hubApi.snapshot(from: lfsRepo, matching: "x.bin") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -726,7 +726,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 1) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(lfsRepo)")) - + let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: lfsRepo) XCTAssertEqual(Set(downloadedFilenames), Set(["x.bin"])) @@ -744,7 +744,7 @@ class SnapshotDownloadTests: XCTestCase { let metadataFile = metadataDestination.appendingPathComponent("x.bin.metadata") try "a".write(to: metadataFile, atomically: true, encoding: .utf8) - + let _ = try await hubApi.snapshot(from: lfsRepo, matching: "x.bin") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -753,7 +753,7 @@ class SnapshotDownloadTests: XCTestCase { attributes = try FileManager.default.attributesOfItem(atPath: filePath.path) let secondDownloadTimestamp = attributes[.modificationDate] as! Date - + // File will not be downloaded again thus last modified date will remain unchanged XCTAssertTrue(originalTimestamp == secondDownloadTimestamp) XCTAssertTrue(FileManager.default.fileExists(atPath: metadataDestination.path)) @@ -767,7 +767,7 @@ class SnapshotDownloadTests: XCTestCase { func testNonLFSFileRedownload() async throws { let hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil - + let downloadedTo = try await hubApi.snapshot(from: repo, matching: "config.json") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -777,7 +777,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 1) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)")) - + let downloadedFilenames = getRelativeFiles(url: downloadDestination, repo: repo) XCTAssertEqual(Set(downloadedFilenames), Set(["config.json"])) @@ -804,7 +804,7 @@ class SnapshotDownloadTests: XCTestCase { attributes = try FileManager.default.attributesOfItem(atPath: filePath.path) let secondDownloadTimestamp = attributes[.modificationDate] as! Date - + // File will be downloaded again thus last modified date will change XCTAssertTrue(originalTimestamp != secondDownloadTimestamp) XCTAssertTrue(FileManager.default.fileExists(atPath: metadataDestination.path)) @@ -818,7 +818,7 @@ class SnapshotDownloadTests: XCTestCase { func testOfflineModeReturnsDestination() async throws { var hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil - + var downloadedTo = try await hubApi.snapshot(from: repo, matching: "*.json") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -831,7 +831,7 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)")) hubApi = HubApi(downloadBase: downloadDestination, useOfflineMode: true) - + downloadedTo = try await hubApi.snapshot(from: repo, matching: "*.json") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -845,7 +845,7 @@ class SnapshotDownloadTests: XCTestCase { func testOfflineModeThrowsError() async throws { let hubApi = HubApi(downloadBase: downloadDestination, useOfflineMode: true) - + do { try await hubApi.snapshot(from: repo, matching: "*.json") XCTFail("Expected an error to be thrown") @@ -901,7 +901,7 @@ class SnapshotDownloadTests: XCTestCase { func testOfflineModeWithCorruptedLFSMetadata() async throws { var hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil - + let downloadedTo = try await hubApi.snapshot(from: lfsRepo, matching: "*") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -918,7 +918,7 @@ class SnapshotDownloadTests: XCTestCase { try "77b984598d90af6143d73d5a2d6214b23eba7e27\n98ea6e4f216f2ab4b69fff9b3a44842c38686ca685f3f55dc48c5d3fb1107be4\n0\n".write(to: metadataDestination, atomically: true, encoding: .utf8) hubApi = HubApi(downloadBase: downloadDestination, useOfflineMode: true) - + do { try await hubApi.snapshot(from: lfsRepo, matching: "*") XCTFail("Expected an error to be thrown") @@ -937,7 +937,7 @@ class SnapshotDownloadTests: XCTestCase { func testOfflineModeWithNoFiles() async throws { var hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil - + let downloadedTo = try await hubApi.snapshot(from: lfsRepo, matching: "x.bin") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -953,7 +953,7 @@ class SnapshotDownloadTests: XCTestCase { try FileManager.default.removeItem(at: fileDestination) hubApi = HubApi(downloadBase: downloadDestination, useOfflineMode: true) - + do { try await hubApi.snapshot(from: lfsRepo, matching: "x.bin") XCTFail("Expected an error to be thrown") @@ -972,13 +972,15 @@ class SnapshotDownloadTests: XCTestCase { func testResumeDownloadFromEmptyIncomplete() async throws { let hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil - var downloadedTo = FileManager.default.homeDirectoryForCurrentUser - .appendingPathComponent("Library/Caches/huggingface-tests/models/coreml-projects/Llama-2-7b-chat-coreml") + var downloadedTo = FileManager.default.homeDirectoryForCurrentUser.appendingPathComponent("Library/Caches/huggingface-tests/models/coreml-projects/Llama-2-7b-chat-coreml") let metadataDestination = downloadedTo.appending(component: ".cache/huggingface/download") + let url = URL(string: "https://huggingface.co/coreml-projects/Llama-2-7b-chat-coreml/resolve/main/config.json")! + let etag = try await Hub.getFileMetadata(fileURL: url).etag! + try FileManager.default.createDirectory(at: metadataDestination, withIntermediateDirectories: true, attributes: nil) - try "".write(to: metadataDestination.appendingPathComponent("config.json.incomplete"), atomically: true, encoding: .utf8) + try "".write(to: metadataDestination.appendingPathComponent("config.json.\(etag).incomplete"), atomically: true, encoding: .utf8) downloadedTo = try await hubApi.snapshot(from: repo, matching: "config.json") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -1001,7 +1003,7 @@ class SnapshotDownloadTests: XCTestCase { "pad_token_id": 0, "vocab_size": 32000 } - + """ XCTAssertTrue(fileContents.contains(expected)) } @@ -1010,12 +1012,15 @@ class SnapshotDownloadTests: XCTestCase { let hubApi = HubApi(downloadBase: downloadDestination) var lastProgress: Progress? = nil var downloadedTo = FileManager.default.homeDirectoryForCurrentUser - .appendingPathComponent("Library/Caches/huggingface-tests/models/coreml-projects/Llama-2-7b-chat-coreml") + .appendingPathComponent("Library/Caches/huggingface-tests/models/coreml-projects/Llama-2-7b-chat-coreml") let metadataDestination = downloadedTo.appending(component: ".cache/huggingface/download") + let url = URL(string: "https://huggingface.co/coreml-projects/Llama-2-7b-chat-coreml/resolve/main/config.json")! + let etag = try await Hub.getFileMetadata(fileURL: url).etag! + try FileManager.default.createDirectory(at: metadataDestination, withIntermediateDirectories: true, attributes: nil) - try "X".write(to: metadataDestination.appendingPathComponent("config.json.incomplete"), atomically: true, encoding: .utf8) + try "X".write(to: metadataDestination.appendingPathComponent("config.json.\(etag).incomplete"), atomically: true, encoding: .utf8) downloadedTo = try await hubApi.snapshot(from: repo, matching: "config.json") { progress in print("Total Progress: \(progress.fractionCompleted)") print("Files Completed: \(progress.completedUnitCount) of \(progress.totalUnitCount)") @@ -1024,9 +1029,9 @@ class SnapshotDownloadTests: XCTestCase { XCTAssertEqual(lastProgress?.fractionCompleted, 1) XCTAssertEqual(lastProgress?.completedUnitCount, 1) XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)")) - + let fileContents = try String(contentsOfFile: downloadedTo.appendingPathComponent("config.json").path) - + let expected = """ X "architectures": [ @@ -1038,11 +1043,11 @@ class SnapshotDownloadTests: XCTestCase { "pad_token_id": 0, "vocab_size": 32000 } - + """ XCTAssertTrue(fileContents.contains(expected)) } - + func testRealDownloadInterruptionAndResumption() async throws { // Use the DepthPro model weights file let targetFile = "DepthProNormalizedInverseDepth.mlpackage/Data/com.apple.CoreML/weights/weight.bin" @@ -1059,8 +1064,12 @@ class SnapshotDownloadTests: XCTestCase { try await Task.sleep(nanoseconds: 15_000_000_000) downloadTask.cancel() - try await hubApi.snapshot(from: repo, matching: targetFile) { progress in + let downloadedTo = try await hubApi.snapshot(from: repo, matching: targetFile) { progress in print("Progress reached 2 \(progress.fractionCompleted * 100)%") } + + let filePath = downloadedTo.appendingPathComponent(targetFile) + XCTAssertTrue(FileManager.default.fileExists(atPath: filePath.path), + "Downloaded file should exist at \(filePath.path)") } } From 1967624f560fe99d75a4c49cb9ce0bdc3842e216 Mon Sep 17 00:00:00 2001 From: Arda Atahan Ibis Date: Fri, 25 Apr 2025 10:39:51 -0700 Subject: [PATCH 6/7] properly cancel downloader task and fix formatting --- Sources/Hub/Downloader.swift | 41 +++++++++++++++------------- Sources/Hub/HubApi.swift | 2 +- Tests/HubTests/DownloaderTests.swift | 4 +-- Tests/HubTests/HubApiTests.swift | 21 ++++++++++---- 4 files changed, 41 insertions(+), 27 deletions(-) diff --git a/Sources/Hub/Downloader.swift b/Sources/Hub/Downloader.swift index 6b893cd..bbdb2f4 100644 --- a/Sources/Hub/Downloader.swift +++ b/Sources/Hub/Downloader.swift @@ -34,21 +34,9 @@ class Downloader: NSObject, ObservableObject { private(set) var expectedSize: Int? private(set) var downloadedSize: Int = 0 - internal var session: URLSession? = nil + var session: URLSession? = nil + var downloadTask: Task? = nil - /// Check if an incomplete file exists for the destination and returns its size - /// - Parameter destination: The destination URL for the download - /// - Returns: Size of the incomplete file if it exists, otherwise 0 - static func incompleteFileSize(at incompletePath: URL) -> Int { - if FileManager.default.fileExists(atPath: incompletePath.path) { - if let attributes = try? FileManager.default.attributesOfItem(atPath: incompletePath.path), let fileSize = attributes[.size] as? Int { - return fileSize - } - } - - return 0 - } - init( from url: URL, to destination: URL, @@ -64,7 +52,7 @@ class Downloader: NSObject, ObservableObject { self.expectedSize = expectedSize // Create incomplete file path based on destination - self.tempFilePath = incompleteDestination + tempFilePath = incompleteDestination // If resume size wasn't specified, check for an existing incomplete file let resumeSize = Self.incompleteFileSize(at: incompleteDestination) @@ -84,6 +72,19 @@ class Downloader: NSObject, ObservableObject { setUpDownload(from: url, with: authToken, resumeSize: resumeSize, headers: headers, expectedSize: expectedSize, timeout: timeout, numRetries: numRetries) } + /// Check if an incomplete file exists for the destination and returns its size + /// - Parameter destination: The destination URL for the download + /// - Returns: Size of the incomplete file if it exists, otherwise 0 + static func incompleteFileSize(at incompletePath: URL) -> Int { + if FileManager.default.fileExists(atPath: incompletePath.path) { + if let attributes = try? FileManager.default.attributesOfItem(atPath: incompletePath.path), let fileSize = attributes[.size] as? Int { + return fileSize + } + } + + return 0 + } + /// Sets up and initiates a file download operation /// /// - Parameters: @@ -119,7 +120,7 @@ class Downloader: NSObject, ObservableObject { } } - Task { + self.downloadTask = Task { do { // Set up the request with appropriate headers var request = URLRequest(url: url) @@ -157,11 +158,12 @@ class Downloader: NSObject, ObservableObject { try tempFile.seekToEnd() } - defer { tempFile.closeFile() } try await self.httpGet(request: request, tempFile: tempFile, resumeSize: self.downloadedSize, numRetries: numRetries, expectedSize: expectedSize) // Clean up and move the completed download to its final destination tempFile.closeFile() + + try Task.checkCancellation() try FileManager.default.moveDownloadedFile(from: self.tempFilePath, to: self.destination) self.downloadState.value = .completed(self.destination) } catch { @@ -189,7 +191,7 @@ class Downloader: NSObject, ObservableObject { numRetries: Int, expectedSize: Int? ) async throws { - guard let session = session else { + guard let session else { throw DownloadError.unexpectedError } @@ -283,6 +285,7 @@ class Downloader: NSObject, ObservableObject { func cancel() { session?.invalidateAndCancel() + downloadTask?.cancel() downloadState.value = .failed(URLError(.cancelled)) } } @@ -315,7 +318,7 @@ extension Downloader: URLSessionDownloadDelegate { extension FileManager { func moveDownloadedFile(from srcURL: URL, to dstURL: URL) throws { - if fileExists(atPath: dstURL.path) { + if fileExists(atPath: dstURL.path()) { try removeItem(at: dstURL) } diff --git a/Sources/Hub/HubApi.swift b/Sources/Hub/HubApi.swift index 9a98139..a9d5020 100644 --- a/Sources/Hub/HubApi.swift +++ b/Sources/Hub/HubApi.swift @@ -366,7 +366,7 @@ public extension HubApi { FileManager.default.fileExists(atPath: destination.path) } - // We're using incomplete destination to prepare cache destination because incomplete files include lfs + non-lfs files (vs only lfs for metadata files) + /// We're using incomplete destination to prepare cache destination because incomplete files include lfs + non-lfs files (vs only lfs for metadata files) func prepareCacheDestination(_ incompleteDestination: URL) throws { let directoryURL = incompleteDestination.deletingLastPathComponent() try FileManager.default.createDirectory(at: directoryURL, withIntermediateDirectories: true, attributes: nil) diff --git a/Tests/HubTests/DownloaderTests.swift b/Tests/HubTests/DownloaderTests.swift index 89ff611..0f8d106 100644 --- a/Tests/HubTests/DownloaderTests.swift +++ b/Tests/HubTests/DownloaderTests.swift @@ -24,9 +24,9 @@ enum DownloadError: LocalizedError { } } -fileprivate extension Downloader { +private extension Downloader { func interruptDownload() { - self.session?.invalidateAndCancel() + session?.invalidateAndCancel() } } diff --git a/Tests/HubTests/HubApiTests.swift b/Tests/HubTests/HubApiTests.swift index 84182cc..5465918 100644 --- a/Tests/HubTests/HubApiTests.swift +++ b/Tests/HubTests/HubApiTests.swift @@ -1003,7 +1003,7 @@ class SnapshotDownloadTests: XCTestCase { "pad_token_id": 0, "vocab_size": 32000 } - + """ XCTAssertTrue(fileContents.contains(expected)) } @@ -1043,27 +1043,38 @@ class SnapshotDownloadTests: XCTestCase { "pad_token_id": 0, "vocab_size": 32000 } - + """ XCTAssertTrue(fileContents.contains(expected)) } func testRealDownloadInterruptionAndResumption() async throws { // Use the DepthPro model weights file - let targetFile = "DepthProNormalizedInverseDepth.mlpackage/Data/com.apple.CoreML/weights/weight.bin" - let repo = "coreml-projects/DepthPro-coreml-normalized-inverse-depth" + let targetFile = "SAM 2 Studio 1.1.zip" + let repo = "coreml-projects/sam-2-studio" let hubApi = HubApi(downloadBase: downloadDestination) + // Create expectation for first progress update + let progressExpectation = expectation(description: "First progress update received") + // Create a task for the download let downloadTask = Task { try await hubApi.snapshot(from: repo, matching: targetFile) { progress in print("Progress reached 1 \(progress.fractionCompleted * 100)%") + if progress.fractionCompleted > 0 { + progressExpectation.fulfill() + } } } - try await Task.sleep(nanoseconds: 15_000_000_000) + // Wait for the first progress update + await fulfillment(of: [progressExpectation], timeout: 30.0) + + // Cancel the download once we've seen progress downloadTask.cancel() + try await Task.sleep(nanoseconds: 5_000_000_000) + // Resume download with a new task let downloadedTo = try await hubApi.snapshot(from: repo, matching: targetFile) { progress in print("Progress reached 2 \(progress.fractionCompleted * 100)%") } From ccb411cdc13b052e9f0fff73787345037ec0e985 Mon Sep 17 00:00:00 2001 From: Arda Atahan Ibis Date: Mon, 28 Apr 2025 11:54:38 -0700 Subject: [PATCH 7/7] remove unnecessary tabs in test cases --- Tests/HubTests/HubApiTests.swift | 2 -- 1 file changed, 2 deletions(-) diff --git a/Tests/HubTests/HubApiTests.swift b/Tests/HubTests/HubApiTests.swift index 5465918..6b0d150 100644 --- a/Tests/HubTests/HubApiTests.swift +++ b/Tests/HubTests/HubApiTests.swift @@ -1003,7 +1003,6 @@ class SnapshotDownloadTests: XCTestCase { "pad_token_id": 0, "vocab_size": 32000 } - """ XCTAssertTrue(fileContents.contains(expected)) } @@ -1043,7 +1042,6 @@ class SnapshotDownloadTests: XCTestCase { "pad_token_id": 0, "vocab_size": 32000 } - """ XCTAssertTrue(fileContents.contains(expected)) }