name: Self-hosted runner (benchmark)

on:
  push:
    branches: [main]
  pull_request:
    types: [ opened, labeled, reopened, synchronize ]

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
  cancel-in-progress: true

env:
  HF_HOME: /mnt/cache
  DATASET_ID: hf-benchmarks/transformers
  MODEL_ID: meta-llama/Llama-3.1-8B-Instruct

jobs:
  benchmark:
    name: Benchmark
    strategy:
      matrix:
        # group: [aws-g5-4xlarge-cache, aws-p4d-24xlarge-plus] (A100 runner is not enabled)
        group: [aws-g5-4xlarge-cache]
    runs-on:
      group: ${{ matrix.group }}
    if: |
      (github.event_name == 'pull_request' && contains( github.event.pull_request.labels.*.name, 'run-benchmark') )||
      (github.event_name == 'push' && github.ref == 'refs/heads/main')
    container:
      image: huggingface/transformers-all-latest-gpu
      options: --gpus all --privileged --ipc host
    steps:
      - name: Get repo
        uses: actions/checkout@v5
        with:
          fetch-depth: 1

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

      - name: Reinstall transformers in edit mode (remove the one installed during docker image build)
        run: python3 -m pip uninstall -y transformers && python3 -m pip install -e ".[torch]"

      - name: Run benchmark
        run: |
          git config --global --add safe.directory /__w/transformers/transformers
          if [ "$GITHUB_EVENT_NAME" = "pull_request" ]; then
            commit_id=$(echo "${{ github.event.pull_request.head.sha }}")
          elif [ "$GITHUB_EVENT_NAME" = "push" ]; then
            commit_id=$GITHUB_SHA
          fi
          commit_msg=$(git show -s --format=%s | cut -c1-70)
          python3 benchmark_v2/run_benchmarks.py -b 32 -s 128 -n 256 --level 2 --branch-name "$BRANCH_NAME" --commit-id "$commit_id" --commit-message "$commit_msg" --model-id "$MODEL_ID" --log-level INFO --push-result-to-dataset "$DATASET_ID"
        env:
          HF_TOKEN: ${{ secrets.HF_HUB_READ_TOKEN }}
          PUSH_TO_HUB_TOKEN: ${{ secrets.PUSH_TO_HUB_TOKEN }}
          # Enable this to see debug logs
          # HF_HUB_VERBOSITY: debug
          # TRANSFORMERS_VERBOSITY: debug
          BRANCH_NAME: ${{ github.head_ref || github.ref_name }}