CI Cleanup and Parallel Testing (#3)

.github/workflows/kernel-bench.yml

@@ -1,6 +1,7 @@
 name: Triton Kernel Benchmarking
 on:
   pull_request:
+    types: [opened, synchronize]
   workflow_dispatch:
   schedule:
   # Runs at 12:00 PM UTC, which is 5:00 AM PST
@@ -17,6 +18,8 @@ concurrency:
 jobs:
   kernel-bench:
     runs-on: amdgpu-mi250-x86-64
+    env:
+      KERNEL_VENV_DIR: /groups/aig_sharks/triton_bench
     steps:
       - name: Checkout Repo
         uses: actions/checkout@v4
@@ -26,19 +29,24 @@ jobs:
         with:
           python-version: "3.11"
 
+      - name : Set Test File
+        run: |
+          python3.11 ci-tools/parse.py --test "${{ github.event.pull_request.body }}"
+
       - name: Setup Triton Python Venv
         run: |
-          python3.11 -m venv triton_bench
-          source triton_bench/bin/activate
+          python3.11 -m venv ${KERNEL_VENV_DIR}
+          source ${KERNEL_VENV_DIR}/bin/activate
           pip install --upgrade pip
-          pip install --pre pytorch-triton-rocm torch --index-url https://download.pytorch.org/whl/nightly/rocm6.2
+          pip install --pre pytorch-triton-rocm==3.1.0+cf34004b8a torch==2.6.0.dev20241023+rocm6.2 --index-url https://download.pytorch.org/whl/nightly/rocm6.2
+          pip install pytest-xdist
           pip install -r requirements.txt
 
       - name: Run Kernel Benchmarking
         run: |
-          source triton_bench/bin/activate
+          source ${KERNEL_VENV_DIR}/bin/activate
           pip freeze
-          pytest -s kernels
+          pytest kernels/${{ env.TEST_FILE }} -s -n 4
       
       - uses: actions/upload-artifact@master
         with:

README.md

@@ -14,6 +14,9 @@ This repo is configured with a custom Github runner donated by AMD. You can queu
 The main things you need to run your own benchmark
 1. In `kernels/` create a new file that must start with the name `test_`. This is because we use `pytest` to discover your kernel
 2. If you want your benchmark results to persist in a Github Artifact, we recommend using the builtin Triton `benchmark.run(save_path="./perf-artifacts/your_kernel", show_plots=True, print_data=True)`
+3. In your PR, if you don't want to run testing on all the kernels, you can specify a specific kernel you want
+to test by adding a line like the following to your PR description: `ci-exactly: <test-file-name.py>` as seen
+in this PR: [Example](https://github.com/gpu-mode/amd-cluster/pull/3)
 
 Have fun! We intend for this to be a social repo, if you have any other requests for things we could do better please let us know!
 

ci-tools/parse.py

@@ -0,0 +1,19 @@
+import argparse
+import os
+
+def set_env_variable_from_string(input_string):
+    if "ci-exactly:" in input_string:
+        env_file = os.getenv('GITHUB_ENV')
+        test_name = input_string.split("ci-exactly:")[1].strip()
+        with open(env_file, "a") as myfile:
+            myfile.write(f"TEST_FILE={test_name}")
+        print(f'Successfully set TEST_FILE to: {test_name}')
+    else:
+        print('The PR body does not contain "ci-exactly:" so running all tests')
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Set TEST_FILE environment variable from PR body")
+    parser.add_argument("--test", required=True, help="PR body containing 'ci-exactly:'")
+
+    args = parser.parse_args()
+    set_env_variable_from_string(args.test)