From 64456844e9bdccd05d2a431c431c4d1b80235a60 Mon Sep 17 00:00:00 2001 From: Justin Larkin Date: Fri, 24 Jul 2026 17:49:41 -0400 Subject: [PATCH] feat(wheels): add configurable build tag hook for wheel filenames Implement the accepted proposal from docs/proposals/wheel-build-tag-hook.md (issue #1059, tracking issue #1181). Add a `wheels.build_tag_hook` option in global settings that lets downstream projects append environment-specific suffixes (OS, accelerator, torch ABI) to wheel build tags via a user-defined callable. The hook receives ctx, req, version, and wheel_tags and returns suffix segments joined with `_`. - Add `WheelSettings` model with `build_tag_hook: ImportString` to settings - Add `get_build_tag()` and `_validate_build_tag_segments()` to wheels.py - Update `add_extra_metadata_to_wheels()`, bootstrapper cache checks, and `_is_wheel_built()` to use computed build tags - Minimal finder update to match suffixed build tag filenames - Validate hook output: reject single strings, invalid chars, non-strings - No behavior change when hook is not configured Closes: #1181 Co-Authored-By: Claude Opus 4.6 Signed-off-by: Justin Larkin Co-authored-by: Cursor Signed-off-by: Justin Larkin Co-authored-by: Cursor --- src/fromager/bootstrapper/_cache.py | 36 +++-- src/fromager/commands/build.py | 9 +- src/fromager/finders.py | 43 +++--- src/fromager/packagesettings/__init__.py | 2 + src/fromager/packagesettings/_models.py | 33 +++++ src/fromager/packagesettings/_settings.py | 17 ++- src/fromager/wheels.py | 63 ++++++++- tests/test_finders.py | 45 +++++++ tests/test_packagesettings.py | 50 +++++++ tests/test_wheels.py | 154 +++++++++++++++++++++- 10 files changed, 408 insertions(+), 44 deletions(-) diff --git a/src/fromager/bootstrapper/_cache.py b/src/fromager/bootstrapper/_cache.py index e8a3c6830..09d5f7e52 100644 --- a/src/fromager/bootstrapper/_cache.py +++ b/src/fromager/bootstrapper/_cache.py @@ -86,22 +86,29 @@ def _look_for_existing_wheel( search_in: pathlib.Path, ) -> tuple[pathlib.Path | None, pathlib.Path | None]: pbi = ctx.package_build_info(req) - expected_build_tag = pbi.build_tag(resolved_version) + base_build_tag = pbi.build_tag(resolved_version) logger.info( - f"looking for existing wheel for version {resolved_version} with build tag {expected_build_tag} in {search_in}" + f"looking for existing wheel for version {resolved_version} with build tag {base_build_tag} in {search_in}" ) wheel_filename = finders.find_wheel( downloads_dir=search_in, req=req, dist_version=str(resolved_version), - build_tag=expected_build_tag, + build_tag=base_build_tag, ) if not wheel_filename: return None, None - _, _, build_tag, _ = wheels.extract_info_from_wheel_file(req, wheel_filename) - if expected_build_tag and expected_build_tag != build_tag: + _, _, actual_build_tag, wheel_tags = wheels.extract_info_from_wheel_file( + req, wheel_filename + ) + expected_build_tag = wheels.get_build_tag( + ctx=ctx, req=req, version=resolved_version, wheel_tags=wheel_tags + ) + expected = expected_build_tag or (0, "") + actual = actual_build_tag or (0, "") + if expected != actual: logger.info( - f"found wheel for {resolved_version} in {wheel_filename} but build tag does not match. Got {build_tag} but expected {expected_build_tag}" + f"found wheel for {resolved_version} in {wheel_filename} but build tag does not match. Got {actual_build_tag} but expected {expected_build_tag}" ) return None, None logger.info(f"found existing wheel {wheel_filename}") @@ -129,16 +136,19 @@ def _download_wheel_from_cache( results = resolver.find_all_matching_from_provider(provider, pinned_req) wheel_url, _ = results[0] wheelfile_name = pathlib.Path(urlparse(wheel_url).path) - pbi = ctx.package_build_info(req) - expected_build_tag = pbi.build_tag(resolved_version) + _, _, actual_build_tag, wheel_tags = wheels.extract_info_from_wheel_file( + req, wheelfile_name + ) + expected_build_tag = wheels.get_build_tag( + ctx=ctx, req=req, version=resolved_version, wheel_tags=wheel_tags + ) logger.info(f"has expected build tag {expected_build_tag}") - changelogs = pbi.get_changelog(resolved_version) - logger.debug(f"has change logs {changelogs}") - _, _, build_tag, _ = wheels.extract_info_from_wheel_file(req, wheelfile_name) - if expected_build_tag and expected_build_tag != build_tag: + expected = expected_build_tag or (0, "") + actual = actual_build_tag or (0, "") + if expected != actual: logger.info( - f"found wheel for {resolved_version} in cache but build tag does not match. Got {build_tag} but expected {expected_build_tag}" + f"found wheel for {resolved_version} in cache but build tag does not match. Got {actual_build_tag} but expected {expected_build_tag}" ) return None, None diff --git a/src/fromager/commands/build.py b/src/fromager/commands/build.py index 2e33f8f0e..adeccaaa9 100644 --- a/src/fromager/commands/build.py +++ b/src/fromager/commands/build.py @@ -485,11 +485,12 @@ def _is_wheel_built( wheel_server_urls=wheel_server_urls, ) logger.info("found candidate wheel %s", url) - pbi = wkctx.package_build_info(req) - build_tag_from_settings = pbi.build_tag(resolved_version) - build_tag = build_tag_from_settings if build_tag_from_settings else (0, "") wheel_basename = downloads.extract_filename_from_url(url) - _, _, build_tag_from_name, _ = parse_wheel_filename(wheel_basename) + _, _, build_tag_from_name, wheel_tags = parse_wheel_filename(wheel_basename) + expected_tag = wheels.get_build_tag( + ctx=wkctx, req=req, version=resolved_version, wheel_tags=wheel_tags + ) + build_tag = expected_tag if expected_tag else (0, "") existing_build_tag = build_tag_from_name if build_tag_from_name else (0, "") if ( existing_build_tag[0] > build_tag[0] diff --git a/src/fromager/finders.py b/src/fromager/finders.py index 69c31e331..57aa1b45a 100644 --- a/src/fromager/finders.py +++ b/src/fromager/finders.py @@ -162,30 +162,27 @@ def find_wheel( """ filename_prefix = _dist_name_to_filename(req.name) canonical_name = canonicalize_name(req.name) - # if build tag is 0 then we can ignore to handle non tagged wheels for backward compatibility - candidate_bases_build_tag = f"{build_tag[0]}{build_tag[1]}-" if build_tag else "" - candidate_bases = set( - [ - # First check if the file is there using the canonically - # transformed name. - f"{filename_prefix}-{dist_version}-{candidate_bases_build_tag}", - # If that didn't work, try the canonical dist name. That's not - # "correct" but we do see it. (charset-normalizer-3.3.2- - # and setuptools-scm-8.0.4-) for example - f"{canonical_name}-{dist_version}-{candidate_bases_build_tag}", - # If *that* didn't work, try the dist name we've been - # given as a dependency. That's not "correct", either but we do - # see it. (oslo.messaging-14.7.0-) for example - f"{req.name}-{dist_version}-{candidate_bases_build_tag}", - # Sometimes the sdist uses '.' instead of '-' in the - # package name portion. - f"{req.name.replace('-', '.')}-{dist_version}-{candidate_bases_build_tag}", - ] - ) - # Case-insensitive globbing was added to Python 3.12, but we - # have to run with older versions, too, so do our own name - # comparison. + build_tag_prefixes: list[str] = [] + if build_tag: + build_tag_prefixes.append(f"{build_tag[0]}{build_tag[1]}-") + if not build_tag[1]: + build_tag_prefixes.append(f"{build_tag[0]}_") + else: + build_tag_prefixes.append("") + + name_variants = [ + filename_prefix, + canonical_name, + req.name, + req.name.replace("-", "."), + ] + + candidate_bases: list[str] = [] + for name in name_variants: + for btp in build_tag_prefixes: + candidate_bases.append(f"{name}-{dist_version}-{btp}") + for base in candidate_bases: logger.debug('looking for wheel as "%s"', base) for filename in downloads_dir.glob("*.whl"): diff --git a/src/fromager/packagesettings/__init__.py b/src/fromager/packagesettings/__init__.py index 6289d742b..cee70a2bf 100644 --- a/src/fromager/packagesettings/__init__.py +++ b/src/fromager/packagesettings/__init__.py @@ -11,6 +11,7 @@ ResolverDist, SbomSettings, VariantInfo, + WheelSettings, ) from ._pbi import PackageBuildInfo from ._resolver import ( @@ -86,6 +87,7 @@ "Variant", "VariantChangelog", "VariantInfo", + "WheelSettings", "default_update_extra_environ", "get_extra_environ", "pep440_tag_matcher", diff --git a/src/fromager/packagesettings/_models.py b/src/fromager/packagesettings/_models.py index 4e982e186..ad70f4549 100644 --- a/src/fromager/packagesettings/_models.py +++ b/src/fromager/packagesettings/_models.py @@ -32,6 +32,39 @@ logger = logging.getLogger(__name__) +class WheelSettings(pydantic.BaseModel): + """Global wheel build settings + + :: + + wheels: + build_tag_hook: "mypackage.hooks:build_tag_hook" + + .. versionadded:: 0.92.0 + """ + + model_config = MODEL_CONFIG + + build_tag_hook: pydantic.ImportString[typing.Callable[..., typing.Any]] | None = ( + None + ) + """Callable that returns suffix segments for the wheel build tag. + + The callable receives keyword-only arguments ``ctx``, ``req``, + ``version``, and ``wheel_tags`` and returns + ``Sequence[str]`` of suffix segments. + + Only invoked when the package already has a non-empty build tag + from its changelog entry for the given version; otherwise the hook + is skipped and no build tag is added. The callable must be + deterministic and independent of wheel contents, build environment, + or ELF metadata so fresh builds and cache lookups compute the same + tag. + + .. versionadded:: 0.92.0 + """ + + class SbomSettings(pydantic.BaseModel): """Global SBOM generation settings diff --git a/src/fromager/packagesettings/_settings.py b/src/fromager/packagesettings/_settings.py index 1d76c948a..6eb86cf0b 100644 --- a/src/fromager/packagesettings/_settings.py +++ b/src/fromager/packagesettings/_settings.py @@ -13,7 +13,7 @@ from pydantic import Field from .. import overrides -from ._models import PackageSettings, SbomSettings +from ._models import PackageSettings, SbomSettings, WheelSettings from ._pbi import PackageBuildInfo from ._typedefs import MODEL_CONFIG, GlobalChangelog, Package, Variant @@ -45,6 +45,14 @@ class SettingsFile(pydantic.BaseModel): are generated. """ + wheels: WheelSettings | None = None + """Wheel build settings + + Configures wheel build tag hooks and other wheel-specific options. + + .. versionadded:: 0.92.0 + """ + @classmethod def from_string( cls, @@ -175,6 +183,13 @@ def sbom_settings(self) -> SbomSettings | None: """Get global SBOM settings, or None if SBOM generation is disabled.""" return self._settings.sbom + @property + def build_tag_hook(self) -> typing.Callable[..., typing.Any] | None: + """Get the wheel build tag hook callable, or None if not configured.""" + if self._settings.wheels is None: + return None + return self._settings.wheels.build_tag_hook + def variant_changelog(self) -> list[str]: """Get global changelog for current variant""" return list(self._settings.changelog.get(self.variant, [])) diff --git a/src/fromager/wheels.py b/src/fromager/wheels.py index dc9bd5241..a699fb45d 100644 --- a/src/fromager/wheels.py +++ b/src/fromager/wheels.py @@ -4,6 +4,7 @@ import logging import os import pathlib +import re import shutil import sys import tempfile @@ -38,12 +39,67 @@ logger = logging.getLogger(__name__) +_BUILD_TAG_SEGMENT_RE = re.compile(r"^[a-zA-Z0-9.]+$") + FROMAGER_BUILD_SETTINGS = "fromager-build-settings" FROMAGER_ELF_PROVIDES = "fromager-elf-provides.txt" FROMAGER_ELF_REQUIRES = "fromager-elf-requires.txt" FROMAGER_BUILD_REQ_PREFIX = "fromager" +def _validate_build_tag_segments(segments: list[str]) -> None: + """Validate that each segment matches ``[a-zA-Z0-9.]``.""" + for seg in segments: + if not isinstance(seg, str): + raise ValueError( + f"build_tag_hook must return strings, got {type(seg).__name__}" + ) + if not _BUILD_TAG_SEGMENT_RE.match(seg): + raise ValueError( + f"build tag hook returned invalid segment {seg!r}: " + "each segment must match [a-zA-Z0-9.]" + ) + + +def get_build_tag( + *, + ctx: context.WorkContext, + req: Requirement, + version: Version, + wheel_tags: frozenset[Tag], +) -> BuildTag: + """Compute the full build tag including any hook-provided suffix. + + Calls ``pbi.build_tag(version)`` for the numeric base, then invokes + the configured ``build_tag_hook`` (if any) to append environment + suffix segments. + + .. versionadded:: 0.92.0 + """ + pbi = ctx.package_build_info(req) + base_tag = pbi.build_tag(version) + if not base_tag: + return base_tag + + hook = ctx.settings.build_tag_hook + if hook is None: + return base_tag + + raw = hook(ctx=ctx, req=req, version=version, wheel_tags=wheel_tags) + if isinstance(raw, (str, bytes)): + raise ValueError( + "build_tag_hook must return a sequence of strings, not a single string" + ) + segments = list(raw) + _validate_build_tag_segments(segments) + + if not segments: + return base_tag + + suffix = base_tag[1] + "_" + "_".join(segments) + return (base_tag[0], suffix) + + def _log_existing_sboms( req: Requirement, dist_info_dir: pathlib.Path, @@ -264,8 +320,11 @@ def add_extra_metadata_to_wheels( ) sbom.write_sbom(sbom=sbom_doc, dist_info_dir=dist_info_dir) - build_tag_from_settings = pbi.build_tag(version) - build_tag = build_tag_from_settings if build_tag_from_settings else (0, "") + build_tag = get_build_tag( + ctx=ctx, req=req, version=version, wheel_tags=wheel_tags + ) + if not build_tag: + build_tag = (0, "") cmd = [ "wheel", diff --git a/tests/test_finders.py b/tests/test_finders.py index 110ccfa1f..691fdfebf 100644 --- a/tests/test_finders.py +++ b/tests/test_finders.py @@ -146,3 +146,48 @@ def test_pypi_cache_provider() -> None: finders.PyPICacheProvider( cache_server_url=url, include_sdists=False, include_wheels=False ) + + +class TestFindWheelBuildTagSuffix: + """Tests for ``find_wheel`` with build tag suffixes.""" + + def test_find_wheel_with_exact_build_tag(self, tmp_path: pathlib.Path) -> None: + """Plain build tag matches a plain-tagged wheel.""" + downloads = tmp_path / "downloads" + downloads.mkdir() + wheel = downloads / "mypkg-1.0.0-2-py3-none-any.whl" + wheel.write_text("not-empty") + result = finders.find_wheel(downloads, Requirement("mypkg"), "1.0.0", (2, "")) + assert result == wheel + + def test_find_wheel_matches_suffixed_with_base_tag( + self, tmp_path: pathlib.Path + ) -> None: + """Base tag (2, '') matches a suffixed wheel when no plain one exists.""" + downloads = tmp_path / "downloads" + downloads.mkdir() + wheel = downloads / "mypkg-1.0.0-2_el9.6-cp312-cp312-linux_x86_64.whl" + wheel.write_text("not-empty") + result = finders.find_wheel(downloads, Requirement("mypkg"), "1.0.0", (2, "")) + assert result == wheel + + def test_find_wheel_no_false_positive_on_higher_number( + self, tmp_path: pathlib.Path + ) -> None: + """Build tag 2 does not match build tag 20.""" + downloads = tmp_path / "downloads" + downloads.mkdir() + (downloads / "mypkg-1.0.0-20-py3-none-any.whl").write_text("not-empty") + result = finders.find_wheel(downloads, Requirement("mypkg"), "1.0.0", (2, "")) + assert result is None + + def test_find_wheel_with_full_suffix(self, tmp_path: pathlib.Path) -> None: + """Exact suffixed build tag matches the right wheel.""" + downloads = tmp_path / "downloads" + downloads.mkdir() + wheel = downloads / "mypkg-1.0.0-2_el9.6_cuda13.0-cp312-cp312-linux_x86_64.whl" + wheel.write_text("not-empty") + result = finders.find_wheel( + downloads, Requirement("mypkg"), "1.0.0", (2, "_el9.6_cuda13.0") + ) + assert result == wheel diff --git a/tests/test_packagesettings.py b/tests/test_packagesettings.py index d16bba9ae..724e51803 100644 --- a/tests/test_packagesettings.py +++ b/tests/test_packagesettings.py @@ -1,3 +1,4 @@ +import os import pathlib import typing from unittest.mock import Mock, patch @@ -931,3 +932,52 @@ def test_version_none_no_reference( result = pbi.get_extra_environ(template_env={}, version=None) assert result["FOO"] == "bar" assert "__version__" not in result + + +class TestWheelSettings: + """Tests for ``WheelSettings`` and ``SettingsFile.wheels``.""" + + def test_settings_file_parses_wheels_section(self) -> None: + """wheels.build_tag_hook is loaded from YAML.""" + sf = SettingsFile.from_string( + """ +wheels: + build_tag_hook: "os.path:join" +""" + ) + assert sf.wheels is not None + assert sf.wheels.build_tag_hook is os.path.join + + def test_settings_file_defaults_to_no_wheels(self) -> None: + """When wheels section is absent, wheels is None.""" + sf = SettingsFile.from_string("") + assert sf.wheels is None + + def test_settings_build_tag_hook_property(self) -> None: + """Settings.build_tag_hook exposes the hook callable.""" + sf = SettingsFile.from_string( + """ +wheels: + build_tag_hook: "os.path:join" +""" + ) + settings = Settings( + settings=sf, + package_settings=[], + variant="cpu", + patches_dir=pathlib.Path("/tmp"), + max_jobs=None, + ) + assert settings.build_tag_hook is os.path.join + + def test_settings_build_tag_hook_none_when_unset(self) -> None: + """Settings.build_tag_hook is None when not configured.""" + sf = SettingsFile.from_string("") + settings = Settings( + settings=sf, + package_settings=[], + variant="cpu", + patches_dir=pathlib.Path("/tmp"), + max_jobs=None, + ) + assert settings.build_tag_hook is None diff --git a/tests/test_wheels.py b/tests/test_wheels.py index 1c589d012..fcbe899f1 100644 --- a/tests/test_wheels.py +++ b/tests/test_wheels.py @@ -5,10 +5,11 @@ import pytest from conftest import make_sbom_ctx from packaging.requirements import Requirement +from packaging.tags import Tag from packaging.version import Version from fromager import build_environment, context, downloads, wheels -from fromager.packagesettings import SbomSettings +from fromager.packagesettings import SbomSettings, Settings, SettingsFile, WheelSettings @patch("pyproject_hooks.BuildBackendHookCaller.build_wheel") @@ -337,3 +338,154 @@ def test_validate_wheel_file( else: with pytest.raises(ValueError): wheels.validate_wheel_filename(req, version, wheel_file) + + +def _ctx_with_hook( + tmp_path: pathlib.Path, + hook: object | None = None, +) -> context.WorkContext: + """Create a WorkContext with an optional build_tag_hook.""" + sf = SettingsFile.from_string("") + if hook is not None: + sf = sf.model_copy(update={"wheels": WheelSettings(build_tag_hook=hook)}) + settings = Settings( + settings=sf, + package_settings=[], + variant="cpu", + patches_dir=tmp_path / "patches", + max_jobs=None, + ) + ctx = context.WorkContext( + active_settings=settings, + patches_dir=tmp_path / "patches", + sdists_repo=tmp_path / "sdists-repo", + wheels_repo=tmp_path / "wheels-repo", + work_dir=tmp_path / "work-dir", + variant="cpu", + ) + ctx.setup() + return ctx + + +class TestGetBuildTag: + """Tests for ``wheels.get_build_tag()``.""" + + def test_no_hook_returns_base_tag(self, tmp_path: pathlib.Path) -> None: + """Without a hook, get_build_tag returns pbi.build_tag() unchanged.""" + ctx = _ctx_with_hook(tmp_path) + req = Requirement("mypkg") + version = Version("1.0") + tags = frozenset({Tag("cp312", "cp312", "linux_x86_64")}) + result = wheels.get_build_tag( + ctx=ctx, req=req, version=version, wheel_tags=tags + ) + pbi = ctx.package_build_info(req) + assert result == pbi.build_tag(version) + + def test_hook_appends_suffix_segments( + self, testdata_context: context.WorkContext + ) -> None: + """Hook-provided segments are joined and appended to the base tag.""" + + def hook(**kwargs: object) -> list[str]: + return ["el9.6", "rocm7.1"] + + testdata_context.settings._settings = ( + testdata_context.settings._settings.model_copy( + update={"wheels": WheelSettings(build_tag_hook=hook)} + ) + ) + req = Requirement("test-pkg") + version = Version("1.0.1") + tags = frozenset({Tag("cp312", "cp312", "linux_x86_64")}) + pbi = testdata_context.package_build_info(req) + base = pbi.build_tag(version) + assert base, "test-pkg must have a changelog entry for 1.0.1" + result = wheels.get_build_tag( + ctx=testdata_context, req=req, version=version, wheel_tags=tags + ) + assert result[0] == base[0] + assert result[1] == base[1] + "_el9.6_rocm7.1" + + def test_hook_empty_segments_returns_base( + self, testdata_context: context.WorkContext + ) -> None: + """When hook returns empty list, base tag is returned.""" + + def hook(**kwargs: object) -> list[str]: + return [] + + testdata_context.settings._settings = ( + testdata_context.settings._settings.model_copy( + update={"wheels": WheelSettings(build_tag_hook=hook)} + ) + ) + req = Requirement("test-pkg") + version = Version("1.0.1") + tags = frozenset({Tag("py3", "none", "any")}) + pbi = testdata_context.package_build_info(req) + base = pbi.build_tag(version) + result = wheels.get_build_tag( + ctx=testdata_context, req=req, version=version, wheel_tags=tags + ) + assert result == base + + def test_hook_returning_string_raises( + self, testdata_context: context.WorkContext + ) -> None: + """Single string return is rejected (would be iterated as chars).""" + + def hook(**kwargs: object) -> str: + return "el9.6" + + testdata_context.settings._settings = ( + testdata_context.settings._settings.model_copy( + update={"wheels": WheelSettings(build_tag_hook=hook)} + ) + ) + req = Requirement("test-pkg") + version = Version("1.0.1") + tags = frozenset({Tag("cp312", "cp312", "linux_x86_64")}) + with pytest.raises(ValueError, match="sequence of strings"): + wheels.get_build_tag( + ctx=testdata_context, req=req, version=version, wheel_tags=tags + ) + + def test_hook_invalid_segment_chars_raises( + self, testdata_context: context.WorkContext + ) -> None: + """Segments with invalid characters are rejected.""" + + def hook(**kwargs: object) -> list[str]: + return ["el9.6", "bad-char"] + + testdata_context.settings._settings = ( + testdata_context.settings._settings.model_copy( + update={"wheels": WheelSettings(build_tag_hook=hook)} + ) + ) + req = Requirement("test-pkg") + version = Version("1.0.1") + tags = frozenset({Tag("cp312", "cp312", "linux_x86_64")}) + with pytest.raises(ValueError, match="invalid segment"): + wheels.get_build_tag( + ctx=testdata_context, req=req, version=version, wheel_tags=tags + ) + + def test_hook_not_called_without_base_tag(self, tmp_path: pathlib.Path) -> None: + """Hook is skipped when the package has no changelog build tag.""" + calls: list[dict[str, object]] = [] + + def hook(**kwargs: object) -> list[str]: + calls.append(kwargs) + return ["el9.6"] + + ctx = _ctx_with_hook(tmp_path, hook=hook) + req = Requirement("mypkg") + version = Version("1.0") + tags = frozenset({Tag("cp312", "cp312", "linux_x86_64")}) + result = wheels.get_build_tag( + ctx=ctx, req=req, version=version, wheel_tags=tags + ) + assert result == () + assert calls == []