Python
/
uv
mirror of https://github.com/astral-sh/uv
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add naive caching for variant provider running

This commit is contained in:
konstin 2025-07-25 16:00:49 +02:00
parent 9d3f508798
commit 8007d79ad3
7 changed files with 175 additions and 107 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
crates/uv-distribution-types/src/lib.rs
View File

@ -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};

41
crates/uv-distribution/src/distribution_database.rs
View File

@ -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<FxHashMap<ResourceId, ResolvedVariants>>,
}
impl VariantProviderCache {
pub fn get(&self, resource_id: &ResourceId) -> Option<ResolvedVariants> {
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<Context: BuildContext>(
resolution: Resolution,
distribution_database: DistributionDatabase<'_, Context>,
variants_cache: Arc<VariantProviderCache>,
) -> anyhow::Result<Resolution> {
// 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<Context: BuildContext>(
.filter_map(|node| extract_variants(node)),
)
.map(async |(variants_json, dist)| {
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

3
crates/uv-distribution/src/lib.rs
View File

@ -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;

192
crates/uv-resolver/src/lock/installable.rs
View File

@ -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<VariantProviderCache>,
) -> Result<Resolution, LockError> {
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<Vec<(String, String, String)>> = 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<usize>)> = 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,
let variant_properties = determine_properties(
package,
self.install_path(),
&distribution_database,
&variants_cache,
)
})
.collect();
Some(known_properties)
} else {
None
}
} else {
None
}
} else {
None
};
.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<Context: BuildContext>(
package: &Package,
workspace_root: &Path,
distribution_database: &DistributionDatabase<'_, Context>,
variants_cache: &VariantProviderCache,
) -> Result<Option<Vec<(String, String, String)>>, 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<usize>)> = 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()))
}
}

2
crates/uv-variants/src/resolved_variants.rs
View File

@ -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<VariantNamespace, VariantProviderOutput>,

6
crates/uv/src/commands/project/sync.rs
View File

@ -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)?;

18
test_variants.sh
View File

@ -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 )