From 8007d79ad33d013be60df6781c2ea53f09f78b84 Mon Sep 17 00:00:00 2001 From: konstin Date: Fri, 25 Jul 2025 16:00:49 +0200 Subject: [PATCH] Add naive caching for variant provider running --- crates/uv-distribution-types/src/lib.rs | 10 + .../src/distribution_database.rs | 49 ++++- crates/uv-distribution/src/lib.rs | 3 +- crates/uv-resolver/src/lock/installable.rs | 194 ++++++++++-------- crates/uv-variants/src/resolved_variants.rs | 2 +- crates/uv/src/commands/project/sync.rs | 6 +- test_variants.sh | 18 +- 7 files changed, 175 insertions(+), 107 deletions(-) diff --git a/crates/uv-distribution-types/src/lib.rs b/crates/uv-distribution-types/src/lib.rs index a9e43d2c8..d8d8ea675 100644 --- a/crates/uv-distribution-types/src/lib.rs +++ b/crates/uv-distribution-types/src/lib.rs @@ -1481,6 +1481,16 @@ pub struct RegistryVariantsJson { pub index: IndexUrl, } +impl Identifier for RegistryVariantsJson { + fn distribution_id(&self) -> DistributionId { + self.file.distribution_id() + } + + fn resource_id(&self) -> ResourceId { + self.file.resource_id() + } +} + #[cfg(test)] mod test { use crate::{BuiltDist, Dist, RemoteSource, SourceDist, UrlString}; diff --git a/crates/uv-distribution/src/distribution_database.rs b/crates/uv-distribution/src/distribution_database.rs index 593fcb488..3bbaca2c0 100644 --- a/crates/uv-distribution/src/distribution_database.rs +++ b/crates/uv-distribution/src/distribution_database.rs @@ -3,7 +3,7 @@ use std::future::Future; use std::io; use std::path::Path; use std::pin::Pin; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; use anyhow::anyhow; @@ -29,7 +29,7 @@ use uv_distribution_filename::WheelFilename; use uv_distribution_types::{ BuildableSource, BuiltDist, Dist, DistributionId, HashPolicy, Hashed, Identifier, IndexUrl, InstalledDist, Name, Node, RegistryBuiltDist, RegistryVariantsJson, Resolution, ResolvedDist, - SourceDist, + ResourceId, SourceDist, }; use uv_extract::hash::Hasher; use uv_fs::write_atomic; @@ -1358,10 +1358,37 @@ impl LocalArchivePointer { } } +/// A very simple variants provider cache +/// +/// TODO(konsti): Cache by provider provider plugin and `requires` compatibility +/// TODO(konsti): Cache to disk +#[derive(Debug, Default)] +pub struct VariantProviderCache { + cache: Mutex>, +} + +impl VariantProviderCache { + pub fn get(&self, resource_id: &ResourceId) -> Option { + self.cache + .lock() + .expect("there was a panic in another thread") + .get(resource_id) + .cloned() + } + + pub fn insert(&self, resource_id: ResourceId, resolved_variants: ResolvedVariants) { + self.cache + .lock() + .expect("there was a panic in another thread") + .insert(resource_id, resolved_variants); + } +} + /// TODO(konsti): Find a better home for those functions pub async fn resolve_variants( resolution: Resolution, distribution_database: DistributionDatabase<'_, Context>, + variants_cache: Arc, ) -> anyhow::Result { // Fetch variants.json and then query providers, running in parallel for all dists // TODO(konsti): Reuse in memory index from resolve phase and merge equivalent requests. @@ -1373,10 +1400,20 @@ pub async fn resolve_variants( .filter_map(|node| extract_variants(node)), ) .map(async |(variants_json, dist)| { - // Fetch variants_json and run providers - let resolved_variants = distribution_database - .fetch_and_query_variants(variants_json) - .await?; + let resolved_variants = + if let Some(resolved_variants) = variants_cache.get(&variants_json.resource_id()) { + resolved_variants.clone() + } else { + // Fetch variants_json and run providers + let resolved_variants = distribution_database + .fetch_and_query_variants(variants_json) + .await?; + + variants_cache.insert(variants_json.resource_id(), resolved_variants.clone()); + + resolved_variants + }; + Ok::<_, anyhow::Error>((dist.distribution_id(), resolved_variants)) }) // TODO(konsti): Buffer size diff --git a/crates/uv-distribution/src/lib.rs b/crates/uv-distribution/src/lib.rs index f30fd227c..744d07d5e 100644 --- a/crates/uv-distribution/src/lib.rs +++ b/crates/uv-distribution/src/lib.rs @@ -1,5 +1,6 @@ pub use distribution_database::{ - DistributionDatabase, HttpArchivePointer, LocalArchivePointer, resolve_variants, + DistributionDatabase, HttpArchivePointer, LocalArchivePointer, VariantProviderCache, + resolve_variants, }; pub use download::LocalWheel; pub use error::Error; diff --git a/crates/uv-resolver/src/lock/installable.rs b/crates/uv-resolver/src/lock/installable.rs index 38579a6b2..aece2ec3c 100644 --- a/crates/uv-resolver/src/lock/installable.rs +++ b/crates/uv-resolver/src/lock/installable.rs @@ -11,8 +11,8 @@ use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; use uv_configuration::ExtrasSpecificationWithDefaults; use uv_configuration::{BuildOptions, DependencyGroupsWithDefaults, InstallOptions}; -use uv_distribution::DistributionDatabase; -use uv_distribution_types::{Edge, Node, Resolution, ResolvedDist}; +use uv_distribution::{DistributionDatabase, VariantProviderCache}; +use uv_distribution_types::{Edge, Identifier, Node, Resolution, ResolvedDist}; use uv_normalize::{ExtraName, GroupName, PackageName}; use uv_platform_tags::Tags; use uv_pypi_types::ResolverMarkerEnvironment; @@ -47,6 +47,7 @@ pub trait Installable<'lock> { build_options: &BuildOptions, install_options: &InstallOptions, distribution_database: DistributionDatabase<'_, Context>, + variants_cache: Arc, ) -> Result { let size_guess = self.lock().packages.len(); let mut petgraph = Graph::with_capacity(size_guess, size_guess); @@ -463,97 +464,18 @@ pub trait Installable<'lock> { Either::Right(package.dependencies.iter()) }; - let known_properties: Option> = if deps - .clone() - .any(|dep| dep.complexified_marker.combined().has_variant_expression()) - { - if let Some(variants_json) = package - .to_registry_variants_json(self.install_path()) - .expect("TODO(konsti)") - { - let resolved_variants = distribution_database - .fetch_and_query_variants(&variants_json) - .await - .expect("TODO(konsti)"); - - // Select best wheel - let mut highest_priority_variant_wheel: Option<(_, Vec)> = None; - for wheel in &package.wheels { - let Some(variant) = wheel.filename.variant() else { - // The non-variant wheel is already supported - continue; - }; - - let Some(variants_properties) = - resolved_variants.variants_json.variants.get(variant) - else { - // TODO(konsti): For production, this should be a warning. - panic!("Variant {variant} not found in variants.json"); - }; - - let Some(scores) = score_variant( - &resolved_variants.variants_json.default_priorities, - &resolved_variants.target_variants, - variants_properties, - ) else { - // The wheel is not compatible. - continue; - }; - - if let Some((_, old_scores)) = &highest_priority_variant_wheel { - if &scores > old_scores { - highest_priority_variant_wheel = Some((variant, scores)); - } - } else { - highest_priority_variant_wheel = Some((variant, scores)); - } - } - - if let Some((best_variant, _)) = highest_priority_variant_wheel { - // TODO(konsti): Use the proper type everywhere, we shouldn't need to do - // this conversion at all. - let mut known_properties = BTreeSet::new(); - for (namespace, features) in resolved_variants - .variants_json - .variants - .get(best_variant) - .expect("TODO(konsti): error handling") - { - for (feature, values) in features { - for value in values { - known_properties.insert(VariantPropertyType { - namespace: namespace.clone(), - feature: feature.clone(), - value: value.clone(), - }); - } - } - } - let known_properties: Vec<(String, String, String)> = known_properties - .into_iter() - .map(|known_property| { - ( - known_property.namespace, - known_property.feature, - known_property.value, - ) - }) - .collect(); - Some(known_properties) - } else { - None - } - } else { - None - } - } else { - None - }; + let variant_properties = determine_properties( + package, + self.install_path(), + &distribution_database, + &variants_cache, + ) + .await?; for dep in deps { if !dep.complexified_marker.evaluate( marker_env, - known_properties.as_deref(), + variant_properties.as_deref(), activated_extras.iter().copied(), activated_groups.iter().copied(), ) { @@ -672,3 +594,97 @@ pub trait Installable<'lock> { } } } + +async fn determine_properties( + package: &Package, + workspace_root: &Path, + distribution_database: &DistributionDatabase<'_, Context>, + variants_cache: &VariantProviderCache, +) -> Result>, LockError> { + let Some(variants_json) = package.to_registry_variants_json(workspace_root)? else { + return Ok(None); + }; + let resolved_variants = + if let Some(resolved_variants) = variants_cache.get(&variants_json.resource_id()) { + resolved_variants.clone() + } else { + // Fetch variants_json and run providers + let resolved_variants = distribution_database + .fetch_and_query_variants(&variants_json) + .await + .expect("TODO(konsti)"); + + variants_cache.insert(variants_json.resource_id(), resolved_variants.clone()); + + resolved_variants + }; + + // Select best wheel + let mut highest_priority_variant_wheel: Option<(_, Vec)> = None; + for wheel in &package.wheels { + let Some(variant) = wheel.filename.variant() else { + // The non-variant wheel is already supported + continue; + }; + + let Some(variants_properties) = resolved_variants.variants_json.variants.get(variant) + else { + // TODO(konsti): For production, this should be a warning. + panic!("Variant {variant} not found in variants.json"); + }; + + let Some(scores) = score_variant( + &resolved_variants.variants_json.default_priorities, + &resolved_variants.target_variants, + variants_properties, + ) else { + // The wheel is not compatible. + continue; + }; + + if let Some((_, old_scores)) = &highest_priority_variant_wheel { + if &scores > old_scores { + highest_priority_variant_wheel = Some((variant, scores)); + } + } else { + highest_priority_variant_wheel = Some((variant, scores)); + } + } + + if let Some((best_variant, _)) = highest_priority_variant_wheel { + // TODO(konsti): Use the proper type everywhere, we shouldn't need to do + // this conversion at all. + let mut known_properties = BTreeSet::new(); + for (namespace, features) in resolved_variants + .variants_json + .variants + .get(best_variant) + .expect("TODO(konsti): error handling") + { + for (feature, values) in features { + for value in values { + known_properties.insert(VariantPropertyType { + namespace: namespace.clone(), + feature: feature.clone(), + value: value.clone(), + }); + } + } + } + let known_properties: Vec<(String, String, String)> = known_properties + .into_iter() + .map(|known_property| { + ( + known_property.namespace, + known_property.feature, + known_property.value, + ) + }) + .collect(); + Ok(Some(known_properties)) + } else { + // When selecting the non-variant wheel, all variant markers evaluate to + // false. + Ok(Some(Vec::new())) + } +} diff --git a/crates/uv-variants/src/resolved_variants.rs b/crates/uv-variants/src/resolved_variants.rs index efa1efd9b..984e2c524 100644 --- a/crates/uv-variants/src/resolved_variants.rs +++ b/crates/uv-variants/src/resolved_variants.rs @@ -2,7 +2,7 @@ use crate::VariantProviderOutput; use crate::variants_json::{VariantNamespace, VariantsJsonContent}; use rustc_hash::FxHashMap; -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct ResolvedVariants { pub variants_json: VariantsJsonContent, pub target_variants: FxHashMap, diff --git a/crates/uv/src/commands/project/sync.rs b/crates/uv/src/commands/project/sync.rs index 38ffc341b..b38d3fea7 100644 --- a/crates/uv/src/commands/project/sync.rs +++ b/crates/uv/src/commands/project/sync.rs @@ -18,7 +18,7 @@ use uv_configuration::{ Preview, PreviewFeatures, TargetTriple, }; use uv_dispatch::BuildDispatch; -use uv_distribution::{DistributionDatabase, resolve_variants}; +use uv_distribution::{DistributionDatabase, VariantProviderCache, resolve_variants}; use uv_distribution_types::{ DirectorySourceDist, Dist, Index, Requirement, Resolution, ResolvedDist, SourceDist, }; @@ -711,6 +711,7 @@ pub(super) async fn do_sync( DistributionDatabase::new(&client, &build_dispatch, concurrency.downloads); // Read the lockfile. + let variants_cache = Arc::new(VariantProviderCache::default()); let resolution = target .to_resolution( &marker_env, @@ -720,6 +721,7 @@ pub(super) async fn do_sync( build_options, &install_options, distribution_database, + variants_cache.clone(), ) .await?; @@ -772,7 +774,7 @@ pub(super) async fn do_sync( // TODO(konsti): Pass this into operations::install let distribution_database = DistributionDatabase::new(&client, &build_dispatch, concurrency.downloads); - let resolution = resolve_variants(resolution, distribution_database).await?; + let resolution = resolve_variants(resolution, distribution_database, variants_cache).await?; let site_packages = SitePackages::from_environment(venv)?; diff --git a/test_variants.sh b/test_variants.sh index 1e691bef0..a31d39abe 100755 --- a/test_variants.sh +++ b/test_variants.sh @@ -6,19 +6,21 @@ cargo build unset VIRTUAL_ENV export RUST_LOG=uv_distribution_types=debug,uv_distribution::distribution_database=debug -# No matching variant wheel, no non-variant wheel or sdist +echo "# No matching variant wheel, no non-variant wheel or sdist" uv venv -c -q && ( ( UV_CPU_LEVEL_OVERRIDE=0 cargo run -q pip install built-by-uv --no-index --no-cache --no-progress --find-links ./files && exit 1 ) || exit 0 ) -# No matching variant wheel, but a non-variant wheel +echo "# No matching variant wheel, but a non-variant wheel" uv venv -c -q && UV_CPU_LEVEL_OVERRIDE=0 cargo run -q pip install built-by-uv --no-index --no-cache --no-progress --find-links ./files --find-links ./files_wheel -# No matching variant wheel, but a non-variant sdist +echo "# No matching variant wheel, but a non-variant sdist" uv venv -c -q && UV_CPU_LEVEL_OVERRIDE=0 cargo run -q pip install built-by-uv --no-index --no-cache --no-progress --find-links ./files --find-links ./files_sdist -# Matching cpu2 variant wheel, to be preferred over the non-variant wheel +echo "# Matching cpu2 variant wheel, to be preferred over the non-variant wheel" uv venv -c -q && UV_CPU_LEVEL_OVERRIDE=2 cargo run -q pip install built-by-uv --no-index --no-cache --no-progress --find-links ./files --find-links ./files_wheel -# Matching cpu2 variant wheel, to be preferred over the non-variant wheel and the sdist +echo "# Matching cpu2 variant wheel, to be preferred over the non-variant wheel and the sdist" uv venv -c -q && UV_CPU_LEVEL_OVERRIDE=2 RUST_LOG=uv_distribution_types=debug cargo run -q pip install built-by-uv --no-index --no-cache --no-progress --find-links ./files --find-links ./files_wheel --find-links ./files_sdist -# sync without a compatible variant wheel -( cd scripts/packages/cpu_user && rm -f uv.lock && ( ( uv venv -c -q && UV_CPU_LEVEL_OVERRIDE=0 cargo run -q sync && exit 1 ) || exit 0 ) && ( ( uv venv -c -q && UV_CPU_LEVEL_OVERRIDE=0 cargo run -q sync && exit 1 ) || exit 0 ) ) -# sync with a compatible variant wheel +echo "# sync without a compatible variant wheel (fresh)" +( cd scripts/packages/cpu_user && rm -f uv.lock && ( ( uv venv -c -q && UV_CPU_LEVEL_OVERRIDE=0 cargo run -q sync && exit 1 ) || exit 0 ) ) +echo "# sync without a compatible variant wheel (existing lockfile)" +( cd scripts/packages/cpu_user && rm -f uv.lock && ( ( uv venv -c -q && UV_CPU_LEVEL_OVERRIDE=0 cargo run -q sync && exit 1 ) || exit 0 ) ) +echo "# sync with a compatible variant wheel" # TODO(konsti): Selecting a different level currently selects the right wheel, but doesn't reinstall. ( cd scripts/packages/cpu_user && rm -f uv.lock && uv venv -c -q && UV_CPU_LEVEL_OVERRIDE=2 cargo run -q sync && UV_CPU_LEVEL_OVERRIDE=2 cargo run -q sync && UV_CPU_LEVEL_OVERRIDE=3 cargo run -q sync )