name: Benchmark v2 Framework

on:
  workflow_dispatch:
    inputs:
      runner:
        description: 'Runner to use for the benchmark job'
        required: true
        type: string
      container_image:
        description: 'Container image to use'
        required: true
        type: string
      container_options:
        description: 'Container options'
        required: false
        type: string
        default: ''
      commit_sha:
        description: 'Commit SHA to benchmark'
        required: false
        type: string
      run_id:
        description: 'Run ID for tracking'
        required: false
        type: string
      benchmark_repo_id:
        description: 'Repository ID to push benchmark results'
        required: true
        type: string

env:
  HF_HOME: /mnt/cache
  TRANSFORMERS_IS_CI: yes
  # For gated repositories, we still need to agree to share information on the Hub repo. page in order to get access.
  # This token is created under the bot `hf-transformers-bot`.
  HF_TOKEN: ${{ secrets.HF_HUB_READ_TOKEN }}

permissions:
  contents: read

jobs:
  benchmark-v2:
    name: Benchmark v2
    runs-on: ${{ inputs.runner }}
    if: |
      (github.event_name == 'pull_request' && contains( github.event.pull_request.labels.*.name, 'run-benchmark')) ||
      (github.event_name == 'schedule')
    container:
      image: ${{ inputs.container_image }}
      options: ${{ inputs.container_options }}
    steps:
      - name: Get repo
        uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
        with:
          ref: ${{ inputs.commit_sha || github.sha }}
          persist-credentials: false

      - name: Install benchmark dependencies
        run: |
          python3 -m pip install -r benchmark_v2/requirements.txt

      - name: Reinstall transformers in edit mode
        run: |
          python3 -m pip uninstall -y transformers
          python3 -m pip install -e ".[torch]"

      - name: Show installed libraries and their versions
        run: |
          python3 -m pip list
          python3 -c "import torch; print(f'PyTorch version: {torch.__version__}')"
          python3 -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}')"
          python3 -c "import torch; print(f'CUDA device count: {torch.cuda.device_count()}')" || true
          nvidia-smi || true

      - name: Run benchmark v2
        working-directory: benchmark_v2
        env:
          HF_TOKEN: ${{ secrets.HF_HUB_READ_TOKEN }}
          COMMIT_ID: ${{ inputs.commit_sha || github.sha }}
          RUN_ID: ${{ inputs.run_id }}
          BENCHMARK_REPO_ID: ${{ inputs.benchmark_repo_id }}
          UPLOAD_TOKEN: ${{ secrets.TRANSFORMERS_CI_RESULTS_UPLOAD_TOKEN }}
        run: |
          echo "Running benchmarks"
          python3 run_benchmarks.py \
          --commit-id "$COMMIT_ID" \
          --run-id "$RUN_ID" \
          --push-to-hub "$BENCHMARK_REPO_ID" \
          --token "$UPLOAD_TOKEN" \
          --log-level INFO