fix: Fix CI-wheels to support arm64 [WIP] #11
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
| name: CI - Build CUDA Wheels | |
| on: | |
| pull_request: | |
| push: | |
| branches: | |
| - main | |
| workflow_dispatch: | |
| jobs: | |
| cuda-wheels: | |
| runs-on: ${{ matrix.runner }} | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| include: | |
| # - runner: ubuntu-latest | |
| # arch: amd64 | |
| - platform: linux/arm64 | |
| runner: vllm-runner-arm | |
| os: ubuntu | |
| base_image_suffix: ubuntu24.04 | |
| image_suffix: "-ubuntu" | |
| steps: | |
| - name: Checkout source | |
| uses: actions/checkout@v4 | |
| - name: Set up Python | |
| uses: actions/setup-python@v5 | |
| with: | |
| python-version: "3.12" | |
| - name: Install system dependencies | |
| run: | | |
| sudo apt-get update | |
| sudo apt-get install -y libnuma-dev | |
| - name: Install CUDA toolkit | |
| shell: bash | |
| run: | | |
| ARCH=$(uname -m) | |
| if [[ "$ARCH" == "aarch64" ]]; then | |
| REPO_ARCH="arm64" | |
| else | |
| REPO_ARCH="x86_64" | |
| fi | |
| wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/${REPO_ARCH}/cuda-keyring_1.1-1_all.deb | |
| sudo dpkg -i cuda-keyring_1.1-1_all.deb | |
| sudo apt-get update | |
| sudo apt-get install -y cuda-toolkit-12-8 | |
| - name: Set CUDA environment | |
| run: | | |
| echo "CUDA_HOME=/usr/local/cuda" >> $GITHUB_ENV | |
| echo "/usr/local/cuda/bin" >> $GITHUB_PATH | |
| - name: Clean build artifacts | |
| working-directory: kv_connectors/llmd_fs_backend | |
| run: | | |
| rm -rf build dist *.egg-info wheels | |
| find . -name "*.so" -delete || true | |
| - name: Install build deps | |
| working-directory: kv_connectors/llmd_fs_backend | |
| run: | | |
| python -m pip install -U pip setuptools wheel build ninja | |
| python -m pip install -U torch==2.9.1 | |
| - name: Build wheel | |
| working-directory: kv_connectors/llmd_fs_backend | |
| env: | |
| TORCH_CUDA_ARCH_LIST: "7.0;7.5;8.0;8.6;8.9;9.0" | |
| CUDA_HOME: /usr/local/cuda | |
| run: | | |
| mkdir -p wheels | |
| make wheel | |
| ls -lh wheels | |
| - name: Upload wheels as workflow artifact | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: cuda-wheels-${{ matrix.arch }} | |
| path: kv_connectors/llmd_fs_backend/wheels/*.whl |