diff --git a/crates/uv-resolver/src/candidate_selector.rs b/crates/uv-resolver/src/candidate_selector.rs index e14a888bbf0a6..a6df1ee506242 100644 --- a/crates/uv-resolver/src/candidate_selector.rs +++ b/crates/uv-resolver/src/candidate_selector.rs @@ -3,6 +3,7 @@ use std::ops::Bound; use either::Either; use itertools::Itertools; +use rustc_hash::FxHashSet; use smallvec::SmallVec; use tracing::{debug, trace}; @@ -29,6 +30,15 @@ pub(crate) struct CandidateSelector { index_strategy: IndexStrategy, } +/// Controls how a batch advances through versions across multiple indexes. +#[derive(Clone, Copy)] +pub(crate) enum CandidateBatchMode { + /// Exhaust the current compatible range before moving to the next index. + Compatible(usize), + /// Move monotonically away from the previously selected version across indexes. + InOrder(usize), +} + impl CandidateSelector { /// Return a [`CandidateSelector`] for the given [`Manifest`]. pub(crate) fn for_resolution( @@ -509,6 +519,120 @@ impl CandidateSelector { } } + /// Select up to `limit` candidates without checking for version preferences. + /// + /// This is equivalent to repeatedly calling [`Self::select_no_preference`] and removing the + /// selected version from `range`, but visits each version map once. + pub(crate) fn select_no_preference_batch<'a>( + &'a self, + package_name: &'a PackageName, + range: &Range, + version_maps: &'a [VersionMap], + env: &ResolverEnvironment, + mode: CandidateBatchMode, + ) -> Vec> { + let (limit, in_order) = match mode { + CandidateBatchMode::Compatible(limit) => (limit, false), + CandidateBatchMode::InOrder(limit) => (limit, true), + }; + if limit == 0 { + return Vec::new(); + } + + let highest = self.use_highest_version(package_name, env); + let allow_prerelease = match self.prerelease_strategy.allows(package_name, env) { + AllowPrerelease::Yes => true, + AllowPrerelease::No => false, + AllowPrerelease::IfNecessary => !version_maps.iter().any(VersionMap::stable), + }; + + let mut selected_versions = FxHashSet::default(); + if self.index_strategy == IndexStrategy::UnsafeBestMatch { + let versions = + if highest { + Either::Left( + version_maps + .iter() + .enumerate() + .map(|(map_index, version_map)| { + version_map + .iter_included(range) + .rev() + .map(move |item| (map_index, item)) + }) + .kmerge_by(|(index1, (version1, _)), (index2, (version2, _))| { + match version1.cmp(version2) { + std::cmp::Ordering::Equal => index1 < index2, + std::cmp::Ordering::Less => false, + std::cmp::Ordering::Greater => true, + } + }) + .map(|(_, item)| item), + ) + } else { + Either::Right( + version_maps + .iter() + .enumerate() + .map(|(map_index, version_map)| { + version_map + .iter_included(range) + .map(move |item| (map_index, item)) + }) + .kmerge_by(|(index1, (version1, _)), (index2, (version2, _))| { + match version1.cmp(version2) { + std::cmp::Ordering::Equal => index1 < index2, + std::cmp::Ordering::Less => true, + std::cmp::Ordering::Greater => false, + } + }) + .map(|(_, item)| item), + ) + }; + return Self::select_candidates( + versions, + package_name, + range, + allow_prerelease, + highest, + limit, + &mut selected_versions, + ); + } + + let mut candidates = Vec::with_capacity(limit); + let mut current_range = range.clone(); + for version_map in version_maps { + let versions = if highest { + Either::Left(version_map.iter_included(¤t_range).rev()) + } else { + Either::Right(version_map.iter_included(¤t_range)) + }; + candidates.extend(Self::select_candidates( + versions, + package_name, + ¤t_range, + allow_prerelease, + highest, + limit - candidates.len(), + &mut selected_versions, + )); + if candidates.len() == limit { + break; + } + if in_order && let Some(candidate) = candidates.last() { + let bound = if highest { + Range::strictly_lower_than(candidate.version().clone()) + } else { + Range::strictly_higher_than(candidate.version().clone()) + }; + current_range = current_range.intersection(&bound); + } + } + + candidates + } + /// By default, we select the latest version, but we also allow using the lowest version instead /// to check the lower bounds. pub(crate) fn use_highest_version( @@ -646,6 +770,101 @@ impl CandidateSelector { ); None } + + /// Collect the candidates selected from an ordered set of versions. + /// + /// For versions present in multiple indexes, the first compatible distribution wins. An + /// incompatible distribution is retained only if no index has a compatible distribution for + /// that version, matching [`Self::select_candidate`]. + fn select_candidates<'a>( + versions: impl Iterator)>, + package_name: &'a PackageName, + range: &Range, + allow_prerelease: bool, + highest: bool, + limit: usize, + selected_versions: &mut FxHashSet<&'a Version>, + ) -> Vec> { + let mut candidates = Vec::with_capacity(limit); + let segments = range.iter(); + let segments = if highest { + Either::Left(segments.rev()) + } else { + Either::Right(segments) + }; + let Some(mut cursor) = RangeCursor::new(segments, highest) else { + return candidates; + }; + let mut incompatible: Option = None; + + for (version, maybe_dist) in versions { + if candidates.len() == limit { + break; + } + if selected_versions.contains(version) { + continue; + } + + if incompatible + .as_ref() + .is_some_and(|candidate| candidate.version != version) + { + if let Some(candidate) = incompatible.take() { + selected_versions.insert(candidate.version); + candidates.push(candidate); + } + if candidates.len() == limit { + break; + } + } + + if version.any_prerelease() && !allow_prerelease { + continue; + } + if !cursor.contains(version) { + continue; + } + let Some(dist) = maybe_dist.prioritized_dist() else { + continue; + }; + let candidate = + Candidate::new(package_name, version, dist, VersionChoiceKind::Compatible); + + if matches!( + candidate.dist(), + CandidateDist::Incompatible { + incompatible_dist: IncompatibleDist::Source(IncompatibleSource::ExcludeNewer( + _ + )) | IncompatibleDist::Wheel( + IncompatibleWheel::ExcludeNewer(_) + ), + .. + } + ) { + continue; + } + + if matches!(candidate.dist(), CandidateDist::Incompatible { .. }) { + if incompatible.is_none() { + incompatible = Some(candidate); + } + continue; + } + + selected_versions.insert(version); + incompatible = None; + candidates.push(candidate); + } + + if candidates.len() < limit + && let Some(candidate) = incompatible + { + selected_versions.insert(candidate.version); + candidates.push(candidate); + } + + candidates + } } /// Tracks membership in a range while visiting versions monotonically. @@ -720,6 +939,18 @@ fn is_after(version: &Version, bound: Bound<&Version>) -> bool { #[cfg(test)] mod tests { + use std::collections::BTreeMap; + + use uv_distribution_filename::WheelFilename; + use uv_distribution_types::{ + File, FileLocation, HashComparison, IndexUrl, RegistryBuiltWheel, WheelCompatibility, + }; + use uv_platform_tags::IncompatibleTag; + use uv_pypi_types::HashDigests; + use uv_small_str::SmallString; + + use crate::FlatDistributions; + use super::*; fn version(value: &str) -> Version { @@ -760,6 +991,203 @@ mod tests { fn range_cursor_descending() { assert_range_cursor(true, &["6", "5", "4", "3", "2.5", "2", "1"]); } + + #[derive(Clone, Copy)] + enum TestCompatibility { + Compatible, + Incompatible, + Excluded, + } + + fn version_map(index: usize, entries: &[(&str, TestCompatibility)]) -> VersionMap { + let index_url = format!("https://index{index}.example/simple") + .parse::() + .expect("valid test index URL"); + let base = SmallString::from(format!("https://index{index}.example/files/")); + let distributions = entries + .iter() + .map(|(value, compatibility)| { + let version = version(value); + let filename = format!("example-{version}-py3-none-any.whl") + .parse::() + .expect("valid test wheel filename"); + let file = File { + dist_info_metadata: true, + filename: SmallString::from(filename.to_string()), + hashes: HashDigests::empty(), + requires_python: None, + size: None, + upload_time_utc_ms: None, + url: FileLocation::new(SmallString::from(filename.to_string()), &base), + yanked: None, + zstd: None, + }; + let wheel = RegistryBuiltWheel { + filename, + file: Box::new(file), + index: index_url.clone(), + }; + let compatibility = match compatibility { + TestCompatibility::Compatible => { + WheelCompatibility::Compatible(HashComparison::Matched, None, None) + } + TestCompatibility::Incompatible => WheelCompatibility::Incompatible( + IncompatibleWheel::Tag(IncompatibleTag::Platform), + ), + TestCompatibility::Excluded => { + WheelCompatibility::Incompatible(IncompatibleWheel::ExcludeNewer(Some(1))) + } + }; + ( + version, + PrioritizedDist::from_built(wheel, vec![], compatibility), + ) + }) + .collect::>(); + VersionMap::from(FlatDistributions::from(distributions)) + } + + fn candidate_signature(candidate: &Candidate<'_>) -> (String, bool, String) { + let index = candidate + .prioritized() + .and_then(PrioritizedDist::best_wheel) + .map(|(wheel, _)| wheel.index.to_string()) + .expect("test candidates have wheels"); + ( + candidate.version().to_string(), + candidate.compatible().is_some(), + index, + ) + } + + fn repeated_candidates( + selector: &CandidateSelector, + package_name: &PackageName, + range: &Range, + version_maps: &[VersionMap], + env: &ResolverEnvironment, + limit: usize, + in_order: bool, + ) -> Vec<(String, bool, String)> { + let mut range = range.clone(); + let mut candidates = Vec::with_capacity(limit); + for _ in 0..limit { + let Some(candidate) = + selector.select_no_preference(package_name, &range, version_maps, env) + else { + break; + }; + let selected = Range::singleton(candidate.version().clone()); + range = if in_order { + let bound = if selector.use_highest_version(package_name, env) { + Range::strictly_lower_than(candidate.version().clone()) + } else { + Range::strictly_higher_than(candidate.version().clone()) + }; + range.intersection(&bound) + } else { + range.intersection(&selected.complement()) + }; + candidates.push(candidate_signature(&candidate)); + } + candidates + } + + #[test] + fn batched_selector_matches_repeated_selection() { + use TestCompatibility::{Compatible, Excluded, Incompatible}; + + let version_maps = vec![ + version_map( + 0, + &[ + ("1", Compatible), + ("2", Compatible), + ("3", Incompatible), + ("3.5a1", Compatible), + ("4", Compatible), + ("5", Compatible), + ("6", Incompatible), + ("7a1", Compatible), + ("8", Excluded), + ("9", Compatible), + ], + ), + version_map( + 1, + &[ + ("1", Incompatible), + ("3", Compatible), + ("4", Incompatible), + ("6", Compatible), + ("7a1", Incompatible), + ("8", Compatible), + ("10", Compatible), + ], + ), + ]; + let package_name = "example" + .parse::() + .expect("valid test package name"); + let env = ResolverEnvironment::universal(vec![]); + let range = [ + (Bound::Unbounded, Bound::Excluded(version("2"))), + (Bound::Included(version("3")), Bound::Included(version("4"))), + (Bound::Excluded(version("5")), Bound::Unbounded), + ] + .into_iter() + .collect::>(); + + for highest in [false, true] { + for index_strategy in [ + IndexStrategy::FirstIndex, + IndexStrategy::UnsafeFirstMatch, + IndexStrategy::UnsafeBestMatch, + ] { + for prerelease_strategy in [ + PrereleaseStrategy::Disallow, + PrereleaseStrategy::Allow, + PrereleaseStrategy::IfNecessary, + ] { + let selector = CandidateSelector { + resolution_strategy: if highest { + ResolutionStrategy::Highest + } else { + ResolutionStrategy::Lowest + }, + prerelease_strategy, + index_strategy, + }; + for (in_order, mode) in [ + (false, CandidateBatchMode::Compatible(6)), + (true, CandidateBatchMode::InOrder(6)), + ] { + let expected = repeated_candidates( + &selector, + &package_name, + &range, + &version_maps, + &env, + 6, + in_order, + ); + let actual = selector + .select_no_preference_batch( + &package_name, + &range, + &version_maps, + &env, + mode, + ) + .iter() + .map(candidate_signature) + .collect::>(); + assert_eq!(actual, expected); + } + } + } + } + } } #[derive(Debug, Clone)] diff --git a/crates/uv-resolver/src/resolver/batch_prefetch.rs b/crates/uv-resolver/src/resolver/batch_prefetch.rs index c4401e5eb4511..9d551125deeba 100644 --- a/crates/uv-resolver/src/resolver/batch_prefetch.rs +++ b/crates/uv-resolver/src/resolver/batch_prefetch.rs @@ -7,7 +7,7 @@ use rustc_hash::{FxHashMap, FxHashSet}; use tokio::sync::mpsc::Sender; use tracing::{debug, trace}; -use crate::candidate_selector::CandidateSelector; +use crate::candidate_selector::{CandidateBatchMode, CandidateSelector}; use crate::pubgrub::{PubGrubPackage, PubGrubPackageInner, Range}; use crate::resolver::Request; use crate::{ @@ -18,22 +18,6 @@ use uv_normalize::PackageName; use uv_pep440::Version; use uv_pep508::MarkerTree; -enum BatchPrefetchStrategy { - /// Go through the next versions assuming the existing selection and its constraints - /// remain. - Compatible { - compatible: Range, - previous: Version, - }, - /// We encounter cases (botocore) where the above doesn't work: Say we previously selected - /// a==x.y.z, which depends on b==x.y.z. a==x.y.z is incompatible, but we don't know that - /// yet. We just selected b==x.y.z and want to prefetch, since for all versions of a we try, - /// we have to wait for the matching version of b. The exiting range gives us only one version - /// of b, so the compatible strategy doesn't prefetch any version. Instead, we try the next - /// heuristic where the next version of b will be x.y.(z-1) and so forth. - InOrder { previous: Version }, -} - /// Prefetch a large number of versions if we already unsuccessfully tried many versions. /// /// This is an optimization specifically targeted at cold cache urllib3/boto3/botocore, where we @@ -120,11 +104,6 @@ impl BatchPrefetcher { .map_err(|_| ResolveError::UnregisteredTask(name.to_string()))? }; - let phase = BatchPrefetchStrategy::Compatible { - compatible: current_range.clone(), - previous: version.clone(), - }; - self.last_prefetch.insert(name.clone(), num_tried); self.prefetch_runner.send_prefetch( @@ -132,7 +111,8 @@ impl BatchPrefetcher { unchangeable_constraints, total_prefetch, &versions_response, - phase, + current_range, + version, python_requirement, selector, env, @@ -212,7 +192,8 @@ impl BatchPrefetcherRunner { unchangeable_constraints: Option<&Term>>, total_prefetch: usize, versions_response: &Arc, - mut phase: BatchPrefetchStrategy, + compatible: &Range, + previous: &Version, python_requirement: &PythonRequirement, selector: &CandidateSelector, env: &ResolverEnvironment, @@ -221,62 +202,47 @@ impl BatchPrefetcherRunner { return Ok(()); }; - let mut prefetch_count = 0; - for _ in 0..total_prefetch { - let candidate = match phase { - BatchPrefetchStrategy::Compatible { - compatible, - previous, - } => { - if let Some(candidate) = - selector.select_no_preference(name, &compatible, version_map, env) - { - let compatible = compatible.intersection( - &Range::singleton(candidate.version().clone()).complement(), - ); - phase = BatchPrefetchStrategy::Compatible { - compatible, - previous: candidate.version().clone(), - }; - candidate - } else { - // We exhausted the compatible version, switch to ignoring the existing - // constraints on the package and instead going through versions in order. - phase = BatchPrefetchStrategy::InOrder { previous }; - continue; - } - } - BatchPrefetchStrategy::InOrder { previous } => { - let mut range = if selector.use_highest_version(name, env) { - Range::strictly_lower_than(previous) - } else { - Range::strictly_higher_than(previous) - }; - // If we have constraints from root, don't go beyond those. Example: We are - // prefetching for foo 1.60 and have a dependency for `foo>=1.50`, so we should - // only prefetch 1.60 to 1.50, knowing 1.49 will always be rejected. - if let Some(unchangeable_constraints) = &unchangeable_constraints { - range = match unchangeable_constraints { - Term::Positive(constraints) => range.intersection(constraints), - Term::Negative(negative_constraints) => { - range.intersection(&negative_constraints.complement()) - } - }; - } - if let Some(candidate) = - selector.select_no_preference(name, &range, version_map, env) - { - phase = BatchPrefetchStrategy::InOrder { - previous: candidate.version().clone(), - }; - candidate - } else { - // Both strategies exhausted their candidates. - break; - } - } + // First consider candidates compatible with the current constraints. Selecting the batch + // in one traversal avoids repeatedly intersecting singleton complements into `compatible`. + let mut candidates = selector.select_no_preference_batch( + name, + compatible, + version_map, + env, + CandidateBatchMode::Compatible(total_prefetch), + ); + let compatible_count = candidates.len(); + + // If the compatible strategy is exhausted, use the remaining slots for versions in order. + // One slot is consumed by the transition, matching the original prefetch loop. + if compatible_count < total_prefetch { + let previous = candidates + .last() + .map_or(previous, |candidate| candidate.version()); + let mut range = if selector.use_highest_version(name, env) { + Range::strictly_lower_than(previous.clone()) + } else { + Range::strictly_higher_than(previous.clone()) }; + if let Some(unchangeable_constraints) = unchangeable_constraints { + range = match unchangeable_constraints { + Term::Positive(constraints) => range.intersection(constraints), + Term::Negative(negative_constraints) => { + range.intersection(&negative_constraints.complement()) + } + }; + } + candidates.extend(selector.select_no_preference_batch( + name, + &range, + version_map, + env, + CandidateBatchMode::InOrder(total_prefetch - compatible_count - 1), + )); + } + let mut prefetch_count = 0; + for (candidate_index, candidate) in candidates.into_iter().enumerate() { let Some(dist) = candidate.compatible() else { continue; }; @@ -305,9 +271,10 @@ impl BatchPrefetcherRunner { // Emit a request to fetch the metadata for this version. trace!( "Prefetching {prefetch_count} ({}) {}", - match phase { - BatchPrefetchStrategy::Compatible { .. } => "compatible", - BatchPrefetchStrategy::InOrder { .. } => "in order", + if candidate_index < compatible_count { + "compatible" + } else { + "in order" }, dist );