mirror of
https://github.com/astral-sh/uv.git
synced 2025-11-20 11:56:03 +00:00
Fetch concurrently for non-first-match index strategies (#10432)
## 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:
parent
201726cda5
commit
a0494bb059
1 changed files with 84 additions and 35 deletions
|
|
@ -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,16 +248,16 @@ impl RegistryClient {
|
|||
}
|
||||
|
||||
let mut results = Vec::new();
|
||||
|
||||
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));
|
||||
|
||||
// If we're only using the first match, we can stop here.
|
||||
if self.index_strategy == IndexStrategy::FirstIndex {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(err) => match err.into_kind() {
|
||||
// The package could not be found in the remote index.
|
||||
ErrorKind::WrappedReqwestError(url, err) => match err.status() {
|
||||
|
|
@ -276,10 +277,58 @@ impl RegistryClient {
|
|||
// The package could not be found in the local index.
|
||||
ErrorKind::FileNotFound(_) => {}
|
||||
|
||||
other => return Err(other.into()),
|
||||
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),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if results.is_empty() {
|
||||
return match self.connectivity {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue