build-pytorch-extension-wheels.yml

name: Build PyTorch Extension Wheels

on:
  workflow_call:
    inputs:
      working-directory:
        type: string
        default: "."
      build-command:
        type: string
        default: "python setup.py bdist_wheel"
      python-versions:
        type: string
        default: "[3.7, 3.8, 3.9, '3.10', 3.11]"
      pytorch-versions:
        type: string
        default: "['1.10', 1.11, 1.12, 1.13]"
      cuda-versions:
        type: string
        default: "[10.2, 11.3, 11.6, 11.7]"

jobs:
  setup:
    name: Setup PyTorch Extension Builder
    runs-on: ubuntu-latest
    outputs:
      image-tag: ${{ steps.get-tag.outputs.image-tag }}
    steps:
      - name: Get Tag
        id: get-tag
        run: |
          REQUEST_BASE=https://api.github.com/repos/${{ github.repository }}/actions
          JOBS_REQUEST=$REQUEST_BASE/runs/${{ github.run_id }}/jobs
          AUTH=username:${{ secrets.GITHUB_TOKEN }}
          HEADER="Accept: application/vnd.github.v3+json"
          JOBS=$(curl -s -u $AUTH -H $HEADER $JOBS_REQUEST -L --output -)
          JOB_IDS=$(echo $JOBS | sed 's/\"id\": /\n"id\": /g' | grep ' / Setup PyTorch Extension Builder' | cut -d ',' -f1 | sed 's/\"id\": //g')

          echo "Job ID: $JOB_IDS"

          LOGS_REQUEST=$REQUEST_BASE/jobs/${JOB_IDS[0]}/logs
          LOGS=$(curl -s -u $AUTH -H $HEADER $LOGS_REQUEST -L --output -)
          TAG=$(echo "$LOGS" | grep "Job defined at: charliebudd/torch-extension-builder/.github/workflows/build-pytorch-extension-wheels.yml@" | grep -v "grep" | sed 's:.*/::' | tr -d '[:space:]')

          echo "Tag: $TAG"

          if [[ $TAG == main ]] || [[ $TAG =~ v[0-9] ]]; then
              echo "image-tag=-$TAG" >> $GITHUB_OUTPUT
          else
              echo "image-tag=" >> $GITHUB_OUTPUT
          fi

  build-wheels:
    name: Build Wheels
    needs: setup
    runs-on: ubuntu-latest
    container: 
      image: docker://torchextensionbuilder/manylinux:py${{ matrix.python-version }}-pt${{ matrix.pytorch-version }}-cu${{ matrix.cuda-version }}${{needs.setup.outputs.image-tag}}
    strategy:
      fail-fast: false
      matrix:
        python-version: ${{ fromJson(inputs.python-versions) }}
        pytorch-version: ${{ fromJson(inputs.pytorch-versions) }}
        cuda-version: ${{ fromJson(inputs.cuda-versions) }}
        exclude:
          - pytorch-version: '1.10'
            python-version: '3.10'
          - pytorch-version: '1.10'
            python-version: 3.11
          - pytorch-version: '1.10'
            cuda-version: 11.6
          - pytorch-version: '1.10'
            cuda-version: 11.7
          - pytorch-version: 1.11
            python-version: 3.11
          - pytorch-version: 1.11
            cuda-version: 11.6
          - pytorch-version: 1.11
            cuda-version: 11.7
          - pytorch-version: 1.12
            python-version: 3.11
          - pytorch-version: 1.12
            cuda-version: 11.7
          - pytorch-version: 1.13
            cuda-version: 10.2
          - pytorch-version: 1.13
            cuda-version: 11.3
    steps:
      - name: Checkout Repository
        uses: actions/checkout@v3
        with:
          fetch-depth: 0

      - name: Build Wheels
        working-directory: ${{ inputs.working-directory }}
        run: ${{ inputs.build-command }}
        
      - name: Repair Wheels
        working-directory: ${{ inputs.working-directory }}
        run: auditwheel repair dist/*.whl --plat manylinux_2_17_x86_64 --only-plat
        
      - name: Patch Wheels
        working-directory: ${{ inputs.working-directory }}
        run: python /torch_extension_builder/patch_wheels.py wheelhouse/*manylinux*.whl

      - name: Cache Patched Wheels
        uses: actions/upload-artifact@v3
        with:
          path: ${{ inputs.working-directory }}/patched_wheels
          name: py${{ matrix.python-version }}-wheels
          retention-days: 1
          
      - name: Cache Script
        uses: actions/upload-artifact@v3
        with:
          path: /torch_extension_builder
          name: scripts
          retention-days: 1
  
  combine-wheels:
    name: Combine Wheels
    needs: build-wheels
    runs-on: ubuntu-latest
    strategy:
      fail-fast: false
      matrix:
        python-version: ${{ fromJson(inputs.python-versions) }}

    steps:
      - name: Download Patched Wheels
        uses: actions/download-artifact@v3
        with:
          name: py${{ matrix.python-version }}-wheels
          path: patched_wheels
        
      - name: Download Script
        uses: actions/download-artifact@v3
        with:
          name: scripts

      - name: Combine Wheels
        run: python combine_wheels.py patched_wheels/*
        
      - name: Cache Final Wheels
        uses: actions/upload-artifact@v3
        with:
          path: combined_wheels
          name: final-wheels
          retention-days: 1