Skip to content

Commit

Permalink
Use install concurrency for bytecode compilation too (#11615)
Browse files Browse the repository at this point in the history
Instead of always using all available threads for bytecode compilation,
respect `UV_CONCURRENT_INSTALLS`, so the parallelism is configurable
instead of hardcoded. We reuse the install limit since bytecode
compilation only runs after install.
  • Loading branch information
konstin authored Feb 20, 2025
1 parent ab551ea commit f9b638a
Show file tree
Hide file tree
Showing 6 changed files with 25 additions and 16 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/uv-dev/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ workspace = true
uv-cache = { workspace = true, features = ["clap"] }
uv-cli = { workspace = true }
uv-client = { workspace = true }
uv-configuration = { workspace = true }
uv-distribution-filename = { workspace = true }
uv-distribution-types = { workspace = true }
uv-extract = { workspace = true, optional = true }
Expand Down
2 changes: 2 additions & 0 deletions crates/uv-dev/src/compile.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use clap::Parser;
use tracing::info;

use uv_cache::{Cache, CacheArgs};
use uv_configuration::Concurrency;
use uv_python::{EnvironmentPreference, PythonEnvironment, PythonRequest};

#[derive(Parser)]
Expand Down Expand Up @@ -33,6 +34,7 @@ pub(crate) async fn compile(args: CompileArgs) -> anyhow::Result<()> {
let files = uv_installer::compile_tree(
&fs_err::canonicalize(args.root)?,
&interpreter,
&Concurrency::default(),
cache.root(),
)
.await?;
Expand Down
12 changes: 5 additions & 7 deletions crates/uv-installer/src/compile.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
use std::num::NonZeroUsize;
use std::panic::AssertUnwindSafe;
use std::path::{Path, PathBuf};
use std::process::Stdio;
Expand All @@ -14,6 +13,7 @@ use tokio::sync::oneshot;
use tracing::{debug, instrument};
use walkdir::WalkDir;

use uv_configuration::Concurrency;
use uv_fs::Simplified;
use uv_static::EnvVars;
use uv_warnings::warn_user;
Expand Down Expand Up @@ -71,28 +71,26 @@ pub enum CompileError {
pub async fn compile_tree(
dir: &Path,
python_executable: &Path,
concurrency: &Concurrency,
cache: &Path,
) -> Result<usize, CompileError> {
debug_assert!(
dir.is_absolute(),
"compileall doesn't work with relative paths: `{}`",
dir.display()
);
let worker_count = std::thread::available_parallelism().unwrap_or_else(|err| {
warn_user!("Couldn't determine number of cores, compiling with a single thread: {err}");
NonZeroUsize::MIN
});
let worker_count = concurrency.installs;

// A larger buffer is significantly faster than just 1 or the worker count.
let (sender, receiver) = async_channel::bounded::<PathBuf>(worker_count.get() * 10);
let (sender, receiver) = async_channel::bounded::<PathBuf>(worker_count * 10);

// Running Python with an actual file will produce better error messages.
let tempdir = tempdir_in(cache).map_err(CompileError::TempFile)?;
let pip_compileall_py = tempdir.path().join("pip_compileall.py");

debug!("Starting {} bytecode compilation workers", worker_count);
let mut worker_handles = Vec::new();
for _ in 0..worker_count.get() {
for _ in 0..worker_count {
let (tx, rx) = oneshot::channel();

let worker = worker(
Expand Down
23 changes: 15 additions & 8 deletions crates/uv/src/commands/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ pub(crate) use tool::uninstall::uninstall as tool_uninstall;
pub(crate) use tool::update_shell::update_shell as tool_update_shell;
pub(crate) use tool::upgrade::upgrade as tool_upgrade;
use uv_cache::Cache;
use uv_configuration::Concurrency;
use uv_distribution_types::InstalledMetadata;
use uv_fs::{Simplified, CWD};
use uv_installer::compile_tree;
Expand Down Expand Up @@ -147,21 +148,27 @@ pub(super) struct DryRunEvent<T: Display> {
/// See the `--compile` option on `pip sync` and `pip install`.
pub(super) async fn compile_bytecode(
venv: &PythonEnvironment,
concurrency: &Concurrency,
cache: &Cache,
printer: Printer,
) -> anyhow::Result<()> {
let start = std::time::Instant::now();
let mut files = 0;
for site_packages in venv.site_packages() {
let site_packages = CWD.join(site_packages);
files += compile_tree(&site_packages, venv.python_executable(), cache.root())
.await
.with_context(|| {
format!(
"Failed to bytecode-compile Python file in: {}",
site_packages.user_display()
)
})?;
files += compile_tree(
&site_packages,
venv.python_executable(),
concurrency,
cache.root(),
)
.await
.with_context(|| {
format!(
"Failed to bytecode-compile Python file in: {}",
site_packages.user_display()
)
})?;
}
let s = if files == 1 { "" } else { "s" };
writeln!(
Expand Down
2 changes: 1 addition & 1 deletion crates/uv/src/commands/pip/operations.rs
Original file line number Diff line number Diff line change
Expand Up @@ -554,7 +554,7 @@ pub(crate) async fn install(
}

if compile {
compile_bytecode(venv, cache, printer).await?;
compile_bytecode(venv, &concurrency, cache, printer).await?;
}

// Construct a summary of the changes made to the environment.
Expand Down

0 comments on commit f9b638a

Please sign in to comment.