From e3b77d71a2c78f324a9a5127227f2694a9140243 Mon Sep 17 00:00:00 2001
From: Vladimir Milosevic <vmilosevic@tenstorrent.com>
Date: Fri, 10 Jan 2025 09:47:54 +0000
Subject: [PATCH] Configure test grouping

---
 .github/workflows/build-and-test.yml | 19 ++++++++++++++++---
 .github/workflows/on-nightly.yml     |  1 +
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml
index c847f7241..b4972879a 100644
--- a/.github/workflows/build-and-test.yml
+++ b/.github/workflows/build-and-test.yml
@@ -17,6 +17,11 @@ on:
         description: 'Git SHA of commit in tenstorrent/tt-mlir or branch name'
         required: false
         type: string
+      test_grouping:
+        description: 'Test grouping strategy'
+        required: false
+        default: '{ "group_count": 3, "group_ids": [1,2,3] }'
+        type: string
   workflow_call:
     inputs:
       test_mark:
@@ -28,14 +33,22 @@ on:
         description: 'Git SHA of commit in tenstorrent/tt-mlir or branch name'
         required: false
         type: string
+      test_grouping:
+        description: 'Test grouping strategy'
+        required: false
+        default: '{ "group_count": 3, "group_ids": [1,2,3] }'
+        type: string
+  push:
 
 permissions:
   packages: write
   checks: write
   pull-requests: write # only required if `comment: true` was enabled
 
-jobs:
+env:
+  test_group_count: ${{ fromJson(inputs.test_grouping || inputs.test_grouping.default).group_count }}
 
+jobs:
   docker-build:
     uses: ./.github/workflows/build-image.yml
     secrets: inherit
@@ -48,7 +61,7 @@ jobs:
       matrix:
         build:
           - runs-on: runner
-        test_group_id: [1,2]
+        test_group_id: ${{ fromJson(inputs.test_grouping || inputs.test_grouping.default).group_ids }}
 
     runs-on:
       - in-service
@@ -152,7 +165,7 @@ jobs:
         apt install -y libgl1 libglx-mesa0
         set -o pipefail # Ensures that the exit code reflects the first command that fails
         pip install pytest-split
-        pytest -m push --splits 2 \
+        pytest -m push --splits ${{ env.test_group_count }} \
                --group ${{ matrix.test_group_id }} \
                --splitting-algorithm least_duration \
                -m "${{ inputs.test_mark }}" \
diff --git a/.github/workflows/on-nightly.yml b/.github/workflows/on-nightly.yml
index 28ad70b6e..7bc0a1406 100644
--- a/.github/workflows/on-nightly.yml
+++ b/.github/workflows/on-nightly.yml
@@ -11,3 +11,4 @@ jobs:
     secrets: inherit
     with:
       test_mark: 'nightly'
+      test_grouping: '{ "group_count": 4, "group_ids": [1,2,3,4] }'