From 81ecd8df13a6e058c20bd56d2b72a2788fa7b733 Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 8 Jul 2026 10:03:37 +0900 Subject: [PATCH] Refactor wheel building script --- .github/actions/build-macos/action.yml | 2 +- .github/actions/build-wheel/action.yml | 7 +-- .github/actions/setup/action.yml | 3 +- .github/workflows/release.yml | 4 ++ setup.py | 78 +++++++++++++++----------- 5 files changed, 56 insertions(+), 38 deletions(-) diff --git a/.github/actions/build-macos/action.yml b/.github/actions/build-macos/action.yml index 84055009e9..c3109277fc 100644 --- a/.github/actions/build-macos/action.yml +++ b/.github/actions/build-macos/action.yml @@ -17,7 +17,7 @@ runs: shell: bash run: | echo "::group::Install dependencies" - uv pip install 'build<=1.4.2' setuptools + uv pip install build setuptools echo "::endgroup::" - name: Build wheel diff --git a/.github/actions/build-wheel/action.yml b/.github/actions/build-wheel/action.yml index 96cca9eadc..01244d4677 100644 --- a/.github/actions/build-wheel/action.yml +++ b/.github/actions/build-wheel/action.yml @@ -35,7 +35,7 @@ runs: shell: bash run: | echo "::group::Install dependencies" - uv pip install 'build<=1.4.2' setuptools + uv pip install build setuptools if ${{ runner.os == 'Linux' }} ; then uv pip install auditwheel patchelf fi @@ -50,8 +50,7 @@ runs: MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} run: | echo "::group::Build frontend package" - python setup.py clean --all - MLX_BUILD_STAGE=1 python -m build -w + MLX_BUILD_FRONTEND_PACKAGE=1 python -m build -w echo "::endgroup::" - name: Post-process frontend package @@ -77,8 +76,8 @@ runs: MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} run: | echo "::group::Build backend package" + MLX_BUILD_BACKEND_PACKAGE=1 python -m build -w python setup.py clean --all - MLX_BUILD_STAGE=2 python -m build -w echo "::endgroup::" - name: Post-process backend package diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml index 4f7b14105f..bd0df81b4c 100644 --- a/.github/actions/setup/action.yml +++ b/.github/actions/setup/action.yml @@ -42,7 +42,7 @@ runs: echo "::group::Install common dependencies" sudo apt-get update sudo apt-get install -y --no-install-recommends \ - gdb g++ ninja-build zip \ + gdb g++ ninja-build unzip \ libblas-dev liblapack-dev liblapacke-dev \ openmpi-bin openmpi-common libopenmpi-dev echo "::endgroup::" @@ -52,6 +52,7 @@ runs: shell: bash run: | echo "::group::Install macOS dependencies" + brew trust aws/tap # suppress warning in github actions brew install openmpi xcodebuild -showComponent MetalToolchain echo "::endgroup::" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index bd94af697b..14b463e3a7 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -86,6 +86,7 @@ jobs: with: cmake-args: ${{ steps.setup.outputs.cmake-args }} build-backend: false + - run: unzip -l wheelhouse/*.whl - uses: actions/upload-artifact@v7 with: name: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }} @@ -127,6 +128,7 @@ jobs: with: cmake-args: ${{ steps.setup.outputs.cmake-args }} build-frontend: false + - run: unzip -l wheelhouse/*.whl - uses: actions/upload-artifact@v7 with: name: backend-${{ matrix.toolkit }}-${{ runner.os }}-${{ runner.arch }} @@ -168,6 +170,8 @@ jobs: macos-target: '26.2' cmake-args: ${{ steps.setup.outputs.cmake-args }} build-backend: ${{ matrix.python-version == '3.10' }} + - name: Display content of the wheels + run: for WHEEL in wheelhouse/*.whl; do unzip -l $WHEEL; echo; done - name: Upload frontend packages uses: actions/upload-artifact@v7 with: diff --git a/setup.py b/setup.py index 3c2f138048..f7d0ed8569 100644 --- a/setup.py +++ b/setup.py @@ -53,7 +53,22 @@ def get_version(): return version -build_stage = int(os.environ.get("MLX_BUILD_STAGE", 0)) +# Release builds for PyPi are separated into 2 packages: +# +# Frontend package: +# - Triggered with `MLX_BUILD_FRONTEND_PACKAGE=1` +# - Include everything except backend-specific binaries (e.g. libmlx.so, mlx.metallib, etc) +# - Wheel has Python ABI and platform tags +# - Wheel should be built for the cross-product of python version and platforms +# - Package name is "mlx" and it depends on backend packages (e.g. mlx-metal, mlx-cuda) +# Backend package: +# - Triggered with `MLX_BUILD_BACKEND_PACKAGE=1` +# - Include headers and backend binaries. +# - Wheel has only platform tags +# - Wheel should be built only for different platforms +# - Package name is back-end specific, e.g mlx-metal, mlx-cuda +build_frontend = int(os.environ.get("MLX_BUILD_FRONTEND_PACKAGE", 0)) +build_backend = int(os.environ.get("MLX_BUILD_BACKEND_PACKAGE", 0)) build_macos = platform.system() == "Darwin" build_cuda = "MLX_BUILD_CUDA=ON" in os.environ.get("CMAKE_ARGS", "") @@ -88,17 +103,9 @@ def build_extension(self, ext: CMakeExtension) -> None: if not build_temp.exists(): build_temp.mkdir(parents=True) - install_prefix = extdir - pybind_out_dir = extdir - if build_stage == 1: - # Don't include MLX libraries in the wheel - install_prefix = build_temp - elif build_stage == 2: - # Don't include Python bindings in the wheel - pybind_out_dir = build_temp cmake_args = [ - f"-DCMAKE_INSTALL_PREFIX={install_prefix}", - f"-DMLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY={pybind_out_dir}", + f"-DCMAKE_INSTALL_PREFIX={extdir}", + f"-DMLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY={extdir}", f"-DCMAKE_BUILD_TYPE={cfg}", f"-DPython_EXECUTABLE={sys.executable}", "-DMLX_BUILD_PYTHON_BINDINGS=ON", @@ -113,8 +120,7 @@ def build_extension(self, ext: CMakeExtension) -> None: if "CMAKE_ARGS" in os.environ: cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item] - # For release wheel force building for all supported arches. - if build_stage == 2 and build_cuda: + if build_backend and build_cuda: # Last arch is always real and virtual for forward-compatibility cuda_archs = [ "75-real", @@ -191,11 +197,35 @@ def run(self): class MLXBdistWheel(bdist_wheel): def get_tag(self) -> tuple[str, str, str]: impl, abi, plat_name = super().get_tag() - if build_stage == 2: + if build_backend: impl = self.python_tag abi = "none" return (impl, abi, plat_name) + def write_wheelfile(self, *args, **kwargs) -> None: + super().write_wheelfile(*args, **kwargs) + + mlx_dir = Path(self.bdist_dir, "mlx") + + def is_backend_file(file): + if file.is_relative_to(Path(mlx_dir, "lib")): + return True + if file.is_relative_to(Path(mlx_dir, "include")): + return True + if file.is_relative_to(Path(mlx_dir, "share")): + return True + if file.suffix == ".dll": + return True + return False + + if build_frontend or build_backend: + for file in Path(self.bdist_dir).rglob("*"): + if not file.is_relative_to(mlx_dir) or not file.is_file(): + continue + bf = is_backend_file(file) + if (build_frontend and bf) or (build_backend and not bf): + file.unlink() + # Read the content of README.md with open(Path(__file__).parent / "README.md", encoding="utf-8") as f: @@ -260,24 +290,8 @@ def get_tag(self) -> tuple[str, str, str]: } install_requires = [] - # Release builds for PyPi are in two stages. - # Each stage should be run from a clean build: - # python setup.py clean --all - # - # Stage 1: - # - Triggered with `MLX_BUILD_STAGE=1` - # - Include everything except backend-specific binaries (e.g. libmlx.so, mlx.metallib, etc) - # - Wheel has Python ABI and platform tags - # - Wheel should be built for the cross-product of python version and platforms - # - Package name is mlx and it depends on subpackage in stage 2 (e.g. mlx-metal) - # Stage 2: - # - Triggered with `MLX_BUILD_STAGE=2` - # - Includes only backend-specific binaries (e.g. libmlx.so, mlx.metallib, etc) - # - Wheel has only platform tags - # - Wheel should be built only for different platforms - # - Package name is back-end specific, e.g mlx-metal - if build_stage != 2: - if build_stage == 1: + if not build_backend: + if build_frontend: install_requires.append( f'mlx-metal=={version}; platform_system == "Darwin"' )