Reuse build (virtual) environments across resolution and installation (#14338)

## Summary

The basic idea here is that we can (should) reuse a build environment
across resolution (`prepare_metadata_for_build_wheel`) and installation.
This also happens to solve the build-PyTorch-from-source problem, since
we use a consistent build environment between the invocations.

Since `SourceDistributionBuilder` is stateless, we instead store the
builds on `BuildContext`, and we key them by various properties: the
underlying interpreter, the configuration settings, etc. This just
ensures that if we build the same package twice within a process, we
don't accidentally reuse an incompatible build (virtual) environment.
(Note that still drop build environments at the end of the command, and
don't attempt to reuse them across processes.)

Closes #14269.

Cargo.lock

@@ -5937,9 +5937,8 @@ name = "uv-types"
 version = "0.0.1"
 dependencies = [
  "anyhow",
- "boxcar",
+ "dashmap",
  "rustc-hash",
- "tempfile",
  "thiserror 2.0.12",
  "uv-cache",
  "uv-configuration",

crates/uv-build-frontend/src/lib.rs

@@ -259,8 +259,6 @@ pub struct SourceBuild {
     environment_variables: FxHashMap<OsString, OsString>,
     /// Runner for Python scripts.
     runner: PythonRunner,
-    /// A file lock representing the source tree, currently only used with setuptools.
-    _source_tree_lock: Option<LockedFile>,
 }
 
 impl SourceBuild {
@@ -394,23 +392,6 @@ impl SourceBuild {
             OsString::from(venv.scripts())
         };
 
-        // Depending on the command, setuptools puts `*.egg-info`, `build/`, and `dist/` in the
-        // source tree, and concurrent invocations of setuptools using the same source dir can
-        // stomp on each other. We need to lock something to fix that, but we don't want to dump a
-        // `.lock` file into the source tree that the user will need to .gitignore. Take a global
-        // proxy lock instead.
-        let mut source_tree_lock = None;
-        if pep517_backend.is_setuptools() {
-            debug!("Locking the source tree for setuptools");
-            let canonical_source_path = source_tree.canonicalize()?;
-            let lock_path = std::env::temp_dir().join(format!(
-                "uv-setuptools-{}.lock",
-                cache_digest(&canonical_source_path)
-            ));
-            source_tree_lock =
-                Some(LockedFile::acquire(lock_path, source_tree.to_string_lossy()).await?);
-        }
-
         // Create the PEP 517 build environment. If build isolation is disabled, we assume the build
         // environment is already setup.
         let runner = PythonRunner::new(concurrent_builds, level);
@@ -457,10 +438,30 @@ impl SourceBuild {
             environment_variables,
             modified_path,
             runner,
-            _source_tree_lock: source_tree_lock,
         })
     }
 
+    /// Acquire a lock on the source tree, if necessary.
+    async fn acquire_lock(&self) -> Result<Option<LockedFile>, Error> {
+        // Depending on the command, setuptools puts `*.egg-info`, `build/`, and `dist/` in the
+        // source tree, and concurrent invocations of setuptools using the same source dir can
+        // stomp on each other. We need to lock something to fix that, but we don't want to dump a
+        // `.lock` file into the source tree that the user will need to .gitignore. Take a global
+        // proxy lock instead.
+        let mut source_tree_lock = None;
+        if self.pep517_backend.is_setuptools() {
+            debug!("Locking the source tree for setuptools");
+            let canonical_source_path = self.source_tree.canonicalize()?;
+            let lock_path = env::temp_dir().join(format!(
+                "uv-setuptools-{}.lock",
+                cache_digest(&canonical_source_path)
+            ));
+            source_tree_lock =
+                Some(LockedFile::acquire(lock_path, self.source_tree.to_string_lossy()).await?);
+        }
+        Ok(source_tree_lock)
+    }
+
     async fn get_resolved_requirements(
         build_context: &impl BuildContext,
         source_build_context: SourceBuildContext,
@@ -631,6 +632,9 @@ impl SourceBuild {
             return Ok(Some(metadata_dir.clone()));
         }
 
+        // Lock the source tree, if necessary.
+        let _lock = self.acquire_lock().await?;
+
         // Hatch allows for highly dynamic customization of metadata via hooks. In such cases, Hatch
         // can't uphold the PEP 517 contract, in that the metadata Hatch would return by
         // `prepare_metadata_for_build_wheel` isn't guaranteed to match that of the built wheel.
@@ -749,6 +753,9 @@ impl SourceBuild {
 
     /// Perform a PEP 517 build for a wheel or source distribution (sdist).
     async fn pep517_build(&self, output_dir: &Path) -> Result<String, Error> {
+        // Lock the source tree, if necessary.
+        let _lock = self.acquire_lock().await?;
+
         // Write the hook output to a file so that we can read it back reliably.
         let outfile = self
             .temp_dir
@@ -862,10 +869,6 @@ impl SourceBuild {
 }
 
 impl SourceBuildTrait for SourceBuild {
-    fn into_build_dir(self) -> TempDir {
-        self.temp_dir
-    }
-
     async fn metadata(&mut self) -> Result<Option<PathBuf>, AnyErrorBuild> {
         Ok(self.get_metadata_without_build().await?)
     }

crates/uv-configuration/src/build_options.rs

@@ -4,7 +4,7 @@ use uv_pep508::PackageName;
 
 use crate::{PackageNameSpecifier, PackageNameSpecifiers};
 
-#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
+#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash)]
 pub enum BuildKind {
     /// A PEP 517 wheel build.
     #[default]

crates/uv-configuration/src/sources.rs

@@ -1,4 +1,6 @@
-#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+#[derive(
+    Debug, Default, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize,
+)]
 #[serde(rename_all = "kebab-case", deny_unknown_fields)]
 pub enum SourceStrategy {
     /// Use `tool.uv.sources` when resolving dependencies.

crates/uv-dispatch/src/lib.rs

@@ -2,12 +2,13 @@
 //! [installer][`uv_installer`] and [build][`uv_build`] through [`BuildDispatch`]
 //! implementing [`BuildContext`].
 
+use std::ffi::{OsStr, OsString};
+use std::path::Path;
+
 use anyhow::{Context, Result};
 use futures::FutureExt;
 use itertools::Itertools;
 use rustc_hash::FxHashMap;
-use std::ffi::{OsStr, OsString};
-use std::path::Path;
 use thiserror::Error;
 use tracing::{debug, instrument, trace};
 
@@ -179,7 +180,7 @@ impl BuildContext for BuildDispatch<'_> {
         &self.shared_state.git
     }
 
-    fn build_arena(&self) -> &BuildArena {
+    fn build_arena(&self) -> &BuildArena<SourceBuild> {
         &self.shared_state.build_arena
     }
 
@@ -526,7 +527,7 @@ pub struct SharedState {
     /// The downloaded distributions.
     in_flight: InFlight,
     /// Build directories for any PEP 517 builds executed during resolution or installation.
-    build_arena: BuildArena,
+    build_arena: BuildArena<SourceBuild>,
 }
 
 impl SharedState {
@@ -565,7 +566,7 @@ impl SharedState {
     }
 
     /// Return the [`BuildArena`] used by the [`SharedState`].
-    pub fn build_arena(&self) -> &BuildArena {
+    pub fn build_arena(&self) -> &BuildArena<SourceBuild> {
         &self.build_arena
     }
 }

crates/uv-distribution/src/error.rs

@@ -108,6 +108,8 @@ pub enum Error {
     CacheHeal(String, HashAlgorithm),
     #[error("The source distribution requires Python {0}, but {1} is installed")]
     RequiresPython(VersionSpecifiers, Version),
+    #[error("Failed to identify base Python interpreter")]
+    BaseInterpreter(#[source] std::io::Error),
 
     /// A generic request middleware error happened while making a request.
     /// Refer to the error message for more details.

crates/uv-distribution/src/source/mod.rs

@@ -43,7 +43,7 @@ use uv_normalize::PackageName;
 use uv_pep440::{Version, release_specifiers_to_ranges};
 use uv_platform_tags::Tags;
 use uv_pypi_types::{HashAlgorithm, HashDigest, HashDigests, PyProjectToml, ResolutionMetadata};
-use uv_types::{BuildContext, BuildStack, SourceBuildTrait};
+use uv_types::{BuildContext, BuildKey, BuildStack, SourceBuildTrait};
 use uv_workspace::pyproject::ToolUvSources;
 
 use crate::distribution_database::ManagedClient;
@@ -2297,35 +2297,73 @@ impl<'a, T: BuildContext> SourceDistributionBuilder<'a, T> {
             // In the uv build backend, the normalized filename and the disk filename are the same.
             name.to_string()
         } else {
-            let builder = self
-                .build_context
-                .setup_build(
-                    source_root,
-                    subdirectory,
-                    source_root,
-                    Some(&source.to_string()),
-                    source.as_dist(),
-                    source_strategy,
-                    if source.is_editable() {
-                        BuildKind::Editable
-                    } else {
-                        BuildKind::Wheel
-                    },
-                    BuildOutput::Debug,
-                    self.build_stack.cloned().unwrap_or_default(),
-                )
-                .await
-                .map_err(|err| Error::Build(err.into()))?;
+            // Identify the base Python interpreter to use in the cache key.
+            let base_python = if cfg!(unix) {
+                self.build_context
+                    .interpreter()
+                    .find_base_python()
+                    .map_err(Error::BaseInterpreter)?
+            } else {
+                self.build_context
+                    .interpreter()
+                    .to_base_python()
+                    .map_err(Error::BaseInterpreter)?
+            };
 
-            // Build the wheel.
-            let wheel = builder.wheel(temp_dir.path()).await.map_err(Error::Build)?;
+            let build_kind = if source.is_editable() {
+                BuildKind::Editable
+            } else {
+                BuildKind::Wheel
+            };
 
-            // Store a reference to the build context.
-            self.build_context
-                .build_arena()
-                .push(builder.into_build_dir());
+            let build_key = BuildKey {
+                base_python: base_python.into_boxed_path(),
+                source_root: source_root.to_path_buf().into_boxed_path(),
+                subdirectory: subdirectory
+                    .map(|subdirectory| subdirectory.to_path_buf().into_boxed_path()),
+                source_strategy,
+                build_kind,
+            };
 
-            wheel
+            if let Some(builder) = self.build_context.build_arena().remove(&build_key) {
+                debug!("Creating build environment for: {source}");
+                let wheel = builder.wheel(temp_dir.path()).await.map_err(Error::Build)?;
+
+                // Store the build context.
+                self.build_context.build_arena().insert(build_key, builder);
+
+                wheel
+            } else {
+                debug!("Reusing existing build environment for: {source}");
+
+                let builder = self
+                    .build_context
+                    .setup_build(
+                        source_root,
+                        subdirectory,
+                        source_root,
+                        Some(&source.to_string()),
+                        source.as_dist(),
+                        source_strategy,
+                        if source.is_editable() {
+                            BuildKind::Editable
+                        } else {
+                            BuildKind::Wheel
+                        },
+                        BuildOutput::Debug,
+                        self.build_stack.cloned().unwrap_or_default(),
+                    )
+                    .await
+                    .map_err(|err| Error::Build(err.into()))?;
+
+                // Build the wheel.
+                let wheel = builder.wheel(temp_dir.path()).await.map_err(Error::Build)?;
+
+                // Store the build context.
+                self.build_context.build_arena().insert(build_key, builder);
+
+                wheel
+            }
         };
 
         // Read the metadata from the wheel.
@@ -2380,6 +2418,26 @@ impl<'a, T: BuildContext> SourceDistributionBuilder<'a, T> {
             }
         }
 
+        // Identify the base Python interpreter to use in the cache key.
+        let base_python = if cfg!(unix) {
+            self.build_context
+                .interpreter()
+                .find_base_python()
+                .map_err(Error::BaseInterpreter)?
+        } else {
+            self.build_context
+                .interpreter()
+                .to_base_python()
+                .map_err(Error::BaseInterpreter)?
+        };
+
+        // Determine whether this is an editable or non-editable build.
+        let build_kind = if source.is_editable() {
+            BuildKind::Editable
+        } else {
+            BuildKind::Wheel
+        };
+
         // Set up the builder.
         let mut builder = self
             .build_context
@@ -2390,11 +2448,7 @@ impl<'a, T: BuildContext> SourceDistributionBuilder<'a, T> {
                 Some(&source.to_string()),
                 source.as_dist(),
                 source_strategy,
-                if source.is_editable() {
-                    BuildKind::Editable
-                } else {
-                    BuildKind::Wheel
-                },
+                build_kind,
                 BuildOutput::Debug,
                 self.build_stack.cloned().unwrap_or_default(),
             )
@@ -2403,15 +2457,25 @@ impl<'a, T: BuildContext> SourceDistributionBuilder<'a, T> {
 
         // Build the metadata.
         let dist_info = builder.metadata().await.map_err(Error::Build)?;
+
+        // Store the build context.
+        self.build_context.build_arena().insert(
+            BuildKey {
+                base_python: base_python.into_boxed_path(),
+                source_root: source_root.to_path_buf().into_boxed_path(),
+                subdirectory: subdirectory
+                    .map(|subdirectory| subdirectory.to_path_buf().into_boxed_path()),
+                source_strategy,
+                build_kind,
+            },
+            builder,
+        );
+
+        // Return the `.dist-info` directory, if it exists.
         let Some(dist_info) = dist_info else {
             return Ok(None);
         };
 
-        // Store a reference to the build context.
-        self.build_context
-            .build_arena()
-            .push(builder.into_build_dir());
-
         // Read the metadata from disk.
         debug!("Prepared metadata for: {source}");
         let content = fs::read(dist_info.join("METADATA"))

crates/uv-types/Cargo.toml

@@ -31,9 +31,8 @@ uv-redacted = { workspace = true }
 uv-workspace = { workspace = true }
 
 anyhow = { workspace = true }
-boxcar = { workspace = true }
+dashmap = { workspace = true }
 rustc-hash = { workspace = true }
-tempfile = { workspace = true }
 thiserror = { workspace = true }
 
 [features]

crates/uv-types/src/builds.rs

@@ -1,5 +1,9 @@
+use std::path::Path;
 use std::sync::Arc;
-use tempfile::TempDir;
+
+use dashmap::DashMap;
+
+use uv_configuration::{BuildKind, SourceStrategy};
 use uv_pep508::PackageName;
 use uv_python::PythonEnvironment;
 
@@ -40,13 +44,41 @@ impl BuildIsolation<'_> {
     }
 }
 
-/// An arena of temporary directories used for builds.
-#[derive(Default, Debug, Clone)]
-pub struct BuildArena(Arc<boxcar::Vec<TempDir>>);
+/// A key for the build cache, which includes the interpreter, source root, subdirectory, source
+/// strategy, and build kind.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct BuildKey {
+    pub base_python: Box<Path>,
+    pub source_root: Box<Path>,
+    pub subdirectory: Option<Box<Path>>,
+    pub source_strategy: SourceStrategy,
+    pub build_kind: BuildKind,
+}
 
-impl BuildArena {
-    /// Push a new temporary directory into the arena.
-    pub fn push(&self, temp_dir: TempDir) {
-        self.0.push(temp_dir);
+/// An arena of in-process builds.
+#[derive(Debug)]
+pub struct BuildArena<T>(Arc<DashMap<BuildKey, T>>);
+
+impl<T> Default for BuildArena<T> {
+    fn default() -> Self {
+        Self(Arc::new(DashMap::new()))
+    }
+}
+
+impl<T> Clone for BuildArena<T> {
+    fn clone(&self) -> Self {
+        Self(self.0.clone())
+    }
+}
+
+impl<T> BuildArena<T> {
+    /// Insert a build entry into the arena.
+    pub fn insert(&self, key: BuildKey, value: T) {
+        self.0.insert(key, value);
+    }
+
+    /// Remove a build entry from the arena.
+    pub fn remove(&self, key: &BuildKey) -> Option<T> {
+        self.0.remove(key).map(|entry| entry.1)
     }
 }

crates/uv-types/src/traits.rs

@@ -5,9 +5,7 @@ use std::path::{Path, PathBuf};
 
 use anyhow::Result;
 use rustc_hash::FxHashSet;
-use tempfile::TempDir;
 
-use crate::BuildArena;
 use uv_cache::Cache;
 use uv_configuration::{BuildKind, BuildOptions, BuildOutput, ConfigSettings, SourceStrategy};
 use uv_distribution_filename::DistFilename;
@@ -20,6 +18,8 @@ use uv_pep508::PackageName;
 use uv_python::{Interpreter, PythonEnvironment};
 use uv_workspace::WorkspaceCache;
 
+use crate::BuildArena;
+
 ///  Avoids cyclic crate dependencies between resolver, installer and builder.
 ///
 /// To resolve the dependencies of a packages, we may need to build one or more source
@@ -70,7 +70,7 @@ pub trait BuildContext {
     fn git(&self) -> &GitResolver;
 
     /// Return a reference to the build arena.
-    fn build_arena(&self) -> &BuildArena;
+    fn build_arena(&self) -> &BuildArena<Self::SourceDistBuilder>;
 
     /// Return a reference to the discovered registry capabilities.
     fn capabilities(&self) -> &IndexCapabilities;
@@ -153,9 +153,6 @@ pub trait BuildContext {
 /// You can either call only `wheel()` to build the wheel directly, call only `metadata()` to get
 /// the metadata without performing the actual or first call `metadata()` and then `wheel()`.
 pub trait SourceBuildTrait {
-    /// Return the temporary build directory.
-    fn into_build_dir(self) -> TempDir;
-
     /// A wrapper for `uv_build::SourceBuild::get_metadata_without_build`.
     ///
     /// For PEP 517 builds, this calls `prepare_metadata_for_build_wheel`
@@ -188,13 +185,13 @@ pub trait InstalledPackagesProvider: Clone + Send + Sync + 'static {
 pub struct EmptyInstalledPackages;
 
 impl InstalledPackagesProvider for EmptyInstalledPackages {
-    fn get_packages(&self, _name: &PackageName) -> Vec<&InstalledDist> {
-        Vec::new()
-    }
-
     fn iter(&self) -> impl Iterator<Item = &InstalledDist> {
         std::iter::empty()
     }
+
+    fn get_packages(&self, _name: &PackageName) -> Vec<&InstalledDist> {
+        Vec::new()
+    }
 }
 
 /// [`anyhow::Error`]-like wrapper type for [`BuildDispatch`] method return values, that also makes