Shard CI tests across 4 parallel jobs per backend to reduce wall-clock time

.github/workflows/actions.yml

@@ -16,19 +16,23 @@ concurrency:
 
 jobs:
   run_tests:
-    name: Test the code
+    name: Test the code (${{ matrix.backend }}, ${{ matrix.version }}, shard ${{ matrix.shard }})
     strategy:
       fail-fast: false
       matrix:
         backend: [tensorflow, jax, torch]
         version: [keras-stable]
+        shard: [1, 2, 3, 4]
         include:
           - backend: jax
             version: keras-3.15
+            shard: 1
           - backend: jax
             version: keras-nightly
+            shard: 1
           - backend: openvino
             version: keras-nightly
+            shard: 1
 
     container:
       image: python:3.11-slim
@@ -65,13 +69,14 @@ jobs:
         pip install keras-nightly --no-cache-dir --progress-bar off
     - name: Test with pytest
       run: |
-        pytest keras_hub/ -n auto --dist loadfile --durations=50
+        pytest keras_hub/ -n auto --dist loadfile --durations=50 --splits 4 --group ${{ matrix.shard }} --splitting-algorithm least_duration
     - name: Run integration tests
+      if: ${{ matrix.shard == 1 }}
       run: |
         python pip_build.py --install
         cd integration_tests && pytest . -k "not NoTensorflow"
     - name: Run no tensorflow integration test
-      if: ${{ matrix.backend != 'tensorflow'}}
+      if: ${{ matrix.backend != 'tensorflow' && matrix.shard == 1 }}
       run: |
         pip uninstall -y tensorflow-text tensorflow
         cd integration_tests && pytest . -k "NoTensorflow"
@@ -104,4 +109,4 @@ jobs:
       - name: Install pre-commit
         run: pip install pre-commit && pre-commit install
       - name: Run pre-commit
-        run: pre-commit run --all-files --hook-stage manual
+        run: pre-commit run --all-files --hook-stage manual

requirements-common.txt

@@ -14,6 +14,7 @@ ruff
 pytest
 pytest-cov
 pytest-xdist
+pytest-split
 build
 namex
 # Optional deps.