build_wheels_and_publish.yml

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: © 2023 Aalto University

name: Build_pypi_wheels

on:
  push:
    branches: [ "*stable" ]
    tags: [ "*rc*"]
  pull_request:
    branches: [ "*stable" ]

jobs:
  build_pypi_binary_wheels:
    name: Build wheels on ${{ matrix.os }} for ${{ matrix.platform }} ${{ matrix.arch }}
    runs-on: ${{ matrix.os }}
    strategy:
      matrix:
        os: [ubuntu-20.04, macos-11]
        platform: [manylinux, musllinux, macosx]
        arch: [x86_64, arm64, aarch64, universal2]
        exclude:
          - os: macos-11
            platform: manylinux
          - os: macos-11
            platform: musllinux
          - os: macos-11
            arch: aarch64
          - os: ubuntu-20.04
            platform: macosx
          - os: ubuntu-20.04
            arch: universal2
          - os: ubuntu-20.04
            arch: arm64
    steps:
      - uses: actions/checkout@v3
        with:
          submodules: 'recursive'

      # - uses: actions/setup-python@v4
      #   name: Install Python
      #   with:
      #     python-version: '3.9'

      - name: Set up QEMU
        if: runner.os == 'Linux' && matrix.arch == 'aarch64'
        uses: docker/setup-qemu-action@v2
        with:
          platforms: all

      - name: Build wheels CPU only
        if: ${{ !(matrix.platform == 'manylinux' && matrix.arch == 'x86_64') }}
        uses: pypa/cibuildwheel@v2.11.2
        env:
          CIBW_ARCHS_LINUX: ${{ matrix.arch }}
          CIBW_ARCHS_MACOS: ${{ matrix.arch }}
          CIBW_BUILD: "*-${{ matrix.platform }}*"
          CIBW_BUILD_VERBOSITY: 2
          CIBW_ENVIRONMENT_MACOS: JAX_CHACHA_PRNG_DISABLE_OPENMP=1
          # CIBW_ENVIRONMENT: JAX_CHACHA_PRNG_BUILD="cpu"

      - name: Build wheels CUDA 11
        if: ${{ matrix.platform == 'manylinux' && matrix.arch == 'x86_64' }}
        uses: pypa/cibuildwheel@v2.11.2
        env:
          CIBW_ARCHS_LINUX: ${{ matrix.arch }}
          CIBW_BUILD: "*-${{ matrix.platform }}*"
          CIBW_BUILD_VERBOSITY: 2
          CIBW_BEFORE_ALL_LINUX: >
            curl -o installer.run -s https://developer.download.nvidia.com/compute/cuda/11.2.2/local_installers/cuda_11.2.2_460.32.03_linux.run &&
            /bin/sh ./installer.run --toolkit --silent
          CIBW_ENVIRONMENT: PATH=/usr/local/cuda/bin:$PATH # JAX_CHACHA_PRNG_BUILD="cuda11.2"
          CIBW_ENVIRONMENT_MACOS: JAX_CHACHA_PRNG_DISABLE_OPENMP=1

      - uses: actions/upload-artifact@v3
        with:
          path: ./wheelhouse/*.whl
          name: dist

  build_pypi_source_dist:
    name: Build source distribution package for PyPI
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v3
        with:
          submodules: 'recursive'

      - uses: actions/setup-python@v4
        name: Install Python
        with:
          python-version: '3.9'

      - name: Setup build environment
        run: pip install build

      - name: Build sdist
        run: python -m build --sdist

      - uses: actions/upload-artifact@v3
        with:
          path: ./dist/*.tar.gz
          name: dist

  upload_pypi_wheels:
    name: Upload wheels to PyPI
    runs-on: ubuntu-latest
    if: ${{ contains(github.ref, 'stable') && github.event_name == 'push' }}
    needs:
      - build_pypi_binary_wheels
      - build_pypi_source_dist
    steps:
      - uses: actions/download-artifact@v3
        with:
          name: dist
          path: ./dist/
      - uses: pypa/gh-action-pypi-publish@release/v1
        name: Upload to PyPI
        with:
          password: ${{ secrets.PYPI_API_TOKEN }}
          skip_existing: true

  create_github_release:
    name: Create GitHub Release
    runs-on: ubuntu-latest
    if: ${{ github.event_name == 'push' }}
    needs:
      - build_pypi_binary_wheels
      - build_pypi_source_dist
    permissions:
      contents: write
    steps:
      - uses: actions/checkout@v3
      - uses: actions/download-artifact@v3
        with:
          name: dist
          path: ./dist/
      - name: Publish Pre-Release
        if: startsWith(github.ref, 'refs/tags/')
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          VERSION=$(python -c "import chacha.version; print(chacha.version.VERSION)")
          TAG_NAME=${GITHUB_REF##refs/tags/}
          echo "Version: ${VERSION}"
          echo "Ref: ${GITHUB_REF}"
          echo "Tag: ${TAG_NAME}"
          assets=()
          for asset in $(find ./dist/ -type f); do
            echo "Adding asset: ${asset}"
            assets+=("-a" "$asset")
          done
          hub release create -p "${assets[@]}" -m "$VERSION" -F ChangeLog.txt "$VERSION"
      - name: Publish Release
        if: contains(github.ref, 'stable')
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          VERSION=$(python -c "import chacha.version; print(chacha.version.VERSION)")
          echo "Version: ${VERSION}"
          echo "Ref: ${GITHUB_REF}"
          echo "Commit: ${GITHUB_SHA}"
          assets=()
          for asset in $(find ./dist/ -type f); do
            echo "Adding asset: ${asset}"
            assets+=("-a" "$asset")
          done
          hub release create "${assets[@]}" -m "$VERSION" -F ChangeLog.txt -t "$GITHUB_SHA" "$VERSION"