distribution-filename: speed up is_compatible (#367)

This PR tweaks the representation of `Tags` in order to offer a
faster implementation of `WheelFilename::is_compatible`. We now use a
nested map of tags that lets us avoid looping over every supported
platform tag. As the code comments suggest, that is the essential gain.
We still do not mind looping over the tags in each wheel name since they
tend to be quite small. And pushing our thumb on that side of things can
make things worse overall since it would likely slow down WheelFilename
construction itself.

For micro-benchmarks, we improve considerably for compatibility
checking:

    $ critcmp base test3
group base test3
----- ---- -----
build_platform_tags/burntsushi-archlinux 1.00 46.2±0.28µs ? ?/sec 2.48
114.8±0.45µs ? ?/sec
wheelname_parsing/flyte-long-compatible 1.00 624.8±3.31ns 174.0 MB/sec
1.01 629.4±4.30ns 172.7 MB/sec
wheelname_parsing/flyte-long-incompatible 1.00 743.6±4.23ns 165.4 MB/sec
1.00 746.9±4.62ns 164.7 MB/sec
wheelname_parsing/flyte-short-compatible 1.00 526.7±4.76ns 54.3 MB/sec
1.01 530.2±5.81ns 54.0 MB/sec
wheelname_parsing/flyte-short-incompatible 1.00 540.4±4.93ns 60.0 MB/sec
1.01 545.7±5.31ns 59.4 MB/sec
wheelname_parsing_failure/flyte-long-extension 1.00 13.6±0.13ns 3.2
GB/sec 1.01 13.7±0.14ns 3.2 GB/sec
wheelname_parsing_failure/flyte-short-extension 1.00 14.0±0.20ns 1160.4
MB/sec 1.01 14.1±0.14ns 1146.5 MB/sec
wheelname_tag_compatibility/flyte-long-compatible 11.33 159.8±2.79ns
680.5 MB/sec 1.00 14.1±0.23ns 7.5 GB/sec
wheelname_tag_compatibility/flyte-long-incompatible 237.60
1671.8±37.99ns 73.6 MB/sec 1.00 7.0±0.08ns 17.1 GB/sec
wheelname_tag_compatibility/flyte-short-compatible 16.07 223.5±8.60ns
128.0 MB/sec 1.00 13.9±0.30ns 2.0 GB/sec
wheelname_tag_compatibility/flyte-short-incompatible 149.83 628.3±2.13ns
51.6 MB/sec 1.00 4.2±0.10ns 7.6 GB/sec

We do regress slightly on the time it takes for `Tags::new` to run, but
this is somewhat expected. And in absolute terms, 114us is perfectly
acceptable given that it's only executed ~once for each `puffin`
invocation.

Ad hoc benchmarks indicate an overall 25% perf improvement in `puffin
pip-compile` times. This roughly corresponds with how much time
`is_compatible` was taking. Indeed, profiling confirms that it has
virtually disappeared from the profile.

Fixes #157
This commit is contained in:
Andrew Gallant 2023-11-09 09:01:03 -05:00 committed by GitHub
parent bdb89b4072
commit 33c0901a28
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 571 additions and 28 deletions

View file

@ -0,0 +1,32 @@
---
source: crates/distribution-filename/src/wheel.rs
expression: "WheelFilename::from_str(\"foo-1.2.3-build-python-abi-platform.whl\")"
---
Ok(
WheelFilename {
distribution: PackageName(
"foo",
),
version: Version {
epoch: 0,
release: [
1,
2,
3,
],
pre: None,
post: None,
dev: None,
local: None,
},
python_tag: [
"python",
],
abi_tag: [
"abi",
],
platform_tag: [
"platform",
],
},
)

View file

@ -0,0 +1,39 @@
---
source: crates/distribution-filename/src/wheel.rs
expression: "WheelFilename::from_str(\"foo-1.2.3-ab.cd.ef-gh-ij.kl.mn.op.qr.st.whl\")"
---
Ok(
WheelFilename {
distribution: PackageName(
"foo",
),
version: Version {
epoch: 0,
release: [
1,
2,
3,
],
pre: None,
post: None,
dev: None,
local: None,
},
python_tag: [
"ab",
"cd",
"ef",
],
abi_tag: [
"gh",
],
platform_tag: [
"ij",
"kl",
"mn",
"op",
"qr",
"st",
],
},
)

View file

@ -0,0 +1,32 @@
---
source: crates/distribution-filename/src/wheel.rs
expression: "WheelFilename::from_str(\"foo-1.2.3-foo-bar-baz.whl\")"
---
Ok(
WheelFilename {
distribution: PackageName(
"foo",
),
version: Version {
epoch: 0,
release: [
1,
2,
3,
],
pre: None,
post: None,
dev: None,
local: None,
},
python_tag: [
"foo",
],
abi_tag: [
"bar",
],
platform_tag: [
"baz",
],
},
)

View file

@ -31,14 +31,17 @@ impl FromStr for WheelFilename {
// The wheel filename should contain either five or six entries. If six, then the third
// entry is the build tag. If five, then the third entry is the Python tag.
// https://www.python.org/dev/peps/pep-0427/#file-name-convention
//
// 2023-11-08(burntsushi): It looks like the code below actually drops
// the build tag if one is found. According to PEP 0427, the build tag
// is used to break ties. This might mean that we generate identical
// `WheelName` values for multiple distinct wheels, but it's not clear
// if this is a problem in practice.
let mut parts = basename.split('-');
let Some(distribution) = parts.next() else {
return Err(WheelFilenameError::InvalidWheelFileName(
filename.to_string(),
"Must have a distribution name".to_string(),
));
};
let distribution = parts
.next()
.expect("split always yields 1 or more elements");
let Some(version) = parts.next() else {
return Err(WheelFilenameError::InvalidWheelFileName(
@ -70,6 +73,12 @@ impl FromStr for WheelFilename {
let (distribution, version, python_tag, abi_tag, platform_tag) =
if let Some(platform_tag) = parts.next() {
if parts.next().is_some() {
return Err(WheelFilenameError::InvalidWheelFileName(
filename.to_string(),
"Must have 5 or 6 components, but has more".to_string(),
));
}
(
distribution,
version,
@ -116,15 +125,7 @@ impl Display for WheelFilename {
impl WheelFilename {
/// Returns `true` if the wheel is compatible with the given tags.
pub fn is_compatible(&self, compatible_tags: &Tags) -> bool {
for tag in compatible_tags.iter() {
if self.python_tag.contains(&tag.0)
&& self.abi_tag.contains(&tag.1)
&& self.platform_tag.contains(&tag.2)
{
return true;
}
}
false
compatible_tags.is_compatible(&self.python_tag, &self.abi_tag, &self.platform_tag)
}
/// Get the tag for this wheel.
@ -170,3 +171,82 @@ pub enum WheelFilenameError {
#[error("The wheel filename \"{0}\" has an invalid package name")]
InvalidPackageName(String, InvalidNameError),
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn err_not_whl_extension() {
let err = WheelFilename::from_str("foo.rs").unwrap_err();
insta::assert_display_snapshot!(err, @r###"The wheel filename "foo.rs" is invalid: Must end with .whl"###);
}
#[test]
fn err_1_part_empty() {
let err = WheelFilename::from_str(".whl").unwrap_err();
insta::assert_display_snapshot!(err, @r###"The wheel filename ".whl" is invalid: Must have a version"###);
}
#[test]
fn err_1_part_no_version() {
let err = WheelFilename::from_str("foo.whl").unwrap_err();
insta::assert_display_snapshot!(err, @r###"The wheel filename "foo.whl" is invalid: Must have a version"###);
}
#[test]
fn err_2_part_no_pythontag() {
let err = WheelFilename::from_str("foo-version.whl").unwrap_err();
insta::assert_display_snapshot!(err, @r###"The wheel filename "foo-version.whl" is invalid: Must have a Python tag"###);
}
#[test]
fn err_3_part_no_abitag() {
let err = WheelFilename::from_str("foo-version-python.whl").unwrap_err();
insta::assert_display_snapshot!(err, @r###"The wheel filename "foo-version-python.whl" is invalid: Must have an ABI tag"###);
}
#[test]
fn err_4_part_no_platformtag() {
let err = WheelFilename::from_str("foo-version-python-abi.whl").unwrap_err();
insta::assert_display_snapshot!(err, @r###"The wheel filename "foo-version-python-abi.whl" is invalid: Must have a platform tag"###);
}
#[test]
fn err_too_many_parts() {
let err =
WheelFilename::from_str("foo-1.2.3-build-python-abi-platform-oops.whl").unwrap_err();
insta::assert_display_snapshot!(err, @r###"The wheel filename "foo-1.2.3-build-python-abi-platform-oops.whl" is invalid: Must have 5 or 6 components, but has more"###);
}
#[test]
fn err_invalid_package_name() {
let err = WheelFilename::from_str("f!oo-1.2.3-python-abi-platform.whl").unwrap_err();
insta::assert_display_snapshot!(err, @r###"The wheel filename "f!oo-1.2.3-python-abi-platform.whl" has an invalid package name"###);
}
#[test]
fn err_invalid_version() {
let err = WheelFilename::from_str("foo-x.y.z-python-abi-platform.whl").unwrap_err();
insta::assert_display_snapshot!(err, @r###"The wheel filename "foo-x.y.z-python-abi-platform.whl" has an invalid version part: Version `x.y.z` doesn't match PEP 440 rules"###);
}
#[test]
fn ok_single_tags() {
insta::assert_debug_snapshot!(WheelFilename::from_str("foo-1.2.3-foo-bar-baz.whl"));
}
#[test]
fn ok_multiple_tags() {
insta::assert_debug_snapshot!(WheelFilename::from_str(
"foo-1.2.3-ab.cd.ef-gh-ij.kl.mn.op.qr.st.whl"
));
}
#[test]
fn ok_build_tag() {
insta::assert_debug_snapshot!(WheelFilename::from_str(
"foo-1.2.3-build-python-abi-platform.whl"
));
}
}