Small reads in uv_distribution::distribution_database::download producing v. slow download performance when installing torch

[thundergolfer]

Hey, this is a successor to #1978 as I'm still trying to integrate uv with Modal's internal mirror. But this time I don't yet have a solution to this problem 🙂.

I'm observed extremely degraded install performance when installing pytorch against our mirror. The install command is uv pip install torch.

On my test machine, installing without our mirror is fast, taking around 7 seconds.

But when I get our mirror involved, performance tanks and an install takes over 5 minutes. The command I'm running is this:

RUST_LOG="uv=trace,uv_extrace=trace,async_zip=debug" \
UV_NO_CACHE=true UV_INDEX_URL="http://172.21.0.1:5555/simple" \
uv --verbose pip install torch

Investigating...

Using debug logs on our mirror I see that there is a vast (~64x) difference in read sizes between uv pip install and pip install.

Our mirror serves .whl files off disk like this:

use hyper::{Body, Response};
use tokio_util::codec::{BytesCodec, FramedRead};

// ... 

let file tokio::fs::File::open(&filepath).await?;
debug!("cache hit: serving {} from disk cache", filepath.display());
let size = file.metadata().await.map(|m| Some(m.len())).unwrap_or(None);
let stream = FramedRead::new(file, BytesCodec::new());
let body = Body::wrap_stream(stream);
let mut response = Response::new(body);

By enabling debug logs (RUST_LOG=debug) I can see that when using pip install read chunks are large:

We can see in the screenshot of the debug logs that reads up to 128KiB are happening. Performance of the pip install command is good, the whl download takes around 3 seconds.

The pip command here was

PIP_TRUSTED_HOST="172.21.0.1:5555" \
PIP_NO_CACHE_DIR=off PIP_INDEX_URL="http://172.21.0.1:5555/simple" \
pip install torch

and pip --version is pip 23.3.1 from /home/ubuntu/modal/venv/lib/python3.11/site-packages/pip (python 3.11).

The read sizes of pip contrast strongly with uv pip install, where I'm seeing read sizes between 11 and 1120 bytes.

I believe these relatively tiny read sizes when downloading are contributing to the slow download performance on the 720.55MiB torch .whl file.

Looking at the trace logs this is the region of code I'm in:

uv/crates/uv-distribution/src/distribution_database.rs

Line 161 in 043d726

