Skip to content

Do a PyPi release for cuda on arm #2866

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
awni merged 1 commit into from
Dec 4, 2025
Merged

Do a PyPi release for cuda on arm #2866

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion .github/actions/build-cuda-release/action.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,15 @@
name: 'Build CUDA wheel'
description: 'Build CUDA wheel'

inputs:
arch:
description: 'Platform architecture tag'
required: true
type: choice
options:
- x86_64
- aarch64

runs:
using: "composite"
steps:
Expand All @@ -12,4 +21,4 @@ runs:
pip install auditwheel build patchelf setuptools
python setup.py clean --all
MLX_BUILD_STAGE=2 python -m build -w
bash python/scripts/repair_cuda.sh
bash python/scripts/repair_cuda.sh ${{ inputs.arch }}
1 change: 1 addition & 0 deletions .github/actions/setup-linux/action.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ runs:
using: "composite"
steps:
- name: Use ccache
if: ${{ runner.arch == 'x86_64' }}

@zcbenz zcbenz Dec 12, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@awni This disables ccache completely as x86_64 is not a valid value for runner.arch. Do you remember why it is disabled for arm build? Maybe we should add an input to just disable ccache for release builds?

@awni awni Dec 12, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think because it wasn't installed on the ubuntu-large-arm machine (the one we use for the release build).

uses: hendrikmuhs/ccache-action@v1.2
with:
key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}-py${{ inputs.python-version }}
Expand Down
7 changes: 6 additions & 1 deletion .github/workflows/release.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,10 @@ jobs:

build_cuda_release:
if: github.repository == 'ml-explore/mlx'
runs-on: ubuntu-22-large
strategy:
Comment thread
madrob marked this conversation as resolved.
matrix:
arch: ['x86_64', 'aarch64']
runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }}
Comment on lines +133 to +134

@madrob madrob Dec 4, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would put the conditional logic in the matrix so it's easier to see:

Suggested change
arch: ['x86_64', 'aarch64']
runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }}
include:
- host: ubuntu-22-large
arch: x86_64
- host: ubuntu-22-large-arm
arch: aarch64
runs-on: ${{ matrix.host }}

env:
PYPI_RELEASE: 1
DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
Expand All @@ -139,6 +142,8 @@ jobs:
toolkit: 'cuda-12.9'
- name: Build Python package
uses: ./.github/actions/build-cuda-release
with:
arch: ${{ matrix.arch }}
- name: Upload artifacts
uses: actions/upload-artifact@v5
with:
Expand Down
2 changes: 1 addition & 1 deletion python/scripts/repair_cuda.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
#!/bin/bash

auditwheel repair dist/* \
--plat manylinux_2_35_x86_64 \
--plat manylinux_2_35_${1} \
--exclude libcublas* \
--exclude libnvrtc* \
--exclude libcuda* \
Expand Down
Loading