Skip to content

Download large model shards as parallel HTTP Range requests#1960

Open
jjang-ai wants to merge 1 commit into
mainfrom
perf/chunked-hf-download
Open

Download large model shards as parallel HTTP Range requests#1960
jjang-ai wants to merge 1 commit into
mainfrom
perf/chunked-hf-download

Conversation

@jjang-ai

Copy link
Copy Markdown
Contributor

Problem

A single-shard repo — e.g. mlx-community/Qwen3-8B-4bit, one 4.6 GB model.safetensors — downloads over exactly one connection. Per-connection throughput to the Hugging Face CDN caps well below most links, so no amount of waiting makes it faster. maxConcurrentFiles = 3 only parallelizes across files, so single-shard repos get nothing from it.

Two failure modes compound it, and these hit every repo:

  • On retry, resumeDataForAttempt = nil restarts the file from byte zero. Resume data is in-memory only, so quitting the app throws away the partial file. A 4.6 GB download that dies at 90% starts over, four times, then fails.
  • Integrity is checked by file size alone. A size-correct but corrupt file passes.

Change

Large LFS objects stream as concurrent HTTP Range requests into a pre-allocated .part file. Small files and any server that doesn't advertise Range fall back to the existing single-task path unchanged.

Three measured details drive the design:

  1. URLSession multiplexes concurrent HTTP/2 requests to one host onto a single TCP connection. Sharing a session across chunks collapses them back onto one connection and wins nothing (81 vs 112 MB/s). Each lane owns its own session. Raising httpMaximumConnectionsPerHost does not substitute for this.
  2. A session per chunk is just as bad — every chunk then pays a fresh TLS handshake, redirect, and TCP slow-start ramp. That measured 1.20×. Reusing connections across chunks measured 1.9×. So a fixed set of lanes drain a shared chunk queue.
  3. Lanes are capped app-wide by LaneBudget. Without it, 3 concurrent files × 8 lanes = 24 sockets, which on a thin link just compete for the same bottleneck and add queueing delay. Reservation never blocks, so a small file can't starve behind a large one.

Also folded in, because chunking makes them nearly free:

  • Durable resume. Completed chunks are recorded in a .part.json manifest, so an interrupted transfer resumes instead of restarting — across a failed retry and across an app relaunch. A dead chunk now costs 32 MB, not 4.6 GB.
  • Real integrity. Chunks are pinned to the resolved commit SHA (x-repo-commit), so a push mid-download can't splice two revisions together — a corruption the size-only check would never catch. The assembled file is verified against the sha256 Hugging Face already publishes in x-linked-etag.

Measurements

Lane count, end-to-end on the 4.6 GB shard:

Lanes Throughput
4 106.9 MB/s
6 146.6 MB/s
8 (chosen) 167.4 MB/s
12 180.7 MB/s
16 184.9 MB/s

Eight buys ~90% of the achievable gain at half the sockets of sixteen; past twelve is inside the noise. On the 4.6 GB shard that's ~50s → ~26s, sha256-verified.

Caveat, stated plainly: the single-connection cap is congestion-dependent, not a fixed ceiling — the legacy path measured anywhere from 34 to 97 MB/s across runs on the same machine. The direction of the win is consistent in every rep, but its magnitude varies with conditions. Users who are genuinely link-bound (rather than connection-bound) will see little throughput change; for them the win is the resume behavior.

All numbers are from one machine on one link.

Testing

  • swift test --filter ChunkedFileDownloaderTests — 14 tests. Unit coverage for commit pinning, the sha256 gate, chunkability, and the lane budget.
  • Live transfer tests are gated behind RUN_LIVE_HF_DOWNLOAD=1, because chunked range semantics are exactly what a mock would get wrong. They assert that parallel chunks reassemble to Hugging Face's published sha256, that an interrupted transfer resumes from non-zero bytes, and that pause surfaces as PauseInfo rather than a failure — tearing down the lanes' sessions makes the in-flight chunk throw URLError.cancelled, which without classification would have marked the model .failed.
  • Existing ModelDownloadRetryPolicyTests, ModelDownloadServiceStorageTests, HuggingFaceDownloadPathTests pass unchanged.
  • Full OsaurusCore suite: 4,986 passed. The 5 failures (tool-exposure ×2, chat-session-store flush, gemma4 tokenizer, path resolver) reproduce on clean main at this same commit and are order-dependent / pre-existing; none touch the download path.

Scope

PrivacyFilterModelDownloader and ImageModelDownloadService have their own downloaders and are untouched. The privacy-filter one is still strictly serial.

Follow-up (not in this PR)

