mirrors/uv
1
0
Fork 0
mirror of https://github.com/astral-sh/uv.git synced 2025-11-20 11:56:03 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Fetch concurrently for non-first-match index strategies (#10432)

Browse source 
## Summary

On a basic test, this speeds up cold resolution by about 25%:

```
❯ hyperfine "uv lock --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match --upgrade --no-cache" "../target/release/uv lock --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match --upgrade --no-cache" --warmup 10 --runs 30
Benchmark 1: uv lock --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match --upgrade --no-cache
  Time (mean ± σ):     585.8 ms ±  28.2 ms    [User: 149.7 ms, System: 97.4 ms]
  Range (min … max):   541.5 ms … 654.8 ms    30 runs

Benchmark 2: ../target/release/uv lock --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match --upgrade --no-cache
  Time (mean ± σ):     468.3 ms ±  52.0 ms    [User: 131.7 ms, System: 76.9 ms]
  Range (min … max):   380.2 ms … 607.0 ms    30 runs

Summary
  ../target/release/uv lock --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match --upgrade --no-cache ran
    1.25 ± 0.15 times faster than uv lock --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match --upgrade --no-cache
```

Given:
```toml
[project]
name = "foo"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12.0"
dependencies = [
    "black>=24.10.0",
    "django>=5.1.4",
    "flask>=3.1.0",
    "requests>=2.32.3",
]
```

And:

```shell
hyperfine "uv lock --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match --upgrade --no-cache" "../target/release/uv lock --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match --upgrade --no-cache" --warmup 10 --runs 30
```

Closes https://github.com/astral-sh/uv/issues/10429.
This commit is contained in:
Charlie Marsh 2025-01-09 12:45:20 -05:00 • committed by GitHub
parent 201726cda5
commit a0494bb059
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 84 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

119
crates/uv-client/src/registry_client.rs
View file

@ -1,14 +1,15 @@
use async_http_range_reader::AsyncHttpRangeReader;
use futures::{FutureExt, TryStreamExt};
use http::HeaderMap;
use itertools::Either;
use reqwest::{Client, Response, StatusCode};
use reqwest_middleware::ClientWithMiddleware;
use std::collections::BTreeMap;
use std::fmt::Debug;
use std::path::PathBuf;
use std::str::FromStr;
use std::time::Duration;
use async_http_range_reader::AsyncHttpRangeReader;
use futures::{FutureExt, StreamExt, TryStreamExt};
use http::HeaderMap;
use itertools::Either;
use reqwest::{Client, Response, StatusCode};
use reqwest_middleware::ClientWithMiddleware;
use tracing::{info_span, instrument, trace, warn, Instrument};
use url::Url;
@ -247,38 +248,86 @@ impl RegistryClient {
}
let mut results = Vec::new();
for index in it {
match self.simple_single_index(package_name, index).await {
Ok(metadata) => {
results.push((index, metadata));
// If we're only using the first match, we can stop here.
if self.index_strategy == IndexStrategy::FirstIndex {
break;
match self.index_strategy {
// If we're searching for the first index that contains the package, fetch serially.
IndexStrategy::FirstIndex => {
for index in it {
match self.simple_single_index(package_name, index).await {
Ok(metadata) => {
results.push((index, metadata));
break;
}
Err(err) => match err.into_kind() {
// The package could not be found in the remote index.
ErrorKind::WrappedReqwestError(url, err) => match err.status() {
Some(StatusCode::NOT_FOUND) => {}
Some(StatusCode::UNAUTHORIZED) => {
capabilities.set_unauthorized(index.clone());
}
Some(StatusCode::FORBIDDEN) => {
capabilities.set_forbidden(index.clone());
}
_ => return Err(ErrorKind::WrappedReqwestError(url, err).into()),
},
// The package is unavailable due to a lack of connectivity.
ErrorKind::Offline(_) => {}
// The package could not be found in the local index.
ErrorKind::FileNotFound(_) => {}
err => return Err(err.into()),
},
};
}
}
// Otherwise, fetch concurrently.
IndexStrategy::UnsafeBestMatch | IndexStrategy::UnsafeFirstMatch => {
let fetches = futures::stream::iter(it)
.map(|index| async move {
match self.simple_single_index(package_name, index).await {
Ok(metadata) => Ok(Some((index, metadata))),
Err(err) => match err.into_kind() {
// The package could not be found in the remote index.
ErrorKind::WrappedReqwestError(url, err) => match err.status() {
Some(StatusCode::NOT_FOUND) => Ok(None),
Some(StatusCode::UNAUTHORIZED) => {
capabilities.set_unauthorized(index.clone());
Ok(None)
}
Some(StatusCode::FORBIDDEN) => {
capabilities.set_forbidden(index.clone());
Ok(None)
}
_ => Err(ErrorKind::WrappedReqwestError(url, err).into()),
},
// The package is unavailable due to a lack of connectivity.
ErrorKind::Offline(_) => Ok(None),
// The package could not be found in the local index.
ErrorKind::FileNotFound(_) => Ok(None),
err => Err(err.into()),
},
}
})
.buffered(8);
futures::pin_mut!(fetches);
while let Some(result) = fetches.next().await {
match result {
Ok(Some((index, metadata))) => {
results.push((index, metadata));
}
Ok(None) => continue,
Err(err) => return Err(err),
}
}
Err(err) => match err.into_kind() {
// The package could not be found in the remote index.
ErrorKind::WrappedReqwestError(url, err) => match err.status() {
Some(StatusCode::NOT_FOUND) => {}
Some(StatusCode::UNAUTHORIZED) => {
capabilities.set_unauthorized(index.clone());
}
Some(StatusCode::FORBIDDEN) => {
capabilities.set_forbidden(index.clone());
}
_ => return Err(ErrorKind::WrappedReqwestError(url, err).into()),
},
// The package is unavailable due to a lack of connectivity.
ErrorKind::Offline(_) => {}
// The package could not be found in the local index.
ErrorKind::FileNotFound(_) => {}
other => return Err(other.into()),
},
};
}
}
if results.is_empty() {