name: torch compile tests

on:
  workflow_dispatch:
    inputs:
      branch:
        description: 'Branch to test on'
        required: true
      pytorch_nightly:
        description: 'Whether to use PyTorch nightly (true/false)'
        required: false
        default: false

env:
  RUN_SLOW: "yes"
  IS_GITHUB_CI: "1"
  # To be able to run tests on CUDA 12.2
  NVIDIA_DISABLE_REQUIRE: "1"

permissions: {}

jobs:
  run_tests_with_compile:
    runs-on:
      group: aws-g6-4xlarge-plus
    env:
      PEFT_DEBUG_WITH_TORCH_COMPILE: 1
      CUDA_VISIBLE_DEVICES: "0"
      TEST_TYPE: "single_gpu_huggingface/peft-gpu:latest"
      USE_PYTORCH_NIGHTLY: "${{ github.event.inputs.pytorch_nightly }}"
    container:
      image: "huggingface/peft-gpu:latest"
      options: --gpus all --shm-size "16gb" --ipc host -v /mnt/cache/.cache/huggingface:/mnt/cache/
    defaults:
      run:
        shell: bash
    steps:
      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0  # v7.0.0
        with:
          ref: ${{ github.event.inputs.branch }}
          repository: ${{ github.event.pull_request.head.repo.full_name }}
          persist-credentials: false
      - name: Pip install
        run: |
          source activate peft
          pip install -e . --no-deps
          pip install pytest-cov pytest-reportlog parameterized datasets scipy einops
          pip install "pytest>=7.2.0,<8.0.0" # see: https://github.com/huggingface/transformers/blob/ce4fff0be7f6464d713f7ac3e0bbaafbc6959ae5/setup.py#L148C6-L148C26
          if [ "${USE_PYTORCH_NIGHTLY}" = "true" ]; then
            python -m pip install --upgrade --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
          fi
      - name: Test compile with pytest
        run: |
          source activate peft
          echo "PEFT_DEBUG_WITH_TORCH_COMPILE=$PEFT_DEBUG_WITH_TORCH_COMPILE"
          make tests_torch_compile