Add UV_CONCURRENT_INSTALLS variable in favor of RAYON_NUM_THREADS (#3646)

## Summary

Continuation of https://github.com/astral-sh/uv/pull/3493. This gives us
more flexibility in case we decide to move away from `rayon` in the
future.
This commit is contained in:
Ibraheem Ahmed 2024-05-17 23:12:37 -04:00 committed by GitHub
parent fe2bc079bc
commit 53633392c3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 50 additions and 7 deletions

1
Cargo.lock generated
View file

@ -4518,6 +4518,7 @@ dependencies = [
"platform-tags",
"predicates",
"pypi-types",
"rayon",
"regex",
"requirements-txt",
"reqwest",

View file

@ -560,6 +560,8 @@ uv accepts the following command-line arguments as environment variables:
will perform at any given time.
- `UV_CONCURRENT_BUILDS`: Sets the maximum number of source distributions that `uv` will build
concurrently at any given time.
- `UV_CONCURRENT_INSTALLS`: Used to control the number of threads used when installing and unzipping
packages.
In each case, the corresponding command-line argument takes precedence over an environment variable.
@ -583,8 +585,6 @@ In addition, uv respects the following environment variables:
- `FISH_VERSION`: Used to detect the use of the Fish shell.
- `BASH_VERSION`: Used to detect the use of the Bash shell.
- `ZSH_VERSION`: Used to detect the use of the Zsh shell.
- `RAYON_NUM_THREADS`: Used to control the number of threads used when unzipping and installing
packages. See the [rayon documentation](https://docs.rs/rayon/latest/rayon/) for more.
- `MACOSX_DEPLOYMENT_TARGET`: Used with `--python-platform macos` and related variants to set the
deployment target (i.e., the minimum supported macOS version). Defaults to `12.0`, the
least-recent non-EOL macOS version at time of writing.

View file

@ -11,13 +11,18 @@ pub struct Concurrency {
///
/// Note this value must be non-zero.
pub builds: usize,
/// The maximum number of concurrent installs.
///
/// Note this value must be non-zero.
pub installs: usize,
}
impl Default for Concurrency {
fn default() -> Self {
Concurrency {
downloads: Concurrency::DEFAULT_DOWNLOADS,
builds: Concurrency::default_builds(),
builds: Concurrency::threads(),
installs: Concurrency::threads(),
}
}
}
@ -26,8 +31,8 @@ impl Concurrency {
// The default concurrent downloads limit.
pub const DEFAULT_DOWNLOADS: usize = 50;
// The default concurrent builds limit.
pub fn default_builds() -> usize {
// The default concurrent builds and install limit.
pub fn threads() -> usize {
std::thread::available_parallelism()
.map(NonZeroUsize::get)
.unwrap_or(1)

View file

@ -115,6 +115,7 @@ impl Combine for PipOptions {
.concurrent_downloads
.combine(other.concurrent_downloads),
concurrent_builds: self.concurrent_builds.combine(other.concurrent_builds),
concurrent_installs: self.concurrent_installs.combine(other.concurrent_installs),
}
}
}

View file

@ -87,4 +87,5 @@ pub struct PipOptions {
pub require_hashes: Option<bool>,
pub concurrent_downloads: Option<NonZeroUsize>,
pub concurrent_builds: Option<NonZeroUsize>,
pub concurrent_installs: Option<NonZeroUsize>,
}

View file

@ -51,6 +51,7 @@ indicatif = { workspace = true }
itertools = { workspace = true }
miette = { workspace = true, features = ["fancy"] }
owo-colors = { workspace = true }
rayon = { workspace = true }
regex = { workspace = true }
rustc-hash = { workspace = true }
serde = { workspace = true }

View file

@ -176,6 +176,10 @@ async fn run() -> Result<ExitStatus> {
// Resolve the settings from the command-line arguments and workspace configuration.
let args = PipCompileSettings::resolve(args, workspace);
rayon::ThreadPoolBuilder::new()
.num_threads(args.shared.concurrency.installs)
.build_global()
.expect("failed to initialize global rayon pool");
// Initialize the cache.
let cache = cache.init()?.with_refresh(args.refresh);
@ -247,6 +251,10 @@ async fn run() -> Result<ExitStatus> {
// Resolve the settings from the command-line arguments and workspace configuration.
let args = PipSyncSettings::resolve(args, workspace);
rayon::ThreadPoolBuilder::new()
.num_threads(args.shared.concurrency.installs)
.build_global()
.expect("failed to initialize global rayon pool");
// Initialize the cache.
let cache = cache.init()?.with_refresh(args.refresh);
@ -293,6 +301,10 @@ async fn run() -> Result<ExitStatus> {
// Resolve the settings from the command-line arguments and workspace configuration.
let args = PipInstallSettings::resolve(args, workspace);
rayon::ThreadPoolBuilder::new()
.num_threads(args.shared.concurrency.installs)
.build_global()
.expect("failed to initialize global rayon pool");
// Initialize the cache.
let cache = cache.init()?.with_refresh(args.refresh);

View file

@ -307,6 +307,7 @@ impl PipCompileSettings {
link_mode,
concurrent_builds: env(env::CONCURRENT_BUILDS),
concurrent_downloads: env(env::CONCURRENT_DOWNLOADS),
concurrent_installs: env(env::CONCURRENT_INSTALLS),
..PipOptions::default()
},
workspace,
@ -415,6 +416,7 @@ impl PipSyncSettings {
require_hashes: flag(require_hashes, no_require_hashes),
concurrent_builds: env(env::CONCURRENT_BUILDS),
concurrent_downloads: env(env::CONCURRENT_DOWNLOADS),
concurrent_installs: env(env::CONCURRENT_INSTALLS),
..PipOptions::default()
},
workspace,
@ -567,6 +569,7 @@ impl PipInstallSettings {
require_hashes: flag(require_hashes, no_require_hashes),
concurrent_builds: env(env::CONCURRENT_BUILDS),
concurrent_downloads: env(env::CONCURRENT_DOWNLOADS),
concurrent_installs: env(env::CONCURRENT_INSTALLS),
..PipOptions::default()
},
workspace,
@ -955,6 +958,7 @@ impl PipSharedSettings {
require_hashes,
concurrent_builds,
concurrent_downloads,
concurrent_installs,
} = workspace
.and_then(|workspace| workspace.options.pip)
.unwrap_or_default();
@ -1044,11 +1048,18 @@ impl PipSharedSettings {
downloads: args
.concurrent_downloads
.or(concurrent_downloads)
.map_or(Concurrency::DEFAULT_DOWNLOADS, NonZeroUsize::get),
.map(NonZeroUsize::get)
.unwrap_or(Concurrency::DEFAULT_DOWNLOADS),
builds: args
.concurrent_builds
.or(concurrent_builds)
.map_or_else(Concurrency::default_builds, NonZeroUsize::get),
.map(NonZeroUsize::get)
.unwrap_or_else(Concurrency::threads),
installs: args
.concurrent_installs
.or(concurrent_installs)
.map(NonZeroUsize::get)
.unwrap_or_else(Concurrency::threads),
},
}
}
@ -1061,6 +1072,9 @@ mod env {
pub(super) const CONCURRENT_BUILDS: (&str, &str) =
("UV_CONCURRENT_BUILDS", "a non-zero integer");
pub(super) const CONCURRENT_INSTALLS: (&str, &str) =
("UV_CONCURRENT_INSTALLS", "a non-zero integer");
}
/// Attempt to load and parse an environment variable with the given name.

8
uv.schema.json generated
View file

@ -263,6 +263,14 @@
"format": "uint",
"minimum": 1.0
},
"concurrent-installs": {
"type": [
"integer",
"null"
],
"format": "uint",
"minimum": 1.0
},
"config-settings": {
"anyOf": [
{