Add support for per-project build-time environment variables (#15095)

## Summary

E.g., you can now do:

```toml
[tool.uv.extra-build-variables]
flash-attn = { FLASH_ATTENTION_SKIP_CUDA_BUILD = "TRUE" }
```

Author: charliermarsh

crates/uv-bench/benches/uv.rs

@@ -1,6 +1,6 @@
+use std::hint::black_box;
 use std::str::FromStr;
 
-use std::hint::black_box;
 use uv_bench::criterion::{Criterion, criterion_group, criterion_main, measurement::WallTime};
 use uv_cache::Cache;
 use uv_client::RegistryClientBuilder;
@@ -92,7 +92,7 @@ mod resolver {
     use uv_dispatch::{BuildDispatch, SharedState};
     use uv_distribution::DistributionDatabase;
     use uv_distribution_types::{
-        DependencyMetadata, ExtraBuildRequires, IndexLocations, RequiresPython,
+        DependencyMetadata, ExtraBuildRequires, ExtraBuildVariables, IndexLocations, RequiresPython,
     };
     use uv_install_wheel::LinkMode;
     use uv_pep440::Version;
@@ -144,6 +144,7 @@ mod resolver {
     ) -> Result<ResolverOutput> {
         let build_isolation = BuildIsolation::default();
         let extra_build_requires = ExtraBuildRequires::default();
+        let extra_build_variables = ExtraBuildVariables::default();
         let build_options = BuildOptions::default();
         let concurrency = Concurrency::default();
         let config_settings = ConfigSettings::default();
@@ -193,6 +194,7 @@ mod resolver {
             &config_settings_package,
             build_isolation,
             &extra_build_requires,
+            &extra_build_variables,
             LinkMode::default(),
             &build_options,
             &hashes,

crates/uv-cli/src/options.rs

@@ -355,6 +355,7 @@ pub fn resolver_options(
         no_build_isolation: flag(no_build_isolation, build_isolation, "build-isolation"),
         no_build_isolation_package: Some(no_build_isolation_package),
         extra_build_dependencies: None,
+        extra_build_variables: None,
         exclude_newer: ExcludeNewer::from_args(
             exclude_newer,
             exclude_newer_package.unwrap_or_default(),
@@ -477,6 +478,7 @@ pub fn resolver_installer_options(
             Some(no_build_isolation_package)
         },
         extra_build_dependencies: None,
+        extra_build_variables: None,
         exclude_newer,
         exclude_newer_package: exclude_newer_package.map(ExcludeNewerPackage::from_iter),
         link_mode,

crates/uv-dispatch/src/lib.rs

@@ -24,9 +24,9 @@ use uv_configuration::{BuildOutput, Concurrency};
 use uv_distribution::DistributionDatabase;
 use uv_distribution_filename::DistFilename;
 use uv_distribution_types::{
-    CachedDist, DependencyMetadata, ExtraBuildRequires, Identifier, IndexCapabilities,
-    IndexLocations, IsBuildBackendError, Name, Requirement, Resolution, SourceDist,
-    VersionOrUrlRef,
+    CachedDist, DependencyMetadata, ExtraBuildRequires, ExtraBuildVariables, Identifier,
+    IndexCapabilities, IndexLocations, IsBuildBackendError, Name, Requirement, Resolution,
+    SourceDist, VersionOrUrlRef,
 };
 use uv_git::GitResolver;
 use uv_installer::{Installer, Plan, Planner, Preparer, SitePackages};
@@ -90,6 +90,7 @@ pub struct BuildDispatch<'a> {
     dependency_metadata: &'a DependencyMetadata,
     build_isolation: BuildIsolation<'a>,
     extra_build_requires: &'a ExtraBuildRequires,
+    extra_build_variables: &'a ExtraBuildVariables,
     link_mode: uv_install_wheel::LinkMode,
     build_options: &'a BuildOptions,
     config_settings: &'a ConfigSettings,
@@ -119,6 +120,7 @@ impl<'a> BuildDispatch<'a> {
         config_settings_package: &'a PackageConfigSettings,
         build_isolation: BuildIsolation<'a>,
         extra_build_requires: &'a ExtraBuildRequires,
+        extra_build_variables: &'a ExtraBuildVariables,
         link_mode: uv_install_wheel::LinkMode,
         build_options: &'a BuildOptions,
         hasher: &'a HashStrategy,
@@ -142,6 +144,7 @@ impl<'a> BuildDispatch<'a> {
             config_settings_package,
             build_isolation,
             extra_build_requires,
+            extra_build_variables,
             link_mode,
             build_options,
             hasher,
@@ -227,6 +230,10 @@ impl BuildContext for BuildDispatch<'_> {
         self.extra_build_requires
     }
 
+    fn extra_build_variables(&self) -> &ExtraBuildVariables {
+        self.extra_build_variables
+    }
+
     async fn resolve<'data>(
         &'data self,
         requirements: &'data [Requirement],
@@ -312,6 +319,7 @@ impl BuildContext for BuildDispatch<'_> {
             self.config_settings,
             self.config_settings_package,
             self.extra_build_requires(),
+            self.extra_build_variables,
             self.cache(),
             venv,
             tags,
@@ -446,6 +454,18 @@ impl BuildContext for BuildDispatch<'_> {
             self.config_settings.clone()
         };
 
+        // Get package-specific environment variables if available.
+        let mut environment_variables = self.build_extra_env_vars.clone();
+        if let Some(name) = dist_name {
+            if let Some(package_vars) = self.extra_build_variables.get(name) {
+                environment_variables.extend(
+                    package_vars
+                        .iter()
+                        .map(|(key, value)| (OsString::from(key), OsString::from(value))),
+                );
+            }
+        }
+
         let builder = SourceBuild::setup(
             source,
             subdirectory,
@@ -464,7 +484,7 @@ impl BuildContext for BuildDispatch<'_> {
             self.extra_build_requires(),
             &build_stack,
             build_kind,
-            self.build_extra_env_vars.clone(),
+            environment_variables,
             build_output,
             self.concurrency.builds,
             self.preview,

crates/uv-distribution-types/src/build_requires.rs

@@ -1,5 +1,7 @@
 use std::collections::BTreeMap;
 
+use serde::{Deserialize, Serialize};
+
 use uv_cache_key::{CacheKey, CacheKeyHasher};
 use uv_normalize::PackageName;
 
@@ -104,3 +106,52 @@ impl ExtraBuildRequires {
             .collect::<Result<Self, _>>()
     }
 }
+
+/// A map of extra build variables, from variable name to value.
+pub type BuildVariables = BTreeMap<String, String>;
+
+/// Extra environment variables to set during builds, on a per-package basis.
+#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub struct ExtraBuildVariables(BTreeMap<PackageName, BuildVariables>);
+
+impl std::ops::Deref for ExtraBuildVariables {
+    type Target = BTreeMap<PackageName, BuildVariables>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl std::ops::DerefMut for ExtraBuildVariables {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl IntoIterator for ExtraBuildVariables {
+    type Item = (PackageName, BuildVariables);
+    type IntoIter = std::collections::btree_map::IntoIter<PackageName, BuildVariables>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.into_iter()
+    }
+}
+
+impl FromIterator<(PackageName, BuildVariables)> for ExtraBuildVariables {
+    fn from_iter<T: IntoIterator<Item = (PackageName, BuildVariables)>>(iter: T) -> Self {
+        Self(iter.into_iter().collect())
+    }
+}
+
+impl CacheKey for ExtraBuildVariables {
+    fn cache_key(&self, state: &mut CacheKeyHasher) {
+        for (package, vars) in &self.0 {
+            package.as_str().cache_key(state);
+            for (key, value) in vars {
+                key.cache_key(state);
+                value.cache_key(state);
+            }
+        }
+    }
+}

crates/uv-distribution/src/index/built_wheel_index.rs

@@ -5,8 +5,8 @@ use uv_cache_info::CacheInfo;
 use uv_cache_key::cache_digest;
 use uv_configuration::{ConfigSettings, PackageConfigSettings};
 use uv_distribution_types::{
-    DirectUrlSourceDist, DirectorySourceDist, ExtraBuildRequirement, ExtraBuildRequires,
-    GitSourceDist, Hashed, PathSourceDist,
+    BuildVariables, DirectUrlSourceDist, DirectorySourceDist, ExtraBuildRequirement,
+    ExtraBuildRequires, ExtraBuildVariables, GitSourceDist, Hashed, PathSourceDist,
 };
 use uv_normalize::PackageName;
 use uv_platform_tags::Tags;
@@ -25,6 +25,7 @@ pub struct BuiltWheelIndex<'a> {
     config_settings: &'a ConfigSettings,
     config_settings_package: &'a PackageConfigSettings,
     extra_build_requires: &'a ExtraBuildRequires,
+    extra_build_variables: &'a ExtraBuildVariables,
 }
 
 impl<'a> BuiltWheelIndex<'a> {
@@ -36,6 +37,7 @@ impl<'a> BuiltWheelIndex<'a> {
         config_settings: &'a ConfigSettings,
         config_settings_package: &'a PackageConfigSettings,
         extra_build_requires: &'a ExtraBuildRequires,
+        extra_build_variables: &'a ExtraBuildVariables,
     ) -> Self {
         Self {
             cache,
@@ -44,6 +46,7 @@ impl<'a> BuiltWheelIndex<'a> {
             config_settings,
             config_settings_package,
             extra_build_requires,
+            extra_build_variables,
         }
     }
 
@@ -75,10 +78,18 @@ impl<'a> BuiltWheelIndex<'a> {
         // If there are build settings, we need to scope to a cache shard.
         let config_settings = self.config_settings_for(&source_dist.name);
         let extra_build_deps = self.extra_build_requires_for(&source_dist.name);
-        let cache_shard = if config_settings.is_empty() && extra_build_deps.is_empty() {
+        let extra_build_vars = self.extra_build_variables_for(&source_dist.name);
+        let cache_shard = if config_settings.is_empty()
+            && extra_build_deps.is_empty()
+            && extra_build_vars.is_none()
+        {
             cache_shard
         } else {
-            cache_shard.shard(cache_digest(&(&config_settings, extra_build_deps)))
+            cache_shard.shard(cache_digest(&(
+                &config_settings,
+                extra_build_deps,
+                extra_build_vars,
+            )))
         };
 
         Ok(self.find(&cache_shard))
@@ -114,10 +125,18 @@ impl<'a> BuiltWheelIndex<'a> {
         // If there are build settings, we need to scope to a cache shard.
         let config_settings = self.config_settings_for(&source_dist.name);
         let extra_build_deps = self.extra_build_requires_for(&source_dist.name);
-        let cache_shard = if config_settings.is_empty() && extra_build_deps.is_empty() {
+        let extra_build_vars = self.extra_build_variables_for(&source_dist.name);
+        let cache_shard = if config_settings.is_empty()
+            && extra_build_deps.is_empty()
+            && extra_build_vars.is_none()
+        {
             cache_shard
         } else {
-            cache_shard.shard(cache_digest(&(&config_settings, extra_build_deps)))
+            cache_shard.shard(cache_digest(&(
+                &config_settings,
+                extra_build_deps,
+                extra_build_vars,
+            )))
         };
 
         Ok(self
@@ -164,10 +183,18 @@ impl<'a> BuiltWheelIndex<'a> {
         // If there are build settings, we need to scope to a cache shard.
         let config_settings = self.config_settings_for(&source_dist.name);
         let extra_build_deps = self.extra_build_requires_for(&source_dist.name);
-        let cache_shard = if config_settings.is_empty() && extra_build_deps.is_empty() {
+        let extra_build_vars = self.extra_build_variables_for(&source_dist.name);
+        let cache_shard = if config_settings.is_empty()
+            && extra_build_deps.is_empty()
+            && extra_build_vars.is_none()
+        {
             cache_shard
         } else {
-            cache_shard.shard(cache_digest(&(&config_settings, extra_build_deps)))
+            cache_shard.shard(cache_digest(&(
+                &config_settings,
+                extra_build_deps,
+                extra_build_vars,
+            )))
         };
 
         Ok(self
@@ -192,10 +219,18 @@ impl<'a> BuiltWheelIndex<'a> {
         // If there are build settings, we need to scope to a cache shard.
         let config_settings = self.config_settings_for(&source_dist.name);
         let extra_build_deps = self.extra_build_requires_for(&source_dist.name);
-        let cache_shard = if config_settings.is_empty() && extra_build_deps.is_empty() {
+        let extra_build_vars = self.extra_build_variables_for(&source_dist.name);
+        let cache_shard = if config_settings.is_empty()
+            && extra_build_deps.is_empty()
+            && extra_build_vars.is_none()
+        {
             cache_shard
         } else {
-            cache_shard.shard(cache_digest(&(&config_settings, extra_build_deps)))
+            cache_shard.shard(cache_digest(&(
+                &config_settings,
+                extra_build_deps,
+                extra_build_vars,
+            )))
         };
 
         self.find(&cache_shard)
@@ -274,4 +309,9 @@ impl<'a> BuiltWheelIndex<'a> {
             .map(Vec::as_slice)
             .unwrap_or(&[])
     }
+
+    /// Determine the extra build variables for the given package name.
+    fn extra_build_variables_for(&self, name: &PackageName) -> Option<&BuildVariables> {
+        self.extra_build_variables.get(name)
+    }
 }