diff --git a/.github/actions/build-macos/action.yml b/.github/actions/build-macos/action.yml index b7a69aaca0..7bcd4e0d09 100644 --- a/.github/actions/build-macos/action.yml +++ b/.github/actions/build-macos/action.yml @@ -15,7 +15,7 @@ runs: steps: - name: Install dependencies shell: bash - run: uv pip install 'build<=1.4.2' setuptools + run: uv pip install build setuptools - name: Build wheel shell: bash diff --git a/.github/actions/build-wheel/action.yml b/.github/actions/build-wheel/action.yml index 6177bf1849..2e15efd93f 100644 --- a/.github/actions/build-wheel/action.yml +++ b/.github/actions/build-wheel/action.yml @@ -34,7 +34,7 @@ runs: - name: Install dependencies shell: bash run: | - uv pip install 'build<=1.4.2' setuptools + uv pip install build setuptools if ${{ runner.os == 'Linux' }} ; then uv pip install auditwheel patchelf fi @@ -48,7 +48,7 @@ runs: MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} run: | python setup.py clean --all - MLX_BUILD_STAGE=1 python -m build -w + MLX_BUILD_FRONTEND_PACKAGE=1 python -m build -w - name: Post-process frontend package if: inputs.build-frontend == 'true' @@ -71,7 +71,7 @@ runs: MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} run: | python setup.py clean --all - MLX_BUILD_STAGE=2 python -m build -w + MLX_BUILD_BACKEND_PACKAGE=1 python -m build -w - name: Post-process backend package if: inputs.build-backend == 'true' diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml index 3c71b13d72..947c0ff1f0 100644 --- a/.github/actions/setup/action.yml +++ b/.github/actions/setup/action.yml @@ -41,7 +41,7 @@ runs: run: | 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 @@ -50,6 +50,7 @@ runs: shell: bash run: | brew update + brew trust aws/tap # suppress warning in github actions brew install openmpi xcodebuild -showComponent MetalToolchain sysctl -a | grep machdep.cpu 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..756bf85f08 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", "") @@ -77,9 +92,7 @@ def finalize_options(self) -> None: self.build_temp = os.path.dirname(self.build_temp) def build_extension(self, ext: CMakeExtension) -> None: - # Must be in this form due to bug in .resolve() only fixed in Python 3.10+ - ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name) # type: ignore[no-untyped-call] - extdir = ext_fullpath.parent.resolve() + extdir = self._get_ext_dir(ext) debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug cfg = "Debug" if debug else "Release" @@ -88,17 +101,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 +118,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", @@ -186,16 +190,50 @@ def run(self): ["cmake", "--install", build_temp, "--component", "core_stub"], check=True, ) + # Copy the type stubs to extdir so they are included in wheels. + stubs_dir = Path("python/mlx/core") + if stubs_dir.exists(): + extdir = self._get_ext_dir(ext) + self.copy_tree(stubs_dir, extdir / "core") + + def _get_ext_dir(self, ext): + # Must be in this form due to bug in .resolve() only fixed in Python 3.10+ + ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name) # type: ignore[no-untyped-call] + return ext_fullpath.parent.resolve() 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 +298,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"' )