From 85162d111152eab2d4282bb2b8b929fa5c9b7825 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 10 Oct 2023 23:46:30 -0400 Subject: [PATCH] Parallelize wheel installations with Rayon (#84) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit It looks like using _either_ async Rust with a `JoinSet` _or_ parallelizing a fixed threadpool with Rayon provide about a ~5% speed-up over our current serial approach: ```console ❯ hyperfine --runs 30 --warmup 5 --prepare "./target/release/puffin venv .venv" \ "./target/release/rayon sync ./scripts/benchmarks/requirements-large.txt" \ "./target/release/async sync ./scripts/benchmarks/requirements-large.txt" \ "./target/release/main sync ./scripts/benchmarks/requirements-large.txt" Benchmark 1: ./target/release/rayon sync ./scripts/benchmarks/requirements-large.txt Time (mean ± σ): 295.7 ms ± 16.9 ms [User: 28.6 ms, System: 263.3 ms] Range (min … max): 249.2 ms … 315.9 ms 30 runs Benchmark 2: ./target/release/async sync ./scripts/benchmarks/requirements-large.txt Time (mean ± σ): 296.2 ms ± 20.2 ms [User: 36.1 ms, System: 340.1 ms] Range (min … max): 258.0 ms … 359.4 ms 30 runs Benchmark 3: ./target/release/main sync ./scripts/benchmarks/requirements-large.txt Time (mean ± σ): 306.6 ms ± 19.5 ms [User: 25.3 ms, System: 220.5 ms] Range (min … max): 269.6 ms … 332.2 ms 30 runs Summary './target/release/rayon sync ./scripts/benchmarks/requirements-large.txt' ran 1.00 ± 0.09 times faster than './target/release/async sync ./scripts/benchmarks/requirements-large.txt' 1.04 ± 0.09 times faster than './target/release/main sync ./scripts/benchmarks/requirements-large.txt' ``` It's much easier to just parallelize with Rayon and avoid async in the underlying wheel code, so this PR takes that approach for now. --- .../install-wheel-rs/src/install_location.rs | 2 +- crates/install-wheel-rs/src/unpacked.rs | 64 +++++++++++++------ crates/install-wheel-rs/src/wheel.rs | 10 +-- crates/puffin-installer/src/installer.rs | 33 +++++----- 4 files changed, 65 insertions(+), 44 deletions(-) diff --git a/crates/install-wheel-rs/src/install_location.rs b/crates/install-wheel-rs/src/install_location.rs index 429cd7ae5..cd03c051d 100644 --- a/crates/install-wheel-rs/src/install_location.rs +++ b/crates/install-wheel-rs/src/install_location.rs @@ -73,7 +73,7 @@ impl AsRef for LockedDir { /// We use a lockfile to prevent multiple instance writing stuff on the same time /// As of pip 22.0, e.g. `pip install numpy; pip install numpy; pip install numpy` will /// non-deterministically fail. -pub struct InstallLocation> { +pub struct InstallLocation { /// absolute path venv_base: T, python_version: (u8, u8), diff --git a/crates/install-wheel-rs/src/unpacked.rs b/crates/install-wheel-rs/src/unpacked.rs index e0dbaa007..807211014 100644 --- a/crates/install-wheel-rs/src/unpacked.rs +++ b/crates/install-wheel-rs/src/unpacked.rs @@ -10,7 +10,7 @@ use fs_err::File; use mailparse::MailHeaderMap; use tracing::{debug, span, Level}; -use crate::install_location::{InstallLocation, LockedDir}; +use crate::install_location::InstallLocation; use crate::wheel::{ extra_dist_info, install_data, parse_wheel_version, read_scripts_from_section, write_script_entrypoints, @@ -24,7 +24,10 @@ use crate::{read_record_file, Error, Script}; /// /// /// Wheel 1.0: -pub fn install_wheel(location: &InstallLocation, wheel: &Path) -> Result<(), Error> { +pub fn install_wheel( + location: &InstallLocation>, + wheel: impl AsRef, +) -> Result<(), Error> { let base_location = location.venv_base(); // TODO(charlie): Pass this in. @@ -43,8 +46,8 @@ pub fn install_wheel(location: &InstallLocation, wheel: &Path) -> Res .join("site-packages") }; - let dist_info_prefix = find_dist_info(wheel)?; - let (name, _version) = read_metadata(&dist_info_prefix, wheel)?; + let dist_info_prefix = find_dist_info(&wheel)?; + let (name, _version) = read_metadata(&dist_info_prefix, &wheel)?; let _my_span = span!(Level::DEBUG, "install_wheel", name); @@ -52,7 +55,9 @@ pub fn install_wheel(location: &InstallLocation, wheel: &Path) -> Res // https://packaging.python.org/en/latest/specifications/binary-distribution-format/#installing-a-wheel-distribution-1-0-py32-none-any-whl // > 1.a Parse distribution-1.0.dist-info/WHEEL. // > 1.b Check that installer is compatible with Wheel-Version. Warn if minor version is greater, abort if major version is greater. - let wheel_file_path = wheel.join(format!("{dist_info_prefix}.dist-info/WHEEL")); + let wheel_file_path = wheel + .as_ref() + .join(format!("{dist_info_prefix}.dist-info/WHEEL")); let wheel_text = std::fs::read_to_string(&wheel_file_path)?; parse_wheel_version(&wheel_text)?; @@ -60,15 +65,19 @@ pub fn install_wheel(location: &InstallLocation, wheel: &Path) -> Res // > 1.d Else unpack archive into platlib (site-packages). // We always install in the same virtualenv site packages debug!(name, "Extracting file"); - let num_unpacked = unpack_wheel_files(&site_packages, wheel)?; + let num_unpacked = unpack_wheel_files(&site_packages, &wheel)?; debug!(name, "Extracted {num_unpacked} files"); // Read the RECORD file. - let mut record_file = File::open(&wheel.join(format!("{dist_info_prefix}.dist-info/RECORD")))?; + let mut record_file = File::open( + wheel + .as_ref() + .join(format!("{dist_info_prefix}.dist-info/RECORD")), + )?; let mut record = read_record_file(&mut record_file)?; debug!(name, "Writing entrypoints"); - let (console_scripts, gui_scripts) = parse_scripts(wheel, &dist_info_prefix, None)?; + let (console_scripts, gui_scripts) = parse_scripts(&wheel, &dist_info_prefix, None)?; write_script_entrypoints(&site_packages, location, &console_scripts, &mut record)?; write_script_entrypoints(&site_packages, location, &gui_scripts, &mut record)?; @@ -117,7 +126,7 @@ pub fn install_wheel(location: &InstallLocation, wheel: &Path) -> Res /// Either way, we just search the wheel for the name /// /// -fn find_dist_info(path: &Path) -> Result { +fn find_dist_info(path: impl AsRef) -> Result { // Iterate over `path` to find the `.dist-info` directory. It should be at the top-level. let Some(dist_info) = std::fs::read_dir(path)?.find_map(|entry| { let entry = entry.ok()?; @@ -147,8 +156,13 @@ fn find_dist_info(path: &Path) -> Result { } /// -fn read_metadata(dist_info_prefix: &str, wheel: &Path) -> Result<(String, String), Error> { - let metadata_file = wheel.join(format!("{dist_info_prefix}.dist-info/METADATA")); +fn read_metadata( + dist_info_prefix: &str, + wheel: impl AsRef, +) -> Result<(String, String), Error> { + let metadata_file = wheel + .as_ref() + .join(format!("{dist_info_prefix}.dist-info/METADATA")); // Read into a buffer. let mut content = Vec::new(); @@ -197,11 +211,13 @@ fn read_metadata(dist_info_prefix: &str, wheel: &Path) -> Result<(String, String /// /// Extras are supposed to be ignored, which happens if you pass None for extras fn parse_scripts( - wheel: &Path, + wheel: impl AsRef, dist_info_prefix: &str, extras: Option<&[String]>, ) -> Result<(Vec