Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Avoid reusing cached downloaded binaries with --no-binary #7772

Merged
merged 1 commit into from
Sep 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 69 additions & 59 deletions crates/uv-distribution/src/index/registry_wheel_index.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,8 @@
use std::collections::hash_map::Entry;
use std::collections::BTreeMap;

use rustc_hash::FxHashMap;

use distribution_types::{CachedRegistryDist, Hashed, IndexLocations, IndexUrl};
use pep440_rs::Version;
use platform_tags::Tags;
use uv_cache::{Cache, CacheBucket, WheelCache};
use uv_fs::{directories, files, symlinks};
Expand All @@ -14,14 +12,23 @@ use uv_types::HashStrategy;
use crate::index::cached_wheel::CachedWheel;
use crate::source::{HttpRevisionPointer, LocalRevisionPointer, HTTP_REVISION, LOCAL_REVISION};

/// An entry in the [`RegistryWheelIndex`].
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct IndexEntry {
/// The cached distribution.
pub dist: CachedRegistryDist,
/// Whether the wheel was built from source (true), or downloaded from the registry directly (false).
pub built: bool,
}

/// A local index of distributions that originate from a registry, like `PyPI`.
#[derive(Debug)]
pub struct RegistryWheelIndex<'a> {
cache: &'a Cache,
tags: &'a Tags,
index_locations: &'a IndexLocations,
hasher: &'a HashStrategy,
index: FxHashMap<&'a PackageName, BTreeMap<Version, CachedRegistryDist>>,
index: FxHashMap<&'a PackageName, Vec<IndexEntry>>,
}