Authenticated requests 302 to cas-bridge.xethub.hf.co (Xet) instead of us.aws.cdn.hf.co (CloudFront), and measured ~half the single-connection throughput in 3/3 reps (37.9 vs 68.7 MB/s median). Users who add an HF token may be getting slower downloads on the current single-connection path. Chunking largely washes the difference out, so it's out of scope here, but it's worth re-testing from another network and filing separately.

A single-shard repo (one 4.6 GB model.safetensors) transferred over exactly
one connection, and per-connection throughput to the Hugging Face CDN caps
well below most links. maxConcurrentFiles only parallelizes across files, so
those repos got no benefit from it at all.

Large LFS objects now stream as concurrent Range requests into a pre-allocated
.part file. Measured end-to-end on a 4.6 GB shard: ~50s -> ~26s.

Three details drive the design:

  - URLSession multiplexes concurrent HTTP/2 requests to one host onto a
    single TCP connection, so sharing a session across chunks collapses them
    back onto one connection. Each lane owns its own session; raising
    httpMaximumConnectionsPerHost does not substitute for this.
  - A session per chunk is just as bad -- every chunk then pays a fresh TLS
    handshake, redirect, and TCP slow-start ramp, which measured 1.20x versus
    1.9x for reused connections. So a fixed set of lanes drain a shared chunk
    queue.
  - Lanes are capped app-wide by LaneBudget. Without it, three concurrent
    files would open 24 sockets and just add queueing delay on a thin link.

Also fixes two long-standing behaviors on this path:

  - A failed retry restarted the file from byte zero, and resume data never
    survived a relaunch. Durable chunks are now recorded in a .part.json
    manifest, so an interrupted 4.6 GB transfer resumes instead of restarting.
  - Integrity was checked by file size alone. Chunks are pinned to the
    resolved commit SHA so a push mid-download cannot splice revisions, and
    the assembled file is verified against the sha256 Hugging Face already
    publishes in x-linked-etag.

Small files and any server that does not advertise Range fall back to the
existing single-task path unchanged.

Live transfer tests are gated behind RUN_LIVE_HF_DOWNLOAD=1: chunked range
semantics are exactly what a mock would get wrong.
@RaajeevChandran

Copy link
Copy Markdown
Contributor

1. HF token leaks to CDN hosts (blocker)

  • TransferLane attaches the user's HF token via HuggingFaceAuth.authorize but has
    no willPerformHTTPRedirection delegate. When a chunk request to
    huggingface.co/resolve/<sha>/... gets a 302 to the CDN, URLSession re-sends the
    Authorization header to the CDN host. DirectDownloader strips it on cross-host
    redirects for exactly this reason (see the comment there referencing huggingface_hub
    behavior), and we recently host-gated the token in the single-connection path too.
  • ChunkedFileDownloader.probe also calls authorize on whatever URL it receives,
    which is not always a huggingface.co URL (see point 2).

The fix should be small: only attach the token when the request host is
huggingface.co, and add the same Authorization-stripping redirect delegate to
TransferLane.

2. Chunking silently never activates for onboarding downloads

Onboarding model downloads now go through the model download proxy (merged recently),
so DirectDownloader receives a presigned CDN URL for those. The probe blocks
redirects to read x-linked-size and x-linked-etag off the huggingface.co 3xx, but
presigned CDN URLs answer 206 directly without those headers, so the probe returns nil
and onboarding stays single-connection. That flow is the one most sensitive to download
speed, so this matters.

The proxy's resolve response already returns size, etag, and commit, which is
exactly the metadata the chunker needs. Suggestion: let the proxy route hand a
HuggingFaceFileMetadata directly to ChunkedFileDownloader and skip the probe. If
you want to keep that out of scope for this PR, at minimum skip the probe for
non-huggingface.co hosts so we do not waste a request and leak the token there.

Minor, non-blocking

  • The disk-space preflight does not account for a leftover .part file's remaining
    growth, so there is a small over-admission window.
  • pinnedURL assumes the /resolve/main/ path shape. It degrades to unpinned when the
    shape changes, which is fine, just noting it.

@mimeding

Copy link
Copy Markdown
Contributor

Ran this live before reading anything else, because chunked range semantics are exactly the thing a desk review would get wrong. Both arms compiled verbatim from their branches into a minimal harness (DirectDownloader from main vs DirectDownloader + ChunkedFileDownloader from this branch; only GlobalProxySettings.makeSession and HuggingFaceAuth.authorize stubbed to their no-proxy/anonymous forms, no service-level retry in either arm). Artifact: mlx-community/Qwen2.5-7B-Instruct-4bit model.safetensors, 4,284,346,255 bytes, alternating runs into an isolated directory, same residential network throughout, AC power; one extra run that spanned a machine sleep was discarded. CDN was us.aws.cdn.hf.co (xet-bridge presigned).

