Batch prefetch per fork (#10029)

Previously, the batch prefetcher was part of the solver loop, used
across forks. This would lead to each preference in a fork being counted
as a tried version, so that after 5 forks with the identical version, we
would start batch prefetching. The reported numbers of tried versions
are also reported. By tracking the batch prefetcher on the fork the
numbers are corrected.

An alternative would be tracking the actually tried versions, but that
would mean more overhead in the top level solver loop when the current
heuristic works.

In `ecosystem/transformers`:

```
$ hyperfine --runs 10 --prepare "rm -f uv.lock" "../../target/release/uv lock --exclude-newer 2024-08-08T00:00:00Z" "uv lock --exclude-newer 2024-08-08T00:00:00Z"
Benchmark 1: ../../target/release/uv lock --exclude-newer 2024-08-08T00:00:00Z
  Time (mean ± σ):     386.2 ms ±   6.1 ms    [User: 396.0 ms, System: 144.5 ms]
  Range (min … max):   378.5 ms … 397.9 ms    10 runs

Benchmark 2: uv lock --exclude-newer 2024-08-08T00:00:00Z
  Time (mean ± σ):     422.0 ms ±   5.5 ms    [User: 459.6 ms, System: 190.3 ms]
  Range (min … max):   415.0 ms … 430.5 ms    10 runs

Summary
  ../../target/release/uv lock --exclude-newer 2024-08-08T00:00:00Z ran
    1.09 ± 0.02 times faster than uv lock --exclude-newer 2024-08-08T00:00:00Z
```
This commit is contained in:
konsti 2024-12-19 15:47:01 +01:00 committed by GitHub
parent dd442450b0
commit ac348eecdf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 17 additions and 5 deletions

View File

@ -39,6 +39,7 @@ enum BatchPrefetchStrategy {
/// have to fetch the metadata for a lot of versions. /// have to fetch the metadata for a lot of versions.
/// ///
/// Note that these all heuristics that could totally prefetch lots of irrelevant versions. /// Note that these all heuristics that could totally prefetch lots of irrelevant versions.
#[derive(Clone)]
pub(crate) struct BatchPrefetcher { pub(crate) struct BatchPrefetcher {
// Internal types. // Internal types.
tried_versions: FxHashMap<PackageName, usize>, tried_versions: FxHashMap<PackageName, usize>,

View File

@ -310,12 +310,17 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
let root = PubGrubPackage::from(PubGrubPackageInner::Root(self.project.clone())); let root = PubGrubPackage::from(PubGrubPackageInner::Root(self.project.clone()));
let pubgrub = State::init(root.clone(), MIN_VERSION.clone()); let pubgrub = State::init(root.clone(), MIN_VERSION.clone());
let mut prefetcher = BatchPrefetcher::new( let prefetcher = BatchPrefetcher::new(
self.capabilities.clone(), self.capabilities.clone(),
self.index.clone(), self.index.clone(),
request_sink.clone(), request_sink.clone(),
); );
let state = ForkState::new(pubgrub, self.env.clone(), self.python_requirement.clone()); let state = ForkState::new(
pubgrub,
self.env.clone(),
self.python_requirement.clone(),
prefetcher,
);
let mut preferences = self.preferences.clone(); let mut preferences = self.preferences.clone();
let mut forked_states = self.env.initial_forked_states(state); let mut forked_states = self.env.initial_forked_states(state);
let mut resolutions = vec![]; let mut resolutions = vec![];
@ -393,7 +398,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
else { else {
// All packages have been assigned, the fork has been successfully resolved // All packages have been assigned, the fork has been successfully resolved
if tracing::enabled!(Level::DEBUG) { if tracing::enabled!(Level::DEBUG) {
prefetcher.log_tried_versions(); state.prefetcher.log_tried_versions();
} }
debug!( debug!(
"{} resolution took {:.3}s", "{} resolution took {:.3}s",
@ -472,7 +477,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
// (idempotent due to caching). // (idempotent due to caching).
self.request_package(next_package, url, index, &request_sink)?; self.request_package(next_package, url, index, &request_sink)?;
prefetcher.version_tried(next_package); state.prefetcher.version_tried(next_package);
let term_intersection = state let term_intersection = state
.pubgrub .pubgrub
@ -545,7 +550,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
// Only consider registry packages for prefetch. // Only consider registry packages for prefetch.
if url.is_none() { if url.is_none() {
prefetcher.prefetch_batches( state.prefetcher.prefetch_batches(
next_package, next_package,
index, index,
&version, &version,
@ -2259,6 +2264,10 @@ pub(crate) struct ForkState {
/// solution that omits Python 3.8 support. /// solution that omits Python 3.8 support.
python_requirement: PythonRequirement, python_requirement: PythonRequirement,
conflict_tracker: ConflictTracker, conflict_tracker: ConflictTracker,
/// Prefetch package versions for packages with many rejected versions.
///
/// Tracked on the fork state to avoid counting each identical version between forks as new try.
prefetcher: BatchPrefetcher,
} }
impl ForkState { impl ForkState {
@ -2266,6 +2275,7 @@ impl ForkState {
pubgrub: State<UvDependencyProvider>, pubgrub: State<UvDependencyProvider>,
env: ResolverEnvironment, env: ResolverEnvironment,
python_requirement: PythonRequirement, python_requirement: PythonRequirement,
prefetcher: BatchPrefetcher,
) -> Self { ) -> Self {
Self { Self {
initial: None, initial: None,
@ -2279,6 +2289,7 @@ impl ForkState {
env, env,
python_requirement, python_requirement,
conflict_tracker: ConflictTracker::default(), conflict_tracker: ConflictTracker::default(),
prefetcher,
} }
} }