From 5c18bc656bb65e98832166f563336123782f6b33 Mon Sep 17 00:00:00 2001 From: Sean Pryor Date: Fri, 6 Feb 2026 17:41:38 -0500 Subject: [PATCH] AIPCC-10079: Allow overriding extract_info_from_wheel_file This happens due to an unfortuante way that TensorFlow is built. Since it uses Bazel to build and package the wheel, it doesn't have the usual infrastructure one uses to build variant wheels. Due to that, a single repo is used to build wheels with different names. In Fromager's case, this means we specify the req as `tensorflow` but in reality will get `tensorflow_cpu` or `tensorflow_rocm`, etc. This is how they create them upstream, and to match the requirements from notebooks that the wheel naming convention remains the same, we have to override the check that the wheel name matches the req in extract_info_from_wheel_file(). To achieve this, I made it overridable like other functions and provided a default_extract_info_from_wheel_file() implemetation to preserve existing behavior. I expect TensorFlow will be one of the only users of this override. Signed-off-by: Sean Pryor --- src/fromager/wheels.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/src/fromager/wheels.py b/src/fromager/wheels.py index 88a2a829..b3610eb4 100644 --- a/src/fromager/wheels.py +++ b/src/fromager/wheels.py @@ -117,7 +117,7 @@ def _extra_metadata_elfdeps( return elfinfos -def extract_info_from_wheel_file( +def default_extract_info_from_wheel_file( req: Requirement, wheel_file: pathlib.Path ) -> tuple[str, Version, BuildTag, frozenset[Tag]]: # parse_wheel_filename normalizes the dist name, however the dist-info @@ -133,6 +133,23 @@ def extract_info_from_wheel_file( return (dist_name, dist_version, build_tag, wheel_tags) +def extract_info_from_wheel_file( + req: Requirement, wheel_file: pathlib.Path +) -> tuple[str, Version, BuildTag, frozenset[Tag]]: + + wheel_info: tuple[str, Version, BuildTag, frozenset[Tag]] = ( + overrides.find_and_invoke( + req.name, + "extract_info_from_wheel_file", + default_extract_info_from_wheel_file, + req=req, + wheel_file=wheel_file, + ) + ) + + return wheel_info + + def default_add_extra_metadata_to_wheels( ctx: context.WorkContext, req: Requirement,