name: Check CUDA Extension Build Before Merge

on:
  workflow_dispatch:
  pull_request:
    paths:
      - 'version.txt'

jobs:
  matrix_preparation:
    name: Prepare Container List
    if: github.repository == 'hpcaitech/ColossalAI'
    runs-on: [self-hosted, ubuntu-latest]
    outputs:
      matrix: ${{ steps.set-matrix.outputs.matrix }}
    steps:
      - uses: actions/checkout@v3

      - id: set-matrix
        run: |
          cuda_ext=$(cat .cuda_ext.json | tr '\n' ' ')
          echo "matrix=${cuda_ext}" >> $GITHUB_OUTPUT

  build:
    name: Release bdist wheels
    needs: matrix_preparation
    runs-on: [self-hosted, ubuntu-latest]
    strategy:
      fail-fast: false
      matrix: ${{fromJson(needs.matrix_preparation.outputs.matrix)}}
    container:
      image: ${{ matrix.build.cuda_image }}
      options: --gpus all --rm
    steps:
      - uses: actions/checkout@v2

      - name: Install PyTorch
        run: eval ${{ matrix.build.torch_command }}

      - name: Download cub for CUDA 10.2
        run: |
          CUDA_VERSION=$(nvcc -V | awk -F ',| ' '/release/{print $6}')

          # check if it is CUDA 10.2
          # download cub
          if [ "$CUDA_VERSION" = "10.2" ]; then
            wget https://github.com/NVIDIA/cub/archive/refs/tags/1.8.0.zip
            unzip 1.8.0.zip
            cp -r cub-1.8.0/cub/ colossalai/kernel/cuda_native/csrc/kernels/include/
          fi

      - name: Build
        run: |
          BUILD_EXT=1 pip install -v -e .