build: bind native extension to isolated Python

Author: apstenku123

cppmega_mlx/training/native_optim/CMakeLists.txt

@@ -12,16 +12,22 @@ find_package(
   Python 3.11
   COMPONENTS Interpreter Development.Module
   REQUIRED)
-execute_process(
-  COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
-  OUTPUT_STRIP_TRAILING_WHITESPACE
-  OUTPUT_VARIABLE nanobind_ROOT)
+if(NOT DEFINED nanobind_DIR)
+  execute_process(
+    COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    OUTPUT_VARIABLE nanobind_ROOT
+    COMMAND_ERROR_IS_FATAL ANY)
+endif()
 find_package(nanobind CONFIG REQUIRED)
 
-execute_process(
-  COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir
-  OUTPUT_STRIP_TRAILING_WHITESPACE
-  OUTPUT_VARIABLE MLX_ROOT)
+if(NOT DEFINED MLX_ROOT)
+  execute_process(
+    COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    OUTPUT_VARIABLE MLX_ROOT
+    COMMAND_ERROR_IS_FATAL ANY)
+endif()
 find_package(MLX CONFIG REQUIRED)
 
 add_library(cppmega_native_optim)

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=68", "mlx>=0.31"]
+requires = ["setuptools>=68", "mlx>=0.31", "nanobind>=2.4"]
 build-backend = "setuptools.build_meta"
 
 [project]

setup.py

@@ -1,8 +1,34 @@
+import os
+import sys
+
 from setuptools import setup
 
+import nanobind
 from mlx import extension
 
 
+class CppmegaCMakeBuild(extension.CMakeBuild):
+    """Build native extensions with the interpreter that invoked setuptools."""
+
+    def build_extension(self, ext: extension.CMakeExtension) -> None:
+        previous_args = os.environ.get("CMAKE_ARGS")
+        build_contract_args = (
+            f"-DPython_EXECUTABLE={sys.executable}",
+            f"-Dnanobind_DIR={nanobind.cmake_dir()}",
+            f"-DMLX_ROOT={extension._MLX_PATH}",
+        )
+        os.environ["CMAKE_ARGS"] = " ".join(
+            part for part in (previous_args, *build_contract_args) if part
+        )
+        try:
+            super().build_extension(ext)
+        finally:
+            if previous_args is None:
+                os.environ.pop("CMAKE_ARGS", None)
+            else:
+                os.environ["CMAKE_ARGS"] = previous_args
+
+
 if __name__ == "__main__":
     setup(
         ext_modules=[
@@ -11,7 +37,7 @@
                 sourcedir="cppmega_mlx/training/native_optim",
             )
         ],
-        cmdclass={"build_ext": extension.CMakeBuild},
+        cmdclass={"build_ext": CppmegaCMakeBuild},
         package_data={
             "cppmega_mlx.training.native_optim": [
                 "*.so",

tests/test_setup_native_build.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import os
+import sys
+import tomllib
+from pathlib import Path
+
+from setuptools import Distribution
+
+
+def test_native_build_dependencies_are_declared() -> None:
+    pyproject = tomllib.loads(
+        (Path(__file__).resolve().parents[1] / "pyproject.toml").read_text()
+    )
+    requirements = pyproject["build-system"]["requires"]
+
+    assert any(requirement.startswith("mlx>=") for requirement in requirements)
+    assert any(requirement.startswith("nanobind>=") for requirement in requirements)
+
+
+def test_native_build_pins_cmake_to_invoking_python(monkeypatch) -> None:
+    repo_root = Path(__file__).resolve().parents[1]
+    monkeypatch.syspath_prepend(str(repo_root))
+
+    import setup as cppmega_setup
+
+    captured: dict[str, str | None] = {}
+
+    def capture_build_extension(_command, _extension) -> None:
+        captured["cmake_args"] = os.environ.get("CMAKE_ARGS")
+
+    monkeypatch.setattr(
+        cppmega_setup.extension.CMakeBuild,
+        "build_extension",
+        capture_build_extension,
+    )
+    monkeypatch.setenv("CMAKE_ARGS", "-DEXISTING_OPTION=ON")
+
+    command = cppmega_setup.CppmegaCMakeBuild(Distribution())
+    command.build_extension(object())
+
+    cmake_args = captured["cmake_args"]
+    assert cmake_args is not None
+    assert cmake_args.startswith(
+        f"-DEXISTING_OPTION=ON -DPython_EXECUTABLE={sys.executable}"
+    )
+    assert "-Dnanobind_DIR=" in cmake_args
+    assert "-DMLX_ROOT=" in cmake_args
+    assert os.environ["CMAKE_ARGS"] == "-DEXISTING_OPTION=ON"