name: ~Build wheel template

on:
  workflow_call:
    inputs:
      runs-on:
        description: "The runner to use for the build"
        required: true
        type: string
      python-version:
        description: "The Python version to use for the build"
        required: true
        type: string
      cuda-version:
        description: "The CUDA version to use for the build"
        required: true
        type: string
      torch-version:
        description: "The PyTorch version to use for the build"
        required: true
        type: string
      cxx11_abi:
        description: "The C++11 ABI to use for the build"
        required: true
        type: string
      upload-to-release:
        description: "Upload wheel to this release"
        required: false
        type: boolean
        default: false
      release-version:
        description: "Upload wheel to this release"
        required: false
        type: string

defaults:
  run:
    shell: bash -x -e -u -o pipefail {0}

jobs:
  build-wheel:
    runs-on: ${{ inputs.runs-on }}
    name: Build wheel (${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }})
    steps:
      - name: Checkout
        uses: actions/checkout@v4
        with:
          ref: ${{ inputs.release-version }}
          submodules: recursive
      
      - name: Checkout build scripts
        uses: actions/checkout@v4
        with:
          path: build-scripts/

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: ${{ inputs.python-version }}

      - name: Set CUDA and PyTorch versions
        run: |
          echo "MATRIX_CUDA_VERSION=$(echo ${{ inputs.cuda-version }} | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
          echo "MATRIX_TORCH_VERSION=$(echo ${{ inputs.torch-version }} | awk -F \. {'print $1 "." $2'})" >> $GITHUB_ENV
          echo "WHEEL_CUDA_VERSION=$(echo ${{ inputs.cuda-version }} | awk -F \. {'print $1'})" >> $GITHUB_ENV
          echo "MATRIX_PYTHON_VERSION=$(echo ${{ inputs.python-version }} | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV

      - name: Free up disk space
        if: ${{ runner.os == 'Linux' }}
        # https://github.com/easimon/maximize-build-space/blob/master/action.yml
        # https://github.com/easimon/maximize-build-space/tree/test-report
        run: |
          sudo rm -rf /usr/share/dotnet
          sudo rm -rf /opt/ghc
          sudo rm -rf /opt/hostedtoolcache/CodeQL

      - name: Set up swap space
        if: runner.os == 'Linux'
        uses: pierotofy/set-swap-space@v1.0
        with:
          swap-size-gb: 10

      - name: Install CUDA ${{ inputs.cuda-version }}
        if: ${{ inputs.cuda-version != 'cpu' }}
        uses: Jimver/cuda-toolkit@v0.2.30
        id: cuda-toolkit
        with:
          cuda: ${{ inputs.cuda-version }}
          linux-local-args: '["--toolkit"]'
          # default method is "local", and we're hitting some error with caching for CUDA 11.8 and 12.1
          # method: ${{ (inputs.cuda-version == '11.8.0' || inputs.cuda-version == '12.1.0') && 'network' || 'local' }}
          method: "network"
          sub-packages: '["nvcc"]'

      - name: Install PyTorch ${{ inputs.torch-version }}+cu${{ inputs.cuda-version }}
        run: |
          pip install --upgrade pip
          # For some reason torch 2.2.0 on python 3.12 errors saying no setuptools
          pip install setuptools==68.0.0
          # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
          # AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
          pip install typing-extensions==4.12.2
          # Pick the highest available PyTorch wheel CUDA version that doesn't exceed system CUDA
          export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; \
            available = { \
              '2.6': [118, 124, 126], \
              '2.7': [118, 126, 128], \
              '2.8': [126, 128, 129], \
              '2.9': [126, 128, 130], \
              '2.10': [126, 128, 130], \
            }[env['MATRIX_TORCH_VERSION']]; \
            sys_cuda = int(env['MATRIX_CUDA_VERSION']); \
            print(max(v for v in available if v <= sys_cuda))" \
          )
          if [[ ${{ inputs.torch-version }} == *"dev"* ]]; then
            # pip install --no-cache-dir --pre torch==${{ inputs.torch-version }} --index-url https://download.pytorch.org/whl/nightly/cu${TORCH_CUDA_VERSION}
            # Can't use --no-deps because we need cudnn etc.
            # Hard-coding this version of pytorch-triton for torch 2.6.0.dev20241001
            pip install jinja2
            pip install https://download.pytorch.org/whl/nightly/pytorch_triton-3.1.0%2Bcf34004b8a-cp${MATRIX_PYTHON_VERSION}-cp${MATRIX_PYTHON_VERSION}-linux_x86_64.whl
            pip install --no-cache-dir --pre https://download.pytorch.org/whl/nightly/cu${TORCH_CUDA_VERSION}/torch-${{ inputs.torch-version }}%2Bcu${TORCH_CUDA_VERSION}-cp${MATRIX_PYTHON_VERSION}-cp${MATRIX_PYTHON_VERSION}-linux_x86_64.whl
          else
            pip install --no-cache-dir torch==${{ inputs.torch-version }} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
          fi
          nvcc --version
          python --version
          python -c "import torch; print('PyTorch:', torch.__version__)"
          python -c "import torch; print('CUDA:', torch.version.cuda)"
          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
        shell: bash

      - name: Build wheel
        id: build_wheel
        env: 
          CXX11_ABI: ${{ inputs.cxx11_abi }}
          MATRIX_TORCH_VERSION: ${{ env.MATRIX_TORCH_VERSION}}
          WHEEL_CUDA_VERSION: ${{ env.WHEEL_CUDA_VERSION }}
          MATRIX_PYTHON_VERSION: ${{ env.MATRIX_PYTHON_VERSION }}
        run: |
          EXIT_CODE=$(bash build-scripts/.github/scripts/build.sh | tail -n 1)

          # Store exit code in GitHub env for later steps
          echo "build_exit_code=$EXIT_CODE" | tee -a "$GITHUB_OUTPUT"

          exit $EXIT_CODE

      - name: Log Built Wheels
        run: |
          ls dist

      - name: Get Release with tag
        id: get_current_release
        uses: joutvhu/get-release@v1
        with:
          tag_name: ${{ inputs.release-version }}
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

      - name: Upload Release Asset
        id: upload_release_asset
        if: inputs.upload-to-release
        uses: actions/upload-release-asset@v1
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        with:
          upload_url: ${{ steps.get_current_release.outputs.upload_url }}
          asset_path: ./dist/${{env.wheel_name}}
          asset_name: ${{env.wheel_name}}
          asset_content_type: application/*