impl<'a> RegistryWheelIndex<'a> {
Expand All @@ -44,26 +51,12 @@ impl<'a> RegistryWheelIndex<'a> {
/// Return an iterator over available wheels for a given package.
///
/// If the package is not yet indexed, this will index the package by reading from the cache.
pub fn get(
&mut self,
name: &'a PackageName,
) -> impl Iterator<Item = (&Version, &CachedRegistryDist)> {
pub fn get(&mut self, name: &'a PackageName) -> impl Iterator<Item = &IndexEntry> {
self.get_impl(name).iter().rev()
}

/// Get the best wheel for the given package name and version.
///
/// If the package is not yet indexed, this will index the package by reading from the cache.
pub fn get_version(
&mut self,
name: &'a PackageName,
version: &Version,
) -> Option<&CachedRegistryDist> {
self.get_impl(name).get(version)
}

/// Get an entry in the index.
fn get_impl(&mut self, name: &'a PackageName) -> &BTreeMap<Version, CachedRegistryDist> {
fn get_impl(&mut self, name: &'a PackageName) -> &[IndexEntry] {
let versions = match self.index.entry(name) {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => entry.insert(Self::index(
Expand All @@ -84,8 +77,8 @@ impl<'a> RegistryWheelIndex<'a> {
tags: &Tags,
index_locations: &IndexLocations,
hasher: &HashStrategy,
) -> BTreeMap<Version, CachedRegistryDist> {
let mut versions = BTreeMap::new();
) -> Vec<IndexEntry> {
let mut entries = vec![];

// Collect into owned `IndexUrl`.
let flat_index_urls: Vec<IndexUrl> = index_locations
Expand Down Expand Up @@ -113,12 +106,19 @@ impl<'a> RegistryWheelIndex<'a> {
if let Some(wheel) =
CachedWheel::from_http_pointer(wheel_dir.join(file), cache)
{
// Enforce hash-checking based on the built distribution.
if wheel.satisfies(
hasher
.get_package(&wheel.filename.name, &wheel.filename.version),
) {
Self::add_wheel(wheel, tags, &mut versions);
if wheel.filename.compatibility(tags).is_compatible() {
// Enforce hash-checking based on the built distribution.
if wheel.satisfies(
hasher.get_package(
&wheel.filename.name,
&wheel.filename.version,
),
) {
entries.push(IndexEntry {
dist: wheel.into_registry_dist(),
built: false,
});
}
}
}
}
Expand All @@ -132,12 +132,19 @@ impl<'a> RegistryWheelIndex<'a> {
if let Some(wheel) =
CachedWheel::from_local_pointer(wheel_dir.join(file), cache)
{
// Enforce hash-checking based on the built distribution.
if wheel.satisfies(
hasher
.get_package(&wheel.filename.name, &wheel.filename.version),
) {
Self::add_wheel(wheel, tags, &mut versions);
if wheel.filename.compatibility(tags).is_compatible() {
// Enforce hash-checking based on the built distribution.
if wheel.satisfies(
hasher.get_package(
&wheel.filename.name,
&wheel.filename.version,
),
) {
entries.push(IndexEntry {
dist: wheel.into_registry_dist(),
built: false,
});
}
}
}
}
Expand Down Expand Up @@ -182,38 +189,41 @@ impl<'a> RegistryWheelIndex<'a> {
if let Some(revision) = revision {
for wheel_dir in symlinks(cache_shard.join(revision.id())) {
if let Some(wheel) = CachedWheel::from_built_source(wheel_dir) {
// Enforce hash-checking based on the source distribution.
if revision.satisfies(
hasher.get_package(&wheel.filename.name, &wheel.filename.version),
) {
Self::add_wheel(wheel, tags, &mut versions);
if wheel.filename.compatibility(tags).is_compatible() {
// Enforce hash-checking based on the source distribution.
if revision.satisfies(
hasher
.get_package(&wheel.filename.name, &wheel.filename.version),
) {
entries.push(IndexEntry {
dist: wheel.into_registry_dist(),
built: true,
});
}
}
}
}
}
}
}

versions
}

/// Add the [`CachedWheel`] to the index.
fn add_wheel(
wheel: CachedWheel,
tags: &Tags,
versions: &mut BTreeMap<Version, CachedRegistryDist>,
) {
let dist_info = wheel.into_registry_dist();

// Pick the wheel with the highest priority
let compatibility = dist_info.filename.compatibility(tags);
if let Some(existing) = versions.get_mut(&dist_info.filename.version) {
// Override if we have better compatibility
if compatibility > existing.filename.compatibility(tags) {
*existing = dist_info;
}
} else if compatibility.is_compatible() {
versions.insert(dist_info.filename.version.clone(), dist_info);
}
// Sort the cached distributions by (1) version, (2) compatibility, and (3) build status.
// We want the highest versions, with the greatest compatibility, that were built from source.
// at the end of the list.
entries.sort_unstable_by(|a, b| {
a.dist
.filename
.version
.cmp(&b.dist.filename.version)
.then_with(|| {
a.dist
.filename
.compatibility(tags)
.cmp(&b.dist.filename.compatibility(tags))
.then_with(|| a.built.cmp(&b.built))
})
});

entries
}
}
17 changes: 14 additions & 3 deletions crates/uv-installer/src/plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,7 @@ impl<'a> Planner<'a> {

// Check if installation of a binary version of the package should be allowed.
let no_binary = build_options.no_binary_package(&requirement.name);
let no_build = build_options.no_build_package(&requirement.name);

let installed_dists = site_packages.remove_packages(&requirement.name);

Expand Down Expand Up @@ -143,9 +144,19 @@ impl<'a> Planner<'a> {
// Identify any cached distributions that satisfy the requirement.
match &requirement.source {
RequirementSource::Registry { specifier, .. } => {
if let Some((_version, distribution)) = registry_index
.get(&requirement.name)
.find(|(version, _)| specifier.contains(version))
if let Some(distribution) =
registry_index.get(&requirement.name).find_map(|entry| {
if !specifier.contains(&entry.dist.filename.version) {
return None;
};
if entry.built && no_build {
return None;
}
if !entry.built && no_binary {
return None;
}
Some(&entry.dist)
})
{
debug!("Requirement already cached: {distribution}");
cached.push(CachedDist::Registry(distribution.clone()));
Expand Down
2 changes: 1 addition & 1 deletion crates/uv/tests/lock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7473,7 +7473,7 @@ fn lock_find_links_http_sdist() -> Result<()> {
----- stdout -----

----- stderr -----
Prepared 1 package in [TIME]
Prepared 2 packages in [TIME]
Installed 2 packages in [TIME]
+ packaging==23.2
+ project==0.1.0 (from file://[TEMP_DIR]/)
Expand Down
117 changes: 117 additions & 0 deletions crates/uv/tests/pip_install.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2010,6 +2010,66 @@ fn install_only_binary_all_and_no_binary_all() {
context.assert_command("import anyio").failure();
}

/// Binary dependencies in the cache should be reused when the user provides `--no-build`.
#[test]
fn install_no_binary_cache() {
let context = TestContext::new("3.12");

// Install a binary distribution.
uv_snapshot!(
context.pip_install().arg("idna"),
@r###"
success: true
exit_code: 0
----- stdout -----

----- stderr -----
Resolved 1 package in [TIME]
Prepared 1 package in [TIME]
Installed 1 package in [TIME]
+ idna==3.6
"###
);

// Re-create the virtual environment.
context.venv().assert().success();

// Re-install. The distribution should be installed from the cache.
uv_snapshot!(
context.pip_install().arg("idna"),
@r###"
success: true
exit_code: 0
----- stdout -----

----- stderr -----
Resolved 1 package in [TIME]
Installed 1 package in [TIME]
+ idna==3.6
"###
);

// Re-create the virtual environment.
context.venv().assert().success();

// Install with `--no-binary`. The distribution should be built from source, despite a binary
// distribution being available in the cache.
uv_snapshot!(
context.pip_install().arg("idna").arg("--no-binary").arg(":all:"),
@r###"
success: true
exit_code: 0
----- stdout -----

----- stderr -----
Resolved 1 package in [TIME]
Prepared 1 package in [TIME]
Installed 1 package in [TIME]
+ idna==3.6
"###
);
}

/// Respect `--only-binary` flags in `requirements.txt`
#[test]
fn only_binary_requirements_txt() {
Expand Down Expand Up @@ -2124,6 +2184,63 @@ fn only_binary_editable_setup_py() {
);
}

#[test]
fn cache_priority() {
let context = TestContext::new("3.12");

// Install a specific `idna` version.
uv_snapshot!(
context.pip_install().arg("idna==3.6"),
@r###"
success: true
exit_code: 0
----- stdout -----

----- stderr -----
Resolved 1 package in [TIME]
Prepared 1 package in [TIME]
Installed 1 package in [TIME]
+ idna==3.6
"###
);

// Install a lower `idna` version.
uv_snapshot!(
context.pip_install().arg("idna==3.0"),
@r###"
success: true
exit_code: 0
----- stdout -----

----- stderr -----
Resolved 1 package in [TIME]
Prepared 1 package in [TIME]
Uninstalled 1 package in [TIME]
Installed 1 package in [TIME]
- idna==3.6
+ idna==3.0
"###
);

// Re-create the virtual environment.
context.venv().assert().success();

// Install `idna` without a version specifier.
uv_snapshot!(
context.pip_install().arg("idna"),
@r###"
success: true
exit_code: 0
----- stdout -----

----- stderr -----
Resolved 1 package in [TIME]
Installed 1 package in [TIME]
+ idna==3.6
"###
);
}

/// Install a package into a virtual environment, and ensuring that the executable permissions
/// are retained.
///
Expand Down
Loading