Use fork markers and fork preferences in resolution with lockfile (#5481)

By resolving for each fork from the lockfile individually and by adding
using preferences for the current fork, we solve the instability #5180.
I've tested the locally and will add the packse test scenarios upstack.

Part of
https://github.com/astral-sh/uv/issues/5180#issuecomment-2247696198

crates/uv-resolver/src/candidate_selector.rs

@@ -1,12 +1,12 @@
 use itertools::Itertools;
 use pubgrub::Range;
 use std::fmt::{Display, Formatter};
-use tracing::debug;
+use tracing::{debug, trace};
 
 use distribution_types::{CompatibleDist, IncompatibleDist, IncompatibleSource};
 use distribution_types::{DistributionMetadata, IncompatibleWheel, Name, PrioritizedDist};
 use pep440_rs::Version;
-use pep508_rs::MarkerEnvironment;
+use pep508_rs::{MarkerEnvironment, MarkerTree};
 use uv_configuration::IndexStrategy;
 use uv_normalize::PackageName;
 use uv_types::InstalledPackagesProvider;
@@ -66,9 +66,7 @@ impl CandidateSelector {
     pub(crate) fn index_strategy(&self) -> &IndexStrategy {
         &self.index_strategy
     }
-}
 
-impl CandidateSelector {
     /// Select a [`Candidate`] from a set of candidate versions and files.
     ///
     /// Unless present in the provided [`Exclusions`], local distributions from the
@@ -84,6 +82,8 @@ impl CandidateSelector {
         exclusions: &'a Exclusions,
         markers: &ResolverMarkers,
     ) -> Option<Candidate<'a>> {
+        // 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,
             range,
@@ -93,11 +93,18 @@ impl CandidateSelector {
             exclusions,
             markers,
         ) {
+            trace!("Using preference {} {}", preferred.name, preferred.version,);
             return Some(preferred);
         }
 
+        // Check for a locally installed distribution that satisfies the range and is allowed.
         if !exclusions.contains(package_name) {
             if let Some(installed) = Self::get_installed(package_name, range, installed_packages) {
+                trace!(
+                    "Using preference {} {} from installed package",
+                    installed.name,
+                    installed.version,
+                );
                 return Some(installed);
             }
         }
@@ -105,74 +112,160 @@ impl CandidateSelector {
         self.select_no_preference(package_name, range, version_maps, markers)
     }
 
-    /// Check for a preference (e.g., an existing version from an existing lockfile or
+    /// If the package has a preference, an existing version from an existing lockfile or a version
-    /// from a previous fork) that satisfies the current range.
+    /// from a sibling fork, and the preference satisfies the current range, use that.
+    ///
+    /// We try to find a resolution that, depending on the input, does not diverge from the
+    /// lockfile or matches a sibling fork. We try an exact match for the current markers (fork
+    /// or specific) first, to ensure stability with repeated locking. If that doesn't work, we
+    /// fall back to preferences that don't match in hopes of still resolving different forks into
+    /// the same version; A solution with less different versions is more desirable than one where
+    /// we may have more recent version in some cases, but overall more versions.
     fn get_preferred<'a, InstalledPackages: InstalledPackagesProvider>(
-        &self,
+        &'a self,
         package_name: &'a PackageName,
         range: &Range<Version>,
         version_maps: &'a [VersionMap],
         preferences: &'a Preferences,
         installed_packages: &'a InstalledPackages,
-        exclusions: &'a Exclusions,
+        exclusions: &Exclusions,
-        markers: &ResolverMarkers,
+        resolver_markers: &ResolverMarkers,
-    ) -> Option<Candidate<'a>> {
+    ) -> Option<Candidate> {
-        let version = preferences.version(package_name)?;
+        // In the branches, we "sort" the preferences by marker-matching through an iterator that
-
+        // first has the matching half and then the mismatching half.
-        // Respect the version range for this requirement.
+        match resolver_markers {
-        if !range.contains(version) {
+            ResolverMarkers::SpecificEnvironment(env) => {
-            return None;
+                // We may hit a combination of fork markers preferences with specific environment
-        }
+                // output in the future when adding support for the PEP 665 successor.
-
+                let preferences_match =
-        // Check for a locally installed distribution that matches the preferred version.
+                    preferences.get(package_name).filter(|(marker, _version)| {
-        if !exclusions.contains(package_name) {
+                        // `.unwrap_or(true)` because the universal marker is considered matching.
-            let installed_dists = installed_packages.get_packages(package_name);
+                        marker
-            match installed_dists.as_slice() {
+                            .map(|marker| marker.evaluate(env, &[]))
-                [] => {}
+                            .unwrap_or(true)
-                [dist] => {
+                    });
-                    if dist.version() == version {
+                let preferences_mismatch =
-                        debug!("Found installed version of {dist} that satisfies preference in {range}");
+                    preferences.get(package_name).filter(|(marker, _version)| {
-
+                        marker
-                        return Some(Candidate {
+                            .map(|marker| !marker.evaluate(env, &[]))
-                            name: package_name,
+                            .unwrap_or(false)
-                            version,
+                    });
-                            dist: CandidateDist::Compatible(CompatibleDist::InstalledDist(dist)),
+                self.get_preferred_from_iter(
-                            choice_kind: VersionChoiceKind::Preference,
+                    preferences_match.chain(preferences_mismatch),
-                        });
+                    package_name,
-                    }
+                    range,
-                }
+                    version_maps,
-                // We do not consider installed distributions with multiple versions because
+                    installed_packages,
-                // during installation these must be reinstalled from the remote
+                    exclusions,
-                _ => {
+                    resolver_markers,
-                    debug!("Ignoring installed versions of {package_name}: multiple distributions found");
+                )
-                }
+            }
+            ResolverMarkers::Universal { .. } => {
+                // In universal mode, all preferences are matching.
+                self.get_preferred_from_iter(
+                    preferences.get(package_name),
+                    package_name,
+                    range,
+                    version_maps,
+                    installed_packages,
+                    exclusions,
+                    resolver_markers,
+                )
+            }
+            ResolverMarkers::Fork(fork_markers) => {
+                let preferences_match =
+                    preferences.get(package_name).filter(|(marker, _version)| {
+                        // `.unwrap_or(true)` because the universal marker is considered matching.
+                        marker.map(|marker| marker == fork_markers).unwrap_or(true)
+                    });
+                let preferences_mismatch =
+                    preferences.get(package_name).filter(|(marker, _version)| {
+                        marker.map(|marker| marker != fork_markers).unwrap_or(false)
+                    });
+                self.get_preferred_from_iter(
+                    preferences_match.chain(preferences_mismatch),
+                    package_name,
+                    range,
+                    version_maps,
+                    installed_packages,
+                    exclusions,
+                    resolver_markers,
+                )
             }
         }
+    }
 
-        // Respect the pre-release strategy for this fork.
+    /// Return the first preference that satisfies the current range and is allowed.
-        if version.any_prerelease()
+    fn get_preferred_from_iter<'a, InstalledPackages: InstalledPackagesProvider>(
-            && self.prerelease_strategy.allows(package_name, markers) != AllowPreRelease::Yes
+        &'a self,
-        {
+        preferences: impl Iterator<Item = (Option<&'a MarkerTree>, &'a Version)>,
-            return None;
+        package_name: &'a PackageName,
+        range: &Range<Version>,
+        version_maps: &'a [VersionMap],
+        installed_packages: &'a InstalledPackages,
+        exclusions: &Exclusions,
+        resolver_markers: &ResolverMarkers,
+    ) -> Option<Candidate<'a>> {
+        for (_marker, version) in preferences {
+            // Respect the version range for this requirement.
+            if !range.contains(version) {
+                continue;
+            }
+
+            // Check for a locally installed distribution that matches the preferred version.
+            if !exclusions.contains(package_name) {
+                let installed_dists = installed_packages.get_packages(package_name);
+                match installed_dists.as_slice() {
+                    [] => {}
+                    [dist] => {
+                        if dist.version() == version {
+                            debug!("Found installed version of {dist} that satisfies preference in {range}");
+
+                            return Some(Candidate {
+                                name: package_name,
+                                version,
+                                dist: CandidateDist::Compatible(CompatibleDist::InstalledDist(
+                                    dist,
+                                )),
+                                choice_kind: VersionChoiceKind::Preference,
+                            });
+                        }
+                    }
+                    // We do not consider installed distributions with multiple versions because
+                    // during installation these must be reinstalled from the remote
+                    _ => {
+                        debug!("Ignoring installed versions of {package_name}: multiple distributions found");
+                    }
+                }
+            }
+
+            // Respect the pre-release strategy for this fork.
+            if version.any_prerelease()
+                && self
+                    .prerelease_strategy
+                    .allows(package_name, resolver_markers)
+                    != AllowPreRelease::Yes
+            {
+                continue;
+            }
+
+            // Check for a remote distribution that matches the preferred version
+            if let Some(file) = version_maps
+                .iter()
+                .find_map(|version_map| version_map.get(version))
+            {
+                return Some(Candidate::new(
+                    package_name,
+                    version,
+                    file,
+                    VersionChoiceKind::Preference,
+                ));
+            }
         }
-
-        // Check for a remote distribution that matches the preferred version
-        if let Some(file) = version_maps
-            .iter()
-            .find_map(|version_map| version_map.get(version))
-        {
-            return Some(Candidate::new(
-                package_name,
-                version,
-                file,
-                VersionChoiceKind::Preference,
-            ));
-        }
-
         None
     }
 
-    /// Check for a locally installed distribution that satisfies the range.
+    /// Check for an installed distribution that satisfies the current range and is allowed.
     fn get_installed<'a, InstalledPackages: InstalledPackagesProvider>(
         package_name: &'a PackageName,
         range: &Range<Version>,
@@ -217,8 +310,8 @@ impl CandidateSelector {
         version_maps: &'a [VersionMap],
         markers: &ResolverMarkers,
     ) -> Option<Candidate> {
-        tracing::trace!(
+        trace!(
-            "selecting candidate for package {package_name} with range {range:?} with {} remote versions",
+            "Selecting candidate for {package_name} with range {range} with {} remote versions",
             version_maps.iter().map(VersionMap::len).sum::<usize>(),
         );
         let highest = self.use_highest_version(package_name);
@@ -408,12 +501,9 @@ impl CandidateSelector {
 
             return Some(candidate);
         }
-        tracing::trace!(
+        trace!(
-            "exhausted all candidates for package {:?} with range {:?} \
+            "Exhausted all candidates for package {package_name} with range {range} \
-             after {} steps",
+             after {steps} steps",
-            package_name,
-            range,
-            steps,
         );
         match prerelease {
             None => None,

crates/uv-resolver/src/lock.rs

@@ -1132,6 +1132,10 @@ impl Distribution {
         &self.id.version
     }
 
+    pub fn fork_markers(&self) -> Option<&BTreeSet<MarkerTree>> {
+        self.fork_markers.as_ref()
+    }
+
     /// Returns a [`VersionId`] for this package that can be used for resolution.
     fn version_id(&self, workspace_root: &Path) -> Result<VersionId, LockError> {
         match &self.id.source {

crates/uv-resolver/src/preferences.rs

@@ -1,4 +1,4 @@
-use std::collections::hash_map::Entry;
+use std::collections::BTreeSet;
 use std::str::FromStr;
 
 use rustc_hash::FxHashMap;
@@ -22,7 +22,11 @@ pub enum PreferenceError {
 pub struct Preference {
     name: PackageName,
     version: Version,
+    /// The markers on the requirement itself (those after the semicolon).
     marker: Option<MarkerTree>,
+    /// If coming from a package with diverging versions, the markers of the forks this preference
+    /// is part of, otherwise `None`.
+    fork_markers: Option<BTreeSet<MarkerTree>>,
     hashes: Vec<HashDigest>,
 }
 
@@ -53,6 +57,8 @@ impl Preference {
             name: requirement.name,
             version: specifier.version().clone(),
             marker: requirement.marker,
+            // requirements.txt doesn't have fork annotations.
+            fork_markers: None,
             hashes: entry
                 .hashes
                 .iter()
@@ -72,6 +78,8 @@ impl Preference {
             name: dist.name().clone(),
             version: version.clone(),
             marker: None,
+            // Installed distributions don't have fork annotations.
+            fork_markers: None,
             hashes: Vec::new(),
         }
     }
@@ -82,6 +90,7 @@ impl Preference {
             name: dist.id.name.clone(),
             version: dist.id.version.clone(),
             marker: None,
+            fork_markers: dist.fork_markers().cloned(),
             hashes: Vec::new(),
         }
     }
@@ -98,58 +107,120 @@ impl Preference {
 }
 
 /// 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, Pin>);
+pub struct Preferences(FxHashMap<PackageName, Vec<(Option<MarkerTree>, Pin)>>);
 
 impl Preferences {
     /// Create a map of pinned packages from an iterator of [`Preference`] entries.
-    /// Takes ownership of the [`Preference`] entries.
     ///
-    /// The provided [`MarkerEnvironment`] will be used to filter  the preferences
+    /// The provided [`MarkerEnvironment`] will be used to filter the preferences
     /// to an applicable subset.
     pub fn from_iter<PreferenceIterator: IntoIterator<Item = Preference>>(
         preferences: PreferenceIterator,
         markers: Option<&MarkerEnvironment>,
     ) -> Self {
-        // TODO(zanieb): We should explicitly ensure that when a package name is seen multiple times
+        let mut slf = Self::default();
-        // that the newest or oldest version is preferred depending on the resolution strategy;
+        for preference in preferences {
-        // right now, the order is dependent on the given iterator.
+            // Filter non-matching preferences when resolving for an environment.
-        let preferences = preferences
+            if let Some(markers) = markers {
-            .into_iter()
+                if !preference
-            .filter_map(|preference| {
+                    .marker
-                if preference.marker.as_ref().map_or(true, |marker| {
+                    .as_ref()
-                    marker.evaluate_optional_environment(markers, &[])
+                    .map_or(true, |marker| marker.evaluate(markers, &[]))
-                }) {
+                {
-                    Some((
-                        preference.name,
-                        Pin {
-                            version: preference.version,
-                            hashes: preference.hashes,
-                        },
-                    ))
-                } else {
                     trace!("Excluding {preference} from preferences due to unmatched markers");
-                    None
+                    continue;
                 }
-            })
-            .collect();
 
-        Self(preferences)
+                if !preference
+                    .fork_markers
+                    .as_ref()
+                    .map(|fork_markers| {
+                        fork_markers
+                            .iter()
+                            .any(|marker| marker.evaluate(markers, &[]))
+                    })
+                    .unwrap_or(true)
+                {
+                    trace!("Excluding {preference} from preferences due to unmatched fork markers");
+                    continue;
+                }
+            }
+
+            // Flatten the list of markers into individual entries.
+            if let Some(fork_markers) = &preference.fork_markers {
+                for fork_marker in fork_markers {
+                    slf.insert(
+                        preference.name.clone(),
+                        Some(fork_marker.clone()),
+                        Pin {
+                            version: preference.version.clone(),
+                            hashes: preference.hashes.clone(),
+                        },
+                    );
+                }
+            } else {
+                slf.insert(
+                    preference.name,
+                    None,
+                    Pin {
+                        version: preference.version,
+                        hashes: preference.hashes,
+                    },
+                );
+            }
+        }
+
+        slf
     }
 
-    /// Return the [`Entry`] for a package in the preferences.
+    /// Insert a preference at the back.
-    pub fn entry(&mut self, package_name: PackageName) -> Entry<PackageName, Pin> {
+    pub(crate) fn insert(
-        self.0.entry(package_name)
+        &mut self,
+        package_name: PackageName,
+        markers: Option<MarkerTree>,
+        pin: impl Into<Pin>,
+    ) {
+        self.0
+            .entry(package_name)
+            .or_default()
+            .push((markers, pin.into()));
     }
 
     /// Returns an iterator over the preferences.
-    pub fn iter(&self) -> impl Iterator<Item = (&PackageName, &Version)> {
+    pub fn iter(
-        self.0.iter().map(|(name, pin)| (name, pin.version()))
+        &self,
+    ) -> impl Iterator<
+        Item = (
+            &PackageName,
+            impl Iterator<Item = (Option<&MarkerTree>, &Version)>,
+        ),
+    > {
+        self.0.iter().map(|(name, preferences)| {
+            (
+                name,
+                preferences
+                    .iter()
+                    .map(|(markers, pin)| (markers.as_ref(), pin.version())),
+            )
+        })
     }
 
     /// Return the pinned version for a package, if any.
-    pub(crate) fn version(&self, package_name: &PackageName) -> Option<&Version> {
+    pub(crate) fn get(
-        self.0.get(package_name).map(Pin::version)
+        &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()))
     }
 
     /// Return the hashes for a package, if the version matches that of the pin.
@@ -160,8 +231,10 @@ impl Preferences {
     ) -> Option<&[HashDigest]> {
         self.0
             .get(package_name)
-            .filter(|pin| pin.version() == version)
+            .into_iter()
-            .map(Pin::hashes)
+            .flatten()
+            .find(|(_markers, pin)| pin.version() == version)
+            .map(|(_markers, pin)| pin.hashes())
     }
 }
 
@@ -173,7 +246,7 @@ impl std::fmt::Display for Preference {
 
 /// The pinned data associated with a package in a locked `requirements.txt` file (e.g., `flask==1.2.3`).
 #[derive(Debug, Clone)]
-pub struct Pin {
+pub(crate) struct Pin {
     version: Version,
     hashes: Vec<HashDigest>,
 }

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

@@ -2,7 +2,6 @@
 
 use std::borrow::Cow;
 use std::cmp::Ordering;
-use std::collections::hash_map::Entry;
 use std::collections::{BTreeMap, BTreeSet, VecDeque};
 use std::fmt::{Display, Formatter, Write};
 use std::ops::Bound;
@@ -326,7 +325,17 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
             self.requires_python.clone(),
         );
         let mut preferences = self.preferences.clone();
-        let mut forked_states = vec![state];
+        let mut forked_states = if let ResolverMarkers::Universal {
+            fork_preferences: Some(fork_preferences),
+        } = &self.markers
+        {
+            fork_preferences
+                .iter()
+                .map(|fork_preference| state.clone().with_markers(fork_preference.clone()))
+                .collect()
+        } else {
+            vec![state]
+        };
         let mut resolutions = vec![];
 
         'FORK: while let Some(mut state) = forked_states.pop() {
@@ -381,11 +390,15 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
 
                     let resolution = state.into_resolution();
 
-                    // Walk over the selected versions, and mark them as preferences.
+                    // Walk over the selected versions, and mark them as preferences. We have to
+                    // add forks back as to not override the preferences from the lockfile for
+                    // the next fork
                     for (package, version) in &resolution.nodes {
-                        if let Entry::Vacant(entry) = preferences.entry(package.name.clone()) {
+                        preferences.insert(
-                            entry.insert(version.clone().into());
+                            package.name.clone(),
-                        }
+                            resolution.markers.fork_markers().cloned(),
+                            version.clone(),
+                        );
                     }
 
                     // If another fork had the same resolution, merge into that fork instead.

crates/uv-resolver/src/yanks.rs

@@ -41,11 +41,11 @@ impl AllowedYanks {
         }
 
         // Allow yanks for any packages that are already pinned in the lockfile.
-        for (name, version) in manifest.preferences.iter() {
+        for (name, preferences) in manifest.preferences.iter() {
             allowed_yanks
                 .entry(name.clone())
                 .or_default()
-                .insert(version.clone());
+                .extend(preferences.map(|(_markers, version)| version.clone()));
         }
 
         Self(Arc::new(allowed_yanks))

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

@@ -440,7 +440,7 @@ async fn do_lock(
                 &Reinstall::default(),
                 upgrade,
                 None,
-                resolver_markers,
+                resolver_markers.clone(),
                 python_requirement.clone(),
                 &client,
                 &flat_index,
@@ -516,7 +516,7 @@ async fn do_lock(
                 &Reinstall::default(),
                 upgrade,
                 None,
-                ResolverMarkers::universal(None),
+                resolver_markers,
                 python_requirement,
                 &client,
                 &flat_index,

crates/uv/tests/lock.rs

@@ -1580,6 +1580,10 @@ fn lock_upgrade_log_multi_version() -> Result<()> {
             lock, @r###"
         version = 1
         requires-python = ">=3.12"
+        environment-markers = [
+            "sys_platform == 'win32'",
+            "sys_platform != 'win32'",
+        ]
         exclude-newer = "2024-03-25 00:00:00 UTC"
 
         [[distribution]]