From b456fa2939f0fb7420207a2b0ae65ed6fd0f3e98 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 17 Apr 2024 10:21:42 -0400 Subject: [PATCH] Incorporate heuristics to improve package prioritization (#3087) See: https://github.com/astral-sh/uv/issues/3078 --- crates/uv-resolver/src/dependency_provider.rs | 2 +- crates/uv-resolver/src/pubgrub/priority.rs | 122 +++++++++++++++--- crates/uv-resolver/src/resolver/mod.rs | 32 ++--- crates/uv/tests/pip_install_scenarios.rs | 56 +++++--- 4 files changed, 163 insertions(+), 49 deletions(-) diff --git a/crates/uv-resolver/src/dependency_provider.rs b/crates/uv-resolver/src/dependency_provider.rs index 40f26d1d8..148eca5fe 100644 --- a/crates/uv-resolver/src/dependency_provider.rs +++ b/crates/uv-resolver/src/dependency_provider.rs @@ -18,7 +18,7 @@ impl DependencyProvider for UvDependencyProvider { fn prioritize(&self, _package: &Self::P, _range: &Self::VS) -> Self::Priority { unimplemented!() } - type Priority = PubGrubPriority; + type Priority = Option; type Err = Infallible; diff --git a/crates/uv-resolver/src/pubgrub/priority.rs b/crates/uv-resolver/src/pubgrub/priority.rs index d8fc3f210..bb6b1e329 100644 --- a/crates/uv-resolver/src/pubgrub/priority.rs +++ b/crates/uv-resolver/src/pubgrub/priority.rs @@ -1,34 +1,126 @@ use std::cmp::Reverse; +use pubgrub::range::Range; use rustc_hash::FxHashMap; +use pep440_rs::Version; use uv_normalize::PackageName; use crate::pubgrub::package::PubGrubPackage; +/// A prioritization map to guide the `PubGrub` resolution process. +/// +/// During resolution, `PubGrub` needs to decide which package to consider next. The priorities +/// encoded here are used to guide that decision. +/// +/// Like `pip`, we prefer packages that are pinned to direct URLs over packages pinned to a single +/// version over packages that are constrained in some way over packages that are unconstrained. +/// +/// See: #[derive(Debug, Default)] -pub(crate) struct PubGrubPriorities(FxHashMap); +pub(crate) struct PubGrubPriorities(FxHashMap); impl PubGrubPriorities { - /// Add a package to the priority map. - pub(crate) fn add(&mut self, package: PackageName) { - let priority = self.0.len(); - self.0.entry(package).or_insert(priority); + /// Add a [`PubGrubPackage`] to the priority map. + pub(crate) fn insert(&mut self, package: &PubGrubPackage, version: &Range) { + let next = self.0.len(); + match package { + PubGrubPackage::Root(_) => {} + PubGrubPackage::Python(_) => {} + PubGrubPackage::Package(name, _, None) => { + match self.0.entry(name.clone()) { + std::collections::hash_map::Entry::Occupied(mut entry) => { + // Preserve the original index. + let index = match entry.get() { + PubGrubPriority::Singleton(Reverse(index)) => *index, + PubGrubPriority::Unconstrained(Reverse(index)) => *index, + PubGrubPriority::Constrained(Reverse(index)) => *index, + PubGrubPriority::DirectUrl(Reverse(index)) => *index, + PubGrubPriority::Root => next, + }; + + // Compute the priority. + let priority = if version.as_singleton().is_some() { + PubGrubPriority::Singleton(Reverse(index)) + } else if version == &Range::full() { + PubGrubPriority::Unconstrained(Reverse(index)) + } else { + PubGrubPriority::Constrained(Reverse(index)) + }; + + // Take the maximum of the new and existing priorities. + if priority > *entry.get() { + entry.insert(priority); + } + } + std::collections::hash_map::Entry::Vacant(entry) => { + // Insert the priority. + entry.insert(if version.as_singleton().is_some() { + PubGrubPriority::Singleton(Reverse(next)) + } else if version == &Range::full() { + PubGrubPriority::Unconstrained(Reverse(next)) + } else { + PubGrubPriority::Constrained(Reverse(next)) + }); + } + } + } + PubGrubPackage::Package(name, _, Some(_)) => { + match self.0.entry(name.clone()) { + std::collections::hash_map::Entry::Occupied(mut entry) => { + // Preserve the original index. + let index = match entry.get() { + PubGrubPriority::Singleton(Reverse(index)) => *index, + PubGrubPriority::Unconstrained(Reverse(index)) => *index, + PubGrubPriority::Constrained(Reverse(index)) => *index, + PubGrubPriority::DirectUrl(Reverse(index)) => *index, + PubGrubPriority::Root => next, + }; + + // Compute the priority. + let priority = PubGrubPriority::DirectUrl(Reverse(index)); + + // Take the maximum of the new and existing priorities. + if priority > *entry.get() { + entry.insert(priority); + } + } + std::collections::hash_map::Entry::Vacant(entry) => { + // Insert the priority. + entry.insert(PubGrubPriority::DirectUrl(Reverse(next))); + } + } + } + } } - /// Return the priority of the given package, if it exists. + /// Return the [`PubGrubPriority`] of the given package, if it exists. pub(crate) fn get(&self, package: &PubGrubPackage) -> Option { match package { - PubGrubPackage::Root(_) => Some(Reverse(0)), - PubGrubPackage::Python(_) => Some(Reverse(0)), - PubGrubPackage::Package(name, _, _) => self - .0 - .get(name) - .copied() - .map(|priority| priority + 1) - .map(Reverse), + PubGrubPackage::Root(_) => Some(PubGrubPriority::Root), + PubGrubPackage::Python(_) => Some(PubGrubPriority::Root), + PubGrubPackage::Package(name, _, _) => self.0.get(name).copied(), } } } -pub(crate) type PubGrubPriority = Reverse; +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub(crate) enum PubGrubPriority { + /// The package has no specific priority. + /// + /// As such, its priority is based on the order in which the packages were added (FIFO), such + /// that the first package we visit is prioritized over subsequent packages. + Unconstrained(Reverse), + + /// The version range is constrained in some way (e.g., with a `<=` or `>` operator). + Constrained(Reverse), + + /// The version range is constrained to a single version (e.g., with the `==` operator). + Singleton(Reverse), + + /// The package was specified via a direct URL. + DirectUrl(Reverse), + + /// The package is the root package. + Root, +} diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs index 9832c6dd0..e7a998cf2 100644 --- a/crates/uv-resolver/src/resolver/mod.rs +++ b/crates/uv-resolver/src/resolver/mod.rs @@ -18,7 +18,7 @@ use tracing::{debug, enabled, info_span, instrument, trace, warn, Instrument, Le use distribution_types::{ BuiltDist, Dist, DistributionMetadata, IncompatibleDist, IncompatibleSource, IncompatibleWheel, - InstalledDist, Name, RemoteSource, ResolvedDist, ResolvedDistRef, SourceDist, VersionOrUrl, + InstalledDist, RemoteSource, ResolvedDist, ResolvedDistRef, SourceDist, VersionOrUrl, }; pub(crate) use locals::Locals; use pep440_rs::{Version, MIN_VERSION}; @@ -316,12 +316,9 @@ impl< } // Choose a package version. - let Some(highest_priority_pkg) = - state - .partial_solution - .pick_highest_priority_pkg(|package, _range| { - priorities.get(package).unwrap_or_default() - }) + let Some(highest_priority_pkg) = state + .partial_solution + .pick_highest_priority_pkg(|package, _range| priorities.get(package)) else { if enabled!(Level::DEBUG) { prefetcher.log_tried_versions(); @@ -523,7 +520,6 @@ impl< async fn visit_package( &self, package: &PubGrubPackage, - priorities: &mut PubGrubPriorities, request_sink: &tokio::sync::mpsc::Sender, ) -> Result<(), ResolveError> { match package { @@ -537,7 +533,6 @@ impl< // Emit a request to fetch the metadata for this package. if self.index.packages.register(name.clone()) { - priorities.add(name.clone()); request_sink.send(Request::Package(name.clone())).await?; } } @@ -550,7 +545,6 @@ impl< // Emit a request to fetch the metadata for this distribution. let dist = Dist::from_url(name.clone(), url.clone())?; if self.index.distributions.register(dist.version_id()) { - priorities.add(dist.name().clone()); request_sink.send(Request::Dist(dist)).await?; } } @@ -845,9 +839,11 @@ impl< for (package, version) in constraints.iter() { debug!("Adding direct dependency: {package}{version}"); + // Update the package priorities. + priorities.insert(package, version); + // Emit a request to fetch the metadata for this package. - self.visit_package(package, priorities, request_sink) - .await?; + self.visit_package(package, request_sink).await?; } // Add a dependency on each editable. @@ -914,9 +910,11 @@ impl< for (dep_package, dep_version) in constraints.iter() { debug!("Adding transitive dependency for {package}{version}: {dep_package}{dep_version}"); + // Update the package priorities. + priorities.insert(dep_package, dep_version); + // Emit a request to fetch the metadata for this package. - self.visit_package(dep_package, priorities, request_sink) - .await?; + self.visit_package(dep_package, request_sink).await?; } // If a package has an extra, insert a constraint on the base package. @@ -1029,9 +1027,11 @@ impl< for (package, version) in constraints.iter() { debug!("Adding transitive dependency: {package}{version}"); + // Update the package priorities. + priorities.insert(package, version); + // Emit a request to fetch the metadata for this package. - self.visit_package(package, priorities, request_sink) - .await?; + self.visit_package(package, request_sink).await?; } // If a package has an extra, insert a constraint on the base package. diff --git a/crates/uv/tests/pip_install_scenarios.rs b/crates/uv/tests/pip_install_scenarios.rs index 43c399e1b..9a80cd7eb 100644 --- a/crates/uv/tests/pip_install_scenarios.rs +++ b/crates/uv/tests/pip_install_scenarios.rs @@ -457,10 +457,17 @@ fn dependency_excludes_range_of_compatible_versions() { ----- stderr ----- × No solution found when resolving dependencies: - ╰─▶ Because only the following versions of package-a are available: + ╰─▶ Because package-a==1.0.0 depends on package-b==1.0.0 and only the following versions of package-a are available: package-a==1.0.0 package-a>=2.0.0,<=3.0.0 - and package-a==1.0.0 depends on package-b==1.0.0, we can conclude that package-a<2.0.0 depends on package-b==1.0.0. (1) + we can conclude that package-a<2.0.0 depends on package-b==1.0.0. + And because package-a==3.0.0 depends on package-b==3.0.0, we can conclude that any of: + package-a<2.0.0 + package-a>=3.0.0 + depends on one of: + package-b<=1.0.0 + package-b>=3.0.0 + (1) Because only the following versions of package-c are available: package-c==1.0.0 @@ -470,8 +477,13 @@ fn dependency_excludes_range_of_compatible_versions() { package-a<2.0.0 package-a>=3.0.0 - And because we know from (1) that package-a<2.0.0 depends on package-b==1.0.0, we can conclude that package-a!=3.0.0, all versions of package-c, package-b!=1.0.0 are incompatible. - And because package-a==3.0.0 depends on package-b==3.0.0, we can conclude that all versions of package-c depend on one of: + And because we know from (1) that any of: + package-a<2.0.0 + package-a>=3.0.0 + depends on one of: + package-b<=1.0.0 + package-b>=3.0.0 + we can conclude that all versions of package-c depend on one of: package-b<=1.0.0 package-b>=3.0.0 @@ -582,10 +594,17 @@ fn dependency_excludes_non_contiguous_range_of_compatible_versions() { ----- stderr ----- × No solution found when resolving dependencies: - ╰─▶ Because only the following versions of package-a are available: + ╰─▶ Because package-a==1.0.0 depends on package-b==1.0.0 and only the following versions of package-a are available: package-a==1.0.0 package-a>=2.0.0,<=3.0.0 - and package-a==1.0.0 depends on package-b==1.0.0, we can conclude that package-a<2.0.0 depends on package-b==1.0.0. (1) + we can conclude that package-a<2.0.0 depends on package-b==1.0.0. + And because package-a==3.0.0 depends on package-b==3.0.0, we can conclude that any of: + package-a<2.0.0 + package-a>=3.0.0 + depends on one of: + package-b<=1.0.0 + package-b>=3.0.0 + (1) Because only the following versions of package-c are available: package-c==1.0.0 @@ -595,8 +614,13 @@ fn dependency_excludes_non_contiguous_range_of_compatible_versions() { package-a<2.0.0 package-a>=3.0.0 - And because we know from (1) that package-a<2.0.0 depends on package-b==1.0.0, we can conclude that all versions of package-c, package-a!=3.0.0, package-b!=1.0.0 are incompatible. - And because package-a==3.0.0 depends on package-b==3.0.0, we can conclude that all versions of package-c depend on one of: + And because we know from (1) that any of: + package-a<2.0.0 + package-a>=3.0.0 + depends on one of: + package-b<=1.0.0 + package-b>=3.0.0 + we can conclude that all versions of package-c depend on one of: package-b<=1.0.0 package-b>=3.0.0 @@ -1023,7 +1047,7 @@ fn extra_incompatible_with_root() { ----- stderr ----- × No solution found when resolving dependencies: - ╰─▶ Because only package-a[extra]==1.0.0 is available and package-a[extra]==1.0.0 depends on package-b==1.0.0, we can conclude that all versions of package-a[extra] depend on package-b==1.0.0. + ╰─▶ Because package-a[extra]==1.0.0 depends on package-b==1.0.0 and only package-a[extra]==1.0.0 is available, we can conclude that all versions of package-a[extra] depend on package-b==1.0.0. And because you require package-a[extra] and you require package-b==2.0.0, we can conclude that the requirements are unsatisfiable. "###); @@ -1183,7 +1207,7 @@ fn transitive_incompatible_with_root_version() { ----- stderr ----- × No solution found when resolving dependencies: - ╰─▶ Because only package-a==1.0.0 is available and package-a==1.0.0 depends on package-b==2.0.0, we can conclude that all versions of package-a depend on package-b==2.0.0. + ╰─▶ Because package-a==1.0.0 depends on package-b==2.0.0 and only package-a==1.0.0 is available, we can conclude that all versions of package-a depend on package-b==2.0.0. And because you require package-a and you require package-b==1.0.0, we can conclude that the requirements are unsatisfiable. "###); @@ -1527,7 +1551,7 @@ fn local_transitive_greater_than() { ----- stderr ----- × No solution found when resolving dependencies: - ╰─▶ Because only package-a==1.0.0 is available and package-a==1.0.0 depends on package-b>2.0.0, we can conclude that all versions of package-a depend on package-b>2.0.0. + ╰─▶ Because package-a==1.0.0 depends on package-b>2.0.0 and only package-a==1.0.0 is available, we can conclude that all versions of package-a depend on package-b>2.0.0. And because you require package-a and you require package-b==2.0.0+foo, we can conclude that the requirements are unsatisfiable. "###); @@ -1638,7 +1662,7 @@ fn local_transitive_less_than() { ----- stderr ----- × No solution found when resolving dependencies: - ╰─▶ Because only package-a==1.0.0 is available and package-a==1.0.0 depends on package-b<2.0.0, we can conclude that all versions of package-a depend on package-b<2.0.0. + ╰─▶ Because package-a==1.0.0 depends on package-b<2.0.0 and only package-a==1.0.0 is available, we can conclude that all versions of package-a depend on package-b<2.0.0. And because you require package-a and you require package-b==2.0.0+foo, we can conclude that the requirements are unsatisfiable. "###); @@ -1797,7 +1821,7 @@ fn local_transitive_conflicting() { ----- stderr ----- × No solution found when resolving dependencies: - ╰─▶ Because only package-a==1.0.0 is available and package-a==1.0.0 depends on package-b==2.0.0+bar, we can conclude that all versions of package-a depend on package-b==2.0.0+bar. + ╰─▶ Because package-a==1.0.0 depends on package-b==2.0.0+bar and only package-a==1.0.0 is available, we can conclude that all versions of package-a depend on package-b==2.0.0+bar. And because you require package-a and you require package-b==2.0.0+foo, we can conclude that the requirements are unsatisfiable. "###); @@ -3331,10 +3355,8 @@ fn transitive_prerelease_and_stable_dependency_many_versions() { ----- stderr ----- × No solution found when resolving dependencies: - ╰─▶ Because only package-a==1.0.0 is available and package-a==1.0.0 depends on package-c>=2.0.0b1, we can conclude that all versions of package-a depend on package-c>=2.0.0b1. - And because only package-c<2.0.0b1 is available, we can conclude that all versions of package-a depend on package-c>3.0.0. - And because package-b==1.0.0 depends on package-c and only package-b==1.0.0 is available, we can conclude that all versions of package-b and all versions of package-a are incompatible. - And because you require package-a and you require package-b, we can conclude that the requirements are unsatisfiable. + ╰─▶ Because only package-c<2.0.0b1 is available and package-a==1.0.0 depends on package-c>=2.0.0b1, we can conclude that package-a==1.0.0 cannot be used. + And because only package-a==1.0.0 is available and you require package-a, we can conclude that the requirements are unsatisfiable. hint: package-c was requested with a pre-release marker (e.g., package-c>=2.0.0b1), but pre-releases weren't enabled (try: `--prerelease=allow`) "###);