ci: isolate MLX native ABI contract

Author: apstenku123

.github/workflows/e2e-matrix.yml

@@ -50,13 +50,21 @@ jobs:
           base_python=/Volumes/external/sources/cppmega.mlx/.venv/bin/python
           test -x "$base_python"
           job_venv="$(mktemp -d "$RUNNER_TEMP/cppmega-mlx-e2e-${GITHUB_RUN_ID}-${GITHUB_JOB}.XXXXXX")"
-          "$base_python" -m venv --system-site-packages "$job_venv"
+          "$base_python" -m venv "$job_venv"
           echo "VBGUI_E2E_PYTHON=$job_venv/bin/python" >> "$GITHUB_ENV"
           echo "VBGUI_E2E_VENV=$job_venv" >> "$GITHUB_ENV"
+          echo "PYTHONPATH=" >> "$GITHUB_ENV"
+          echo "PYTHONNOUSERSITE=1" >> "$GITHUB_ENV"
       - name: Python deps
         run: |
           "$VBGUI_E2E_PYTHON" -m pip install --disable-pip-version-check \
             -e ".[gui,parquet,widget]"
+          "$VBGUI_E2E_PYTHON" - <<'PY'
+          from cppmega_mlx.training.native_optim import status
+          native_status = status()
+          if not native_status.get("available", False):
+              raise RuntimeError(native_status.get("reason", "native optimizer extension unavailable"))
+          PY
       - name: Generate fixtures
         run: |
           "$VBGUI_E2E_PYTHON" tests/fixtures/build_e2e_matrix.py
@@ -109,12 +117,20 @@ jobs:
           base_python=/Volumes/external/sources/cppmega.mlx/.venv/bin/python
           test -x "$base_python"
           job_venv="$(mktemp -d "$RUNNER_TEMP/cppmega-mlx-e2e-${GITHUB_RUN_ID}-${GITHUB_JOB}.XXXXXX")"
-          "$base_python" -m venv --system-site-packages "$job_venv"
+          "$base_python" -m venv "$job_venv"
           echo "VBGUI_E2E_PYTHON=$job_venv/bin/python" >> "$GITHUB_ENV"
           echo "VBGUI_E2E_VENV=$job_venv" >> "$GITHUB_ENV"
+          echo "PYTHONPATH=" >> "$GITHUB_ENV"
+          echo "PYTHONNOUSERSITE=1" >> "$GITHUB_ENV"
       - run: |
           "$VBGUI_E2E_PYTHON" -m pip install --disable-pip-version-check \
             -e ".[gui,parquet,widget]"
+          "$VBGUI_E2E_PYTHON" - <<'PY'
+          from cppmega_mlx.training.native_optim import status
+          native_status = status()
+          if not native_status.get("available", False):
+              raise RuntimeError(native_status.get("reason", "native optimizer extension unavailable"))
+          PY
       - run: |
           "$VBGUI_E2E_PYTHON" tests/fixtures/build_e2e_matrix.py
       - working-directory: vbgui
@@ -162,12 +178,20 @@ jobs:
           base_python=/Volumes/external/sources/cppmega.mlx/.venv/bin/python
           test -x "$base_python"
           job_venv="$(mktemp -d "$RUNNER_TEMP/cppmega-mlx-e2e-${GITHUB_RUN_ID}-${GITHUB_JOB}.XXXXXX")"
-          "$base_python" -m venv --system-site-packages "$job_venv"
+          "$base_python" -m venv "$job_venv"
           echo "VBGUI_E2E_PYTHON=$job_venv/bin/python" >> "$GITHUB_ENV"
           echo "VBGUI_E2E_VENV=$job_venv" >> "$GITHUB_ENV"
+          echo "PYTHONPATH=" >> "$GITHUB_ENV"
+          echo "PYTHONNOUSERSITE=1" >> "$GITHUB_ENV"
       - run: |
           "$VBGUI_E2E_PYTHON" -m pip install --disable-pip-version-check \
             -e ".[gui,parquet,widget]"
+          "$VBGUI_E2E_PYTHON" - <<'PY'
+          from cppmega_mlx.training.native_optim import status
+          native_status = status()
+          if not native_status.get("available", False):
+              raise RuntimeError(native_status.get("reason", "native optimizer extension unavailable"))
+          PY
       - run: |
           "$VBGUI_E2E_PYTHON" tests/fixtures/build_e2e_matrix.py
       - working-directory: vbgui

