fix(ext): fix parallax_ext

.github/workflows/ci.yml

@@ -6,20 +6,13 @@ on:
 
 jobs:
   test:
-    name: Test on ${{ matrix.os }}
+    name: Test on ${{ matrix.os }} / Python ${{ matrix.python-version }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
-        include:
-          - os: ubuntu-latest
-            python-version: '3.11'
-            # Install gpu (sglang) and dev dependencies for Linux
-            extras: 'gpu, dev'
-          - os: macos-latest
-            python-version: '3.11'
-            # Install mac (mlx) and dev dependencies for macOS
-            extras: 'mac, dev'
+        os: [ubuntu-latest, macos-15, macos-26]
+        python-version: ['3.11', '3.12', '3.13']
 
     steps:
     - name: Free Disk Space (Ubuntu)
@@ -43,6 +36,7 @@ jobs:
       with:
         filters: |
           src:
+            - '.github/workflows/ci.yml'
             - 'src/**'
             - 'tests/**'
             - 'pyproject.toml'
@@ -58,17 +52,21 @@ jobs:
       uses: actions/cache@v4
       with:
         path: ~/.cache/pip
-        key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('pyproject.toml') }}
+        key: ${{ matrix.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('pyproject.toml') }}
         restore-keys: |
-          ${{ runner.os }}-pip-${{ matrix.python-version }}-
-          ${{ runner.os }}-pip-
+          ${{ matrix.os }}-pip-${{ matrix.python-version }}-
+          ${{ matrix.os }}-pip-
 
     - name: Install dependencies
       if: steps.changes.outputs.src == 'true'
       run: |
         python -m pip install --upgrade pip
         # Install extras dependencies based on matrix variable
-        pip install -e ".[${{ matrix.extras }}]"
+        if [[ "${{ runner.os }}" == "Linux" ]]; then
+          pip install -e ".[gpu, dev]"
+        else
+          pip install -e ".[mac, dev]"
+        fi
 
     - name: Run Unit Tests
       if: steps.changes.outputs.src == 'true'
@@ -77,15 +75,15 @@ jobs:
         pytest tests/ -v --cov=src/parallax --cov-report=xml
 
     - name: Upload coverage to Codecov
-      if: steps.changes.outputs.src == 'true' && runner.os == 'macOS'
+      if: steps.changes.outputs.src == 'true' && matrix.os == 'macos-15' && matrix.python-version == '3.11'
       uses: codecov/codecov-action@v4
       with:
         file: ./coverage.xml
         fail_ci_if_error: false
         token: ${{ secrets.CODECOV_TOKEN }}
 
     - name: Run E2E tests (macOS only)
-      if: steps.changes.outputs.src == 'true' && runner.os == 'macOS'
+      if: steps.changes.outputs.src == 'true' && matrix.os == 'macos-26' && matrix.python-version == '3.12'
       shell: bash
       env:
         TERM: xterm-256color

src/parallax_extensions/CMakelists.txt

@@ -81,5 +81,13 @@ nanobind_add_module(
 target_link_libraries(_ext PRIVATE parallax_ext)
 
 if(BUILD_SHARED_LIBS)
-  target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path)
+  set(_PARALLAX_MLX_RPATH "@loader_path/../../mlx/lib")
+  set_target_properties(
+    parallax_ext
+    PROPERTIES BUILD_WITH_INSTALL_RPATH TRUE
+               INSTALL_RPATH "${_PARALLAX_MLX_RPATH}")
+  set_target_properties(
+    _ext
+    PROPERTIES BUILD_WITH_INSTALL_RPATH TRUE
+               INSTALL_RPATH "@loader_path;${_PARALLAX_MLX_RPATH}")
 endif()

src/parallax_extensions/ops.py

@@ -1,9 +1,5 @@
 import importlib
-import os
-import shutil
-import subprocess
 import sys
-import sysconfig
 from pathlib import Path
 from types import ModuleType
 from typing import Optional
@@ -38,93 +34,8 @@ def _build_import_error(original_error: Exception) -> ImportError:
     return ImportError(msg)
 
 
