name: "Test self-hosted runner"
on:
  workflow_dispatch:

jobs:
  docker:
    runs-on:
      - self-hosted
    timeout-minutes: 10
    container:
      image: docker:dind
    steps:
      - run: ls
      - run: pwd

  test-build:
    runs-on:
      - self-hosted
    timeout-minutes: 1000
    strategy:
      fail-fast: false
      matrix:
        flash-attn-version: ["2.4.3"]
        python-version: ["3.12"]
        torch-version: ["2.7.0"]
        # https://developer.nvidia.com/cuda-toolkit-archive
        cuda-version: ["12.8.1"]
    container:
      image: ubuntu:22.04
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}
      - uses: Jimver/cuda-toolkit@master
        with:
          cuda: ${{ matrix.cuda-version }}
          linux-local-args: '["--toolkit"]'
          method: "network"

      - name: Set CUDA and PyTorch versions
        run: |
          echo "MATRIX_CUDA_VERSION=$(echo ${{ matrix.cuda-version }} | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
          echo "MATRIX_TORCH_VERSION=$(echo ${{ matrix.torch-version }} | awk -F \. {'print $1 "." $2'})" >> $GITHUB_ENV
          echo "CACHE_KEY=cuda-ext-${{ matrix.flash-attn-version }}-py${{ matrix.python-version }}-torch${{ matrix.torch-version }}-cuda${{ matrix.cuda-version }}" >> $GITHUB_ENV

      - name: Install build dependencies
        run: |
          sudo apt install -y ninja-build clang
          pip install -U pip setuptools==75.8.0 wheel setuptools packaging psutil

      - name: Install PyTorch ${{ matrix.torch-version }}+cu${{ matrix.cuda-version }}
        run: |
          export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; \
              support_cuda_versions = { \
                  '2.0': [117, 118], \
                  '2.1': [118, 121], \
                  '2.2': [118, 121], \
                  '2.3': [118, 121], \
                  '2.4': [118, 121, 124], \
                  '2.5': [118, 121, 124], \
                  '2.6': [118, 124, 126], \
                  '2.7': [118, 126, 128], \
              }; \
              target_cuda_versions = support_cuda_versions[env['MATRIX_TORCH_VERSION']]; \
              cuda_version = int(env['MATRIX_CUDA_VERSION']); \
              closest_version = min(target_cuda_versions, key=lambda x: abs(x - cuda_version)); \
              print(closest_version) \
          ")
          pip install --no-cache-dir torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
          nvcc --version
          python -V
          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)"

      - name: Checkout flash-attn
        run: |
          git clone https://github.com/Dao-AILab/flash-attention.git -b "v${{ matrix.flash-attn-version }}"
          cd flash-attention

      - name: Build wheels
        timeout-minutes: 800
        run: |
          export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
          export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
          export MAX_JOBS=2
          export NVCC_THREADS=2
          export FLASH_ATTENTION_FORCE_BUILD="TRUE"
          cd flash-attention
          python setup.py bdist_wheel --dist-dir=dist
          base_wheel_name=$(basename $(ls dist/*.whl | head -n 1))
          wheel_name=$(echo $base_wheel_name | sed "s/${{ matrix.flash-attn-version }}/${{ matrix.flash-attn-version }}+cu${{ env.MATRIX_CUDA_VERSION }}torch${{ env.MATRIX_TORCH_VERSION }}/")
          mv dist/$base_wheel_name dist/$wheel_name
          echo "wheel_name=$wheel_name" >> $GITHUB_ENV

      - name: Install Test
        run: |
          pip install flash-attention/dist/${{ env.wheel_name }}
          python -c "import flash_attn; print(flash_attn.__version__)"