diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 12dfad1e..b96b3168 100644 --- a/.github/workflows/ci.yml +++ b/.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,7 +75,7 @@ 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 @@ -85,7 +83,7 @@ jobs: 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 diff --git a/src/parallax_extensions/CMakelists.txt b/src/parallax_extensions/CMakelists.txt index e5ceb040..fe39262f 100755 --- a/src/parallax_extensions/CMakelists.txt +++ b/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() diff --git a/src/parallax_extensions/lib/_ext.cpython-311-darwin.so b/src/parallax_extensions/lib/_ext.cpython-311-darwin.so index a2615e87..fe80c883 100755 Binary files a/src/parallax_extensions/lib/_ext.cpython-311-darwin.so and b/src/parallax_extensions/lib/_ext.cpython-311-darwin.so differ diff --git a/src/parallax_extensions/lib/_ext.cpython-312-darwin.so b/src/parallax_extensions/lib/_ext.cpython-312-darwin.so index 3188162c..42b0a688 100755 Binary files a/src/parallax_extensions/lib/_ext.cpython-312-darwin.so and b/src/parallax_extensions/lib/_ext.cpython-312-darwin.so differ diff --git a/src/parallax_extensions/lib/_ext.cpython-313-darwin.so b/src/parallax_extensions/lib/_ext.cpython-313-darwin.so index 67ca7c75..32bc2863 100755 Binary files a/src/parallax_extensions/lib/_ext.cpython-313-darwin.so and b/src/parallax_extensions/lib/_ext.cpython-313-darwin.so differ diff --git a/src/parallax_extensions/lib/libparallax_ext.dylib b/src/parallax_extensions/lib/libparallax_ext.dylib index 09d458cd..3925a2c9 100755 Binary files a/src/parallax_extensions/lib/libparallax_ext.dylib and b/src/parallax_extensions/lib/libparallax_ext.dylib differ diff --git a/src/parallax_extensions/lib/parallax_ext.metallib b/src/parallax_extensions/lib/parallax_ext.metallib index f551bf6a..9584de0e 100644 Binary files a/src/parallax_extensions/lib/parallax_ext.metallib and b/src/parallax_extensions/lib/parallax_ext.metallib differ diff --git a/src/parallax_extensions/ops.py b/src/parallax_extensions/ops.py index 9e2fe155..e5a23902 100644 --- a/src/parallax_extensions/ops.py +++ b/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. diff --git a/src/parallax_extensions/setup.py b/src/parallax_extensions/setup.py index 1463dfaf..8bfd717b 100755 --- a/src/parallax_extensions/setup.py +++ b/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",