Configure test grouping

.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 }}" \

.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] }'