cppmega_mlx/training/native_optim/__init__.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import mlx.core as _mlx_core  # noqa: F401 - loads the matching MLX dylib first.
+
 try:
     from cppmega_mlx.training.native_optim._ext import (  # type: ignore[attr-defined]
         fused_adam8bit_step,

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=68", "mlx>=0.31", "nanobind>=2.4"]
+requires = ["setuptools>=68", "mlx==0.32.0", "nanobind>=2.4"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -8,7 +8,7 @@ version = "0.1.0"
 description = "MLX-native local training port of cppmega model semantics"
 requires-python = ">=3.11"
 dependencies = [
-    "mlx>=0.31",
+    "mlx==0.32.0",
     "mlx-lm>=0.31",
     "numpy>=1.26",
     "safetensors>=0.7",

setup.py

@@ -1,4 +1,5 @@
 import os
+import subprocess
 import sys
 from pathlib import Path
 
@@ -8,6 +9,34 @@
 from mlx import extension
 
 
+def _target_mlx_cmake_root() -> Path:
+    probe = subprocess.run(
+        [sys.executable, "-m", "mlx", "--cmake-dir"],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    if probe.returncode == 0 and probe.stdout.strip():
+        runtime_root = Path(probe.stdout.strip()).resolve()
+        runtime_config = runtime_root / "share" / "cmake" / "MLX" / "MLXConfig.cmake"
+        if runtime_config.is_file():
+            return runtime_root
+
+    build_root = Path(extension.__file__).resolve().parent
+    build_config = build_root / "share" / "cmake" / "MLX" / "MLXConfig.cmake"
+    if not build_config.is_file():
+        raise RuntimeError(
+            "neither the target interpreter nor the isolated build dependency "
+            "provides MLXConfig.cmake"
+        )
+    print(
+        "cppmega build: target interpreter has no MLX CMake package; "
+        f"using isolated build dependency at {build_root}",
+        file=sys.stderr,
+    )
+    return build_root
+
+
 class CppmegaCMakeBuild(extension.CMakeBuild):
     """Build native extensions with the interpreter that invoked setuptools."""
 
@@ -16,7 +45,7 @@ def build_extension(self, ext: extension.CMakeExtension) -> None:
         build_contract_args = (
             f"-DPython_EXECUTABLE={sys.executable}",
             f"-Dnanobind_DIR={nanobind.cmake_dir()}",
-            f"-DMLX_ROOT={Path(extension.__file__).resolve().parent}",
+            f"-DMLX_ROOT={_target_mlx_cmake_root()}",
         )
         os.environ["CMAKE_ARGS"] = " ".join(
             part for part in (previous_args, *build_contract_args) if part

tests/test_setup_native_build.py

@@ -4,6 +4,7 @@
 import sys
 import tomllib
 from pathlib import Path
+from types import SimpleNamespace
 
 from setuptools import Distribution
 
@@ -14,8 +15,9 @@ def test_native_build_dependencies_are_declared() -> None:
     )
     requirements = pyproject["build-system"]["requires"]
 
-    assert any(requirement.startswith("mlx>=") for requirement in requirements)
+    assert "mlx==0.32.0" in requirements
     assert any(requirement.startswith("nanobind>=") for requirement in requirements)
+    assert "mlx==0.32.0" in pyproject["project"]["dependencies"]
 
 
 def test_native_build_pins_cmake_to_invoking_python(monkeypatch) -> None:
@@ -47,3 +49,27 @@ def capture_build_extension(_command, _extension) -> None:
     assert "-Dnanobind_DIR=" in cmake_args
     assert "-DMLX_ROOT=" in cmake_args
     assert os.environ["CMAKE_ARGS"] == "-DEXISTING_OPTION=ON"
+
+
+def test_native_build_prefers_target_interpreter_mlx(tmp_path, monkeypatch) -> None:
+    repo_root = Path(__file__).resolve().parents[1]
+    monkeypatch.syspath_prepend(str(repo_root))
+
+    import setup as cppmega_setup
+
+    runtime_root = tmp_path / "runtime-mlx"
+    config = runtime_root / "share" / "cmake" / "MLX" / "MLXConfig.cmake"
+    config.parent.mkdir(parents=True)
+    config.write_text("# fixture\n")
+
+    monkeypatch.setattr(
+        cppmega_setup.subprocess,
+        "run",
+        lambda *args, **kwargs: SimpleNamespace(
+            returncode=0,
+            stdout=f"{runtime_root}\n",
+            stderr="",
+        ),
+    )
+
+    assert cppmega_setup._target_mlx_cmake_root() == runtime_root.resolve()

tests/test_workflow_runner_policy.py

@@ -83,14 +83,19 @@ def test_macos_e2e_uses_an_isolated_job_venv() -> None:
     assert "pip install --upgrade pip" not in workflow
     assert 'job_venv="$RUNNER_TEMP/cppmega-mlx-e2e-venv"' not in workflow
     assert workflow.count('mktemp -d "$RUNNER_TEMP/cppmega-mlx-e2e-') == 3
-    assert workflow.count("-m venv --system-site-packages") == 3
+    assert "--system-site-packages" not in workflow
+    assert workflow.count('"$base_python" -m venv "$job_venv"') == 3
+    assert workflow.count('echo "PYTHONPATH=" >> "$GITHUB_ENV"') == 3
+    assert workflow.count('echo "PYTHONNOUSERSITE=1" >> "$GITHUB_ENV"') == 3
     assert "--no-build-isolation" not in workflow
     assert workflow.count('-e ".[gui,parquet,widget]"') == 3
     assert workflow.count('echo "VBGUI_E2E_PYTHON=$job_venv/bin/python"') == 3
     assert workflow.count('rm -rf "$VBGUI_E2E_VENV"') == 3
+    assert workflow.count("native optimizer extension unavailable") == 3
 
 
 def test_build_backend_declares_mlx_imported_by_setup_py() -> None:
     config = tomllib.loads((REPO_ROOT / "pyproject.toml").read_text(encoding="utf-8"))
 
-    assert "mlx>=0.31" in config["build-system"]["requires"]
+    assert "mlx==0.32.0" in config["build-system"]["requires"]
+    assert "mlx==0.32.0" in config["project"]["dependencies"]