ci: launch GB200 unit tests via launch_on_gb200 marker

.github/workflows/cicd-main.yml

@@ -734,6 +734,95 @@ jobs:
           container-image: ${{ env.container-registry }}/megatron-lm:${{ needs.configure.outputs.sha }}
           sha: ${{ needs.configure.outputs.sha }}
 
+  cicd-parse-unit-tests-gb200:
+    runs-on: ubuntu-latest
+    outputs:
+      unit-tests-gb200: ${{ steps.parse-unit-tests.outputs.unit-tests-gb200 }}
+    needs:
+      - is-not-external-contributor
+      - pre-flight
+      - configure
+      - cicd-wait-in-queue
+      - cicd-container-build
+    if: |
+      needs.pre-flight.result != 'cancelled'
+      && needs.configure.result != 'cancelled'
+      && needs.cicd-wait-in-queue.result != 'cancelled'
+      && needs.cicd-container-build.result != 'cancelled'
+      && needs.is-not-external-contributor.outputs.is_maintainer == 'true'
+      && vars.ENABLE_GB200_TESTING == 'true'
+      && (
+        success()
+        || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || needs.pre-flight.outputs.force_run_all == 'true'
+        || needs.pre-flight.outputs.is_merge_group == 'true'
+      )
+      && !cancelled()
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+        with:
+          ref: ${{ needs.configure.outputs.sha }}
+      - name: Parse unit tests
+        id: parse-unit-tests
+        run: |
+          cat tests/test_utils/recipes/gb200/unit-tests.yaml | yq -o json '[.products[].test_case[] | { "bucket": .}] | sort_by(.model, .test_case)' | jq -c > unit-tests-gb200.json
+          echo "unit-tests-gb200=$(cat unit-tests-gb200.json)" | tee -a $GITHUB_OUTPUT
+
+  cicd-unit-tests-latest-gb200:
+    strategy:
+      fail-fast: false
+      matrix:
+        include: ${{ fromJson(needs.cicd-parse-unit-tests-gb200.outputs.unit-tests-gb200) }}
+    needs:
+      - is-not-external-contributor
+      - pre-flight
+      - configure
+      - cicd-wait-in-queue
+      - cicd-container-build
+      - cicd-parse-unit-tests-gb200
+    runs-on: ${{ needs.is-not-external-contributor.outputs.selected_runner_gb200 }}
+    timeout-minutes: 60
+    name: "${{ matrix.bucket }} - gb200 latest"
+    if: |
+      needs.is-not-external-contributor.result != 'cancelled'
+      && needs.pre-flight.result != 'cancelled'
+      && needs.configure.result != 'cancelled'
+      && needs.cicd-wait-in-queue.result != 'cancelled'
+      && needs.cicd-container-build.result != 'cancelled'
+      && needs.cicd-parse-unit-tests-gb200.result == 'success'
+      && needs.is-not-external-contributor.outputs.is_maintainer == 'true'
+      && vars.ENABLE_GB200_TESTING == 'true'
+      && (
+        success()
+        || needs.pre-flight.outputs.is_ci_workload == 'true'
+        || needs.pre-flight.outputs.force_run_all == 'true'
+        || needs.pre-flight.outputs.is_merge_group == 'true'
+      )
+      && !cancelled()
+    env:
+      PIP_DISABLE_PIP_VERSION_CHECK: 1
+      PIP_NO_PYTHON_VERSION_WARNING: 1
+      PIP_ROOT_USER_ACTION: ignore
+      PIP_DEFAULT_TIMEOUT: 120
+      PIP_RETRIES: 5
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+        with:
+          ref: ${{ needs.configure.outputs.sha }}
+      - name: main
+        uses: ./.github/actions
+        with:
+          test_case: ${{ matrix.bucket }}
+          tag: latest
+          timeout: ${{ matrix.timeout || 30 }}
+          is_unit_test: "true"
+          PAT: ${{ secrets.PAT }}
+          container-image: ${{ env.container-registry-gb200 }}/megatron-lm:${{ needs.configure.outputs.sha }}
+          platform: dgx_gb200
+          sha: ${{ needs.configure.outputs.sha }}
+
   # Single source of truth for "should integration tests run?".
   # Encodes two independent gates:
   #   (A) Approval gate — `cicd-wait-in-queue` must have succeeded
@@ -1001,6 +1090,7 @@ jobs:
       - pre-flight
       - is-not-external-contributor
       - cicd-unit-tests-latest
+      - cicd-unit-tests-latest-gb200
       - cicd-integration-tests-latest-h100
       - cicd-integration-tests-latest-gb200
     if: |
@@ -1032,6 +1122,7 @@ jobs:
           FORCE_RUN_ALL: ${{ needs.pre-flight.outputs.force_run_all }}
           ENABLE_GB200_TESTING: ${{ vars.ENABLE_GB200_TESTING }}
           UNIT_RESULT: ${{ needs.cicd-unit-tests-latest.result }}
+          UNIT_GB200_RESULT: ${{ needs.cicd-unit-tests-latest-gb200.result }}
           H100_RESULT: ${{ needs.cicd-integration-tests-latest-h100.result }}
           GB200_RESULT: ${{ needs.cicd-integration-tests-latest-gb200.result }}
         run: |
@@ -1067,14 +1158,18 @@ jobs:
             FAILED=true
           fi
 
-          # GB200 integration tests are required only when explicitly enabled.
+          # GB200 tests are required only when explicitly enabled.
           if [ "$ENABLE_GB200_TESTING" == "true" ]; then
-            # GB200 integration tests may be skipped only for non-maintainer PRs
+            # GB200 tests may be skipped only for non-maintainer PRs
             # (no GB200 runners available); maintainer runs must always succeed.
             if [ "$GB200_RESULT" == "skipped" ] && [ "$IS_MAINTAINER" == "true" ]; then
               echo "❌ cicd-integration-tests-latest-gb200: skipped unexpectedly for a maintainer run"
               FAILED=true
             fi
+            if [ "$UNIT_GB200_RESULT" == "skipped" ] && [ "$IS_MAINTAINER" == "true" ]; then
+              echo "❌ cicd-unit-tests-latest-gb200: skipped unexpectedly for a maintainer run"
+              FAILED=true
+            fi
           else
             echo "✅ GB200 integration tests disabled by ENABLE_GB200_TESTING"
           fi

pyproject.toml

@@ -267,6 +267,7 @@ markers = [
     "internal: mark a test as a test to private/internal functions.",
     "flaky: mark flaky tests for LTS environment",
     "flaky_in_dev: mark flaky tests for DEV environment",
+    "launch_on_gb200: mark a unit test to be launched on GB200 hardware (4 GPUs/node)",
 ]
 
 [tool.coverage.run]

tests/test_utils/recipes/gb200/unit-tests.yaml

@@ -5,7 +5,7 @@ loggers: [stdout]
 spec:
   name: "{test_case}_{environment}_{platforms}_{tag}"
   model: unit-tests
-  nodes: 2
+  nodes: 1
   build: mcore-pyt-{environment}
   gpus: 4
   platforms: dgx_gb200
@@ -51,105 +51,26 @@ spec:
       --tag $TAG \
       --environment $ENVIRONMENT \
       --bucket $BUCKET \
+      --platform gb200 \
       --unit-test-repeat $UNIT_TEST_REPEAT \
       --log-dir {assets_dir}/logs/1/
 
-    ls -al 
+    ls -al
 
     cd $TEST_PATH
-    /opt/venv/bin/coverage xml 
+    /opt/venv/bin/coverage xml
     cp .coverage {assets_dir}/coverage_report
     cp coverage.xml {assets_dir}
 
+# GB200 unit-test selection is marker-driven: a single catch-all bucket is
+# narrowed to files carrying @pytest.mark.launch_on_gb200 by find_test_cases.py
+# (see tests/unit_tests/run_ci_test.sh --platform gb200). Re-shard into smaller
+# buckets here if the marked set grows large enough to need parallelism.
 products:
-  - test_case: [tests/unit_tests/test_model_configs.py]
-    products:
-      - environment: [dev]
-        tag: [latest, legacy]
-        scope: [unit-tests]
-        n_repeat: [1]
-        time_limit: [1800]
-  - test_case: [tests/unit_tests/test_fp8_param.py]
-    products:
-      - environment: [dev]
-        tag: [latest, legacy]
-        scope: [unit-tests]
-        n_repeat: [1]
-        time_limit: [1800]
-  - test_case: [tests/unit_tests/pipeline_parallel/**/*.py]
-    products:
-      - environment: [dev]
-        tag: [latest, legacy]
-        scope: [unit-tests]
-        n_repeat: [1]
-        time_limit: [1800]
-  - test_case: [tests/unit_tests/models/**/*.py]
-    products:
-      - environment: [dev]
-        tag: [latest, legacy]
-        scope: [unit-tests]
-        n_repeat: [1]
-        time_limit: [1800]
-  - test_case: [tests/unit_tests/data/**/*.py]
-    products:
-      - environment: [dev]
-        tag: [latest, legacy]
-        scope: [unit-tests]
-        n_repeat: [1]
-        time_limit: [1800]
-  - test_case: [tests/unit_tests/dist_checkpointing/test_optimizer.py]
-    products:
-      - environment: [dev]
-        tag: [latest, legacy]
-        scope: [unit-tests]
-        n_repeat: [1]
-        time_limit: [1800]
-  - test_case: [tests/unit_tests/dist_checkpointing/**/*.py]
-    products:
-      - environment: [dev]
-        tag: [latest, legacy]
-        scope: [unit-tests]
-        n_repeat: [1]
-        time_limit: [1800]
-  - test_case: [tests/unit_tests/dist_checkpointing/models/**/*.py]
-    products:
-      - environment: [dev]
-        tag: [latest, legacy]
-        scope: [unit-tests]
-        n_repeat: [1]
-        time_limit: [1800]
-  - test_case: [tests/unit_tests/dist_checkpointing/models/test_moe_experts.py]
-    products:
-      - environment: [dev]
-        tag: [latest, legacy]
-        scope: [unit-tests]
-        n_repeat: [1]
-        time_limit: [1800]
-  - test_case: [tests/unit_tests/transformer/**/*.py]
-    products:
-      - environment: [dev]
-        tag: [latest, legacy]
-        scope: [unit-tests]
-        n_repeat: [1]
-        time_limit: [1800]
-  - test_case: [tests/unit_tests/transformer/moe/**/*.py]
-    products:
-      - environment: [dev]
-        tag: [latest, legacy]
-        scope: [unit-tests]
-        n_repeat: [1]
-        time_limit: [1800]
-  - test_case: [tests/unit_tests/distributed/megatron_fsdp/**/*.py]
-    products:
-      - environment: [dev]
-        tag: [latest]
-        scope: [unit-tests]
-        n_repeat: [1]
-        time_limit: [1800]
   - test_case: [tests/unit_tests/**/*.py]
     products:
       - environment: [dev]
-        tag: [latest, legacy]
+        tag: [latest]
         scope: [unit-tests]
         n_repeat: [1]
         time_limit: [1800]

tests/unit_tests/find_test_cases.py

@@ -5,6 +5,26 @@
 import sys
 from pathlib import Path
 
