Only respect preferences across the same indexes (#9302)

## 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.

Author: charliermarsh

crates/uv-requirements/src/upgrade.rs

@@ -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 {
@@ -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();
 
@@ -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 })
 }

crates/uv-resolver/src/candidate_selector.rs

@@ -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;
@@ -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,
@@ -93,6 +94,7 @@ impl CandidateSelector {
             preferences,
             installed_packages,
             is_excluded,
+            index,
             env,
         ) {
             trace!("Using preference {} {}", preferred.name, preferred.version);
@@ -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,

crates/uv-resolver/src/preferences.rs

@@ -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 {
@@ -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>,
@@ -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()
@@ -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`].
@@ -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();
@@ -152,6 +168,7 @@ impl Preferences {
             if preference.fork_markers.is_empty() {
                 slf.insert(
                     preference.name,
+                    preference.index,
                     None,
                     Pin {
                         version: preference.version,
@@ -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(),
@@ -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.
@@ -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(),
+                    )
+                }),
             )
         })
     }
@@ -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.
@@ -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())
     }
 }
 

crates/uv-resolver/src/resolver/mod.rs

@@ -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(),
                             );
@@ -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
@@ -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.
@@ -1934,6 +1937,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
                     &self.preferences,
                     &self.installed_packages,
                     &self.exclusions,
+                    None,
                     &env,
                 ) else {
                     return Ok(None);

crates/uv-resolver/src/yanks.rs

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

crates/uv/src/commands/project/lock.rs

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