Skip to content

Commit

Permalink
Only respect preferences across the same indexes (#9302)
Browse files Browse the repository at this point in the history
## Summary

The issue here is fairly complex. Consider the following:

```toml
[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.12.0"
dependencies = []

[project.optional-dependencies]
cpu = [
  "torch>=2.5.1",
  "torchvision>=0.20.1",
]
cu124 = [
  "torch>=2.5.1",
  "torchvision>=0.20.1",
]

[tool.uv]
conflicts = [
  [
    { extra = "cpu" },
    { extra = "cu124" },
  ],
]

[tool.uv.sources]
torch = [
  { index = "pytorch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'" },
]
torchvision = [
  { index = "pytorch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'" },
]

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
```

When solving this project, we first pick a PyTorch version from PyPI, to
solve the `cu124` extra, selecting `2.5.1`.

Later, we try to solve the `cpu` extra. In solving that extra, we look
at the PyTorch CPU index. Ideally, we'd select `2.5.1+cpu`... But
`2.5.1` is already a preference. So we choose that.

Now, we only respect preferences for explicit indexes if they came from
the same index.

Closes #9295.
  • Loading branch information
charliermarsh authored Nov 21, 2024
1 parent c6482dd commit 5e48819
Show file tree
Hide file tree
Showing 9 changed files with 129 additions and 64 deletions.
12 changes: 8 additions & 4 deletions crates/uv-requirements/src/upgrade.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use uv_configuration::Upgrade;
use uv_fs::CWD;
use uv_git::ResolvedRepositoryReference;
use uv_requirements_txt::RequirementsTxt;
use uv_resolver::{Lock, Preference, PreferenceError};
use uv_resolver::{Lock, LockError, Preference, PreferenceError};

#[derive(Debug, Default)]
pub struct LockedRequirements {
Expand Down Expand Up @@ -63,7 +63,11 @@ pub async fn read_requirements_txt(
}

/// Load the preferred requirements from an existing lockfile, applying the upgrade strategy.
pub fn read_lock_requirements(lock: &Lock, upgrade: &Upgrade) -> LockedRequirements {
pub fn read_lock_requirements(
lock: &Lock,
install_path: &Path,
upgrade: &Upgrade,
) -> Result<LockedRequirements, LockError> {
let mut preferences = Vec::new();
let mut git = Vec::new();

Expand All @@ -74,13 +78,13 @@ pub fn read_lock_requirements(lock: &Lock, upgrade: &Upgrade) -> LockedRequireme
}

// Map each entry in the lockfile to a preference.
preferences.push(Preference::from_lock(package));
preferences.push(Preference::from_lock(package, install_path)?);

// Map each entry in the lockfile to a Git SHA.
if let Some(git_ref) = package.as_git_ref() {
git.push(git_ref);
}
}

LockedRequirements { preferences, git }
Ok(LockedRequirements { preferences, git })
}
46 changes: 32 additions & 14 deletions crates/uv-resolver/src/candidate_selector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use std::fmt::{Display, Formatter};
use tracing::{debug, trace};

use uv_configuration::IndexStrategy;
use uv_distribution_types::{CompatibleDist, IncompatibleDist, IncompatibleSource};
use uv_distribution_types::{CompatibleDist, IncompatibleDist, IncompatibleSource, IndexUrl};
use uv_distribution_types::{DistributionMetadata, IncompatibleWheel, Name, PrioritizedDist};
use uv_normalize::PackageName;
use uv_pep440::Version;
Expand Down Expand Up @@ -80,11 +80,12 @@ impl CandidateSelector {
preferences: &'a Preferences,
installed_packages: &'a InstalledPackages,
exclusions: &'a Exclusions,
index: Option<&'a IndexUrl>,
env: &ResolverEnvironment,
) -> Option<Candidate<'a>> {
let is_excluded = exclusions.contains(package_name);

// Check for a preference from a lockfile or a previous fork that satisfies the range and
// Check for a preference from a lockfile or a previous fork that satisfies the range and
// is allowed.
if let Some(preferred) = self.get_preferred(
package_name,
Expand All @@ -93,6 +94,7 @@ impl CandidateSelector {
preferences,
installed_packages,
is_excluded,
index,
env,
) {
trace!("Using preference {} {}", preferred.name, preferred.version);
Expand Down Expand Up @@ -131,23 +133,39 @@ impl CandidateSelector {
preferences: &'a Preferences,
installed_packages: &'a InstalledPackages,
is_excluded: bool,
index: Option<&'a IndexUrl>,
env: &ResolverEnvironment,
) -> Option<Candidate> {
// In the branches, we "sort" the preferences by marker-matching through an iterator that
// first has the matching half and then the mismatching half.
let preferences_match = preferences.get(package_name).filter(|(marker, _version)| {
// `.unwrap_or(true)` because the universal marker is considered matching.
marker
.map(|marker| env.included_by_marker(marker))
.unwrap_or(true)
});
let preferences_mismatch = preferences.get(package_name).filter(|(marker, _version)| {
marker
.map(|marker| !env.included_by_marker(marker))
.unwrap_or(false)
});
let preferences_match =
preferences
.get(package_name)
.filter(|(marker, _index, _version)| {
// `.unwrap_or(true)` because the universal marker is considered matching.
marker
.map(|marker| env.included_by_marker(marker))
.unwrap_or(true)
});
let preferences_mismatch =
preferences
.get(package_name)
.filter(|(marker, _index, _version)| {
marker
.map(|marker| !env.included_by_marker(marker))
.unwrap_or(false)
});
let preferences = preferences_match.chain(preferences_mismatch).filter_map(
|(marker, source, version)| {
// If the package is mapped to an explicit index, only consider preferences that
// match the index.
index
.map_or(true, |index| source == Some(index))
.then_some((marker, version))
},
);
self.get_preferred_from_iter(
preferences_match.chain(preferences_mismatch),
preferences,
package_name,
range,
version_maps,
Expand Down
74 changes: 50 additions & 24 deletions crates/uv-resolver/src/preferences.rs
Original file line number Diff line number Diff line change
@@ -1,16 +1,17 @@
use std::path::Path;
use std::str::FromStr;

use rustc_hash::FxHashMap;
use tracing::trace;

use uv_distribution_types::{InstalledDist, InstalledMetadata, InstalledVersion, Name};
use uv_distribution_types::{IndexUrl, InstalledDist, InstalledMetadata, InstalledVersion, Name};
use uv_normalize::PackageName;
use uv_pep440::{Operator, Version};
use uv_pep508::{MarkerTree, VersionOrUrl};
use uv_pypi_types::{HashDigest, HashError};
use uv_requirements_txt::{RequirementEntry, RequirementsTxtRequirement};

use crate::ResolverEnvironment;
use crate::{LockError, ResolverEnvironment};

#[derive(thiserror::Error, Debug)]
pub enum PreferenceError {
Expand All @@ -25,6 +26,8 @@ pub struct Preference {
version: Version,
/// The markers on the requirement itself (those after the semicolon).
marker: MarkerTree,
/// The index URL of the package, if any.
index: Option<IndexUrl>,
/// If coming from a package with diverging versions, the markers of the forks this preference
/// is part of, otherwise `None`.
fork_markers: Vec<MarkerTree>,
Expand Down Expand Up @@ -60,6 +63,7 @@ impl Preference {
marker: requirement.marker,
// requirements.txt doesn't have fork annotations.
fork_markers: vec![],
index: None,
hashes: entry
.hashes
.iter()
Expand All @@ -79,21 +83,26 @@ impl Preference {
name: dist.name().clone(),
version: version.clone(),
marker: MarkerTree::TRUE,
index: None,
// Installed distributions don't have fork annotations.
fork_markers: vec![],
hashes: Vec::new(),
}
}

/// Create a [`Preference`] from a locked distribution.
pub fn from_lock(package: &crate::lock::Package) -> Self {
Self {
pub fn from_lock(
package: &crate::lock::Package,
install_path: &Path,
) -> Result<Self, LockError> {
Ok(Self {
name: package.id.name.clone(),
version: package.id.version.clone(),
marker: MarkerTree::TRUE,
index: package.index(install_path)?,
fork_markers: package.fork_markers().to_vec(),
hashes: Vec::new(),
}
})
}

/// Return the [`PackageName`] of the package for this [`Preference`].
Expand All @@ -107,22 +116,29 @@ impl Preference {
}
}

#[derive(Debug, Clone)]
struct Entry {
marker: Option<MarkerTree>,
index: Option<IndexUrl>,
pin: Pin,
}

/// A set of pinned packages that should be preserved during resolution, if possible.
///
/// The marker is the marker of the fork that resolved to the pin, if any.
///
/// Preferences should be prioritized first by whether their marker matches and then by the order
/// they are stored, so that a lockfile has higher precedence than sibling forks.
#[derive(Debug, Clone, Default)]
pub struct Preferences(FxHashMap<PackageName, Vec<(Option<MarkerTree>, Pin)>>);
pub struct Preferences(FxHashMap<PackageName, Vec<Entry>>);

impl Preferences {
/// Create a map of pinned packages from an iterator of [`Preference`] entries.
///
/// The provided [`ResolverEnvironment`] will be used to filter the preferences
/// to an applicable subset.
pub fn from_iter<PreferenceIterator: IntoIterator<Item = Preference>>(
preferences: PreferenceIterator,
pub fn from_iter(
preferences: impl IntoIterator<Item = Preference>,
env: &ResolverEnvironment,
) -> Self {
let mut slf = Self::default();
Expand Down Expand Up @@ -152,6 +168,7 @@ impl Preferences {
if preference.fork_markers.is_empty() {
slf.insert(
preference.name,
preference.index,
None,
Pin {
version: preference.version,
Expand All @@ -162,6 +179,7 @@ impl Preferences {
for fork_marker in preference.fork_markers {
slf.insert(
preference.name.clone(),
preference.index.clone(),
Some(fork_marker),
Pin {
version: preference.version.clone(),
Expand All @@ -179,13 +197,15 @@ impl Preferences {
pub(crate) fn insert(
&mut self,
package_name: PackageName,
index: Option<IndexUrl>,
markers: Option<MarkerTree>,
pin: impl Into<Pin>,
) {
self.0
.entry(package_name)
.or_default()
.push((markers, pin.into()));
self.0.entry(package_name).or_default().push(Entry {
marker: markers,
index,
pin: pin.into(),
});
}

/// Returns an iterator over the preferences.
Expand All @@ -194,15 +214,19 @@ impl Preferences {
) -> impl Iterator<
Item = (
&PackageName,
impl Iterator<Item = (Option<&MarkerTree>, &Version)>,
impl Iterator<Item = (Option<&MarkerTree>, Option<&IndexUrl>, &Version)>,
),
> {
self.0.iter().map(|(name, preferences)| {
(
name,
preferences
.iter()
.map(|(markers, pin)| (markers.as_ref(), pin.version())),
preferences.iter().map(|entry| {
(
entry.marker.as_ref(),
entry.index.as_ref(),
entry.pin.version(),
)
}),
)
})
}
Expand All @@ -211,12 +235,14 @@ impl Preferences {
pub(crate) fn get(
&self,
package_name: &PackageName,
) -> impl Iterator<Item = (Option<&MarkerTree>, &Version)> {
self.0
.get(package_name)
.into_iter()
.flatten()
.map(|(markers, pin)| (markers.as_ref(), pin.version()))
) -> impl Iterator<Item = (Option<&MarkerTree>, Option<&IndexUrl>, &Version)> {
self.0.get(package_name).into_iter().flatten().map(|entry| {
(
entry.marker.as_ref(),
entry.index.as_ref(),
entry.pin.version(),
)
})
}

/// Return the hashes for a package, if the version matches that of the pin.
Expand All @@ -229,8 +255,8 @@ impl Preferences {
.get(package_name)
.into_iter()
.flatten()
.find(|(_markers, pin)| pin.version() == version)
.map(|(_markers, pin)| pin.hashes())
.find(|entry| entry.pin.version() == version)
.map(|entry| entry.pin.hashes())
}
}

Expand Down
8 changes: 6 additions & 2 deletions crates/uv-resolver/src/resolver/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -381,6 +381,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
for (package, version) in &resolution.nodes {
preferences.insert(
package.name.clone(),
package.index.clone(),
resolution.env.try_markers().cloned(),
version.clone(),
);
Expand Down Expand Up @@ -669,14 +670,15 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
diverging_packages: &'a [PackageName],
) -> impl Iterator<Item = Result<ForkState, ResolveError>> + 'a {
debug!(
"Splitting resolution on {}=={} over {} into {} resolution with separate markers",
"Splitting resolution on {}=={} over {} into {} resolution{} with separate markers",
current_state.next,
version,
diverging_packages
.iter()
.map(ToString::to_string)
.join(", "),
forks.len()
forks.len(),
if forks.len() == 1 { "" } else { "s" }
);
assert!(forks.len() >= 2);
// This is a somewhat tortured technique to ensure
Expand Down Expand Up @@ -1075,6 +1077,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
preferences,
&self.installed_packages,
&self.exclusions,
index,
env,
) else {
// Short circuit: we couldn't find _any_ versions for a package.
Expand Down Expand Up @@ -1934,6 +1937,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
&self.preferences,
&self.installed_packages,
&self.exclusions,
None,
&env,
) else {
return Ok(None);
Expand Down
2 changes: 1 addition & 1 deletion crates/uv-resolver/src/yanks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ impl AllowedYanks {
allowed_yanks
.entry(name.clone())
.or_default()
.extend(preferences.map(|(_markers, version)| version.clone()));
.extend(preferences.map(|(.., version)| version.clone()));
}

Self(Arc::new(allowed_yanks))
Expand Down
3 changes: 2 additions & 1 deletion crates/uv/src/commands/project/lock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -563,7 +563,8 @@ async fn do_lock(

// If an existing lockfile exists, build up a set of preferences.
let LockedRequirements { preferences, git } = versions_lock
.map(|lock| read_lock_requirements(lock, upgrade))
.map(|lock| read_lock_requirements(lock, workspace.install_path(), upgrade))
.transpose()?
.unwrap_or_default();

// Populate the Git resolver.
Expand Down
Loading

0 comments on commit 5e48819

Please sign in to comment.