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)