diff --git a/.github/actions/build-linux-release/action.yml b/.github/actions/build-linux-release/action.yml index 85a0125e..c8dec00b 100644 --- a/.github/actions/build-linux-release/action.yml +++ b/.github/actions/build-linux-release/action.yml @@ -18,12 +18,9 @@ inputs: runs: using: "composite" steps: - - name: Generate package stubs + - name: Build MLX shell: bash - run: | - pip install -e ".[dev]" -v - pip install typing_extensions - python setup.py generate_stubs + run: pip install -e . -v - name: Build Python package shell: bash diff --git a/.github/actions/build-linux/action.yml b/.github/actions/build-linux/action.yml index 337b94a2..21ec6304 100644 --- a/.github/actions/build-linux/action.yml +++ b/.github/actions/build-linux/action.yml @@ -21,19 +21,13 @@ runs: if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then # There is no GPU in arm64 runner, use a common arch. CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=90a" - # Can not build tests when the built executables can not run. - CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF" + # Can not build tests and stubs when the built executables can not run. + CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF -DMLX_BUILD_PYTHON_STUBS=OFF" fi pip install --no-build-isolation -e ".[dev]" -v # Pass the CMAKE_ARGS to following steps. echo CMAKE_ARGS="$CMAKE_ARGS" >> $GITHUB_OUTPUT - - name: Generate package stubs - shell: sh - run: | - pip install typing_extensions - python setup.py generate_stubs - - name: Build CPP only shell: bash run: | diff --git a/.github/actions/build-macos/action.yml b/.github/actions/build-macos/action.yml index 87bac593..067fb8d5 100644 --- a/.github/actions/build-macos/action.yml +++ b/.github/actions/build-macos/action.yml @@ -11,15 +11,9 @@ runs: shell: bash -l {0} run: | pip install --upgrade pip - pip install cmake setuptools nanobind==2.10.2 + pip install cmake setuptools typing_extensions pip install -e . -v - - name: Generate package stubs - shell: bash -l {0} - run: | - pip install typing_extensions - python setup.py generate_stubs - - name: Install tests dependencies shell: bash -l {0} run: | diff --git a/.github/actions/setup-linux/action.yml b/.github/actions/setup-linux/action.yml index ab685154..08bb528f 100644 --- a/.github/actions/setup-linux/action.yml +++ b/.github/actions/setup-linux/action.yml @@ -48,9 +48,9 @@ runs: echo "::group::Setup Python venv" python -m venv .venv source .venv/bin/activate - pip install setuptools cmake nanobind==2.10.2 + pip install setuptools cmake typing_extensions echo PATH=$PATH >> $GITHUB_ENV - # Make cmake search .venv for nanobind + # Search python packages in .venv echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV echo "::endgroup::" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index b5ff734f..18db54c8 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -98,13 +98,8 @@ jobs: shell: bash -l {0} run: | pip install --upgrade pip - pip install cmake setuptools nanobind==2.10.2 + pip install cmake setuptools typing_extensions pip install -e . -v - - name: Generate package stubs - shell: bash -l {0} - run: | - pip install typing_extensions - python setup.py generate_stubs - name: Build macOS 14 package uses: ./.github/actions/build-macos-release with: diff --git a/CMakeLists.txt b/CMakeLists.txt index 705e15a1..0ed30932 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -41,6 +41,7 @@ option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF) option(MLX_BUILD_GGUF "Include support for GGUF format" ON) option(MLX_BUILD_SAFETENSORS "Include support for safetensors format" ON) option(MLX_BUILD_BLAS_FROM_SOURCE "Build OpenBLAS from source code" OFF) +option(MLX_BUILD_PYTHON_STUBS "Build stub files for python bindings" ON) option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF) option(MLX_USE_CCACHE "Use CCache for compilation cache when available" ON) option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF) @@ -338,11 +339,13 @@ if(MLX_BUILD_PYTHON_BINDINGS) Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED) - execute_process( - COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir - OUTPUT_STRIP_TRAILING_WHITESPACE - OUTPUT_VARIABLE nanobind_ROOT) - find_package(nanobind CONFIG REQUIRED) + FetchContent_Declare( + nanobind + GIT_REPOSITORY https://github.com/wjakob/nanobind.git + GIT_TAG v2.10.2 + GIT_SHALLOW TRUE + EXCLUDE_FROM_ALL) + FetchContent_MakeAvailable(nanobind) add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/python/src) endif() diff --git a/docs/src/install.rst b/docs/src/install.rst index 9c72fdb4..bb7ef1f1 100644 --- a/docs/src/install.rst +++ b/docs/src/install.rst @@ -128,13 +128,6 @@ Run the tests with: python -m unittest discover python/tests -Optional: Install stubs to enable auto completions and type checking from your -IDE: - -.. code-block:: shell - - python setup.py generate_stubs - C++ API ^^^^^^^ diff --git a/pyproject.toml b/pyproject.toml index 60b5f646..45675ae8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [build-system] requires = [ "setuptools>=80", - "nanobind==2.10.2", "cmake>=3.25", + "typing_extensions", ] build-backend = "setuptools.build_meta" diff --git a/python/src/CMakeLists.txt b/python/src/CMakeLists.txt index 7053e61c..1a70bf4a 100644 --- a/python/src/CMakeLists.txt +++ b/python/src/CMakeLists.txt @@ -29,6 +29,37 @@ nanobind_add_module( ${CMAKE_CURRENT_SOURCE_DIR}/trees.cpp ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp) +if(MLX_BUILD_PYTHON_STUBS) + nanobind_add_stub( + core_stub + # Run stubgen -m mlx.core -i python -p _stub_patterns.txt -o python/mlx + RECURSIVE + MODULE + "mlx.core" + PYTHON_PATH + "${CMAKE_CURRENT_SOURCE_DIR}/.." + PATTERN_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/_stub_patterns.txt" + OUTPUT_PATH + "${CMAKE_CURRENT_SOURCE_DIR}/../mlx" + # Note that the list is passed to cmake for dependency managment and not + # used by stubgen. + OUTPUT + "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/__init__.pyi" + "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/cuda.pyi" + "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/distributed.pyi" + "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/fast.pyi" + "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/fft.pyi" + "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/linalg.pyi" + "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/metal.pyi" + "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/random.pyi" + # Make this an optional installable component. + EXCLUDE_FROM_ALL + INSTALL_TIME + COMPONENT + core_stub) +endif() + if(NOT MLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY) if(NOT CMAKE_LIBRARY_OUTPUT_DIRECTORY) set(MLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}) diff --git a/setup.py b/setup.py index 61aae881..00c75144 100644 --- a/setup.py +++ b/setup.py @@ -8,7 +8,7 @@ import subprocess from functools import partial from pathlib import Path -from setuptools import Command, Extension, find_namespace_packages, setup +from setuptools import Extension, find_namespace_packages, setup from setuptools.command.bdist_wheel import bdist_wheel from setuptools.command.build_ext import build_ext @@ -155,43 +155,25 @@ class CMakeBuild(build_ext): def run(self): super().run() + ext = next(ext for ext in self.extensions if ext.name == "mlx.core") + # Based on https://github.com/pypa/setuptools/blob/main/setuptools/command/build_ext.py#L102 if self.inplace: - for ext in self.extensions: - if ext.name == "mlx.core": - # Resolve inplace package dir - build_py = self.get_finalized_command("build_py") - inplace_file, regular_file = self._get_inplace_equivalent( - build_py, ext - ) + # Resolve inplace package dir + build_py = self.get_finalized_command("build_py") + inplace_file, regular_file = self._get_inplace_equivalent(build_py, ext) - inplace_dir = str(Path(inplace_file).parent.resolve()) - regular_dir = str(Path(regular_file).parent.resolve()) + inplace_dir = str(Path(inplace_file).parent.resolve()) + regular_dir = str(Path(regular_file).parent.resolve()) - self.copy_tree(regular_dir, inplace_dir) + self.copy_tree(regular_dir, inplace_dir) - -class GenerateStubs(Command): - user_options = [] - - def initialize_options(self): - pass - - def finalize_options(self): - pass - - def run(self) -> None: - out_path = "python/mlx" - stub_cmd = [ - "python", - "-m", - "nanobind.stubgen", - "-m", - "mlx.core", - "-p", - "python/mlx/_stub_patterns.txt", - ] - subprocess.run(stub_cmd + ["-r", "-O", out_path]) + # Build type stubs. + build_temp = Path(self.build_temp) / ext.name + subprocess.run( + ["cmake", "--install", build_temp, "--component", "core_stub"], + check=True, + ) class MLXBdistWheel(bdist_wheel): @@ -243,7 +225,6 @@ if __name__ == "__main__": ext_modules=[CMakeExtension("mlx.core")], cmdclass={ "build_ext": CMakeBuild, - "generate_stubs": GenerateStubs, "bdist_wheel": MLXBdistWheel, }, )