diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 17258b853..b7b8b70b1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -201,7 +201,7 @@ jobs: steps: - name: Create Dev Drive using ReFS run: | - $Volume = New-VHD -Path C:/uv_dev_drive.vhdx -SizeBytes 12GB | + $Volume = New-VHD -Path C:/uv_dev_drive.vhdx -SizeBytes 14GB | Mount-VHD -Passthru | Initialize-Disk -Passthru | New-Partition -AssignDriveLetter -UseMaximumSize | diff --git a/Cargo.toml b/Cargo.toml index 02250b5c5..4f7ff7b96 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -132,7 +132,7 @@ tempfile = { version = "3.9.0" } textwrap = { version = "0.16.1" } thiserror = { version = "1.0.56" } tl = { version = "0.7.7" } -tokio = { version = "1.35.1", features = ["fs", "io-util", "macros", "process", "rt-multi-thread", "sync"] } +tokio = { version = "1.35.1", features = ["fs", "io-util", "macros", "process", "sync"] } tokio-stream = { version = "0.1.14" } tokio-tar = { version = "0.3.1" } tokio-util = { version = "0.7.10", features = ["compat"] } diff --git a/crates/bench/benches/uv.rs b/crates/bench/benches/uv.rs index 7af692eb8..a7f9aaf35 100644 --- a/crates/bench/benches/uv.rs +++ b/crates/bench/benches/uv.rs @@ -36,7 +36,7 @@ fn resolve_warm_jupyter(c: &mut Criterion) { } fn resolve_warm_airflow(c: &mut Criterion) { - let runtime = &tokio::runtime::Builder::new_multi_thread() + let runtime = &tokio::runtime::Builder::new_current_thread() // CodSpeed limits the total number of threads to 500 .max_blocking_threads(256) .enable_all() diff --git a/crates/uv-dev/src/main.rs b/crates/uv-dev/src/main.rs index 4cf109bd0..bacb98081 100644 --- a/crates/uv-dev/src/main.rs +++ b/crates/uv-dev/src/main.rs @@ -84,7 +84,7 @@ async fn run() -> Result<()> { Ok(()) } -#[tokio::main] +#[tokio::main(flavor = "current_thread")] async fn main() -> ExitCode { let (duration_layer, _guard) = if let Ok(location) = env::var("TRACING_DURATIONS_FILE") { let location = PathBuf::from(location); diff --git a/crates/uv-dispatch/src/lib.rs b/crates/uv-dispatch/src/lib.rs index f424b2934..c910c52ab 100644 --- a/crates/uv-dispatch/src/lib.rs +++ b/crates/uv-dispatch/src/lib.rs @@ -278,16 +278,17 @@ impl<'a> BuildContext for BuildDispatch<'a> { } // Install the resolved distributions. - let wheels = wheels.into_iter().chain(cached).collect::>(); + let mut wheels = wheels.into_iter().chain(cached).collect::>(); if !wheels.is_empty() { debug!( "Installing build requirement{}: {}", if wheels.len() == 1 { "" } else { "s" }, wheels.iter().map(ToString::to_string).join(", ") ); - Installer::new(venv) + wheels = Installer::new(venv) .with_link_mode(self.link_mode) - .install(&wheels) + .install(wheels) + .await .context("Failed to install build dependencies")?; } diff --git a/crates/uv-installer/src/installer.rs b/crates/uv-installer/src/installer.rs index 2d0003bd6..a39c5c7a6 100644 --- a/crates/uv-installer/src/installer.rs +++ b/crates/uv-installer/src/installer.rs @@ -1,5 +1,9 @@ +use std::convert; + use anyhow::{Context, Error, Result}; +use install_wheel_rs::{linker::LinkMode, Layout}; use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; +use tokio::sync::oneshot; use tracing::instrument; use distribution_types::CachedDist; @@ -7,7 +11,7 @@ use uv_python::PythonEnvironment; pub struct Installer<'a> { venv: &'a PythonEnvironment, - link_mode: install_wheel_rs::linker::LinkMode, + link_mode: LinkMode, reporter: Option>, installer_name: Option, } @@ -17,7 +21,7 @@ impl<'a> Installer<'a> { pub fn new(venv: &'a PythonEnvironment) -> Self { Self { venv, - link_mode: install_wheel_rs::linker::LinkMode::default(), + link_mode: LinkMode::default(), reporter: None, installer_name: Some("uv".to_string()), } @@ -25,7 +29,7 @@ impl<'a> Installer<'a> { /// Set the [`LinkMode`][`install_wheel_rs::linker::LinkMode`] to use for this installer. #[must_use] - pub fn with_link_mode(self, link_mode: install_wheel_rs::linker::LinkMode) -> Self { + pub fn with_link_mode(self, link_mode: LinkMode) -> Self { Self { link_mode, ..self } } @@ -49,33 +53,73 @@ impl<'a> Installer<'a> { /// Install a set of wheels into a Python virtual environment. #[instrument(skip_all, fields(num_wheels = %wheels.len()))] - pub fn install(self, wheels: &[CachedDist]) -> Result<()> { - let layout = self.venv.interpreter().layout(); - tokio::task::block_in_place(|| { - wheels.par_iter().try_for_each(|wheel| { - install_wheel_rs::linker::install_wheel( - &layout, - wheel.path(), - wheel.filename(), - wheel - .parsed_url()? - .as_ref() - .map(pypi_types::DirectUrl::try_from) - .transpose()? - .as_ref(), - self.installer_name.as_deref(), - self.link_mode, - ) - .with_context(|| format!("Failed to install: {} ({wheel})", wheel.filename()))?; + pub async fn install(self, wheels: Vec) -> Result> { + let (tx, rx) = oneshot::channel(); - if let Some(reporter) = self.reporter.as_ref() { - reporter.on_install_progress(wheel); - } + let Self { + venv, + link_mode, + reporter, + installer_name, + } = self; + let layout = venv.interpreter().layout(); - Ok::<(), Error>(()) - }) - }) + rayon::spawn(move || { + let result = install(wheels, layout, installer_name, link_mode, reporter); + tx.send(result).unwrap(); + }); + + rx.await + .map_err(|_| anyhow::anyhow!("`install_blocking` task panicked")) + .and_then(convert::identity) } + + /// Install a set of wheels into a Python virtual environment synchronously. + #[instrument(skip_all, fields(num_wheels = %wheels.len()))] + pub fn install_blocking(self, wheels: Vec) -> Result> { + install( + wheels, + self.venv.interpreter().layout(), + self.installer_name, + self.link_mode, + self.reporter, + ) + } +} + +/// Install a set of wheels into a Python virtual environment synchronously. +#[instrument(skip_all, fields(num_wheels = %wheels.len()))] +fn install( + wheels: Vec, + layout: Layout, + installer_name: Option, + link_mode: LinkMode, + reporter: Option>, +) -> Result> { + wheels.par_iter().try_for_each(|wheel| { + install_wheel_rs::linker::install_wheel( + &layout, + wheel.path(), + wheel.filename(), + wheel + .parsed_url()? + .as_ref() + .map(pypi_types::DirectUrl::try_from) + .transpose()? + .as_ref(), + installer_name.as_deref(), + link_mode, + ) + .with_context(|| format!("Failed to install: {} ({wheel})", wheel.filename()))?; + + if let Some(reporter) = reporter.as_ref() { + reporter.on_install_progress(wheel); + } + + Ok::<(), Error>(()) + })?; + + Ok(wheels) } pub trait Reporter: Send + Sync { diff --git a/crates/uv/src/commands/pip/operations.rs b/crates/uv/src/commands/pip/operations.rs index 906e8eb5c..0fd4cff54 100644 --- a/crates/uv/src/commands/pip/operations.rs +++ b/crates/uv/src/commands/pip/operations.rs @@ -442,13 +442,16 @@ pub(crate) async fn install( } // Install the resolved distributions. - let wheels = wheels.into_iter().chain(cached).collect::>(); + let mut wheels = wheels.into_iter().chain(cached).collect::>(); if !wheels.is_empty() { let start = std::time::Instant::now(); - uv_installer::Installer::new(venv) + wheels = uv_installer::Installer::new(venv) .with_link_mode(link_mode) .with_reporter(InstallReporter::from(printer).with_length(wheels.len() as u64)) - .install(&wheels)?; + // This technically can block the runtime, but we are on the main thread and + // have no other running tasks at this point, so this lets us avoid spawning a blocking + // task. + .install_blocking(wheels)?; let s = if wheels.len() == 1 { "" } else { "s" }; writeln!( diff --git a/crates/uv/src/main.rs b/crates/uv/src/main.rs index 14b30151c..523bf8006 100644 --- a/crates/uv/src/main.rs +++ b/crates/uv/src/main.rs @@ -1019,7 +1019,7 @@ fn main() -> ExitCode { let result = if let Ok(stack_size) = env::var("UV_STACK_SIZE") { let stack_size = stack_size.parse().expect("Invalid stack size"); let tokio_main = move || { - let runtime = tokio::runtime::Builder::new_multi_thread() + let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .thread_stack_size(stack_size) .build() @@ -1041,7 +1041,7 @@ fn main() -> ExitCode { .join() .expect("Tokio executor failed, was there a panic?") } else { - let runtime = tokio::runtime::Builder::new_multi_thread() + let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .expect("Failed building the Runtime");