Build mlx wheels (riscv64) #4
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # SPDX-FileCopyrightText: 2026 The RISE Project | |
| # SPDX-License-Identifier: MIT | |
| --- | |
| # Mirrors the Linux "cpu" leg of upstream's .github/workflows/release.yml | |
| # (build_frontend + build_backend jobs, via .github/actions/build-wheel). Upstream | |
| # splits every release into a per-interpreter "mlx" frontend wheel (thin nanobind | |
| # bindings) and one per-platform "mlx-cpu" backend wheel (the compiled libmlx.so); | |
| # both are built here since a working install needs both. | |
| name: Build mlx wheels (riscv64) | |
| on: | |
| workflow_dispatch: | |
| inputs: | |
| version: | |
| description: 'mlx version to build (git tag without leading v, e.g. 0.32.2)' | |
| required: true | |
| default: '0.32.2' | |
| pull_request: | |
| paths: | |
| - '.github/workflows/build-mlx.yml' | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ inputs.version || '0.32.2' }}-${{ github.head_ref || github.run_id }} | |
| cancel-in-progress: true | |
| permissions: | |
| contents: read | |
| env: | |
| MLX_VERSION: ${{ inputs.version || '0.32.2' }} | |
| MANYLINUX_RISCV64_IMAGE: quay.io/pypa/manylinux_2_39_riscv64 | |
| jobs: | |
| setup: | |
| uses: $/.github/workflows/_setup.yml | |
| build_backend: | |
| needs: [setup] | |
| name: Build mlx-cpu ${{ inputs.version || '0.32.2' }}-manylinux_riscv64 | |
| runs-on: ubuntu-24.04-riscv | |
| timeout-minutes: 360 | |
| steps: | |
| - name: Checkout ml-explore/mlx v${{ env.MLX_VERSION }} | |
| uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 | |
| with: | |
| repository: ml-explore/mlx | |
| ref: v${{ env.MLX_VERSION }} | |
| path: mlx | |
| submodules: recursive | |
| persist-credentials: false | |
| - name: Build mlx-cpu wheel | |
| run: | | |
| mkdir -p output | |
| docker run -i --rm \ | |
| -v "$PWD/mlx:/mlx" \ | |
| -v "$PWD/output:/output" \ | |
| "$MANYLINUX_RISCV64_IMAGE" bash -s <<'MLX_BUILD_EOF' | |
| #!/usr/bin/env bash | |
| set -euxo pipefail | |
| git config --global --add safe.directory '*' | |
| dnf install -y --setopt=install_weak_deps=False openblas-devel lapack-devel blas-devel ninja-build | |
| # auditwheel vendors libopenblas.so.0 (dynamically links libgfortran, see the | |
| # gpl_sources job below), so its licence has to travel with the wheel too. | |
| cp /usr/share/licenses/openblas/LICENSE /mlx/LICENSE.openblas | |
| PY=/opt/python/cp312-cp312/bin/python | |
| "$PY" -m pip install -U --no-cache-dir build auditwheel patchelf | |
| cd /mlx | |
| # CMAKE_INSTALL_LIBDIR=lib: Rocky's GNUInstallDirs default (lib64) makes | |
| # setup.py's MLXBdistWheel.is_backend_file() (which only recognizes "lib") | |
| # drop the compiled libmlx.so from the mlx-cpu wheel entirely; forcing "lib" | |
| # avoids patching setup.py. | |
| export CMAKE_ARGS="-DMLX_BUILD_METAL=OFF -DMLX_BUILD_CUDA=OFF -DCMAKE_INSTALL_LIBDIR=lib -G Ninja" | |
| export MLX_BUILD_BACKEND_PACKAGE=1 | |
| "$PY" -m build -w | |
| auditwheel repair dist/mlx_cpu-*.whl --plat manylinux_2_39_riscv64 --wheel-dir /output | |
| MLX_BUILD_EOF | |
| - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 | |
| with: | |
| name: mlx_cpu-${{ env.MLX_VERSION }}-manylinux_riscv64 | |
| path: output/*.whl | |
| if-no-files-found: error | |
| build_frontend: | |
| needs: [setup] | |
| name: Build mlx ${{ inputs.version || '0.32.2' }} ${{ matrix.python }}-manylinux_riscv64 | |
| runs-on: ubuntu-24.04-riscv | |
| timeout-minutes: 360 | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| python: ["cp312", "cp313", "cp314", "cp314t"] | |
| steps: | |
| - name: Checkout ml-explore/mlx v${{ env.MLX_VERSION }} | |
| uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 | |
| with: | |
| repository: ml-explore/mlx | |
| ref: v${{ env.MLX_VERSION }} | |
| path: mlx | |
| submodules: recursive | |
| persist-credentials: false | |
| - name: Build mlx wheel | |
| run: | | |
| mkdir -p output | |
| docker run -i --rm \ | |
| -v "$PWD/mlx:/mlx" \ | |
| -v "$PWD/output:/output" \ | |
| -e PYTAG="${{ matrix.python }}" \ | |
| "$MANYLINUX_RISCV64_IMAGE" bash -s <<'MLX_BUILD_EOF' | |
| #!/usr/bin/env bash | |
| set -euxo pipefail | |
| git config --global --add safe.directory '*' | |
| dnf install -y --setopt=install_weak_deps=False openblas-devel lapack-devel blas-devel ninja-build | |
| # The image names the free-threaded 3.14 dir cp314-cp314t, not cp314t-cp314t. | |
| PYDIR="$PYTAG-$PYTAG" | |
| if [ "$PYTAG" = "cp314t" ]; then | |
| PYDIR="cp314-cp314t" | |
| fi | |
| PY="/opt/python/$PYDIR/bin/python" | |
| "$PY" -m pip install -U --no-cache-dir build auditwheel patchelf | |
| cd /mlx | |
| export CMAKE_ARGS="-DMLX_BUILD_METAL=OFF -DMLX_BUILD_CUDA=OFF -DCMAKE_INSTALL_LIBDIR=lib -G Ninja" | |
| export MLX_BUILD_FRONTEND_PACKAGE=1 | |
| "$PY" -m build -w | |
| # libmlx.so ships separately in mlx-cpu, so it (and anything only reachable | |
| # through it, e.g. libopenblas) is excluded rather than vendored here; | |
| # --only-plat then skips the policy check that exclusion would otherwise fail. | |
| auditwheel repair dist/mlx-*.whl --plat manylinux_2_39_riscv64 --exclude 'libmlx.so*' --only-plat --wheel-dir /output | |
| MLX_BUILD_EOF | |
| - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 | |
| with: | |
| name: mlx-${{ env.MLX_VERSION }}-${{ matrix.python }}-manylinux_riscv64 | |
| path: output/*.whl | |
| if-no-files-found: error | |
| test_wheels: | |
| needs: [setup, build_backend, build_frontend] | |
| name: Test mlx ${{ inputs.version || '0.32.2' }} ${{ matrix.python }}-manylinux_riscv64 | |
| runs-on: ubuntu-24.04-riscv | |
| timeout-minutes: 120 | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| python: ["cp312", "cp313", "cp314", "cp314t"] | |
| steps: | |
| - name: Checkout ml-explore/mlx v${{ env.MLX_VERSION }} | |
| uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 | |
| with: | |
| repository: ml-explore/mlx | |
| ref: v${{ env.MLX_VERSION }} | |
| path: mlx | |
| persist-credentials: false | |
| - name: Download mlx-cpu wheel | |
| uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 | |
| with: | |
| name: mlx_cpu-${{ env.MLX_VERSION }}-manylinux_riscv64 | |
| path: wheels | |
| - name: Download mlx wheel | |
| uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 | |
| with: | |
| name: mlx-${{ env.MLX_VERSION }}-${{ matrix.python }}-manylinux_riscv64 | |
| path: wheels | |
| - name: Test mlx wheel | |
| run: | | |
| docker run -i --rm \ | |
| -v "$PWD/mlx:/mlx" \ | |
| -v "$PWD/wheels:/wheels" \ | |
| -e PYTAG="${{ matrix.python }}" \ | |
| "$MANYLINUX_RISCV64_IMAGE" bash -s <<'MLX_TEST_EOF' | |
| #!/usr/bin/env bash | |
| set -euxo pipefail | |
| PYDIR="$PYTAG-$PYTAG" | |
| if [ "$PYTAG" = "cp314t" ]; then | |
| PYDIR="cp314-cp314t" | |
| fi | |
| PY="/opt/python/$PYDIR/bin/python" | |
| # torch/ml_dtypes are optional, import-guarded test deps upstream also | |
| # runs without on any CI leg that lacks them; only numpy is required. | |
| "$PY" -m pip install -U --no-cache-dir \ | |
| --extra-index-url https://pypi.riseproject.dev/simple/ --only-binary numpy numpy | |
| "$PY" -m pip install /wheels/mlx_cpu-*.whl /wheels/mlx-*.whl | |
| cd /mlx/python/tests | |
| if [ "$PYTAG" = "cp314t" ]; then | |
| # test_compile_release_on_another_thread traces a function on one thread | |
| # and asserts that dropping its last reference (gc.collect()) on a | |
| # different thread makes a later same-id compile re-trace. MLX 0.32 made | |
| # its compiler cache thread_local with a weak_ptr cross-thread erase path | |
| # specifically for this (ml-explore/mlx#4377); the assertion assumes that | |
| # erase always completes before the next compile call observes it, which | |
| # cp314t's true parallelism does not guarantee the way GIL-serialized | |
| # builds do -- a free-threading test-design gap (reproduced twice, | |
| # deterministic on this leg), not a riscv64 or correctness issue. | |
| "$PY" - <<'PYEOF' | |
| import sys | |
| import unittest | |
| DESELECTED = { | |
| "test_compile.TestCompile.test_compile_release_on_another_thread", | |
| } | |
| def prune(suite): | |
| kept = unittest.TestSuite() | |
| for t in suite: | |
| if isinstance(t, unittest.TestSuite): | |
| sub = prune(t) | |
| if sub.countTestCases(): | |
| kept.addTest(sub) | |
| elif t.id() not in DESELECTED: | |
| kept.addTest(t) | |
| return kept | |
| suite = prune(unittest.TestLoader().discover(".")) | |
| result = unittest.TextTestRunner(verbosity=2, failfast=True).run(suite) | |
| sys.exit(0 if result.wasSuccessful() else 1) | |
| PYEOF | |
| else | |
| "$PY" run.py -v --failfast | |
| fi | |
| MLX_TEST_EOF | |
| gpl_sources: | |
| needs: [setup] | |
| name: Collect GPL sources | |
| runs-on: ubuntu-24.04-riscv | |
| steps: | |
| - name: Checkout python-wheels | |
| uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 | |
| with: | |
| persist-credentials: false | |
| # OpenBLAS is compiled with gfortran, so auditwheel vendors the image's | |
| # libgfortran into mlx_cpu.libs/. | |
| - uses: ./actions/collect-gpl-sources | |
| with: | |
| image: ${{ env.MANYLINUX_RISCV64_IMAGE }} | |
| packages: gcc | |
| output: gpl-sources.tar | |
| - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 | |
| with: | |
| name: mlx-${{ env.MLX_VERSION }}-gpl-sources | |
| path: gpl-sources.tar | |
| if-no-files-found: error | |
| publish_frontend: | |
| name: Publish mlx ${{ inputs.version || '0.32.2' }} | |
| # No GPL sources here: --only-plat repair vendors nothing into the frontend wheel. | |
| needs: [setup, build_frontend, test_wheels] | |
| permissions: | |
| contents: write | |
| pull-requests: write | |
| uses: $/.github/workflows/_publish-wheel.yml | |
| with: | |
| artifact-pattern: mlx-${{ inputs.version || '0.32.2' }}-cp3*-manylinux_riscv64 | |
| publish_backend: | |
| name: Publish mlx-cpu ${{ inputs.version || '0.32.2' }} | |
| needs: [setup, build_backend, gpl_sources, test_wheels] | |
| permissions: | |
| contents: write | |
| pull-requests: write | |
| uses: $/.github/workflows/_publish-wheel.yml | |
| with: | |
| artifact-pattern: mlx_cpu-${{ inputs.version || '0.32.2' }}-manylinux_riscv64 | |
| gpl-sources-artifact: mlx-${{ inputs.version || '0.32.2' }}-gpl-sources | |
| gpl-sources-description: gcc |