Add UV_COMPILE_BYTECODE_TIMEOUT environment variable

crates/uv-installer/src/compile.rs

@@ -2,7 +2,7 @@ use std::panic::AssertUnwindSafe;
 use std::path::{Path, PathBuf};
 use std::process::Stdio;
 use std::time::Duration;
-use std::{io, panic};
+use std::{env, io, panic};
 
 use async_channel::{Receiver, SendError};
 use tempfile::tempdir_in;
@@ -20,7 +20,7 @@ use uv_warnings::warn_user;
 
 const COMPILEALL_SCRIPT: &str = include_str!("pip_compileall.py");
 /// This is longer than any compilation should ever take.
-const COMPILE_TIMEOUT: Duration = Duration::from_secs(60);
+const DEFAULT_COMPILE_TIMEOUT: Duration = Duration::from_secs(60);
 
 #[derive(Debug, Error)]
 pub enum CompileError {
@@ -55,6 +55,8 @@ pub enum CompileError {
     },
     #[error("Python startup timed out ({}s)", _0.as_secs_f32())]
     StartupTimeout(Duration),
+    #[error("Got invalid value from environment for {var}: {message}.")]
+    EnvironmentError { var: &'static str, message: String },
 }
 
 /// Bytecode compile all file in `dir` using a pool of Python interpreters running a Python script
@@ -88,6 +90,29 @@ pub async fn compile_tree(
     let tempdir = tempdir_in(cache).map_err(CompileError::TempFile)?;
     let pip_compileall_py = tempdir.path().join("pip_compileall.py");
 
+    let timeout: Option<Duration> = match env::var(EnvVars::UV_COMPILE_BYTECODE_TIMEOUT) {
+        Ok(value) => {
+            if value == "0" {
+                debug!("Disabling bytecode compilation timeout");
+                None
+            } else {
+                if let Ok(duration) = value.parse::<u64>().map(Duration::from_secs) {
+                    debug!(
+                        "Using bytecode compilation timeout of {}s",
+                        duration.as_secs()
+                    );
+                    Some(duration)
+                } else {
+                    return Err(CompileError::EnvironmentError {
+                        var: "UV_COMPILE_BYTECODE_TIMEOUT",
+                        message: format!("Expected an integer number of seconds, got \"{value}\""),
+                    });
+                }
+            }
+        }
+        Err(_) => Some(DEFAULT_COMPILE_TIMEOUT),
+    };
+
     debug!("Starting {} bytecode compilation workers", worker_count);
     let mut worker_handles = Vec::new();
     for _ in 0..worker_count {
@@ -98,6 +123,7 @@ pub async fn compile_tree(
             python_executable.to_path_buf(),
             pip_compileall_py.clone(),
             receiver.clone(),
+            timeout,
         );
 
         // Spawn each worker on a dedicated thread.
@@ -189,6 +215,7 @@ async fn worker(
     interpreter: PathBuf,
     pip_compileall_py: PathBuf,
     receiver: Receiver<PathBuf>,
+    timeout: Option<Duration>,
 ) -> Result<(), CompileError> {
     fs_err::tokio::write(&pip_compileall_py, COMPILEALL_SCRIPT)
         .await
@@ -208,12 +235,17 @@ async fn worker(
             }
         }
     };
+
     // Handle a broken `python` by using a timeout, one that's higher than any compilation
     // should ever take.
     let (mut bytecode_compiler, child_stdin, mut child_stdout, mut child_stderr) =
-        tokio::time::timeout(COMPILE_TIMEOUT, wait_until_ready)
-            .await
-            .map_err(|_| CompileError::StartupTimeout(COMPILE_TIMEOUT))??;
+        if let Some(duration) = timeout {
+            tokio::time::timeout(duration, wait_until_ready)
+                .await
+                .map_err(|_| CompileError::StartupTimeout(timeout.unwrap()))??
+        } else {
+            wait_until_ready.await?
+        };
 
     let stderr_reader = tokio::task::spawn(async move {
         let mut child_stderr_collected: Vec<u8> = Vec::new();
@@ -223,7 +255,7 @@ async fn worker(
         Ok(child_stderr_collected)
     });
 
-    let result = worker_main_loop(receiver, child_stdin, &mut child_stdout).await;
+    let result = worker_main_loop(receiver, child_stdin, &mut child_stdout, timeout).await;
     // Reap the process to avoid zombies.
     let _ = bytecode_compiler.kill().await;
 
@@ -340,6 +372,7 @@ async fn worker_main_loop(
     receiver: Receiver<PathBuf>,
     mut child_stdin: ChildStdin,
     child_stdout: &mut BufReader<ChildStdout>,
+    timeout: Option<Duration>,
 ) -> Result<(), CompileError> {
     let mut out_line = String::new();
     while let Ok(source_file) = receiver.recv().await {
@@ -372,12 +405,16 @@ async fn worker_main_loop(
 
         // Handle a broken `python` by using a timeout, one that's higher than any compilation
         // should ever take.
-        tokio::time::timeout(COMPILE_TIMEOUT, python_handle)
-            .await
-            .map_err(|_| CompileError::CompileTimeout {
-                elapsed: COMPILE_TIMEOUT,
-                source_file: source_file.clone(),
-            })??;
+        if let Some(duration) = timeout {
+            tokio::time::timeout(duration, python_handle)
+                .await
+                .map_err(|_| CompileError::CompileTimeout {
+                    elapsed: duration,
+                    source_file: source_file.clone(),
+                })??;
+        } else {
+            python_handle.await?;
+        }
 
         // This is a sanity check, if we don't get the path back something has gone wrong, e.g.
         // we're not actually running a python interpreter.

crates/uv-static/src/env_vars.rs

@@ -162,6 +162,9 @@ impl EnvVars {
     /// will compile Python source files to bytecode after installation.
     pub const UV_COMPILE_BYTECODE: &'static str = "UV_COMPILE_BYTECODE";
 
+    /// Timeout (in seconds) for bytecode compilation.
+    pub const UV_COMPILE_BYTECODE_TIMEOUT: &'static str = "UV_COMPILE_BYTECODE_TIMEOUT";
+
     /// Equivalent to the `--no-editable` command-line argument. If set, uv
     /// installs any editable dependencies, including the project and any workspace members, as
     /// non-editable

docs/reference/environment.md

@@ -26,6 +26,10 @@ directory for caching instead of the default cache directory.
 Equivalent to the `--compile-bytecode` command-line argument. If set, uv
 will compile Python source files to bytecode after installation.
 
+### `UV_COMPILE_BYTECODE_TIMEOUT`
+
+Timeout (in seconds) for bytecode compilation.
+
 ### `UV_CONCURRENT_BUILDS`
 
 Sets the maximum number of source distributions that uv will build