From 1c5cdcd70aa0550be6ac691ece6619a76d98960d Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sun, 29 Oct 2023 17:48:36 -0700 Subject: [PATCH] Prioritize packages in visited order (#222) --- crates/puffin-cli/tests/pip_compile.rs | 8 +-- crates/puffin-resolver/src/resolver.rs | 71 ++++++++++++++++++-------- scripts/benchmarks/requirements.in | 22 ++++++++ 3 files changed, 77 insertions(+), 24 deletions(-) diff --git a/crates/puffin-cli/tests/pip_compile.rs b/crates/puffin-cli/tests/pip_compile.rs index e7da7943e..eeb8a3933 100644 --- a/crates/puffin-cli/tests/pip_compile.rs +++ b/crates/puffin-cli/tests/pip_compile.rs @@ -70,7 +70,7 @@ fn compile_requirements_in() -> Result<()> { insta::with_settings!({ filters => vec![ - (r"\d+ms", "[TIME]"), + (r"\d+(ms|s)", "[TIME]"), (r"# .* pip-compile", "# [BIN_PATH] pip-compile"), (r"--cache-dir .*", "--cache-dir [CACHE_DIR]"), ] @@ -120,7 +120,7 @@ dependencies = [ insta::with_settings!({ filters => vec![ - (r"\d+ms", "[TIME]"), + (r"\d+(ms|s)", "[TIME]"), (r"# .* pip-compile", "# [BIN_PATH] pip-compile"), (r"--cache-dir .*", "--cache-dir [CACHE_DIR]"), ] @@ -164,7 +164,7 @@ fn compile_constraints_txt() -> Result<()> { insta::with_settings!({ filters => vec![ - (r"\d+ms", "[TIME]"), + (r"\d+(ms|s)", "[TIME]"), (r"# .* pip-compile", "# [BIN_PATH] pip-compile"), (r"--cache-dir .*", "--cache-dir [CACHE_DIR]"), ] @@ -211,7 +211,7 @@ fn compile_constraints_inline() -> Result<()> { insta::with_settings!({ filters => vec![ - (r"\d+ms", "[TIME]"), + (r"\d+(ms|s)", "[TIME]"), (r"# .* pip-compile", "# [BIN_PATH] pip-compile"), (r"--cache-dir .*", "--cache-dir [CACHE_DIR]"), ] diff --git a/crates/puffin-resolver/src/resolver.rs b/crates/puffin-resolver/src/resolver.rs index d9830fe92..89bc12c7d 100644 --- a/crates/puffin-resolver/src/resolver.rs +++ b/crates/puffin-resolver/src/resolver.rs @@ -52,6 +52,28 @@ pub struct Resolver<'a, Context: BuildContext + Sync> { build_context: &'a Context, } +#[derive(Debug, Default)] +struct Priorities(FxHashMap); + +impl Priorities { + fn add(&mut self, package: PackageName) { + let priority = self.0.len(); + self.0.entry(package).or_insert(priority); + } + + fn get(&self, package: &PubGrubPackage) -> Option { + match package { + PubGrubPackage::Root => Some(Reverse(0)), + PubGrubPackage::Package(name, _) => self + .0 + .get(name) + .copied() + .map(|priority| priority + 1) + .map(Reverse), + } + } +} + impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> { /// Initialize a new resolver. pub fn new( @@ -79,13 +101,6 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> { // metadata (e.g., given `flask==1.0.0`, fetch the metadata for that version). let (request_sink, request_stream) = futures::channel::mpsc::unbounded(); - // Push all the requirements into the package sink. - for requirement in &self.requirements { - debug!("Adding root dependency: {}", requirement); - let package_name = PackageName::normalize(&requirement.name); - request_sink.unbounded_send(Request::Package(package_name))?; - } - // Run the fetcher. let requests_fut = self.fetch(request_stream); @@ -120,6 +135,17 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> { let mut requested_packages = FxHashSet::default(); let mut requested_versions = FxHashSet::default(); let mut pins = FxHashMap::default(); + let mut priorities = Priorities::default(); + + // Push all the requirements into the package sink. + for requirement in &self.requirements { + debug!("Adding root dependency: {}", requirement); + let package_name = PackageName::normalize(&requirement.name); + if requested_packages.insert(package_name.clone()) { + priorities.add(package_name.clone()); + request_sink.unbounded_send(Request::Package(package_name))?; + } + } // Start the solve. let mut state = State::init(root.clone(), MIN_VERSION.clone()); @@ -139,9 +165,12 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> { )?; // Choose a package version. - let Some(highest_priority_pkg) = state - .partial_solution - .pick_highest_priority_pkg(|package, range| self.prioritize(package, range)) + let Some(highest_priority_pkg) = + state + .partial_solution + .pick_highest_priority_pkg(|package, _range| { + priorities.get(package).unwrap_or_default() + }) else { let selection = state.partial_solution.extract_solution(); return Ok(Graph::from_state(&selection, &pins, &state)); @@ -193,6 +222,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> { package, &version, &mut pins, + &mut priorities, &mut requested_packages, request_sink, ) @@ -236,16 +266,6 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> { } } - #[allow(clippy::unused_self)] - fn prioritize( - &self, - _package: &PubGrubPackage, - _range: &Range, - ) -> 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( @@ -369,6 +389,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> { package: &PubGrubPackage, version: &PubGrubVersion, pins: &mut FxHashMap>, + priorities: &mut Priorities, requested_packages: &mut FxHashSet, request_sink: &futures::channel::mpsc::UnboundedSender, ) -> Result { @@ -381,6 +402,15 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> { for (package, version) in iter_requirements(self.requirements.iter(), None, self.markers) { + // Emit a request to fetch the metadata for this package. + if let PubGrubPackage::Package(package_name, None) = &package { + if requested_packages.insert(package_name.clone()) { + priorities.add(package_name.clone()); + request_sink.unbounded_send(Request::Package(package_name.clone()))?; + } + } + + // Add it to the constraints. match constraints.entry(package) { Entry::Occupied(mut entry) => { entry.insert(entry.get().intersection(&version)); @@ -431,6 +461,7 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> { // Emit a request to fetch the metadata for this package. if let PubGrubPackage::Package(package_name, None) = &package { if requested_packages.insert(package_name.clone()) { + priorities.add(package_name.clone()); request_sink.unbounded_send(Request::Package(package_name.clone()))?; } } diff --git a/scripts/benchmarks/requirements.in b/scripts/benchmarks/requirements.in index 1a871e8d6..87f6e1d6f 100644 --- a/scripts/benchmarks/requirements.in +++ b/scripts/benchmarks/requirements.in @@ -6,3 +6,25 @@ 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