add metadata and resumable download support with tests #2

ardaatahan · 2024-12-29T04:42:10Z

This PR aims to bring over some of huggingface_hub Python library's capabilities to swift-transformers:

Disables automatic redirection during head requests to fix getFileMetadata, which currently doesn't return expected commit hash and etag for lfs files.
Adds resumable downloads by downloading HuggingFace files in async bytes, stores them in 10MB buffers and writes buffers to memory.
Adds retry logic to downloads to mitigate connection errors during downloading.
Adds file metadata support to prevent unnecessary redownloads and ensure file integrity.
Adds additional tests to verify metadata and download logic.

BlaiseMuhirwa · 2025-01-02T17:30:06Z

Sources/Hub/Downloader.swift

+        from url: URL,
+        to destination: URL,
+        metadataDirURL: URL,
+        using authToken: String? = nil,


Why are we allowing this to potentially be nil? Is there a case where you can download from HF without providing an auth token?

it was made optional by huggingface. you don't need and auth token to download from public repos

BlaiseMuhirwa · 2025-01-02T17:38:33Z

Sources/Hub/Downloader.swift

+            if newNumRetries <= 0 {
+                throw error
+            }
+            try await Task.sleep(nanoseconds: 1_000_000_000)


Why not just specify one second? It's less confusing and doesn't require someone to do some quick math to figure out the conversion.

BlaiseMuhirwa · 2025-01-02T17:54:54Z

Sources/Hub/Downloader.swift

+            try await downloadWithStreaming(
+                request: request,
+                resumeSize: downloadedSize,
+                numRetries: newNumRetries - 1,
+                expectedSize: expectedSize
+            )


This means we would be retrying the request on every error, but not every error is retryable. For instance, if the HF API returns 400, 401 or 403, we shouldn't retry that request because it's never going to succeed. We should only retry the request if the response from HF is in the [500, 599] range, which is a server-side error, which can be transient (hence, retryable).

Also, generally recommend against implementing this with recursion since it's slightly less readable compared to a simple iterative solution (where we try/catch the error and retry the request until you reach the preset number of retries).

checking for error ranges makes sense. I primarily did it this way to follow the design decisions made in the python library, which uses recursion for retries.

BlaiseMuhirwa · 2025-01-02T18:44:34Z

Sources/Hub/HubApi.swift

+                let contents = try String(contentsOf: metadataPath, encoding: .utf8)
+                let lines = contents.components(separatedBy: .newlines)
+
+                guard lines.count == 4 else {


Why are we hard-coding 4 here? Would be useful to add a comment at the top of the function explaining why we expect the line count to be exactly 4 in the metadata file. Otherwise, if the line count can be something other than 4, maybe make it a parameter?

BlaiseMuhirwa · 2025-01-02T18:57:35Z

Sources/Hub/HubApi.swift

+        }
+
+        func isValidSHA256(_ hash: String) -> Bool {
+            let sha256Pattern = "^[0-9a-f]{64}$"


Sha-256 hashes can also contain upper case letters. We should modify the regex to account for them so we don't have false negatives.

afaik, sha256 hashes are usually lower case only. at least that's the format huggingface uses to serve the hashes for large files, so I reused their pattern from the python library.

Ah, never mind, I forgot that this is a hexadecimal encoding and not a base64 encoding. With hex, 0xab37 is not different from 0xAB37, so we're good.

BlaiseMuhirwa · 2025-01-02T18:58:15Z

Sources/Hub/HubApi.swift

+            let localCommitHash = localMetadata?.commitHash ?? ""
+            let remoteCommitHash = remoteMetadata.commitHash ?? ""
+
+            // Local file exists + metadata exists + commit_hash matches => return file


BlaiseMuhirwa · 2025-01-02T19:02:55Z

Tests/HubTests/HubApiTests.swift

+    func testDownloadFileMetadata() 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)")
+            lastProgress = progress
+        }
+        XCTAssertEqual(lastProgress?.fractionCompleted, 1)
+        XCTAssertEqual(lastProgress?.completedUnitCount, 6)
+        XCTAssertEqual(downloadedTo, downloadDestination.appending(path: "models/\(repo)"))
+
+        let downloadedFilenames = getRelativeFiles(url: downloadDestination)
+        XCTAssertEqual(
+            Set(downloadedFilenames),
+            Set([
+                "config.json", "tokenizer.json", "tokenizer_config.json",
+                "llama-2-7b-chat.mlpackage/Manifest.json",
+                "llama-2-7b-chat.mlpackage/Data/com.apple.CoreML/FeatureDescriptions.json",
+                "llama-2-7b-chat.mlpackage/Data/com.apple.CoreML/Metadata.json",
+            ])
+        )
+
+        let metadataDestination = downloadedTo.appending(component: ".cache/huggingface")
+        let downloadedMetadataFilenames = getRelativeFiles(url: metadataDestination)
+        XCTAssertEqual(
+            Set(downloadedMetadataFilenames),
+            Set([
+                ".cache/huggingface/config.metadata",


This is a very good test.

BlaiseMuhirwa · 2025-01-06T23:58:44Z

Sources/Hub/HubApi.swift

+            let remoteCommitHash = remoteMetadata.commitHash ?? ""
+
+            // Local file exists + metadata exists + commit_hash matches => return file
+            if isValidSHA256(remoteCommitHash) && downloaded && localMetadata != nil && localCommitHash == remoteCommitHash {


Suggested change

if isValidSHA256(remoteCommitHash) && downloaded && localMetadata != nil && localCommitHash == remoteCommitHash {

if downloaded && localMetadata != nil && localCommitHash == remoteCommitHash {

Here isValidSHA256 is actually immaterial since we're also checking whether the local commit hash is the same as the remote one. There is no case where the hash is not a real SHA hash and the condition localCommitHash == remoteCommitHash is true.

ZachNagengast

Looks good to go

ZachNagengast · 2025-01-15T23:35:04Z

Approved, closing in favor of upstream PR

ardaatahan added 6 commits December 28, 2024 20:33

add metadata and resumable download support with tests

02d2571

fix download progress handling

f590932

fix httpHead by disabling redirection for lfs

22b6892

remove unnecessary test case

bedfc7a

fix metadata location for lfs files and relative redirects

af26e60

add test case for large file metadata

b4e1c49

ardaatahan requested review from ZachNagengast, a2they and BlaiseMuhirwa January 2, 2025 17:14

BlaiseMuhirwa reviewed Jan 2, 2025

View reviewed changes

preserve original repo file structure for metadata files

26707b8

BlaiseMuhirwa reviewed Jan 6, 2025

View reviewed changes

ardaatahan and others added 4 commits January 14, 2025 16:43

add test case for lfs and more comments

5839d33

Minor cleanup

9d39cf1

Merge branch 'main' into metadata

97b6163

Remove resumable downloads for future PR

fe2f32b

ZachNagengast approved these changes Jan 15, 2025

View reviewed changes

Cleanup unused variable

30adb75

ZachNagengast closed this Jan 15, 2025

	if isValidSHA256(remoteCommitHash) && downloaded && localMetadata != nil && localCommitHash == remoteCommitHash {
	if downloaded && localMetadata != nil && localCommitHash == remoteCommitHash {

add metadata and resumable download support with tests #2

add metadata and resumable download support with tests #2

Uh oh!

Conversation

ardaatahan commented Dec 29, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ZachNagengast left a comment

Choose a reason for hiding this comment

Uh oh!

ZachNagengast commented Jan 15, 2025

Uh oh!

Uh oh!

ardaatahan commented Dec 29, 2024 •

edited

Loading