# #########################################################
# Build wheels with self-hosted runner
# #########################################################

name: "[Linux x86_64, self-hosted] 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
      runner:
        description: "Runner type"
        required: false
        type: string
        default: '["self-hosted"]'
      is-upload:
        description: "Whether to upload the release asset"
        required: false
        type: boolean
        default: true
      use-container:
        description: "Whether to use container"
        required: false
        type: boolean
        default: true
      container-image:
        description: "Container image"
        required: false
        type: string
        default: "ubuntu:22.04"

jobs:
  build_wheels_self_hosted:
    if: ${{ inputs.use-container }}
    name: Build wheels and Upload (Linux x86_64, self-hosted runner)
    runs-on: ${{ fromjson(inputs.runner) }}
    container:
      image: ${{ inputs.container-image }}
    defaults:
      run:
        shell: bash
    env:
      DEBIAN_FRONTEND: noninteractive
      TERM: xterm-256color
    timeout-minutes: 2160

    steps:
      - name: Install tools
        shell: bash
        run: |
          apt-get update && apt-get install -y --no-install-recommends \
            curl \
            ca-certificates \
            sudo \
            software-properties-common \
            wget \
            unzip \
            zip \
            git \
            build-essential \
            gcc \
            g++ \
            clang \
            ninja-build \
            keyboard-configuration \
            time

      - name: Install gh
        shell: bash
        run: |
          sudo mkdir -p -m 755 /etc/apt/keyrings
          out=$(mktemp)
          wget -nv -O$out https://cli.github.com/packages/githubcli-archive-keyring.gpg
          cat $out | sudo tee /etc/apt/keyrings/githubcli-archive-keyring.gpg > /dev/null
          sudo chmod go+r /etc/apt/keyrings/githubcli-archive-keyring.gpg
          echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" | sudo tee /etc/apt/sources.list.d/github-cli.list > /dev/null
          sudo apt update
          sudo apt install gh -y

      - uses: actions/checkout@v4

      - name: Configure Git safe directory
        shell: bash
        run: |
          git config --global --add safe.directory $(pwd)

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

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

      - name: Install build dependencies
        shell: bash
        run: |
          pip install -U pip setuptools==75.8.0 wheel packaging psutil

      - name: Build wheels
        timeout-minutes: 2160
        id: build_wheels
        shell: bash
        run: |
          chmod +x build_linux.sh
          ./build_linux.sh ${{ inputs.flash-attn-version }} ${{ inputs.python-version }} ${{ inputs.torch-version }} ${{ inputs.cuda-version }}
          wheel_name=$(basename $(ls flash-attention/dist/*.whl | head -n 1))
          echo "WHEEL_NAME=$wheel_name" >> $GITHUB_OUTPUT

      - name: Install Test
        shell: bash
        run: |
          pip uninstall -y flash-attn > /dev/null 2>&1
          pip install --no-cache-dir flash-attention/dist/${{ steps.build_wheels.outputs.WHEEL_NAME }}
          python -c "import flash_attn; print(flash_attn.__version__)"

      - name: Upload Release Asset
        if: ${{ inputs.is-upload }}
        shell: bash
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          tag_name=${{ github.ref_name }}
          wheel_path="flash-attention/dist/${{ steps.build_wheels.outputs.WHEEL_NAME }}"

          # Check if the file exists
          if [ ! -f "$wheel_path" ]; then
            echo "Error: Wheel file not found at $wheel_path"
            exit 1
          fi

          # Upload the release asset using GitHub CLI
          gh release upload "$tag_name" "$wheel_path" --clobber

      # - name: Apply auditwheel repair
      #   continue-on-error: true
      #   run: |
      #     pip install auditwheel
      #     auditwheel show flash-attention/dist/${{ steps.build_wheels.outputs.WHEEL_NAME }}
      #     auditwheel repair flash-attention/dist/${{ steps.build_wheels.outputs.WHEEL_NAME }} -w flash-attention/dist_manylinux
      #     wheel_name=$(basename $(ls flash-attention/dist_manylinux/*.whl | head -n 1))
      #     echo "WHEEL_NAME_MANYLINUX=$wheel_name" >> $GITHUB_OUTPUT

      # - name: Apply auditwheel repair
      #   if: ${{ inputs.is-upload }}
      #   continue-on-error: true
      #   env:
      #     GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
      #   run: |
      #     pip uninstall -y flash-attn > /dev/null 2>&1
      #     pip install --no-cache-dir flash-attention/dist_manylinux/${{ steps.build_wheels.outputs.WHEEL_NAME_MANYLINUX }}
      #     python -c "import flash_attn; print(flash_attn.__version__)"

      #     wheel_path_manylinux="flash-attention/dist_manylinux/${{ steps.build_wheels.outputs.WHEEL_NAME_MANYLINUX }}"
      #     if [ ! -f "$wheel_path_manylinux" ]; then
      #       echo "Error: Wheel file not found at $wheel_path_manylinux"
      #       exit 1
      #     fi

      #     # Upload the release asset using GitHub CLI
      #     gh release upload "$tag_name" "$wheel_path_manylinux" --clobber

      - name: Clean up
        shell: bash
        if: always()
        run: |
          rm -rf /opt/hostedtoolcache/Python

  build_wheels_self_hosted_no_container:
    if: ${{ !inputs.use-container }}
    name: Build wheels and Upload (Linux x86_64, self-hosted runner)
    runs-on: ${{ fromjson(inputs.runner) }}
    defaults:
      run:
        shell: bash
    env:
      DEBIAN_FRONTEND: noninteractive
      TERM: xterm-256color
    timeout-minutes: 2160
    steps:
      - name: Check environment
        shell: bash
        run: |
          cat /etc/os-release || true
          echo "-----------------------------"
          cat /etc/lsb-release || true
          echo "-----------------------------"
          lscpu || true
          echo "-----------------------------"
          df -h || true
          echo "-----------------------------"
          free -h || true

      - name: Install tools
        shell: bash
        run: |
          apt-get update && apt-get install -y --no-install-recommends \
            curl \
            ca-certificates \
            sudo \
            software-properties-common \
            wget \
            unzip \
            zip \
            git \
            build-essential \
            gcc \
            g++ \
            clang \
            ninja-build \
            keyboard-configuration \
            time

      - name: Install gh
        shell: bash
        run: |
          sudo mkdir -p -m 755 /etc/apt/keyrings
          out=$(mktemp)
          wget -nv -O$out https://cli.github.com/packages/githubcli-archive-keyring.gpg
          cat $out | sudo tee /etc/apt/keyrings/githubcli-archive-keyring.gpg > /dev/null
          sudo chmod go+r /etc/apt/keyrings/githubcli-archive-keyring.gpg
          echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" | sudo tee /etc/apt/sources.list.d/github-cli.list > /dev/null
          sudo apt update
          sudo apt install gh -y

      - uses: actions/checkout@v4

      - name: Configure Git safe directory
        shell: bash
        run: |
          git config --global --add safe.directory $(pwd)

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

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

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

      - name: Build wheels
        timeout-minutes: 2160
        id: build_wheels
        shell: bash
        run: |
          chmod +x build_linux.sh
          ./build_linux.sh ${{ inputs.flash-attn-version }} ${{ inputs.python-version }} ${{ inputs.torch-version }} ${{ inputs.cuda-version }}
          wheel_name=$(basename $(ls flash-attention/dist/*.whl | head -n 1))
          echo "WHEEL_NAME=$wheel_name" >> $GITHUB_OUTPUT

      - name: Install Test
        shell: bash
        run: |
          pip uninstall -y flash-attn > /dev/null 2>&1
          pip install --no-cache-dir flash-attention/dist/${{ steps.build_wheels.outputs.WHEEL_NAME }}
          python -c "import flash_attn; print(flash_attn.__version__)"

      - name: Upload Release Asset
        if: ${{ inputs.is-upload }}
        shell: bash
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          tag_name=${{ github.ref_name }}
          wheel_path="flash-attention/dist/${{ steps.build_wheels.outputs.WHEEL_NAME }}"

          # Check if the file exists
          if [ ! -f "$wheel_path" ]; then
            echo "Error: Wheel file not found at $wheel_path"
            exit 1
          fi

          # Upload the release asset using GitHub CLI
          gh release upload "$tag_name" "$wheel_path" --clobber

      # - name: Apply auditwheel repair
      #   continue-on-error: true
      #   run: |
      #     pip install auditwheel
      #     auditwheel show flash-attention/dist/${{ steps.build_wheels.outputs.WHEEL_NAME }}
      #     auditwheel repair flash-attention/dist/${{ steps.build_wheels.outputs.WHEEL_NAME }} -w flash-attention/dist_manylinux
      #     wheel_name=$(basename $(ls flash-attention/dist_manylinux/*.whl | head -n 1))
      #     echo "WHEEL_NAME_MANYLINUX=$wheel_name" >> $GITHUB_OUTPUT

      # - name: Apply auditwheel repair
      #   if: ${{ inputs.is-upload }}
      #   continue-on-error: true
      #   env:
      #     GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
      #   run: |
      #     pip uninstall -y flash-attn > /dev/null 2>&1
      #     pip install --no-cache-dir flash-attention/dist_manylinux/${{ steps.build_wheels.outputs.WHEEL_NAME_MANYLINUX }}
      #     python -c "import flash_attn; print(flash_attn.__version__)"

      #     wheel_path_manylinux="flash-attention/dist_manylinux/${{ steps.build_wheels.outputs.WHEEL_NAME_MANYLINUX }}"
      #     if [ ! -f "$wheel_path_manylinux" ]; then
      #       echo "Error: Wheel file not found at $wheel_path_manylinux"
      #       exit 1
      #     fi

      #     # Upload the release asset using GitHub CLI
      #     gh release upload "$tag_name" "$wheel_path_manylinux" --clobber

      - name: Clean up
        shell: bash
        if: always()
        run: |
          rm -rf /opt/hostedtoolcache/Python