From e89334c39cc1239903bed15e90bf3613b91d7e67 Mon Sep 17 00:00:00 2001 From: Tyler Yankee Date: Thu, 3 Sep 2026 17:58:26 -0400 Subject: [PATCH] [wheel] Test ABI3 wheels across Python versions By definition, an ABI3 wheel can be installed and run under newer versions of Python than what it was built with. Amend 0c490ac612 by adapting the wheel builder to handle testing a wheel on multiple versions of Python within a given target (i.e., version with which it was built). Co-authored-by: Matthew Woehlke --- tools/wheel/wheel_builder/common.py | 5 + tools/wheel/wheel_builder/linux.py | 195 +++++++++++------------ tools/wheel/wheel_builder/linux_types.py | 65 +++++--- tools/wheel/wheel_builder/macos.py | 33 ++-- tools/wheel/wheel_builder/macos_types.py | 8 +- 5 files changed, 174 insertions(+), 132 deletions(-) diff --git a/tools/wheel/wheel_builder/common.py b/tools/wheel/wheel_builder/common.py index 7f813bd8daf7..2df84c62ddcd 100644 --- a/tools/wheel/wheel_builder/common.py +++ b/tools/wheel/wheel_builder/common.py @@ -66,6 +66,11 @@ def __init__(self, *version_parts: int): self.version = ".".join(pv_parts[:2]) self.tag = "".join(pv_parts[:2]) + def validate(self, n_components: int) -> None: + """Asserts that the number of version components belonging to `self` + is equal to `n_components`.""" + assert len(self.version_tuple) == n_components, self.version_tuple + def gripe(message): """ diff --git a/tools/wheel/wheel_builder/linux.py b/tools/wheel/wheel_builder/linux.py index 6483afb2faa1..510163ca3f94 100644 --- a/tools/wheel/wheel_builder/linux.py +++ b/tools/wheel/wheel_builder/linux.py @@ -22,16 +22,23 @@ wheel_name, wheelhouse, ) -from .linux_types import BUILD, TEST, Platform, PythonManager, Role, Target +from .linux_types import Platform, Target, TestCase # Artifacts that need to be cleaned up. DO NOT MODIFY outside of this file. -_files_to_remove = [] -_images_to_remove = [] +_files_to_remove = set() +_images_to_remove = set() tag_base = "pip-drake" ARCH = platform.machine() +# Supported platforms on which every wheel is tested. +_TEST_PLATFORMS = ( + Platform("amazonlinux", "2023", "AL2023"), + Platform("ubuntu", "24.04", "noble"), + Platform("ubuntu", "26.04", "resolute"), +) + # This is the complete set of defined targets (i.e. potential wheels). By # default, all targets matching the currently running architecture are built, # but the user may down-select from this set. The platform alias is used for @@ -52,89 +59,69 @@ "x86_64": ( Target( python_binder=PythonBinder.NANOBIND, - python=PythonTarget(3, 12, 13), build_platform=Platform("amd64/almalinux", "9", "almalinux9"), - test_platforms=( - Platform("amazonlinux", "2023", "AL2023"), - Platform("ubuntu", "24.04", "noble"), - Platform("ubuntu", "26.04", "resolute", PythonManager.UV), - # TODO(jwnimmer-tri) We should test this same abi3 wheel on all - # newer Python versions (so 3.13, 3.14, etc.). + build_python=PythonTarget(3, 12, 13), + test_platforms=_TEST_PLATFORMS, + test_pythons=( + PythonTarget(3, 12), + PythonTarget(3, 13), + PythonTarget(3, 14), ), ), Target( python_binder=PythonBinder.PYBIND11, - python=PythonTarget(3, 12, 13), build_platform=Platform("amd64/almalinux", "9", "almalinux9"), - test_platforms=( - Platform("amazonlinux", "2023", "AL2023"), - Platform("ubuntu", "24.04", "noble"), - Platform("ubuntu", "26.04", "resolute", PythonManager.UV), - ), + build_python=PythonTarget(3, 12, 13), + test_platforms=_TEST_PLATFORMS, + test_pythons=(PythonTarget(3, 12),), ), Target( python_binder=PythonBinder.PYBIND11, - python=PythonTarget(3, 13, 15), build_platform=Platform("amd64/almalinux", "9", "almalinux9"), - test_platforms=( - Platform("amazonlinux", "2023", "AL2023"), - Platform("ubuntu", "24.04", "noble", PythonManager.UV), - Platform("ubuntu", "26.04", "resolute", PythonManager.UV), - ), + build_python=PythonTarget(3, 13, 15), + test_platforms=_TEST_PLATFORMS, + test_pythons=(PythonTarget(3, 13),), ), Target( python_binder=PythonBinder.PYBIND11, - python=PythonTarget(3, 14, 7), build_platform=Platform("amd64/almalinux", "9", "almalinux9"), - test_platforms=( - Platform("amazonlinux", "2023", "AL2023"), - Platform("ubuntu", "24.04", "noble", PythonManager.UV), - Platform("ubuntu", "26.04", "resolute"), - ), + build_python=PythonTarget(3, 14, 7), + test_platforms=_TEST_PLATFORMS, + test_pythons=(PythonTarget(3, 14),), ), ), "aarch64": ( Target( python_binder=PythonBinder.NANOBIND, - python=PythonTarget(3, 12, 13), build_platform=Platform("arm64v8/almalinux", "9", "almalinux9"), - test_platforms=( - Platform("amazonlinux", "2023", "AL2023"), - Platform("ubuntu", "24.04", "noble"), - Platform("ubuntu", "26.04", "resolute", PythonManager.UV), - # TODO(jwnimmer-tri) We should test this same abi3 wheel on all - # newer Python versions (so 3.13, 3.14, etc.). + build_python=PythonTarget(3, 12, 13), + test_platforms=_TEST_PLATFORMS, + test_pythons=( + PythonTarget(3, 12), + PythonTarget(3, 13), + PythonTarget(3, 14), ), ), Target( python_binder=PythonBinder.PYBIND11, - python=PythonTarget(3, 12, 13), build_platform=Platform("arm64v8/almalinux", "9", "almalinux9"), - test_platforms=( - Platform("amazonlinux", "2023", "AL2023"), - Platform("ubuntu", "24.04", "noble"), - Platform("ubuntu", "26.04", "resolute", PythonManager.UV), - ), + build_python=PythonTarget(3, 12, 13), + test_platforms=_TEST_PLATFORMS, + test_pythons=(PythonTarget(3, 12),), ), Target( python_binder=PythonBinder.PYBIND11, - python=PythonTarget(3, 13, 15), build_platform=Platform("arm64v8/almalinux", "9", "almalinux9"), - test_platforms=( - Platform("amazonlinux", "2023", "AL2023"), - Platform("ubuntu", "24.04", "noble", PythonManager.UV), - Platform("ubuntu", "26.04", "resolute", PythonManager.UV), - ), + build_python=PythonTarget(3, 13, 15), + test_platforms=_TEST_PLATFORMS, + test_pythons=(PythonTarget(3, 13),), ), Target( python_binder=PythonBinder.PYBIND11, - python=PythonTarget(3, 14, 7), build_platform=Platform("arm64v8/almalinux", "9", "almalinux9"), - test_platforms=( - Platform("amazonlinux", "2023", "AL2023"), - Platform("ubuntu", "24.04", "noble", PythonManager.UV), - Platform("ubuntu", "26.04", "resolute"), - ), + build_python=PythonTarget(3, 14, 7), + test_platforms=_TEST_PLATFORMS, + test_pythons=(PythonTarget(3, 14),), ), ), }[ARCH] @@ -238,26 +225,33 @@ def _create_source_tar(path): out.close() -def _tagname( - target: Target, role: Role, tag_prefix: str, test_index: int | None = None -): +def _build_tagname(target: Target, tag_prefix: str) -> str: """ - Generates a Docker tag name for a target and tag prefix. - Iff the role is the TEST role, then the test_index must be provided. + Generates a Docker tag name for a build-role target and tag prefix. """ - platform = target.platform(role, test_index).alias - python_tag = target.python.tag + platform = target.build_platform.alias + python_tag = target.build_python.tag python_binder = target.python_binder.value return f"{tag_base}:{tag_prefix}-{platform}-py{python_tag}-{python_binder}" +def _test_tagname(test_case: TestCase, tag_prefix: str) -> str: + """ + Generates a Docker tag name for a test-role TestCase and tag prefix. + """ + platform = test_case.platform.alias + manager = test_case.python_manager.value + python_tag = test_case.python.tag + return f"{tag_base}:{tag_prefix}-{platform}-py{python_tag}-{manager}" + + def _build_stage(target, args, tag_prefix, stage=None): """ Runs a Docker build and return the build tag. """ # Generate canonical tag from target. - tag = _tagname(target, BUILD, tag_prefix) + tag = _build_tagname(target, tag_prefix) # Generate extra arguments to specify what stage to build. if stage is not None: @@ -272,32 +266,34 @@ def _build_stage(target, args, tag_prefix, stage=None): return tag -def _target_args(target: Target, role: Role, test_index: int | None = None): +def _target_args(platform: Platform, python: PythonTarget): + return [ + "--build-arg", f"PLATFORM={platform.name}:{platform.version}", + "--build-arg", f"PYTHON={python.version_full}", + ] # fmt: skip + + +def _build_target_args(target: Target): """ - Returns the Docker build arguments for the specified platform target. - Iff the role is the TEST role, then the test_index must be provided. + Returns the Docker build arguments for the build Dockerfile. """ - platform = target.platform(role, test_index) + return _target_args(target.build_platform, target.build_python) + [ + "--build-arg", f"DRAKE_PYTHON_BINDER={target.python_binder.value}", + ] # fmt: skip - if role == BUILD: - python_args = [ - "--build-arg", f"PYTHON={target.python.version_full}", - "--build-arg", f"DRAKE_PYTHON_BINDER={target.python_binder.value}", - ] # fmt: skip - else: - python_args = [ - "--build-arg", f"PYTHON={target.python.version}", - "--build-arg", f"PYTHON_MANAGER={platform.python_manager.value}", - ] # fmt: skip - return [ - "--build-arg", f"PLATFORM={platform.name}:{platform.version}", - ] + python_args # fmt: skip +def _test_target_args(test_case: TestCase): + """ + Returns the Docker build arguments for the test Dockerfile. + """ + return _target_args(test_case.platform, test_case.python) + [ + "--build-arg", f"PYTHON_MANAGER={test_case.python_manager.value}", + ] # fmt: skip def _build_image(target, identifier, version, options): """ - Runs the build for a target and (optionally) extract the wheel. + Runs the build for a target and (optionally) extracts the wheel. """ drake_is_abi3_wheel = ( "1" if target.python_binder == PythonBinder.NANOBIND else "0" @@ -306,7 +302,7 @@ def _build_image(target, identifier, version, options): "--build-arg", f"DRAKE_VERSION={version}", "--build-arg", f"DRAKE_GIT_SHA={_git_sha(resource_root)}", "--build-arg", f"DRAKE_IS_ABI3_WHEEL={drake_is_abi3_wheel}", - ] + _target_args(target, BUILD) # fmt: skip + ] + _build_target_args(target) # fmt: skip if not options.keep_containers: args.append("--force-rm") @@ -323,7 +319,7 @@ def _build_image(target, identifier, version, options): ) else: tag = _build_stage(target, args, tag_prefix=identifier) - _images_to_remove.append(tag) + _images_to_remove.add(tag) # Extract the wheel (if requested). if options.extract: @@ -352,40 +348,43 @@ def _test_wheel(target, identifier, version, options): """ Runs the test script for the wheel matching the specified target. """ - glibc = glibc_versions[target.platform(BUILD).alias] + glibc = glibc_versions[target.build_platform.alias] wheel = wheel_name( python_binder=target.python_binder, - python_version=target.python.tag, + python_version=target.build_python.tag, wheel_version=version, wheel_platform=f"manylinux_{glibc}_{ARCH}", ) + test_dir = os.path.join(resource_root, "test") - for test_index, test_platform in enumerate(target.test_platforms): - print(f"[-] Testing on {test_platform.alias} ...") - test_image = _tagname(target, TEST, f"test-{identifier}", test_index) + for test_case in target.test_cases(): + print( + f"[-] Testing on {test_case.platform.alias}" + f" (Python {test_case.python.version}) ..." + ) + test_image = _test_tagname(test_case, f"test-{identifier}") test_container = test_image.replace(":", "__") if options.tag_stages: - base_image = _tagname(target, TEST, "test", test_index) + base_image = _test_tagname(test_case, "test") else: base_image = test_image - test_dir = os.path.join(resource_root, "test") # Build the test base image. _docker( "build", "-t", base_image, - *_target_args(target, TEST, test_index), + *_test_target_args(test_case), test_dir, ) if not options.tag_stages: - _images_to_remove.append(base_image) + _images_to_remove.add(base_image) # Install the wheel. install_command = [ "/test/install-wheel.sh", os.path.join(wheelhouse, wheel), - test_platform.python_manager.value, + test_case.python_manager.value, ] # fmt: skip _docker( "run", "-t", f"--name={test_container}", @@ -398,7 +397,7 @@ def _test_wheel(target, identifier, version, options): _docker("commit", test_container, test_image) _docker("container", "rm", test_container) if options.tag_stages: - _images_to_remove.append(test_image) + _images_to_remove.add(test_image) # Run individual tests. test_script = "/test/test-wheel.sh" @@ -423,8 +422,8 @@ def build(options): targets_to_build = [] for t in targets: if ( - t.platform(BUILD).name in options.platforms - and t.python.tag in options.python_versions + t.build_platform.name in options.platforms + and t.build_python.tag in options.python_versions ): targets_to_build.append(t) @@ -442,12 +441,12 @@ def build(options): # Provide the SNOPT source archive as a dependency. snopt_tgz = os.path.join(resource_root, "image", "snopt.tar.gz") - _files_to_remove.append(snopt_tgz) + _files_to_remove.add(snopt_tgz) create_snopt_tgz(snopt_path=options.snopt_path, output=snopt_tgz) # Generate the Drake repository source archive. source_tar = os.path.join(resource_root, "image", "drake-src.tar") - _files_to_remove.append(source_tar) + _files_to_remove.add(source_tar) _create_source_tar(source_tar) # Build the requested wheels. @@ -487,14 +486,14 @@ def add_selection_arguments(parser): parser.add_argument( "--platform", dest="platforms", - default=",".join({t.platform(BUILD).name for t in targets}), + default=",".join({t.build_platform.name for t in targets}), help="platform(s) to build; separate with ',' (default: %(default)s)", ) parser.add_argument( "--python", dest="python_versions", metavar="VERSIONS", - default=",".join(sorted({t.python.tag for t in targets})), + default=",".join(sorted({t.build_python.tag for t in targets})), help=( "python version(s) to build; separate with ','" " (default: %(default)s)" diff --git a/tools/wheel/wheel_builder/linux_types.py b/tools/wheel/wheel_builder/linux_types.py index 0d2cfd284125..4781b711a2a0 100644 --- a/tools/wheel/wheel_builder/linux_types.py +++ b/tools/wheel/wheel_builder/linux_types.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from enum import Enum +import itertools from .common import PythonBinder, PythonTarget @@ -14,9 +15,12 @@ class PythonManager(Enum): UV = "uv" -@dataclass -class Role: - name: str +# Python versions available via each test platform's system package manager. +_DISTRO_PYTHONS: dict[str, set[tuple[int, int]]] = { + "AL2023": {(3, 12), (3, 13), (3, 14)}, + "noble": {(3, 12)}, + "resolute": {(3, 14)}, +} @dataclass @@ -24,32 +28,49 @@ class Platform: name: str version: str alias: str - python_manager: PythonManager = PythonManager.PIP + + +@dataclass +class TestCase: + """A (platform, python) combination with which to test. + + python_manager is selected as PIP if the platform's system package manager + natively provides the requested Python version, or UV otherwise. + """ + + platform: Platform + python: PythonTarget + + def __post_init__(self): + platform_pythons = _DISTRO_PYTHONS[self.platform.alias] + self.python_manager = ( + PythonManager.PIP + if self.python.version_tuple in platform_pythons + else PythonManager.UV + ) @dataclass class Target: python_binder: PythonBinder - python: PythonTarget build_platform: Platform - test_platforms: tuple[Platform] + build_python: PythonTarget + test_platforms: tuple[Platform, ...] + test_pythons: tuple[PythonTarget, ...] def __post_init__(self): - assert len(self.python.version_tuple) == 3, self.python.version_tuple + self.build_python.validate(n_components=3) assert isinstance(self.test_platforms, tuple) + assert isinstance(self.test_pythons, tuple) + for test_python in self.test_pythons: + test_python.validate(n_components=2) - def platform(self, role: Role, test_index: int | None = None) -> Platform: - """Returns the Platform for the given `role`. For the test role, the - `test_index` into the `self.test_platforms` tuple is required. For the - build role, the `test_index` must be None.""" - if role.name == "build": - assert test_index is None - return self.build_platform - if role.name == "test": - assert test_index is not None - return self.test_platforms[test_index] - raise NotImplementedError(role.name) - - -BUILD = Role("build") -TEST = Role("test") + def test_cases(self) -> tuple[TestCase, ...]: + """Returns the Cartesian product of `test_platforms` and + `test_pythons` as a tuple of `TestCase` instances.""" + return tuple( + itertools.starmap( + TestCase, + itertools.product(self.test_platforms, self.test_pythons), + ) + ) diff --git a/tools/wheel/wheel_builder/macos.py b/tools/wheel/wheel_builder/macos.py index 4781b751ed07..1d6e2e98c7f4 100644 --- a/tools/wheel/wheel_builder/macos.py +++ b/tools/wheel/wheel_builder/macos.py @@ -41,9 +41,21 @@ # * the Python versions supported by MOSEK, in tools/wheel/setup.py. If # there is any Python version supported by Drake, but not MOSEK, a note # should be added to the aforementioned installation documentation. - Target(PythonBinder.NANOBIND, PythonTarget(3, 13)), - Target(PythonBinder.PYBIND11, PythonTarget(3, 13)), - Target(PythonBinder.PYBIND11, PythonTarget(3, 14)), + Target( + python_binder=PythonBinder.NANOBIND, + build_python=PythonTarget(3, 13), + test_pythons=(PythonTarget(3, 13), PythonTarget(3, 14)), + ), + Target( + python_binder=PythonBinder.PYBIND11, + build_python=PythonTarget(3, 13), + test_pythons=(PythonTarget(3, 13),), + ), + Target( + python_binder=PythonBinder.PYBIND11, + build_python=PythonTarget(3, 14), + test_pythons=(PythonTarget(3, 14),), + ), ) @@ -55,7 +67,7 @@ def _find_wheel(path, version, target): """ pattern = wheel_name( python_binder=target.python_binder, - python_version=target.python.tag, + python_version=target.build_python.tag, wheel_version=version, wheel_platform="*", ) @@ -80,7 +92,7 @@ def _assert_isdir(path, name): die(f"{name} '{path}' is not a valid directory") -def _test_wheel(wheel, target, env): +def _test_wheel(wheel, python_target, env): """ Runs the test script on `wheel`. """ @@ -89,7 +101,7 @@ def _test_wheel(wheel, target, env): resource_root, "macos", "provision-test-python.sh" ) subprocess.check_call( - ["bash", setup_script, target.python.version], env=env + ["bash", setup_script, python_target.version], env=env ) # Install the wheel. @@ -114,7 +126,7 @@ def build(options): # Collect set of wheels to be built. targets_to_build = [] for t in TARGETS: - if t.python.tag in options.python_versions: + if t.build_python.tag in options.python_versions: targets_to_build.append(t) # Check if there is anything to do. @@ -176,7 +188,7 @@ def build(options): build_script = os.path.join(resource_root, "macos", "build-wheel.sh") build_command = ["bash", build_script] build_command.append(version) - build_command.append(target.python.version) + build_command.append(target.build_python.version) subprocess.check_call(build_command, env=environment) @@ -188,7 +200,8 @@ def build(options): ) if options.test: - _test_wheel(wheel, target=target, env=environment) + for python_target in target.test_pythons: + _test_wheel(wheel, python_target, env=environment) if options.extract: shutil.copy2(wheel, options.output_dir) @@ -225,7 +238,7 @@ def add_selection_arguments(parser): "--python", dest="python_versions", metavar="VERSIONS", - default=",".join(sorted({t.python.tag for t in TARGETS})), + default=",".join(sorted({t.build_python.tag for t in TARGETS})), help=( "python version(s) to build; " "separate with ',' (default: %(default)s)" diff --git a/tools/wheel/wheel_builder/macos_types.py b/tools/wheel/wheel_builder/macos_types.py index 6b36c1f41bba..6c0ff009f886 100644 --- a/tools/wheel/wheel_builder/macos_types.py +++ b/tools/wheel/wheel_builder/macos_types.py @@ -21,7 +21,11 @@ class Target: """ python_binder: PythonBinder - python: PythonTarget + build_python: PythonTarget + test_pythons: tuple[PythonTarget] def __post_init__(self): - assert len(self.python.version_tuple) == 2, self.python.version_tuple + self.build_python.validate(n_components=2) + assert isinstance(self.test_pythons, tuple) + for test_python in self.test_pythons: + test_python.validate(n_components=2)