tritonbench-test.yml

name: Tritonbench PR Test on Triton nightly
on:
  pull_request:
    paths:
      - 'torchbenchmark/operators/*'
      - 'torchbenchmark/util/kernels/*'
      - 'torchbenchmark/util/triton_op.py'
      - 'userbenchmark/triton/*'
      - '.ci/tritonbench/*'
  workflow_dispatch:
  push:
    branches:
      - main
    paths:
      - 'torchbenchmark/operators/*'
      - 'torchbenchmark/util/kernels/*'
      - 'torchbenchmark/util/triton_op.py'
      - 'userbenchmark/triton/*'
      - '.ci/tritonbench/*'

jobs:
  cuda-test:
    # Don't run on forked repos
    if: github.repository_owner == 'pytorch'
    runs-on: [a100-runner]
    timeout-minutes: 240
    environment: docker-s3-upload
    env:
      BASE_CONDA_ENV: "torchbench"
      CONDA_ENV: "tritonbench-pr-test-cuda"
      SETUP_SCRIPT: "/workspace/setup_instance.sh"
      TEST_CONFIG: "cuda"
      HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
    steps:
      - name: Checkout TorchBench
        uses: actions/checkout@v3
        with:
          submodules: 'true'
      - name: Tune Nvidia GPU
        run: |
          sudo nvidia-smi -pm 1
          sudo nvidia-smi -ac 1215,1410
          sudo ldconfig
          nvidia-smi
      - name: Install triton-nightly
        run: |
          bash ./.ci/tritonbench/install-triton-nightly.sh
      - name: Test Tritonbench install
        run: |
          bash ./.ci/tritonbench/test-install.sh
      - name: Test Tritonbench operators
        run: |
          bash ./.ci/tritonbench/test-operators.sh
      - name: Clean up Conda env
        if: always()
        run: |
          . "${SETUP_SCRIPT}"
          conda deactivate && conda deactivate
          conda remove -n "${CONDA_ENV}" --all

concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
  cancel-in-progress: true