name: CI
on: [pull_request]

concurrency:
  group: "${{ github.workflow }}-${{ github.ref }}"
  cancel-in-progress: ${{ github.event_name == 'pull_request' }}

env:
  PYTHON_VERSION: '3.12'
  VENV_DIR: tilelang_ci

jobs:
  format-check:
    runs-on: [self-hosted, nvidia]
    permissions:
      contents: write

    steps:
    - name: Checkout repository
      uses: actions/checkout@v5
      with:
        fetch-depth: 0
        submodules: recursive

    - name: Install python via uv
      uses: astral-sh/setup-uv@v6
      with:
        enable-cache: false
        cache-local-path: ${{ runner.tool_cache }}/uv
        activate-environment: true
        python-version: ${{ env.PYTHON_VERSION }}

    - name: Ensure venv (local & persistent)
      run: |
        [[ -f requirements-test.txt ]] && \
          uv pip install -r requirements-test.txt --no-build-isolation
        uv pip install flash_attn==2.5.8 --no-build-isolation

    - name: Run format check
      run: |
        set -ex
        mkdir -p build
        # run cmake to create the build directory with compile_commands.json
        uv pip install cmake
        cd build; USE_CUDA=1 cmake .. -DCMAKE_EXPORT_COMPILE_COMMANDS=ON; cd ..
        if ! output=$(./format.sh 2>&1); then
          echo "------------------------------------"
          echo "message:"
          echo "$output"
          printf '%s\n' "$output" | grep "Please review and stage the changes."
          echo "------------------------------------"
          exit 1
        fi
        rm -rf build

  build-test-nvidia:
    runs-on: [self-hosted, nvidia]
    needs: format-check
    permissions:
      contents: read
    steps:
    - name: Checkout repository
      uses: actions/checkout@v5
      with:
        fetch-depth: 0
        submodules: recursive
        repository: ${{ github.event.pull_request.head.repo.full_name }}
        ref: ${{ github.event.pull_request.head.ref }}

    - name: Install python via uv
      uses: astral-sh/setup-uv@v6
      with:
        enable-cache: false
        cache-local-path: ${{ runner.tool_cache }}/uv
        activate-environment: true
        python-version: ${{ env.PYTHON_VERSION }}

    - name: Setup venv
      run: |
        [[ -f requirements-test.txt ]] && \
          uv pip install -r requirements-test.txt --no-build-isolation
        uv pip install flash_attn==2.5.8 --no-build-isolation

    - name: Install project (wheel form)
      run: |
        uv pip install .

    - name: Run examples
      run: |
        cd examples
        python -m pytest -n 4 **/test*.py -v -r fE --durations=0 --cache-clear

    - name: Run tests
      run: |
        cd testing/python
        python -m pytest -n 4 -v -r fE --durations=0 --cache-clear --timeout=3600