Skip to content

Commit

Permalink
Resolver module improvements (#9773)
Browse files Browse the repository at this point in the history 
Further small refactorings for the resolver.
  • Loading branch information
@konstin
konstin authored Dec 11, 2024
1 parent 441ed3b commit 509dc83
Show file tree
Hide file tree
Showing 2 changed files with 84 additions and 64 deletions.
33 changes: 24 additions & 9 deletions crates/uv-resolver/src/resolver/batch_prefetch.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,13 +39,31 @@ enum BatchPrefetchStrategy {
/// have to fetch the metadata for a lot of versions.
///
/// Note that these all heuristics that could totally prefetch lots of irrelevant versions.
#[derive(Default)]
pub(crate) struct BatchPrefetcher {
// Internal types.
tried_versions: FxHashMap<PackageName, usize>,
last_prefetch: FxHashMap<PackageName, usize>,
// Shared (e.g., `Arc`) types.
capabilities: IndexCapabilities,
index: InMemoryIndex,
request_sink: Sender<Request>,
}

impl BatchPrefetcher {
pub(crate) fn new(
capabilities: IndexCapabilities,
index: InMemoryIndex,
request_sink: Sender<Request>,
) -> Self {
Self {
tried_versions: FxHashMap::default(),
last_prefetch: FxHashMap::default(),
capabilities,
index,
request_sink,
}
}

/// Prefetch a large number of versions if we already unsuccessfully tried many versions.
pub(crate) fn prefetch_batches(
&mut self,
Expand All @@ -55,9 +73,6 @@ impl BatchPrefetcher {
current_range: &Range<Version>,
unchangeable_constraints: Option<&Term<Range<Version>>>,
python_requirement: &PythonRequirement,
request_sink: &Sender<Request>,
in_memory: &InMemoryIndex,
capabilities: &IndexCapabilities,
selector: &CandidateSelector,
env: &ResolverEnvironment,
) -> anyhow::Result<(), ResolveError> {
Expand All @@ -79,12 +94,12 @@ impl BatchPrefetcher {

// This is immediate, we already fetched the version map.
let versions_response = if let Some(index) = index {
in_memory
self.index
.explicit()
.wait_blocking(&(name.clone(), index.clone()))
.ok_or_else(|| ResolveError::UnregisteredTask(name.to_string()))?
} else {
in_memory
self.index
.implicit()
.wait_blocking(name)
.ok_or_else(|| ResolveError::UnregisteredTask(name.to_string()))?
Expand Down Expand Up @@ -166,7 +181,7 @@ impl BatchPrefetcher {
// Avoid prefetching built distributions that don't support _either_ PEP 658 (`.metadata`)
// or range requests.
if !(wheel.file.dist_info_metadata
|| capabilities.supports_range_requests(&wheel.index))
|| self.capabilities.supports_range_requests(&wheel.index))
{
debug!("Abandoning prefetch for {wheel} due to missing registry capabilities");
return Ok(());
Expand Down Expand Up @@ -214,9 +229,9 @@ impl BatchPrefetcher {
);
prefetch_count += 1;

if in_memory.distributions().register(candidate.version_id()) {
if self.index.distributions().register(candidate.version_id()) {
let request = Request::from(dist);
request_sink.blocking_send(request)?;
self.request_sink.blocking_send(request)?;
}
}

Expand Down
115 changes: 60 additions & 55 deletions crates/uv-resolver/src/resolver/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -306,7 +306,11 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag

let root = PubGrubPackage::from(PubGrubPackageInner::Root(self.project.clone()));
let pubgrub = State::init(root.clone(), MIN_VERSION.clone());
let mut prefetcher = BatchPrefetcher::default();
let mut prefetcher = BatchPrefetcher::new(
self.capabilities.clone(),
self.index.clone(),
request_sink.clone(),
);
let state = ForkState::new(pubgrub, self.env.clone(), self.python_requirement.clone());
let mut preferences = self.preferences.clone();
let mut forked_states = self.env.initial_forked_states(state);
Expand Down Expand Up @@ -497,9 +501,6 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
.partial_solution
.unchanging_term_for_package(next_id),
&state.python_requirement,
&request_sink,
&self.index,
&self.capabilities,
&self.selector,
&state.env,
)?;
Expand Down Expand Up @@ -1078,57 +1079,8 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
}
};

// Validate the Python requirement.
let requires_python = match dist {
CompatibleDist::InstalledDist(_) => None,
CompatibleDist::SourceDist { sdist, .. }
| CompatibleDist::IncompatibleWheel { sdist, .. } => {
sdist.file.requires_python.as_ref()
}
CompatibleDist::CompatibleWheel { wheel, .. } => wheel.file.requires_python.as_ref(),
};
let incompatibility = requires_python.and_then(|requires_python| {
if python_requirement.installed() == python_requirement.target() {
if !python_requirement
.installed()
.is_contained_by(requires_python)
{
return if matches!(dist, CompatibleDist::CompatibleWheel { .. }) {
Some(IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython(
requires_python.clone(),
PythonRequirementKind::Installed,
)))
} else {
Some(IncompatibleDist::Source(
IncompatibleSource::RequiresPython(
requires_python.clone(),
PythonRequirementKind::Installed,
),
))
};
}
} else {
if !python_requirement.target().is_contained_by(requires_python) {
return if matches!(dist, CompatibleDist::CompatibleWheel { .. }) {
Some(IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython(
requires_python.clone(),
PythonRequirementKind::Target,
)))
} else {
Some(IncompatibleDist::Source(
IncompatibleSource::RequiresPython(
requires_python.clone(),
PythonRequirementKind::Target,
),
))
};
}
}
None
});

// The version is incompatible due to its Python requirement.
if let Some(incompatibility) = incompatibility {
// Check whether the version is incompatible due to its Python requirement.
if let Some(incompatibility) = Self::check_requires_python(dist, python_requirement) {
return Ok(Some(ResolverVersion::Unavailable(
candidate.version().clone(),
UnavailableVersion::IncompatibleDist(incompatibility),
Expand Down Expand Up @@ -1180,6 +1132,59 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
Ok(Some(ResolverVersion::Available(version)))
}

/// Check if the distribution is incompatible with the Python requirement, and if so, return
/// the incompatibility.
fn check_requires_python(
dist: &CompatibleDist,
python_requirement: &PythonRequirement,
) -> Option<IncompatibleDist> {
let requires_python = match dist {
CompatibleDist::InstalledDist(_) => None,
CompatibleDist::SourceDist { sdist, .. }
| CompatibleDist::IncompatibleWheel { sdist, .. } => {
sdist.file.requires_python.as_ref()
}
CompatibleDist::CompatibleWheel { wheel, .. } => wheel.file.requires_python.as_ref(),
}?;
if python_requirement.installed() == python_requirement.target() {
if !python_requirement
.installed()
.is_contained_by(requires_python)
{
return if matches!(dist, CompatibleDist::CompatibleWheel { .. }) {
Some(IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython(
requires_python.clone(),
PythonRequirementKind::Installed,
)))
} else {
Some(IncompatibleDist::Source(
IncompatibleSource::RequiresPython(
requires_python.clone(),
PythonRequirementKind::Installed,
),
))
};
}
} else {
if !python_requirement.target().is_contained_by(requires_python) {
return if matches!(dist, CompatibleDist::CompatibleWheel { .. }) {
Some(IncompatibleDist::Wheel(IncompatibleWheel::RequiresPython(
requires_python.clone(),
PythonRequirementKind::Target,
)))
} else {
Some(IncompatibleDist::Source(
IncompatibleSource::RequiresPython(
requires_python.clone(),
PythonRequirementKind::Target,
),
))
};
}
}
None
}

/// Given a candidate package and version, return its dependencies.
#[instrument(skip_all, fields(%package, %version))]
fn get_dependencies_forking(
Expand Down

0 comments on commit 509dc83

Please sign in to comment.