-def _build_signature() -> str:
-    try:
-        from importlib.metadata import version
-
-        mlx_version = version("mlx")
-        nanobind_version = version("nanobind")
-    except Exception:
-        mlx_version = "unknown"
-        nanobind_version = "unknown"
-
-    return "|".join(
-        [
-            sys.implementation.cache_tag or "",
-            sys.version.split()[0],
-            mx.__file__,
-            mlx_version,
-            nanobind_version,
-        ]
-    )
-
-
-def _ensure_build_tools() -> None:
-    missing = []
-    try:
-        import setuptools  # noqa: F401
-    except Exception:
-        missing.append("setuptools")
-
-    if shutil.which("cmake") is None:
-        missing.append("cmake")
-    if shutil.which("ninja") is None:
-        missing.append("ninja")
-
-    if missing:
-        subprocess.run(
-            [sys.executable, "-m", "pip", "install", *missing],
-            check=True,
-        )
-
-
-def _rebuild_for_github_actions() -> None:
-    """Build native kernels against the exact Python/MLX used by GitHub macOS CI."""
-    if os.environ.get("GITHUB_ACTIONS") != "true" or sys.platform != "darwin":
-        return
-    if os.environ.get("PARALLAX_SKIP_CI_EXTENSION_REBUILD") == "1":
-        return
-
-    package_dir = Path(__file__).resolve().parent
-    lib_dir = package_dir / "lib"
-    ext_suffix = sysconfig.get_config_var("EXT_SUFFIX") or ".so"
-    expected_ext = lib_dir / f"_ext{ext_suffix}"
-    stamp = lib_dir / f".ci-build-{sys.implementation.cache_tag or 'python'}"
-    signature = _build_signature()
-
-    if expected_ext.exists() and stamp.exists() and stamp.read_text() == signature:
-        return
-
-    _ensure_build_tools()
-
-    env = os.environ.copy()
-    env["DEBUG"] = "0"
-    cmake_args = env.get("CMAKE_ARGS", "")
-    python_arg = f"-DPython_EXECUTABLE={sys.executable}"
-    env["CMAKE_ARGS"] = f"{python_arg} {cmake_args}".strip()
-
-    log_path = Path("/tmp/parallax_ext_build.log")
-    with log_path.open("w") as log:
-        subprocess.run(
-            [sys.executable, "setup.py", "build_ext", "-j8", "--inplace"],
-            cwd=package_dir,
-            env=env,
-            stdout=log,
-            stderr=subprocess.STDOUT,
-            check=True,
-        )
-
-    stamp.write_text(signature)
-    print(f"Rebuilt parallax_extensions native kernels for CI: {expected_ext}")
-
-
 def load_extension_module() -> ModuleType:
     """Load the compiled extension module for the current Python runtime."""
-    try:
-        _rebuild_for_github_actions()
-    except Exception as exc:  # pragma: no cover - GitHub runner dependent
-        raise _build_import_error(exc) from exc
-
     try:
         # Python's import machinery selects the matching ABI-tagged binary
         # (e.g. _ext.cpython-312-*.so) from parallax_extensions/lib.

src/parallax_extensions/setup.py

@@ -1,7 +1,24 @@
+import os
+import sys
+
 from mlx import extension
 from setuptools import setup
 
+
+def _set_macos_build_defaults() -> None:
+    if sys.platform != "darwin":
+        return
+
+    deployment_target = os.environ.setdefault("MACOSX_DEPLOYMENT_TARGET", "15.0")
+    cmake_args = os.environ.get("CMAKE_ARGS", "")
+    if "CMAKE_OSX_DEPLOYMENT_TARGET" not in cmake_args:
+        os.environ["CMAKE_ARGS"] = (
+            f"-DCMAKE_OSX_DEPLOYMENT_TARGET={deployment_target} {cmake_args}"
+        ).strip()
+
+
 if __name__ == "__main__":
+    _set_macos_build_defaults()
     setup(
         name="parallax_extensions",
         version="0.0.1",