Skip to content

Commit

Permalink
Add a proxy layer for extras (#3100)
Browse files Browse the repository at this point in the history 
Given requirements like:

```
black==23.1.0
black[colorama]
```

The resolver will (on `main`) add a dependency on Black, and then try to
use the most recent version of Black to satisfy `black[colorama]`. For
sake of example, assume `black==24.0.0` is the most recent version. Once
the selects this most recent version, it'll fetch the metadata, then
return the dependencies for `black==24.0.0` with the `colorama` extra
enabled. Finally, it will tack on `black==24.0.0` (a dependency on the
base package). The resolver will then detect a conflict between
`black==23.1.0` and `black==24.0.0`, and throw out
`black[colorama]==24.0.0`, trying to next most-recent version.

This is both wasteful and can cause problems, since we're fetching
metadata for versions that will _never_ satisfy the resolver. In the
`apache-airflow[all]` case, I also ran into an issue whereby we were
attempting to build very old versions of `apache-airflow` due to
`apache-airflow[pandas]`, which in turn led to resolution failures.

The solution proposed here is that we create a new proxy package with
exactly two dependencies: one on `black` and one of `black[colorama]`.
Both of these packages must be at the same version as the proxy package,
so the resolver knows much _earlier_ that (in the above example) the
extra variant _must_ match `23.1.0`.
  • Loading branch information
@charliermarsh
charliermarsh committed Apr 19, 2024
1 parent 822ae19 commit 2e88bb6
Show file tree
Hide file tree
Showing 6 changed files with 128 additions and 98 deletions.
53 changes: 48 additions & 5 deletions crates/uv-resolver/src/error.rs
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
use std::collections::{BTreeMap, BTreeSet};
use std::fmt::Formatter;
use std::ops::Deref;
use std::sync::Arc;

use dashmap::{DashMap, DashSet};
use indexmap::IndexMap;
use pubgrub::range::Range;
use pubgrub::report::{DefaultStringReporter, DerivationTree, Reporter};
use pubgrub::report::{DefaultStringReporter, DerivationTree, External, Reporter};
use rustc_hash::FxHashMap;

use distribution_types::{
Expand Down Expand Up @@ -109,6 +110,45 @@ impl<T> From<tokio::sync::mpsc::error::SendError<T>> for ResolveError {
}
}

/// Given a [`DerivationTree`], collapse any [`External::FromDependencyOf`] incompatibilities
/// wrap an [`PubGrubPackage::Extra`] package.
fn collapse_extra_proxies(derivation_tree: &mut DerivationTree<PubGrubPackage, Range<Version>>) {
match derivation_tree {
DerivationTree::External(_) => {}
DerivationTree::Derived(derived) => {
match (
Arc::make_mut(&mut derived.cause1),
Arc::make_mut(&mut derived.cause2),
) {
(
DerivationTree::External(External::FromDependencyOf(
PubGrubPackage::Extra(..),
..,
)),
ref mut cause,
) => {
collapse_extra_proxies(cause);
*derivation_tree = cause.clone();
}
(
ref mut cause,
DerivationTree::External(External::FromDependencyOf(
PubGrubPackage::Extra(..),
..,
)),
) => {
collapse_extra_proxies(cause);
*derivation_tree = cause.clone();
}
_ => {
collapse_extra_proxies(Arc::make_mut(&mut derived.cause1));
collapse_extra_proxies(Arc::make_mut(&mut derived.cause2));
}
}
}
}
}

impl From<pubgrub::error::PubGrubError<UvDependencyProvider>> for ResolveError {
fn from(value: pubgrub::error::PubGrubError<UvDependencyProvider>) -> Self {
match value {
Expand All @@ -119,7 +159,9 @@ impl From<pubgrub::error::PubGrubError<UvDependencyProvider>> for ResolveError {
unreachable!()
}
pubgrub::error::PubGrubError::Failure(inner) => Self::Failure(inner),
pubgrub::error::PubGrubError::NoSolution(derivation_tree) => {
pubgrub::error::PubGrubError::NoSolution(mut derivation_tree) => {
collapse_extra_proxies(&mut derivation_tree);

Self::NoSolution(NoSolutionError {
derivation_tree,
// The following should be populated before display for the best error messages
Expand Down Expand Up @@ -208,7 +250,8 @@ impl NoSolutionError {
BTreeSet::from([python_requirement.target().deref().clone()]),
);
}
PubGrubPackage::Package(name, ..) => {
PubGrubPackage::Extra(_, _, _) => {}
PubGrubPackage::Package(name, _, _) => {
// Avoid including available versions for packages that exist in the derivation
// tree, but were never visited during resolution. We _may_ have metadata for
// these packages, but it's non-deterministic, and omitting them ensures that
Expand Down Expand Up @@ -256,7 +299,7 @@ impl NoSolutionError {
) -> Self {
let mut new = FxHashMap::default();
for package in self.derivation_tree.packages() {
if let PubGrubPackage::Package(name, ..) = package {
if let PubGrubPackage::Package(name, _, _) = package {
if let Some(entry) = unavailable_packages.get(name) {
let reason = entry.value();
new.insert(name.clone(), reason.clone());
Expand All @@ -275,7 +318,7 @@ impl NoSolutionError {
) -> Self {
let mut new = FxHashMap::default();
for package in self.derivation_tree.packages() {
if let PubGrubPackage::Package(name, ..) = package {
if let PubGrubPackage::Package(name, _, _) = package {
if let Some(entry) = incomplete_packages.get(name) {
let versions = entry.value();
for entry in versions {
Expand Down
21 changes: 20 additions & 1 deletion crates/uv-resolver/src/pubgrub/package.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,13 +61,31 @@ pub enum PubGrubPackage {
/// _after_ a registry variant.
Option<VerbatimUrl>,
),
/// A proxy package to represent a dependency with an extra (e.g., `black[colorama]`).
///
/// For a given package `black`, and an extra `colorama`, we create a virtual package
/// with exactly two dependencies: `PubGrubPackage::Package("black", None)` and
/// `PubGrubPackage::Package("black", Some("colorama")`. Both dependencies are pinned to the
/// same version, and the virtual package is discarded at the end of the resolution process.
///
/// The benefit of the proxy package (versus `PubGrubPackage::Package("black", Some("colorama")`
/// on its own) is that it enables us to avoid attempting to retrieve metadata for irrelevant
/// versions the extra variants by making it clear to PubGrub that the extra variant must match
/// the exact same version of the base variant. Without the proxy package, then when provided
/// requirements like `black==23.0.1` and `black[colorama]`, PubGrub may attempt to retrieve
/// metadata for `black[colorama]` versions other than `23.0.1`.
Extra(PackageName, ExtraName, Option<VerbatimUrl>),
}

impl PubGrubPackage {
/// Create a [`PubGrubPackage`] from a package name and optional extra name.
pub(crate) fn from_package(name: PackageName, extra: Option<ExtraName>, urls: &Urls) -> Self {
let url = urls.get(&name).cloned();
Self::Package(name, extra, url)
if let Some(extra) = extra {
Self::Extra(name, extra, url)
} else {
Self::Package(name, extra, url)
}
}
}

Expand All @@ -94,6 +112,7 @@ impl std::fmt::Display for PubGrubPackage {
Self::Package(name, Some(extra), ..) => {
write!(f, "{name}[{extra}]")
}
Self::Extra(name, extra, ..) => write!(f, "{name}[{extra}]"),
}
}
}
6 changes: 4 additions & 2 deletions crates/uv-resolver/src/pubgrub/priority.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,8 @@ impl PubGrubPriorities {
match package {
PubGrubPackage::Root(_) => {}
PubGrubPackage::Python(_) => {}
PubGrubPackage::Package(name, _, None) => {

PubGrubPackage::Extra(name, _, None) | PubGrubPackage::Package(name, _, None) => {
match self.0.entry(name.clone()) {
std::collections::hash_map::Entry::Occupied(mut entry) => {
// Preserve the original index.
Expand Down Expand Up @@ -65,7 +66,7 @@ impl PubGrubPriorities {
}
}
}
PubGrubPackage::Package(name, _, Some(_)) => {
PubGrubPackage::Extra(name, _, Some(_)) | PubGrubPackage::Package(name, _, Some(_)) => {
match self.0.entry(name.clone()) {
std::collections::hash_map::Entry::Occupied(mut entry) => {
// Preserve the original index.
Expand Down Expand Up @@ -99,6 +100,7 @@ impl PubGrubPriorities {
match package {
PubGrubPackage::Root(_) => Some(PubGrubPriority::Root),
PubGrubPackage::Python(_) => Some(PubGrubPriority::Root),
PubGrubPackage::Extra(name, _, _) => self.0.get(name).copied(),
PubGrubPackage::Package(name, _, _) => self.0.get(name).copied(),
}
}
Expand Down
98 changes: 39 additions & 59 deletions crates/uv-resolver/src/resolver/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -525,6 +525,7 @@ impl<
match package {
PubGrubPackage::Root(_) => {}
PubGrubPackage::Python(_) => {}
PubGrubPackage::Extra(_, _, _) => {}
PubGrubPackage::Package(name, _extra, None) => {
// Verify that the package is allowed under the hash-checking policy.
if !self.hasher.allows_package(name) {
Expand Down Expand Up @@ -561,7 +562,7 @@ impl<
// Iterate over the potential packages, and fetch file metadata for any of them. These
// represent our current best guesses for the versions that we _might_ select.
for (package, range) in packages {
let PubGrubPackage::Package(package_name, _extra, None) = package else {
let PubGrubPackage::Package(package_name, None, None) = package else {
continue;
};
request_sink
Expand Down Expand Up @@ -604,16 +605,9 @@ impl<
}
}

PubGrubPackage::Package(package_name, extra, Some(url)) => {
if let Some(extra) = extra {
debug!(
"Searching for a compatible version of {package_name}[{extra}] @ {url} ({range})",
);
} else {
debug!(
"Searching for a compatible version of {package_name} @ {url} ({range})"
);
}
PubGrubPackage::Extra(package_name, _, Some(url))
| PubGrubPackage::Package(package_name, _, Some(url)) => {
debug!("Searching for a compatible version of {package} @ {url} ({range})");

// If the dist is an editable, return the version from the editable metadata.
if let Some((_local, metadata)) = self.editables.get(package_name) {
Expand Down Expand Up @@ -702,7 +696,8 @@ impl<
Ok(Some(ResolverVersion::Available(version.clone())))
}

PubGrubPackage::Package(package_name, extra, None) => {
PubGrubPackage::Extra(package_name, _, None)
| PubGrubPackage::Package(package_name, _, None) => {
// Wait for the metadata to be available.
let versions_response = self
.index
Expand Down Expand Up @@ -732,13 +727,7 @@ impl<
}
};

if let Some(extra) = extra {
debug!(
"Searching for a compatible version of {package_name}[{extra}] ({range})",
);
} else {
debug!("Searching for a compatible version of {package_name} ({range})");
}
debug!("Searching for a compatible version of {package} ({range})");

// Find a version.
let Some(candidate) = self.selector.select(
Expand Down Expand Up @@ -770,22 +759,13 @@ impl<
}
ResolvedDistRef::Installed(_) => Cow::Borrowed("installed"),
};
if let Some(extra) = extra {
debug!(
"Selecting: {}[{}]=={} ({})",
candidate.name(),
extra,
candidate.version(),
filename,
);
} else {
debug!(
"Selecting: {}=={} ({})",
candidate.name(),
candidate.version(),
filename,
);
}

debug!(
"Selecting: {}=={} ({})",
package,
candidate.version(),
filename,
);

// We want to return a package pinned to a specific version; but we _also_ want to
// store the exact file that we selected to satisfy that version.
Expand All @@ -794,13 +774,16 @@ impl<
let version = candidate.version().clone();

// Emit a request to fetch the metadata for this version.
if self.index.distributions.register(candidate.version_id()) {
let request = match dist.for_resolution() {
ResolvedDistRef::Installable(dist) => Request::Dist(dist.clone()),
ResolvedDistRef::Installed(dist) => Request::Installed(dist.clone()),
};
request_sink.send(request).await?;
if matches!(package, PubGrubPackage::Package(_, _, _)) {
if self.index.distributions.register(candidate.version_id()) {
let request = match dist.for_resolution() {
ResolvedDistRef::Installable(dist) => Request::Dist(dist.clone()),
ResolvedDistRef::Installed(dist) => Request::Installed(dist.clone()),
};
request_sink.send(request).await?;
}
}

Ok(Some(ResolverVersion::Available(version)))
}
}
Expand Down Expand Up @@ -896,7 +879,7 @@ impl<

// Determine if the distribution is editable.
if let Some((_local, metadata)) = self.editables.get(package_name) {
let mut constraints = PubGrubDependencies::from_requirements(
let constraints = PubGrubDependencies::from_requirements(
&metadata.requires_dist,
&self.constraints,
&self.overrides,
Expand All @@ -917,14 +900,6 @@ impl<
self.visit_package(dep_package, request_sink).await?;
}

// If a package has an extra, insert a constraint on the base package.
if extra.is_some() {
constraints.push(
PubGrubPackage::Package(package_name.clone(), None, url.clone()),
Range::singleton(version.clone()),
);
}

return Ok(Dependencies::Available(constraints.into()));
}

Expand Down Expand Up @@ -1013,7 +988,7 @@ impl<
}
};

let mut constraints = PubGrubDependencies::from_requirements(
let constraints = PubGrubDependencies::from_requirements(
&metadata.requires_dist,
&self.constraints,
&self.overrides,
Expand All @@ -1034,16 +1009,20 @@ impl<
self.visit_package(package, request_sink).await?;
}

// If a package has an extra, insert a constraint on the base package.
if extra.is_some() {
constraints.push(
PubGrubPackage::Package(package_name.clone(), None, url.clone()),
Range::singleton(version.clone()),
);
}

Ok(Dependencies::Available(constraints.into()))
}

// Add a dependency on both the extra and base package.
PubGrubPackage::Extra(package_name, extra, url) => Ok(Dependencies::Available(vec![
(
PubGrubPackage::Package(package_name.clone(), None, url.clone()),
Range::singleton(version.clone()),
),
(
PubGrubPackage::Package(package_name.clone(), Some(extra.clone()), url.clone()),
Range::singleton(version.clone()),
),
])),
}
}

Expand Down Expand Up @@ -1251,6 +1230,7 @@ impl<
match package {
PubGrubPackage::Root(_) => {}
PubGrubPackage::Python(_) => {}
PubGrubPackage::Extra(_, _, _) => {}
PubGrubPackage::Package(package_name, _extra, Some(url)) => {
reporter.on_progress(package_name, &VersionOrUrl::Url(url));
}
Expand Down
Loading

0 comments on commit 2e88bb6

Please sign in to comment.