Run Arm Outcome Wall Effective CPU u+s Peak RSS
A1 main completed 486.6 s 8.8 MB/s 19.5 s 34 MB
B1 PR completed 615.6 s 7.0 MB/s 52.8 s 4.33 GB
A2 main failed at 47% (-1001 after throughput collapsed) 9.7 s 31 MB
B2 PR failed at 16% (outage outlasted 3 per-chunk attempts, -1005) 8.2 s 44 MB

Integrity: A1, B1, and a resumed PR run all hash to 86110f368236b53c…f071, byte-identical to the published x-linked-etag.

Two caveats on the table before the questions. First, my link tops out around 8–19 MB/s — far below the per-connection cap your lane curve targets — so the throughput claim is simply not testable from here; 7.0 vs 8.8 MB/s is one run each on a noisy link, not a regression signal. Second, what the link's flakiness did surface is the thing the manifest buys: main hard-failed 2 of 3 full-file attempts (one sat at 0 bytes for 928 s before NSURLErrorNetworkConnectionLost while curl pulled 14 MB/s on the same link), and every main failure forfeits the file (resumeDataForAttempt = nil on retry). Meanwhile a PR run interrupted at 109/128 durable chunks — process killed, machine slept — was finished by a fresh process in 43.7 s, first progress report already crediting the durable 3.4 GB, scratch state cleaned up, final hash exact. That's the strongest evidence in this PR's favor and it reproduced on the first try.

Questions, in the order I'd care about them:

  1. hashFile peaks at ~file-size RSS — can it get an autoreleasepool? FileHandle.read(upToCount:) returns autoreleased NSData on Darwin, and the verify loop never drains, so all ~1075 blocks of a 4.3 GB shard coexist. Measured: B1 peaked at 4.33 GB (flat ~44 MB during transfer — the spike is verify); isolated repro with the verbatim loop on a 2 GiB file hits 2062 MB, and the same loop with the body in autoreleasepool { } hits 10 MB with an identical digest. With maxConcurrentFiles = 3 shards finishing close together this sums, and a 16 GB machine with a model resident is jetsam territory. The live test's 335 MB file can't see it.

  2. Can pause() race a lane into dataTask(with:) on an invalidated session? pause()/invalidate() call session.invalidateAndCancel() on the lanes, but a lane that already passed checkPause() at the top of transferChunk goes on to session.dataTask(...).resume() — task creation on an invalidated session raises an ObjC exception Swift can't catch. The window is one instruction wide per chunk boundary, but a 4.3 GB shard offers ~134 boundaries × 8 lanes, and Pause is a button. Would an isInvalidated bit on the lane, checked under the same lock teardown takes, close it without giving up fast teardown? (DirectDownloader.invalidate()chunked.invalidate() has the same exposure and doesn't set pauseRequested first.)

  3. Did the lane measurements include the per-chunk re-resolve? Every chunk goes back through huggingface.co/resolve/<sha>/ and follows a fresh 302 — ~134 hub hits per 4.3 GB shard where main issued 1, which is new 429 surface for anonymous users (the aggressively rate-limited case per HuggingFaceAuth's own comment). Was reusing the probe's presigned Location until its Expires considered and rejected — e.g. because the 1 h expiry I measured on the xet bridge still needs re-resolve-on-403 for slow links anyway? A one-line comment on that decision would save the next reader the same question.

  4. TransferLane has no willPerformHTTPRedirection — intentional? Main's DirectDownloader strips Authorization cross-host explicitly. I verified current CFNetwork (3860.600.21) drops the header on a cross-host 302 by itself — canary token present at hop 1, absent at the redirect target — so no live leak. But the lane path now depends on an undocumented default the neighbouring class refuses to rely on, and S3-presigned endpoints reject requests carrying both query auth and an Authorization header. One delegate method makes it explicit either way.

  5. Small ones: the chunked branch reads pauseRequested but never resets it, so a once-paused DirectDownloader instance throws PauseInfo from any later chunkable download() (latent — transferFile builds fresh instances, but downloadMissingFiles reuses one across a file loop); try? probe folds real hub errors (401/429) into "not chunkable" and costs a doomed legacy request before proper classification; and a chunk completing exactly as pause lands is discarded because the success-path checkPause() runs before the caller records it in the manifest — record first and it's free.

The integrity core held up under everything I threw at it: short-206 refusal, 200-refusal, commit pinning, manifest invalidation on etag/commit change, sync-before-record, and a bit-exact file assembled across a process kill and a machine sleep. #1 and #2 are the only things I'd want resolved before merge.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants