ci: shard core unit tests

Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>

Author: svcnemo-autobot

.github/actions/test-template/action.yml

@@ -59,6 +59,14 @@ inputs:
   runner:
     description: "Runner to use for test"
     required: true
+  shard-id:
+    description: "Zero-based pytest shard index"
+    required: false
+    default: ""
+  num-shards:
+    description: "Total number of pytest shards"
+    required: false
+    default: ""
 
 runs:
   using: "composite"
@@ -135,6 +143,8 @@ runs:
           --env NEMO_HOME=/home/TestData/nemo_home \
           --env RUN_ID=${{ github.run_id }} \
           --env UV_NO_SYNC=1 \
+          --env SHARD_ID=${{ inputs.shard-id }} \
+          --env NUM_SHARDS=${{ inputs.num-shards }} \
           --workdir /opt/Megatron-Bridge \
           $VOLUME_ARGS \
           ${{ inputs.container-image }} \

.github/workflows/cicd-main.yml

@@ -567,6 +567,10 @@ jobs:
               --skip-file /opt/import-check/import_check_skip.txt
 
   cicd-unit-tests-core:
+    strategy:
+      fail-fast: false
+      matrix:
+        shard-id: [0, 1, 2]
     if: |
       (
         success()
@@ -579,7 +583,7 @@ jobs:
       && (github.event.inputs.test_suite == '' || github.event.inputs.test_suite == 'all' || github.event.inputs.test_suite == 'unit-only' || contains('L0 L1 L2', github.event.inputs.test_suite))
     needs: [pre-flight, configure, cicd-container-build]
     runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}
-    name: Launch_Unit_Tests_Core
+    name: Launch_Unit_Tests_Core_${{ matrix.shard-id }}
     env:
       TRANSFORMERS_OFFLINE: "1"
       HF_HUB_OFFLINE: "1"
@@ -603,6 +607,8 @@ jobs:
           script: Launch_Unit_Tests_Core
           timeout: 18
           is_unit_test: "true"
+          shard-id: ${{ matrix.shard-id }}
+          num-shards: 3
           github-token: ${{ steps.app-token.outputs.token }}
           container-image: ${{ needs.pre-flight.outputs.registry }}/megatron-bridge:${{ needs.configure.outputs.merge_sha }}
           test-data-path: ${{ needs.pre-flight.outputs.test_data_path }}

pyproject.toml

@@ -201,6 +201,7 @@ test = [
     "pytest-timeout>=2.4.0",
     "pytest-mock>=3.14.0",
     "pytest-runner>=6.0.1",
+    "pytest-shard>=0.1.2",
     "pygithub",
     "click",
 ]

tests/unit_tests/Launch_Unit_Tests_Core.sh

@@ -17,6 +17,9 @@ set -xeuo pipefail # Exit immediately if a command exits with a non-zero status
 
 REPO_ROOT=$(cd "$(dirname "$0")/../.." && pwd)
 
+: "${SHARD_ID:?SHARD_ID must be set}"
+: "${NUM_SHARDS:?NUM_SHARDS must be set}"
+
 echo "=================================================="
 echo "🧪 UNIT TESTS (core)"
 echo "=================================================="
@@ -31,6 +34,8 @@ CUDA_VISIBLE_DEVICES="0,1" uv run coverage run -a --data-file="${REPO_ROOT}/.cov
     -o log_cli=true \
     -o log_cli_level=INFO \
     --disable-warnings \
+    --shard-id="${SHARD_ID}" \
+    --num-shards="${NUM_SHARDS}" \
     -vs tests/unit_tests \
     --ignore=tests/unit_tests/diffusion \
     -m "not pleasefixme"