From 32aa12263907222442bdfff792f485a2541cf9a8 Mon Sep 17 00:00:00 2001
From: saienduri <77521230+saienduri@users.noreply.github.com>
Date: Tue, 29 Oct 2024 20:41:17 -0700
Subject: [PATCH] CI Cleanup and Parallel Testing (#3)

---
 .github/workflows/kernel-bench.yml | 18 +++++++++++++-----
 README.md                          |  3 +++
 ci-tools/parse.py                  | 19 +++++++++++++++++++
 3 files changed, 35 insertions(+), 5 deletions(-)
 create mode 100644 ci-tools/parse.py

diff --git a/.github/workflows/kernel-bench.yml b/.github/workflows/kernel-bench.yml
index b1f584e..6205af3 100644
--- a/.github/workflows/kernel-bench.yml
+++ b/.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:
diff --git a/README.md b/README.md
index ec5ff55..0b912a6 100644
--- a/README.md
+++ b/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!
 
diff --git a/ci-tools/parse.py b/ci-tools/parse.py
new file mode 100644
index 0000000..a194870
--- /dev/null
+++ b/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)