Add an option to bytecode compile during installation (#2086)

Add a `--compile` option to `pip install` and `pip sync`.

I chose to implement this as a separate pass over the entire venv. If we
wanted to compile during installation, we'd have to make sure that
writing is exclusive, to avoid concurrent processes writing broken
`.pyc` files. Additionally, this ensures that the entire site-packages
are bytecode compiled, even if there are packages that aren't from this
`uv` invocation. The disadvantage is that we do not update RECORD and
rely on this comment from [PEP 491](https://peps.python.org/pep-0491/):

> Uninstallers should be smart enough to remove .pyc even if it is not
mentioned in RECORD.

If this is a problem we can change it to run during installation and
write RECORD entries.

Internally, this is implemented as an async work-stealing subprocess
worker pool. The producer is a directory traversal over site-packages,
sending each `.py` file to a bounded async FIFO queue/channel. Each
worker has a long-running python process. It pops the queue to get a
single path (or exists if the channel is closed), then sends it to
stdin, waits until it's informed that the compilation is done through a
line on stdout, and repeat. This is fast, e.g. installing `jupyter
plotly` on Python 3.12 it processes 15876 files in 319ms with 32 threads
(vs. 3.8s with a single core). The python processes internally calls
`compileall.compile_file`, the same as pip.

Like pip, we ignore and silence all compilation errors
(#1559). There is a 10s timeout to
handle the case when the workers got stuck. For the reviewers, please
check if i missed any spots where we could deadlock, this is the hardest
part of this PR.

I've added `uv-dev compile <dir>` and `uv-dev clear-compile <dir>`
commands, mainly for my own benchmarking. I don't want to expose them in
`uv`, they almost certainly not the correct workflow and we don't want
to support them.

Fixes #1788
Closes #1559
Closes #1928

Cargo.toml

@@ -20,6 +20,7 @@ license = "MIT OR Apache-2.0"
 anstream = { version = "0.6.13" }
 anyhow = { version = "1.0.80" }
 async-compression = { version = "0.4.6" }
+async-channel = { version = "2.2.0" }
 async-trait = { version = "0.1.77" }
 async_http_range_reader = { version = "0.7.0" }
 async_zip = { git = "https://github.com/charliermarsh/rs-async-zip", rev = "d76801da0943de985254fc6255c0e476b57c5836", features = ["deflate"] }

crates/uv-dev/Cargo.toml

@@ -61,6 +61,7 @@ tracing-durations-export = { workspace = true, features = ["plot"] }
 tracing-indicatif = { workspace = true }
 tracing-subscriber = { workspace = true }
 url = { workspace = true }
+walkdir = { workspace = true }
 which = { workspace = true }
 
 [target.'cfg(target_os = "windows")'.dependencies]

crates/uv-dev/src/clear_compile.rs

@@ -0,0 +1,33 @@
+use std::path::PathBuf;
+
+use clap::Parser;
+use tracing::info;
+use walkdir::WalkDir;
+
+#[derive(Parser)]
+pub(crate) struct ClearCompileArgs {
+    /// Compile all `.py` in this or any subdirectory to bytecode
+    root: PathBuf,
+}
+
+pub(crate) fn clear_compile(args: &ClearCompileArgs) -> anyhow::Result<()> {
+    let mut removed_files = 0;
+    let mut removed_directories = 0;
+    for entry in WalkDir::new(&args.root).contents_first(true) {
+        let entry = entry?;
+        let metadata = entry.metadata()?;
+        if metadata.is_file() {
+            if entry.path().extension().is_some_and(|ext| ext == "pyc") {
+                fs_err::remove_file(entry.path())?;
+                removed_files += 1;
+            }
+        } else if metadata.is_dir() {
+            if entry.file_name() == "__pycache__" {
+                fs_err::remove_dir(entry.path())?;
+                removed_directories += 1;
+            }
+        }
+    }
+    info!("Removed {removed_files} files and {removed_directories} directories");
+    Ok(())
+}

crates/uv-dev/src/compile.rs

@@ -0,0 +1,37 @@
+use std::path::PathBuf;
+
+use clap::Parser;
+use platform_host::Platform;
+use tracing::info;
+use uv_cache::{Cache, CacheArgs};
+use uv_interpreter::PythonEnvironment;
+
+#[derive(Parser)]
+pub(crate) struct CompileArgs {
+    /// Compile all `.py` in this or any subdirectory to bytecode
+    root: PathBuf,
+    python: Option<PathBuf>,
+    #[command(flatten)]
+    cache_args: CacheArgs,
+}
+
+pub(crate) async fn compile(args: CompileArgs) -> anyhow::Result<()> {
+    let cache = Cache::try_from(args.cache_args)?;
+
+    let interpreter = if let Some(python) = args.python {
+        python
+    } else {
+        let platform = Platform::current()?;
+        let venv = PythonEnvironment::from_virtualenv(platform, &cache)?;
+        venv.python_executable().to_path_buf()
+    };
+
+    let files = uv_installer::compile_tree(
+        &fs_err::canonicalize(args.root)?,
+        &interpreter,
+        cache.root(),
+    )
+    .await?;
+    info!("Compiled {files} files");
+    Ok(())
+}

crates/uv-dev/src/main.rs

@@ -19,6 +19,8 @@ use tracing_subscriber::EnvFilter;
 use resolve_many::ResolveManyArgs;
 
 use crate::build::{build, BuildArgs};
+use crate::clear_compile::ClearCompileArgs;
+use crate::compile::CompileArgs;
 use crate::install_many::InstallManyArgs;
 use crate::render_benchmarks::RenderBenchmarksArgs;
 use crate::resolve_cli::ResolveCliArgs;
@@ -41,6 +43,8 @@ static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
 mod build;
+mod clear_compile;
+mod compile;
 mod install_many;
 mod render_benchmarks;
 mod resolve_cli;
@@ -67,6 +71,10 @@ enum Cli {
     Resolve(ResolveCliArgs),
     WheelMetadata(WheelMetadataArgs),
     RenderBenchmarks(RenderBenchmarksArgs),
+    /// Compile all `.py` to `.pyc` files in the tree.
+    Compile(CompileArgs),
+    /// Remove all `.pyc` in the tree.
+    ClearCompile(ClearCompileArgs),
 }
 
 #[instrument] // Anchor span to check for overhead
@@ -88,6 +96,8 @@ async fn run() -> Result<()> {
         }
         Cli::WheelMetadata(args) => wheel_metadata::wheel_metadata(args).await?,
         Cli::RenderBenchmarks(args) => render_benchmarks::render_benchmarks(&args)?,
+        Cli::Compile(args) => compile::compile(args).await?,
+        Cli::ClearCompile(args) => clear_compile::clear_compile(&args)?,
     }
     Ok(())
 }

crates/uv-installer/Cargo.toml

@@ -31,8 +31,10 @@ uv-git = { path = "../uv-git", features = ["vendored-openssl"] }
 uv-interpreter = { path = "../uv-interpreter" }
 uv-normalize = { path = "../uv-normalize" }
 uv-traits = { path = "../uv-traits" }
+uv-warnings = { path = "../uv-warnings" }
 
 anyhow = { workspace = true }
+async-channel = { workspace = true }
 fs-err = { workspace = true }
 futures = { workspace = true }
 pyproject-toml = { workspace = true }
@@ -45,3 +47,4 @@ tokio = { workspace = true }
 toml = { workspace = true }
 tracing = { workspace = true }
 url = { workspace = true }
+walkdir = { workspace = true }