Skip to content

Commit

Permalink
Incorporate heuristics to improve package prioritization
Browse files Browse the repository at this point in the history
  • Loading branch information
charliermarsh committed Apr 17, 2024
1 parent d1b07a3 commit 77dd8a9
Show file tree
Hide file tree
Showing 4 changed files with 163 additions and 49 deletions.
2 changes: 1 addition & 1 deletion crates/uv-resolver/src/dependency_provider.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<PubGrubPriority>;

type Err = Infallible;

Expand Down
122 changes: 107 additions & 15 deletions crates/uv-resolver/src/pubgrub/priority.rs
Original file line number Diff line number Diff line change
@@ -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: <https://github.com/pypa/pip/blob/ef78c129b1a966dbbbdb8ebfffc43723e89110d1/src/pip/_internal/resolution/resolvelib/provider.py#L120>
#[derive(Debug, Default)]
pub(crate) struct PubGrubPriorities(FxHashMap<PackageName, usize>);
pub(crate) struct PubGrubPriorities(FxHashMap<PackageName, PubGrubPriority>);

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<Version>) {
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<PubGrubPriority> {
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<usize>;
#[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<usize>),

/// The version range is constrained in some way (e.g., with a `<=` or `>` operator).
Constrained(Reverse<usize>),

/// The version range is constrained to a single version (e.g., with the `==` operator).
Singleton(Reverse<usize>),

/// The package was specified via a direct URL.
DirectUrl(Reverse<usize>),

/// The package is the root package.
Root,
}
32 changes: 16 additions & 16 deletions crates/uv-resolver/src/resolver/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -523,7 +520,6 @@ impl<
async fn visit_package(
&self,
package: &PubGrubPackage,
priorities: &mut PubGrubPriorities,
request_sink: &tokio::sync::mpsc::Sender<Request>,
) -> Result<(), ResolveError> {
match package {
Expand All @@ -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?;
}
}
Expand All @@ -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?;
}
}
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
56 changes: 39 additions & 17 deletions crates/uv/tests/pip_install_scenarios.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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.
"###);

Expand Down Expand Up @@ -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.
"###);

Expand Down Expand Up @@ -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.
"###);

Expand Down Expand Up @@ -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.
"###);

Expand Down Expand Up @@ -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.
"###);

Expand Down Expand Up @@ -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`)
"###);
Expand Down

0 comments on commit 77dd8a9

Please sign in to comment.