let download = |response: reqwest::Response| {

---

uv version

uv --version
uv 0.1.14

OS

uname -a
Linux ip-10-1-1-198 5.15.0-1052-aws #57~20.04.1-Ubuntu SMP Mon Jan 15 17:04:56 UTC 2024 x86_64 x86_64 x86_64 GNU/Linux

[thundergolfer]

I'm happy to look more into what's going on later in the week 🙂. I've read through some of the code in uv and async_zip but haven't figured out what's going on.

[charliermarsh]

Cool, would definitely welcome more exploration / eyes on it and happy to support. We unzip the wheel to disk directly as we stream the download in-memory, it looks like we're making way too many small reads right now?

[thundergolfer]

Yeah, I added BufReader in a couple places which was ineffectual, and see there's already BufReader used in a couple places (in uv and in async_zip) despite seeing reads << than the 8KiB default buffer size.

It'd take me a decent bit more time to unpack why these small reads are being produced and slap a buffered reader in the right place.

[charliermarsh]

Is there any way to reproduce this with publicly available indexes?

[thundergolfer]

Yeh I think that's job 1 so that we can collab better on this! We won't open-source our mirror implementation just yet.

[hmc-cs-mdrissi]

This behavior seems very similar to one I’m seeing as well. For our internal index, proxied public packages have good performance even when wheel is large (PyTorch), but internal packages with large wheels are much slower with uv then pip.

The only thing I can add is I know for our index we don’t support range requests/head so maybe the fallback path there is needed to trigger slow behavior here.

[charliermarsh]

We don't use range requests for the wheel download IIRC, just for fetching wheel metadata, so hopefully not a factor here.

[thundergolfer]

Adding some details of attempting a repro with an open-source mirror, proxpi: docker run -p 5000:5000 epicwink/proxpi.

time UV_NO_CACHE=true UV_INDEX_URL="http://127.0.0.1:5000/index" target/release/uv --verbose pip install torch

Has proxpi log a whl download time of around 7.4s.

2024-03-09 04:37:34 [    INFO] gunicorn.access: 172.18.0.1 "GET /index/torch/torch-2.2.1-cp311-cp311-manylinux1_x86_64.whl HTTP/1.1" 200 0 7431ms

With just pip:

PIP_TRUSTED_HOST="127.0.0.1:5000" PIP_NO_CACHE_DIR=off PIP_INDEX_URL="http://127.0.0.1:5000/index" pip install torch

2024-03-09 04:39:12 [    INFO] gunicorn.access: 172.18.0.1 "GET /index/torch/ HTTP/1.1" 200 24056 7ms
2024-03-09 04:39:14 [    INFO] gunicorn.access: 172.18.0.1 "GET /index/torch/torch-2.2.1-cp311-cp311-manylinux1_x86_64.whl HTTP/1.1" 200 0 2080ms

I get around 2s.

So there's a difference, but nothing like the ~5 minute slowdown seen in our internal mirror. So a negative result on providing an open-source reproduction. I'll keep digging.

[charliermarsh]

I'd really like to fix this. I'll try to play around with it today. The simplest thing we should try:

diff --git a/crates/uv-extract/src/stream.rs b/crates/uv-extract/src/stream.rs
index 8d8dfb11..b9272df5 100644
--- a/crates/uv-extract/src/stream.rs
+++ b/crates/uv-extract/src/stream.rs
@@ -18,7 +18,7 @@ pub async fn unzip<R: tokio::io::AsyncRead + Unpin>(
     target: impl AsRef<Path>,
 ) -> Result<(), Error> {
     let target = target.as_ref();
-    let mut reader = reader.compat();
+    let mut reader = futures::io::BufReader::new( reader.compat());
     let mut zip = async_zip::base::read::stream::ZipFileReader::new(&mut reader);
 
     let mut directories = FxHashSet::default();

(It looks like the async ZIP reader uses a BufReader internally, but only for the entry bodies, and not the headers.)

[charliermarsh]

Not seeing any improvement from this change.

FWIW, on my machine, the difference is much smaller with the open-source proxy.

pip:

gunicorn.access: 172.17.0.1 "GET /index/torch/torch-2.2.1-cp312-none-macosx_11_0_arm64.whl HTTP/1.1" 200 0 888ms

uv:

gunicorn.access: 172.17.0.1 "GET /index/torch/torch-2.2.1-cp312-none-macosx_11_0_arm64.whl HTTP/1.1" 200 0 1374ms

Which is maybe plausible given that we're unzipping and writing to disk -- not sure.

[charliermarsh]

We could also try increasing the buffer size.

[thundergolfer]

Ok from looking at trace logs for our internal mirror I think the encoding strategy difference is key:

pip

2024-03-09T17:41:23.715190Z DEBUG pypi_mirror: cache hit: serving /tmp/pypi-mirror-worker-cache/packages/2c/df/5810707da6f2fd4be57f0cc417987c0fa16a2eecf0b1b71f82ea555dc619/torch-2.2.1-cp311-cp311-manylinux1_x86_64.whl from disk cache
2024-03-09T17:41:23.715200Z TRACE hyper::proto::h1::conn: flushed({role=server}): State { reading: KeepAlive, writing: Init, keep_alive: Busy }
2024-03-09T17:41:23.715260Z TRACE encode_headers: hyper::proto::h1::role: Server::encode status=200, body=Some(Unknown), req_method=Some(GET)
2024-03-09T17:41:23.715329Z DEBUG hyper::proto::h1::io: flushed 83 bytes
2024-03-09T17:41:23.715336Z TRACE hyper::proto::h1::conn: flushed({role=server}): State { reading: KeepAlive, writing: Body(Encoder { kind: Length(755552143), is_last: false }), keep_alive: Busy }
2024-03-09T17:41:23.715355Z TRACE hyper::proto::h1::encode: sized write, len = 8192
2024-03-09T17:41:23.715362Z TRACE hyper::proto::h1::io: buffer.queue self.len=0 buf.len=8192

uv

2024-03-09T17:40:37.658227Z DEBUG pypi_mirror: cache hit: serving /tmp/pypi-mirror-worker-cache/packages/2c/df/5810707da6f2fd4be57f0cc417987c0fa16a2eecf0b1b71f82ea555dc619/torch-2.2.1-cp311-cp311-manylinux1_x86_64.whl from disk cache
2024-03-09T17:40:37.658414Z TRACE hyper::proto::h1::conn: flushed({role=server}): State { reading: KeepAlive, writing: Init, keep_alive: Busy }
2024-03-09T17:40:37.658530Z TRACE encode_headers: hyper::proto::h1::role: Server::encode status=200, body=Some(Unknown), req_method=Some(GET)
2024-03-09T17:40:37.658576Z DEBUG hyper::proto::h1::io: flushed 108 bytes
2024-03-09T17:40:37.658591Z TRACE hyper::proto::h1::conn: flushed({role=server}): State { reading: KeepAlive, writing: Body(Encoder { kind: Chunked, is_last: false }), keep_alive: Busy }
2024-03-09T17:40:37.658755Z TRACE hyper::proto::h1::encode: encoding chunked 10B
2024-03-09T17:40:37.658759Z TRACE hyper::proto::h1::io: buffer.queue self.len=0 buf.len=15

Will work out why this difference occurs.

[charliermarsh]

Ohh nice!

[thundergolfer]

I think this is another content-encoding related issue.

I added some extra trace! logs into hyper when it was running encode_headers and deciding on the encoding strategy.

uv

Server::encode head=MessageHead { version: HTTP/1.1, subject: 200, headers: {"content-encoding": "gzip"}, extensions: Extensions }, status=200, body=Some(Unknown), req_method=Some(GET)

pip

Server::encode head=MessageHead { version: HTTP/1.1, subject: 200, headers: {"content-length": "755552143"}, extensions: Extensions }, status=200, body=Some(Unknown), req_method=Some(GET)

Seeing gzip in the headers cued me to think this was just the same issue as #1978 showing up in a different place.

By doing the same header change that's done in #1978 torch installs much, much faster against the mirror!

A 15.64s second download against the mirror is still slower than what uv gets against PyPI (around 10s) so there may be some additional change I can make to get performance of our mirror to parity with PyPi. It really should be faster as it's serving over the local network.

Checking the pip repository for its history of dealing with content-encoding I see this pypa/pip#1688

[thundergolfer]

Looking at the flushes in our mirror after adding the identity content-encoding header, I can now see they're averaging around 8KiB, much higher than the ~1KiB before:

Still lower than what pip was producing in our mirror, which went up to 128KiB.

[charliermarsh]

Can you try tweaking the buffer size on the BufReader? I think 8KiB might be the exact default.

[thundergolfer]

You mean doing this but with a non-default buffer size? #2220 (comment)

[charliermarsh]

Yeah. Alternatively I can change the buffer used inside async-zip-rs, we use a fork anyway.

[charliermarsh]

(I don’t have much intuition for how nested BufReaders interact.)

[thundergolfer]

Yep with this diff I get comparable performance to pip 🙂

diff --git a/crates/uv-extract/src/stream.rs b/crates/uv-extract/src/stream.rs
index 8d8dfb11..db3c42fb 100644
--- a/crates/uv-extract/src/stream.rs
+++ b/crates/uv-extract/src/stream.rs
@@ -18,7 +18,7 @@ pub async fn unzip<R: tokio::io::AsyncRead + Unpin>(
     target: impl AsRef<Path>,
 ) -> Result<(), Error> {
     let target = target.as_ref();
-    let mut reader = reader.compat();
+    let mut reader = futures::io::BufReader::with_capacity( 128 * 1024, reader.compat());
     let mut zip = async_zip::base::read::stream::ZipFileReader::new(&mut reader);

     let mut directories = FxHashSet::default();

Debug logging on hyper shows flushes up-to 128KiB.

---

[thundergolfer]

Think this is a enough to make a PR. It'd be two things:

1. setting identity, in the spirit of Accept encoding identity pypa/pip#1688
2. increasing buffer from 8KiB to 128KiB.

[charliermarsh]

Sounds great, thank you! The only thing I’ll want to test before merging is whether the buffer size changes benchmarks much / at all against PyPI.

[thundergolfer]

Makes sense. Looks like I could even check this myself by reading BENCHMARKS.md and going from there?

[charliermarsh]

Totally, up to you! I'd just do a release build on main, move it to uv-main, then do a release build on your branch, then run something ike:

python -m scripts.bench \
        --uv-path ./target/release/uv \
        --uv-path ./target/release/uv-main \
        ./scripts/requirements/compiled/trio.txt --benchmark install-cold

[charliermarsh]

By the way, I can get this released pretty quickly.

[hmc-cs-mdrissi]

I tested the fix on my slow large wheel and it works great. Installation time is now fast for that 1 wheel and I guess my index server had same issue here. Thanks to both of you for continuing to improve uv's performance.

[charliermarsh]

All @thundergolfer!

[charliermarsh]

I'll probably cut a release tomorrow with this in it.