From a9d8de94d59acf391719a2f078be5767da1c1fc2 Mon Sep 17 00:00:00 2001 From: konstin Date: Tue, 26 Nov 2024 09:20:06 +0100 Subject: [PATCH] Initialize rayon lazily --- Cargo.lock | 3 ++- crates/uv-configuration/Cargo.toml | 1 + crates/uv-configuration/src/lib.rs | 2 ++ crates/uv-configuration/src/rayon.rs | 21 +++++++++++++++++++++ crates/uv-extract/Cargo.toml | 1 + crates/uv-extract/src/sync.rs | 10 ++++++---- crates/uv-installer/src/installer.rs | 6 ++++++ crates/uv/Cargo.toml | 1 - crates/uv/src/lib.rs | 20 +++++++++----------- 9 files changed, 48 insertions(+), 17 deletions(-) create mode 100644 crates/uv-configuration/src/rayon.rs diff --git a/Cargo.lock b/Cargo.lock index 670f51b5890e..7cdc815a1849 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4430,7 +4430,6 @@ dependencies = [ "owo-colors", "petgraph", "predicates", - "rayon", "regex", "reqwest", "rkyv", @@ -4746,6 +4745,7 @@ dependencies = [ "clap", "either", "fs-err 3.0.0", + "rayon", "rustc-hash", "schemars", "serde", @@ -4958,6 +4958,7 @@ dependencies = [ "tokio", "tokio-util", "tracing", + "uv-configuration", "uv-distribution-filename", "uv-pypi-types", "xz2", diff --git a/crates/uv-configuration/Cargo.toml b/crates/uv-configuration/Cargo.toml index 0fa47d96d5af..83de06f5cc35 100644 --- a/crates/uv-configuration/Cargo.toml +++ b/crates/uv-configuration/Cargo.toml @@ -29,6 +29,7 @@ uv-static = { workspace = true } clap = { workspace = true, features = ["derive"], optional = true } either = { workspace = true } fs-err = { workspace = true } +rayon = { workspace = true } rustc-hash = { workspace = true } schemars = { workspace = true, optional = true } serde = { workspace = true } diff --git a/crates/uv-configuration/src/lib.rs b/crates/uv-configuration/src/lib.rs index bd45cab7753c..6eb47a89a82a 100644 --- a/crates/uv-configuration/src/lib.rs +++ b/crates/uv-configuration/src/lib.rs @@ -15,6 +15,7 @@ pub use overrides::*; pub use package_options::*; pub use preview::*; pub use project_build_backend::*; +pub use rayon::*; pub use sources::*; pub use target_triple::*; pub use trusted_host::*; @@ -38,6 +39,7 @@ mod overrides; mod package_options; mod preview; mod project_build_backend; +mod rayon; mod sources; mod target_triple; mod trusted_host; diff --git a/crates/uv-configuration/src/rayon.rs b/crates/uv-configuration/src/rayon.rs new file mode 100644 index 000000000000..d00877b8c0d1 --- /dev/null +++ b/crates/uv-configuration/src/rayon.rs @@ -0,0 +1,21 @@ +//! Initialize the rayon threadpool once, before we need it. +//! +//! The `uv` crate sets [`RAYON_PARALLELISM`] from the user settings, and the extract and install +//! code initialize the threadpool lazily only if they are actually used by calling +//! `LazyLock::force(&RAYON_INITIALIZE)`. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::LazyLock; + +/// The number of threads for the rayon threadpool. +/// +/// The default of 0 makes rayon use its default. +pub static RAYON_PARALLELISM: AtomicUsize = AtomicUsize::new(0); + +/// Initialize the threadpool lazily. Always call before using rayon the potentially first time. +pub static RAYON_INITIALIZE: LazyLock<()> = LazyLock::new(|| { + rayon::ThreadPoolBuilder::new() + .num_threads(RAYON_PARALLELISM.load(Ordering::SeqCst)) + .build_global() + .expect("failed to initialize global rayon pool"); +}); diff --git a/crates/uv-extract/Cargo.toml b/crates/uv-extract/Cargo.toml index 1ebc4edc69c3..7f3851f46fe7 100644 --- a/crates/uv-extract/Cargo.toml +++ b/crates/uv-extract/Cargo.toml @@ -16,6 +16,7 @@ doctest = false workspace = true [dependencies] +uv-configuration = { workspace = true } uv-distribution-filename = { workspace = true } uv-pypi-types = { workspace = true } diff --git a/crates/uv-extract/src/sync.rs b/crates/uv-extract/src/sync.rs index 4c1ad1a8ac58..af510f39b945 100644 --- a/crates/uv-extract/src/sync.rs +++ b/crates/uv-extract/src/sync.rs @@ -1,14 +1,14 @@ use std::path::{Path, PathBuf}; -use std::sync::Mutex; +use std::sync::{LazyLock, Mutex}; +use crate::vendor::{CloneableSeekableReader, HasLength}; +use crate::Error; use rayon::prelude::*; use rustc_hash::FxHashSet; use tracing::warn; +use uv_configuration::RAYON_INITIALIZE; use zip::ZipArchive; -use crate::vendor::{CloneableSeekableReader, HasLength}; -use crate::Error; - /// Unzip a `.zip` archive into the target directory. pub fn unzip( reader: R, @@ -18,6 +18,8 @@ pub fn unzip( let reader = std::io::BufReader::new(reader); let archive = ZipArchive::new(CloneableSeekableReader::new(reader))?; let directories = Mutex::new(FxHashSet::default()); + // Initialize the threadpool with the user settings. + LazyLock::force(&RAYON_INITIALIZE); (0..archive.len()) .into_par_iter() .map(|file_number| { diff --git a/crates/uv-installer/src/installer.rs b/crates/uv-installer/src/installer.rs index 1b059bb725f7..e2ce95610dc2 100644 --- a/crates/uv-installer/src/installer.rs +++ b/crates/uv-installer/src/installer.rs @@ -1,11 +1,13 @@ use anyhow::{Context, Error, Result}; use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; use std::convert; +use std::sync::LazyLock; use tokio::sync::oneshot; use tracing::instrument; use uv_install_wheel::{linker::LinkMode, Layout}; use uv_cache::Cache; +use uv_configuration::RAYON_INITIALIZE; use uv_distribution_types::CachedDist; use uv_python::PythonEnvironment; @@ -85,6 +87,8 @@ impl<'a> Installer<'a> { let layout = venv.interpreter().layout(); let relocatable = venv.relocatable(); + // Initialize the threadpool with the user settings. + LazyLock::force(&RAYON_INITIALIZE); rayon::spawn(move || { let result = install( wheels, @@ -136,6 +140,8 @@ fn install( reporter: Option>, relocatable: bool, ) -> Result> { + // Initialize the threadpool with the user settings. + LazyLock::force(&RAYON_INITIALIZE); let locks = uv_install_wheel::linker::Locks::default(); wheels.par_iter().try_for_each(|wheel| { uv_install_wheel::linker::install_wheel( diff --git a/crates/uv/Cargo.toml b/crates/uv/Cargo.toml index 68fe1c9042a7..9d0a0c3b47e7 100644 --- a/crates/uv/Cargo.toml +++ b/crates/uv/Cargo.toml @@ -76,7 +76,6 @@ jiff = { workspace = true } miette = { workspace = true, features = ["fancy-no-backtrace"] } owo-colors = { workspace = true } petgraph = { workspace = true } -rayon = { workspace = true } regex = { workspace = true } reqwest = { workspace = true } rkyv = { workspace = true } diff --git a/crates/uv/src/lib.rs b/crates/uv/src/lib.rs index 19c4221304b5..052f71f30526 100644 --- a/crates/uv/src/lib.rs +++ b/crates/uv/src/lib.rs @@ -1,3 +1,9 @@ +use anstream::eprintln; +use anyhow::{bail, Result}; +use clap::error::{ContextKind, ContextValue}; +use clap::{CommandFactory, Parser}; +use owo_colors::OwoColorize; +use settings::PipTreeSettings; use std::borrow::Cow; use std::env; use std::ffi::OsString; @@ -5,13 +11,7 @@ use std::fmt::Write; use std::io::stdout; use std::path::Path; use std::process::ExitCode; - -use anstream::eprintln; -use anyhow::{bail, Result}; -use clap::error::{ContextKind, ContextValue}; -use clap::{CommandFactory, Parser}; -use owo_colors::OwoColorize; -use settings::PipTreeSettings; +use std::sync::atomic::Ordering; use tokio::task::spawn_blocking; use tracing::{debug, instrument}; use uv_cache::{Cache, Refresh}; @@ -254,10 +254,8 @@ async fn run(mut cli: Cli) -> Result { ) }))?; - rayon::ThreadPoolBuilder::new() - .num_threads(globals.concurrency.installs) - .build_global() - .expect("failed to initialize global rayon pool"); + // Don't initialize the rayon threadpool yet, this is too costly when we're doing a noop sync. + uv_configuration::RAYON_PARALLELISM.store(globals.concurrency.installs, Ordering::SeqCst); debug!("uv {}", uv_cli::version::version());