# #########################################################
# Build wheels with AWS CodeBuild on Windows x86_64
# #########################################################

name: "[Windows x86_64, AWS CodeBuild] Build wheels and upload to GitHub Releases"
on:
  workflow_call:
    inputs:
      flash-attn-version:
        description: "Flash-Attention version"
        required: true
        type: string
      python-version:
        description: "Python version"
        required: true
        type: string
      torch-version:
        description: "PyTorch version"
        required: true
        type: string
      cuda-version:
        description: "CUDA version"
        required: true
        type: string
      is-upload:
        description: "Whether to upload the release asset"
        required: false
        type: boolean
        default: true

jobs:
  build_wheels:
    name: Build wheels and Upload (Windows x86_64, AWS CodeBuild)
    timeout-minutes: 2160
    runs-on: codebuild-flash-attention-pre-build-wheel-windows-${{ github.run_id }}-${{ github.run_attempt }}
    # Large Instance
    env:
      MAX_JOBS: 6
      NVCC_THREADS: 2
    steps:
      - uses: actions/checkout@v4

      - name: Enable Git long paths
        shell: pwsh
        run: git config --system core.longpaths true

      # Install Python using uv because setup-python needs newer version of glibc
      - name: Install uv
        uses: astral-sh/setup-uv@v7
      - name: Install Python
        shell: pwsh
        run: |
          uv venv -p ${{ inputs.python-version }}
          uv pip install -U pip setuptools==75.8.0 wheel packaging psutil numpy ninja auditwheel
          $current_dir = (Get-Location).Path
          echo "$current_dir\.venv\Scripts" >> $env:GITHUB_PATH

      - uses: mjun0812/setup-cuda@v1
        with:
          version: ${{ inputs.cuda-version }}

      - name: Install Visual Studio BuildTools (C++/CMake + Windows SDK)
        shell: pwsh
        run: |
          choco upgrade -y visualstudio2022buildtools `
            --params "--add Microsoft.VisualStudio.Component.VC.Tools.x86.x64 `
                         --add Microsoft.VisualStudio.Component.VC.CMake.Project `
                         --add Microsoft.VisualStudio.Component.Windows11SDK.22621 `
                         --includeRecommended --includeOptional"

      - name: Setup MSVC Developer Command Prompt
        uses: TheMrMilchmann/setup-msvc-dev@v3
        with:
          arch: x64

      - name: Add msbuild to PATH
        uses: microsoft/setup-msbuild@v2

      - name: Build wheels
        shell: pwsh
        timeout-minutes: 2160
        run: |
          .\build_windows.ps1 -FlashAttnVersion "${{ inputs.flash-attn-version }}" -PythonVersion "${{ inputs.python-version }}" -TorchVersion "${{ inputs.torch-version }}" -CudaVersion "${{ inputs.cuda-version }}"
          # Script changes to flash-attention directory, wheel is in dist/
          $wheelName = Get-ChildItem -Path "dist\*.whl" | Select-Object -First 1 | ForEach-Object { $_.Name }
          echo "wheel_name=$wheelName" >> $env:GITHUB_ENV

      - name: Install Test
        shell: pwsh
        working-directory: flash-attention
        run: |
          pip install --no-cache-dir dist/$env:wheel_name
          python -c "import flash_attn; print(flash_attn.__version__)"

      - name: Upload Release Asset
        if: ${{ inputs.is-upload }}
        shell: pwsh
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          $tag_name = "${{ github.ref_name }}"
          $wheel_path = "flash-attention/dist/$env:wheel_name"
          gh release upload "$tag_name" "$wheel_path" --clobber