From 277a0f10520c872348102aec6ac56f8aac3d4459 Mon Sep 17 00:00:00 2001
From: adamomainz <amainz@meta.com>
Date: Fri, 23 Aug 2024 10:02:07 -0700
Subject: [PATCH] cuda wasnt working last run so trying with different configs

---
 .github/workflows/tests.yml | 77 ++++++++++++++++++++++++++++++-------
 requirements.txt            |  1 +
 2 files changed, 65 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 72c9f4f..55f62d7 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -1,4 +1,4 @@
-name: pytests
+name: tests
 
 on:
   push:
@@ -7,17 +7,68 @@ on:
     branches: [ main, "*"]
 
 jobs:
-  tests:
+  Runner-Preparation:
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    outputs:
+      matrix-CUDA: ${{ steps.set-matrix.outputs.matrix-CUDA }}
+      matrix-HIP: ${{ steps.set-matrix.outputs.matrix-HIP }}
+      matrix-MACOS: ${{ steps.set-matrix.outputs.matrix-MACOS }}
+    steps:
+      - name: Prepare runner matrix
+        id: set-matrix
+        if: env.enable_integration == 'true'
+        run: |
+          if [ x"${{ github.repository }}" == x"triton-lang/kernels" ]; then
+            echo '::set-output name=matrix-CUDA::[["a100-runner-set"], ["h100-runner-set"]]'
+            echo '::set-output name=matrix-HIP::[["self-hosted", "gfx90a"]]'
+            echo '::set-output name=matrix-MACOS::[["macos-latest"]]'
+          else
+            echo '::set-output name=matrix-CUDA::["ubuntu-latest"]'
+            echo '::set-output name=matrix-HIP::["ubuntu-latest"]'
+            echo '::set-output name=matrix-MACOS::[["macos-latest"]]'
+          fi
+  pre-commit:
+    name: pre-commit (code formatting)
+    needs: Runner-Preparation
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v3
-    - name: Setup Python
-      uses: actions/setup-python@v4
-      with:
-        python-version: '3.12'
-    - name:  Install pip dependencies
-      run: |
-        python3 -m pip install --upgrade pip
-        python3 -m pip install -r requirements.txt
-    - name: Run Python unit tests
-      run: python3  -m pytest
+      - name: Checkout
+        uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+          cache: 'pip'
+      - name: install dependencies
+        run: |
+          python3 -m pip install --upgrade pip
+          python3 -m pip install -r requirements.txt
+      - name: Check pre-commit
+        run: |
+          # TODO: ignore the first yapf failure until https://github.com/google/yapf/issues/1164 is fixed
+          python3 -m pre_commit run --all-files --verbose yapf &> /dev/null || true
+          # If first run of yapf worked and made changes reset the tree to the original state
+          git reset --hard
+          python3 -m pre_commit run --all-files --verbose
+      - name: Print diff of changes if pre-commit failed
+        if: failure()
+        run: |
+          git diff
+  tests:
+    needs: pre-commit
+    if: needs.Runner-Preparation.outputs.matrix-CUDA != ''
+    strategy:
+      matrix:
+        runner: ${{fromJson(needs.Runner-Preparation.outputs.matrix-CUDA)}}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          submodules: "true"
+      - uses: actions/checkout@v3
+      - name: Setup Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'
+      - name: Run Python unit tests
+        run: python3  -m pytest
diff --git a/requirements.txt b/requirements.txt
index af7ed72..d16b448 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,3 +7,4 @@ fairscale
 triton
 pandas
 pytest
+pre-commit