Update

crates/puffin-resolver/src/pubgrub/mod.rs

@@ -10,6 +10,7 @@ use crate::pubgrub::specifier::PubGrubSpecifier;
 use crate::pubgrub::version::{PubGrubVersion, MAX_VERSION};
 
 pub(crate) mod package;
+pub(crate) mod priority;
 mod specifier;
 pub(crate) mod version;
 

crates/puffin-resolver/src/pubgrub/priority.rs

@@ -0,0 +1,3 @@
+use std::cmp::Reverse;
+
+pub(crate) type PubGrubPriority = Reverse<usize>;

crates/puffin-resolver/src/resolution.rs

@@ -1,9 +1,8 @@
+use std::hash::BuildHasherDefault;
+
 use colored::Colorize;
 use fxhash::FxHashMap;
 use petgraph::visit::EdgeRef;
-use std::cmp::Reverse;
-use std::hash::BuildHasherDefault;
-
 use pubgrub::range::Range;
 use pubgrub::solver::{Kind, State};
 use pubgrub::type_aliases::SelectedDependencies;
@@ -14,6 +13,7 @@ use puffin_client::File;
 use puffin_package::package_name::PackageName;
 
 use crate::pubgrub::package::PubGrubPackage;
+use crate::pubgrub::priority::PubGrubPriority;
 use crate::pubgrub::version::PubGrubVersion;
 
 /// A package pinned at a specific version.
@@ -96,7 +96,7 @@ impl Graph {
     pub fn from_state(
         selection: &SelectedDependencies<PubGrubPackage, PubGrubVersion>,
         pins: &FxHashMap<PackageName, FxHashMap<Version, File>>,
-        state: &State<PubGrubPackage, Range<PubGrubVersion>, Reverse<usize>>,
+        state: &State<PubGrubPackage, Range<PubGrubVersion>, PubGrubPriority>,
     ) -> Self {
         // TODO(charlie): petgraph is a really heavy and unnecessary dependency here. We should
         // write our own graph, given that our requirements are so simple.

crates/puffin-resolver/src/resolver.rs

@@ -34,6 +34,7 @@ use crate::distribution::{DistributionFile, SdistFile, WheelFile};
 use crate::error::ResolveError;
 use crate::manifest::Manifest;
 use crate::pubgrub::package::PubGrubPackage;
+use crate::pubgrub::priority::PubGrubPriority;
 use crate::pubgrub::version::{PubGrubVersion, MIN_VERSION};
 use crate::pubgrub::{iter_requirements, version_range};
 use crate::resolution::Graph;
@@ -130,10 +131,17 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
             // Run unit propagation.
             state.unit_propagation(next)?;
 
+            // Pre-visit all candidate packages, to allow metadata to be fetched in parallel.
+            self.pre_visit(
+                state.partial_solution.prioritized_packages(),
+                &mut requested_versions,
+                request_sink,
+            )?;
+
             // Choose a package version.
             let Some(highest_priority_pkg) = state
                 .partial_solution
-                .pick_highest_priority_pkg(|p, r| self.prioritize(p, r))
+                .pick_highest_priority_pkg(|package, range| self.prioritize(package, range))
             else {
                 let selection = state.partial_solution.extract_solution();
                 return Ok(Graph::from_state(&selection, &pins, &state));
@@ -228,11 +236,65 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
         }
     }
 
-    fn prioritize(&self, _package: &PubGrubPackage, _range: &Range<PubGrubVersion>) -> Priority {
+    #[allow(clippy::unused_self)]
+    fn prioritize(
+        &self,
+        _package: &PubGrubPackage,
+        _range: &Range<PubGrubVersion>,
+    ) -> PubGrubPriority {
         // TODO(charlie): Define a priority function.
         Reverse(0)
     }
 
+    /// Visit the set of candidate packages prior to selection. This allows us to fetch metadata for
+    /// all of the packages in parallel.
+    fn pre_visit(
+        &self,
+        packages: impl Iterator<Item = (&'a PubGrubPackage, &'a Range<PubGrubVersion>)>,
+        in_flight: &mut FxHashSet<String>,
+        request_sink: &futures::channel::mpsc::UnboundedSender<Request>,
+    ) -> Result<(), ResolveError> {
+        // Iterate over the potential packages, and fetch file metadata for any of them. These
+        // represent our current best guesses for the versions that we _might_ select.
+        for (package, range) in packages {
+            let PubGrubPackage::Package(package_name, _) = package else {
+                continue;
+            };
+
+            // If we don't have metadata for this package, we can't make an early decision.
+            let Some(entry) = self.index.packages.get(package_name) else {
+                continue;
+            };
+            let version_map = entry.value();
+
+            // Try to find a compatible version. If there aren't any compatible versions,
+            // short-circuit and return `None`.
+            let Some(candidate) = self.selector.select(package_name, range, version_map) else {
+                // Short-circuit: we couldn't find _any_ compatible versions for a package.
+                return Ok(());
+            };
+
+            // Emit a request to fetch the metadata for this version.
+            match candidate.file {
+                DistributionFile::Wheel(file) => {
+                    if in_flight.insert(file.hashes.sha256.clone()) {
+                        request_sink.unbounded_send(Request::Wheel(file.clone()))?;
+                    }
+                }
+                DistributionFile::Sdist(file) => {
+                    if in_flight.insert(file.hashes.sha256.clone()) {
+                        request_sink.unbounded_send(Request::Sdist(
+                            file.clone(),
+                            candidate.package_name.clone(),
+                            candidate.version.clone().into(),
+                        ))?;
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+
     /// Given a set of candidate packages, choose the next package (and version) to add to the
     /// partial solution.
     async fn choose_version(
@@ -535,8 +597,6 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
     }
 }
 
-type Priority = Reverse<usize>;
-
 /// Fetch the metadata for an item
 #[derive(Debug)]
 enum Request {

scripts/benchmarks/requirements.in

@@ -6,25 +6,3 @@ pygls>=1.0.1
 lsprotocol>=2023.0.0a1
 ruff>=0.0.274
 typing_extensions
-scipy
-numpy
-pandas<2.0.0
-matplotlib>=3.0.0
-scikit-learn
-rich
-textual
-jupyter>=1.0.0,<2.0.0
-transformers[torch]
-django<4.0.0
-sqlalchemy
-psycopg2-binary
-trio<0.20
-trio-websocket
-trio-asyncio
-trio-typing
-trio-protocol
-fastapi
-typer
-pydantic
-uvicorn
-traitlets

vendor/pubgrub/src/internal/partial_solution.rs

@@ -246,6 +246,24 @@ impl<P: Package, VS: VersionSet, Priority: Ord + Clone> PartialSolution<P, VS, P
         }
     }
 
+    pub fn prioritized_packages(&self) -> impl Iterator<Item = (&P, &VS)> {
+        let check_all = self.changed_this_decision_level
+            == self.current_decision_level.0.saturating_sub(1) as usize;
+        let current_decision_level = self.current_decision_level;
+        self.package_assignments
+            .get_range(self.changed_this_decision_level..)
+            .unwrap()
+            .iter()
+            .filter(move |(_, pa)| {
+                // We only actually need to update the package if its Been changed
+                // since the last time we called prioritize.
+                // Which means it's highest decision level is the current decision level,
+                // or if we backtracked in the mean time.
+                check_all || pa.highest_decision_level == current_decision_level
+            })
+            .filter_map(|(p, pa)| pa.assignments_intersection.potential_package_filter(p))
+    }
+
     pub fn pick_highest_priority_pkg(
         &mut self,
         prioritizer: impl Fn(&P, &VS) -> Priority,