Refactor setup_rocm_libs.sh to use rocm manylinux image #146
This file contains 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: PyTorch Triton ROCm Wheels | |
| # Build on every branch push or tag push: | |
| on: | |
| push: | |
| pull_request: | |
| jobs: | |
| build_wheels: | |
| name: Build wheels on ${{ matrix.python }}-manylinux_x86_64 | |
| runs-on: ubuntu-latest | |
| strategy: | |
| matrix: | |
| python: ["cp37", "cp38", "cp39", "cp310", "cp311"] | |
| env: | |
| DOCKER_IMAGE: pytorch/manylinux-rocm:5.6 | |
| PY_VERS: ${{ matrix.python }} | |
| steps: | |
| - uses: actions/checkout@v3 | |
| - name: Patch setup.py | |
| run: | | |
| export HASH=$(git rev-parse --short HEAD) | |
| sed -i -r "s/name\=\"(.*)\"/name=\"pytorch-triton-rocm\"/g" python/setup.py | |
| sed -i -r "s/version\=\"(.*)\"/version=\"\1+"$HASH"\"/g" python/setup.py | |
| echo "" >> python/setup.cfg | |
| echo "[build_ext]" >> python/setup.cfg | |
| echo "base-dir=/triton" >> python/setup.cfg | |
| - name: Patch setup.py version if tagged | |
| # patch setup.py version for any tag starting with 'v' | |
| if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/pytorch-triton-rocm-v') | |
| run: | | |
| # strip off initial 'pytorch-triton-rocm-v' | |
| VERSION=${GITHUB_REF_NAME#pytorch-triton-rocm-v} | |
| echo Version: $VERSION | |
| sed -i -r "s/version\=\"(.*)\"/version=\"$VERSION\"/g" python/setup.py | |
| - name: Build Triton wheel | |
| run: | | |
| set -x | |
| mkdir -p "artifacts/" | |
| TRITON_ROCM_DIR=/triton/python/triton/third_party/rocm | |
| container_name=$(docker run \ | |
| --tty \ | |
| --detach \ | |
| -v "${GITHUB_WORKSPACE}:/triton" \ | |
| -v "${GITHUB_WORKSPACE}/artifacts:/artifacts" \ | |
| -w /artifacts/ \ | |
| "${DOCKER_IMAGE}" \ | |
| ) | |
| # Determine python executable for given version | |
| case $PY_VERS in | |
| cp37) | |
| PYTHON_EXECUTABLE=/opt/python/cp37-cp37m/bin/python | |
| ;; | |
| cp38) | |
| PYTHON_EXECUTABLE=/opt/python/cp38-cp38/bin/python | |
| ;; | |
| cp39) | |
| PYTHON_EXECUTABLE=/opt/python/cp39-cp39/bin/python | |
| ;; | |
| cp310) | |
| PYTHON_EXECUTABLE=/opt/python/cp310-cp310/bin/python | |
| ;; | |
| cp311) | |
| PYTHON_EXECUTABLE=/opt/python/cp311-cp311/bin/python | |
| ;; | |
| *) | |
| echo "Unsupported python version ${PY_VERS}" | |
| exit 1 | |
| ;; | |
| esac | |
| docker exec -t "${container_name}" yum install -y zlib-devel python3 zip unzip | |
| docker exec -t "${container_name}" /bin/bash -c "TRITON_ROCM_DIR=${TRITON_ROCM_DIR} /triton/scripts/amd/setup_rocm_libs.sh" | |
| docker exec -t "${container_name}" $PYTHON_EXECUTABLE -m pip install -U setuptools==67.4.0 | |
| docker exec -t "${container_name}" $PYTHON_EXECUTABLE /triton/.github/scripts/build_triton_wheel.py | |
| docker exec -t "${container_name}" /bin/bash -c "MANYLINUX_VERSION='manylinux_2_17_x86_64.manylinux2014_x86_64' \ | |
| /triton/.github/scripts/fix_so.sh" | |
| - uses: actions/upload-artifact@v3 | |
| with: | |
| path: ./artifacts/*.whl | |
| upload_pypi: | |
| needs: build_wheels | |
| runs-on: ubuntu-latest | |
| # upload to PyPI on every tag starting with 'v' | |
| if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/pytorch-triton-rocm-v') | |
| # alternatively, to publish when a GitHub Release is created, use the following rule: | |
| # if: github.event_name == 'release' && github.event.action == 'published' | |
| steps: | |
| - uses: actions/download-artifact@v3 | |
| with: | |
| # unpacks default artifact into dist/ | |
| # if `name: artifact` is omitted, the action will create extra parent dir | |
| name: artifact | |
| path: dist | |
| - uses: pypa/[email protected] | |
| with: | |
| user: __token__ | |
| password: ${{ secrets.PYPI_API_TOKEN }} | |
| # To test: repository_url: https://test.pypi.org/legacy/ |