+# Platforms whose unit-test selection is driven by a pytest marker rather than
+# by the full recipe bucket. Only files carrying the marker are launched.
+PLATFORM_MARKERS = {"gb200": "launch_on_gb200"}
+
+
+def file_has_marker(filepath, marker):
+    """Return True if the test file references the given pytest marker.
+
+    Args:
+        filepath: Path to a Python test file.
+        marker: The pytest marker name to look for (e.g. ``launch_on_gb200``).
+
+    Returns:
+        True if the marker name appears anywhere in the file, else False.
+    """
+    try:
+        return marker in Path(filepath).read_text()
+    except (OSError, UnicodeDecodeError):
+        return False
+
 
 def get_test_cases(yaml_file):
     result = subprocess.run(
@@ -62,6 +82,17 @@ def main():
         if test_case != BUCKET and is_child_of_bucket(test_case, BUCKET):
             files_to_ignore.update(expand_pattern(test_case))
 
+    # On marker-driven platforms, ignore any test file that does not carry the
+    # platform marker so only marked tests are launched. Restrict to pytest test
+    # files (test_*.py) so conftest.py and helper modules stay collectable.
+    marker = PLATFORM_MARKERS.get(GPU_TYPE)
+    if marker:
+        files_to_ignore.update(
+            f
+            for f in bucket_files
+            if Path(f).name.startswith("test_") and not file_has_marker(f, marker)
+        )
+
     # Output files to ignore
     for file in sorted(files_to_ignore & bucket_files):
         print(f"--ignore={file}")

tests/unit_tests/run_ci_test.sh

@@ -3,7 +3,7 @@ set -euxo pipefail
 
 # Parse command line arguments
 usage() {
-    echo "Usage: $0 --tag {latest|legacy} --environment {lts|dev} --bucket BUCKET [--unit-test-repeat N] [--unit-test-timeout N] --log-dir LOG_DIR"
+    echo "Usage: $0 --tag {latest|legacy} --environment {lts|dev} --bucket BUCKET [--platform {h100|gb200}] [--unit-test-repeat N] [--unit-test-timeout N] --log-dir LOG_DIR"
     exit 1
 }
 
@@ -15,6 +15,7 @@ cd $SCRIPT_PATH/../../
 UNIT_TEST_REPEAT=1
 UNIT_TEST_TIMEOUT=10
 LOG_DIR=$(pwd)/logs
+PLATFORM=h100
 
 # Parse arguments
 while [[ $# -gt 0 ]]; do
@@ -34,6 +35,10 @@ while [[ $# -gt 0 ]]; do
         BUCKET="$2"
         shift 2
         ;;
+    --platform)
+        PLATFORM="$2"
+        shift 2
+        ;;
     --unit-test-repeat)
         UNIT_TEST_REPEAT="$2"
         shift 2
@@ -96,6 +101,10 @@ fi
 cd $TEST_PATH
 
 MARKER=()
+if [[ "$PLATFORM" == "gb200" ]]; then
+    MARKER+=("launch_on_gb200")
+fi
+
 if [[ "$TAG" == "legacy" ]]; then
     MARKER+=("not internal")
 fi
@@ -117,7 +126,7 @@ export BUCKET
 IGNORE_ARGS=()
 while IFS= read -r line; do
     [[ -n "$line" ]] && IGNORE_ARGS+=("$line")
-done < <(python tests/unit_tests/find_test_cases.py "$BUCKET" "h100")
+done < <(python tests/unit_tests/find_test_cases.py "$BUCKET" "$PLATFORM")
 
 echo "------ARGUMENTS for SLURM ---"
 MASTER_ADDR=${MASTER_ADDR:-localhost}
@@ -141,6 +150,23 @@ export NCCL_MAX_NCHANNELS=1
 export NCCL_NVLS_ENABLE=0
 export ONE_LOGGER_JOB_CATEGORY=test
 
+# Run a pytest command. On marker-driven platforms a bucket can legitimately
+# contain no matching tests; treat pytest's "no tests collected" (exit 5) as a
+# pass there instead of aborting the job under `set -e`.
+run_test_cmd() {
+    local cmd="$1"
+    local rc=0
+    set +e
+    eval "$cmd"
+    rc=$?
+    set -e
+    if [[ "$rc" -eq 5 && "$PLATFORM" == "gb200" ]]; then
+        echo "No tests collected for this bucket on $PLATFORM (pytest exit 5) — treating as pass."
+        return 0
+    fi
+    return "$rc"
+}
+
 for i in $(seq $UNIT_TEST_REPEAT); do
     echo "Running prod test suite."
     CMD=$(echo uv run --no-sync python -m torch.distributed.run ${DISTRIBUTED_ARGS[@]} \
@@ -151,7 +177,7 @@ for i in $(seq $UNIT_TEST_REPEAT); do
         -vs \
         ${IGNORE_ARGS[@]} \
         -m "'not experimental and ${MARKER_ARG}'" $(echo "$BUCKET" | sed 's|/\*\*/\*\.py$||'))
-    eval "$CMD"
+    run_test_cmd "$CMD"
 
     if [[ "$TAG" == "latest" ]]; then
         CMD=$(echo uv run --no-sync python -m torch.distributed.run ${DISTRIBUTED_ARGS[@]} -m pytest \
@@ -160,7 +186,7 @@ for i in $(seq $UNIT_TEST_REPEAT); do
              ${IGNORE_ARGS[@]} \
             -m "'experimental and ${MARKER_ARG}'" $(echo "$BUCKET" | sed 's|/\*\*/\*\.py$||'))
 
-        eval "$CMD"
+        run_test_cmd "$CMD"
     fi
 
 done