From 53633392c3344dc2d805922f6b4a6999de41b6df Mon Sep 17 00:00:00 2001 From: Ibraheem Ahmed Date: Fri, 17 May 2024 23:12:37 -0400 Subject: [PATCH] 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. --- Cargo.lock | 1 + README.md | 4 ++-- crates/uv-configuration/src/concurrency.rs | 11 ++++++++--- crates/uv-workspace/src/combine.rs | 1 + crates/uv-workspace/src/settings.rs | 1 + crates/uv/Cargo.toml | 1 + crates/uv/src/main.rs | 12 ++++++++++++ crates/uv/src/settings.rs | 18 ++++++++++++++++-- uv.schema.json | 8 ++++++++ 9 files changed, 50 insertions(+), 7 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 92cc3d11a..5f9267be6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4518,6 +4518,7 @@ dependencies = [ "platform-tags", "predicates", "pypi-types", + "rayon", "regex", "requirements-txt", "reqwest", diff --git a/README.md b/README.md index 5ee53a0ef..ee9d17ae0 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/crates/uv-configuration/src/concurrency.rs b/crates/uv-configuration/src/concurrency.rs index 64a86eb58..8e9ef2200 100644 --- a/crates/uv-configuration/src/concurrency.rs +++ b/crates/uv-configuration/src/concurrency.rs @@ -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) diff --git a/crates/uv-workspace/src/combine.rs b/crates/uv-workspace/src/combine.rs index 8091187cb..f53b1ad20 100644 --- a/crates/uv-workspace/src/combine.rs +++ b/crates/uv-workspace/src/combine.rs @@ -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), } } } diff --git a/crates/uv-workspace/src/settings.rs b/crates/uv-workspace/src/settings.rs index 133eca118..f05e51106 100644 --- a/crates/uv-workspace/src/settings.rs +++ b/crates/uv-workspace/src/settings.rs @@ -87,4 +87,5 @@ pub struct PipOptions { pub require_hashes: Option, pub concurrent_downloads: Option, pub concurrent_builds: Option, + pub concurrent_installs: Option, } diff --git a/crates/uv/Cargo.toml b/crates/uv/Cargo.toml index 7a681a372..442f5faa2 100644 --- a/crates/uv/Cargo.toml +++ b/crates/uv/Cargo.toml @@ -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 } diff --git a/crates/uv/src/main.rs b/crates/uv/src/main.rs index e523b1f56..c436127df 100644 --- a/crates/uv/src/main.rs +++ b/crates/uv/src/main.rs @@ -176,6 +176,10 @@ async fn run() -> Result { // 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 { // 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 { // 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); diff --git a/crates/uv/src/settings.rs b/crates/uv/src/settings.rs index 3f449a420..4a71af60d 100644 --- a/crates/uv/src/settings.rs +++ b/crates/uv/src/settings.rs @@ -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. diff --git a/uv.schema.json b/uv.schema.json index f1f7f85e9..3ba5ccbb0 100644 --- a/uv.schema.json +++ b/uv.schema.json @@ -263,6 +263,14 @@ "format": "uint", "minimum": 1.0 }, + "concurrent-installs": { + "type": [ + "integer", + "null" + ], + "format": "uint", + "minimum": 1.0 + }, "config-settings": { "anyOf": [ {