CI: upload Python coverage.json from pytest (non-blocking)

.github/workflows/rocm-ci.yml

@@ -179,6 +179,8 @@ jobs:
           pip install --upgrade hypothesis setuptools
           pip install --no-build-isolation --no-deps "$TE_TORCH_PKG"
           pip install --no-build-isolation --no-deps "$TE_JAX_PKG"
+          # Optional: wrap pytest with coverage.py. Export is non-blocking.
+          pip install 'coverage>=7'
           EOF
           )"
 
@@ -193,11 +195,13 @@ jobs:
           # One subdir per suite so PyTorch/JAX/Core get independent reports and
           # never collide on identically-named XML (e.g. test_sanity_import.auto.xml
           # exists in both the torch and jax suites and runs in parallel).
-          rm -rf test-results && mkdir -p test-results/torch test-results/jax test-results/core
+          rm -rf test-results python-coverage && mkdir -p test-results/torch test-results/jax test-results/core python-coverage
 
           docker exec \
             -e TEST_SGPU=1 \
             -e TEST_LEVEL=${{ env.TEST_LEVEL }} \
+            -e TE_COVERAGE=1 \
+            -e COVERAGE_FILE=/workspace/python-coverage/.coverage \
             -e JUNITXML_PREFIX=/workspace/test-results/ \
             -e JUNITXML_SUFFIX=.xml \
             -e HF_TOKEN="$HF_TOKEN" \
@@ -220,6 +224,20 @@ jobs:
           python3 ci/junit_report.py test-results/core \
             --title "sGPU Core (${{ matrix.arch_label }})"
 
+      - name: Export Python coverage
+        if: always()
+        run: |
+          # JUnit XML is pass/fail of test cases. coverage.json is which
+          # installed transformer_engine Python functions this job executed.
+          # A missing artifact means export failed; do not fail the job.
+          docker exec \
+            -e TE_PATH=/workspace \
+            -e CI_ARCH=${{ matrix.arch_label }} \
+            -e CI_SUITE=sgpu \
+            -e COVERAGE_FILE=/workspace/python-coverage/.coverage \
+            te-runner bash /workspace/ci/export_python_coverage.sh /workspace/python-coverage \
+            || true
+
       - name: Check suite failure status
         if: always()
         run: |
@@ -252,6 +270,17 @@ jobs:
           if-no-files-found: ignore
           retention-days: 5
 
+      - name: Upload Python coverage
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: python-coverage-sgpu-${{ matrix.arch_label }}
+          path: |
+            python-coverage/coverage.json
+            python-coverage/coverage-meta.txt
+          if-no-files-found: ignore
+          retention-days: 14
+
       - name: Cleanup container
         if: always()
         run: docker rm -f te-runner || true
@@ -312,6 +341,7 @@ jobs:
           pip install ninja pybind11[global]
           pip install --upgrade hypothesis setuptools
           pip install --no-build-isolation --no-deps "$TE_FW_PKG"
+          pip install 'coverage>=7'
           EOF
           )"
 
@@ -329,11 +359,13 @@ jobs:
             *)       echo "::error::Unknown framework: ${{ matrix.framework }}"; exit 1 ;;
           esac
 
-          rm -rf test-results && mkdir -p test-results
+          rm -rf test-results python-coverage && mkdir -p test-results python-coverage
 
           docker exec \
             -e TEST_MGPU=1 \
             -e TEST_LEVEL=${{ env.TEST_LEVEL }} \
+            -e TE_COVERAGE=1 \
+            -e COVERAGE_FILE=/workspace/python-coverage/.coverage \
             -e TEST_SCRIPT=$TEST_SCRIPT \
             -e LOG_FILE=$LOG_FILE \
             -e SUITE_NAME=$SUITE_NAME \
@@ -367,6 +399,20 @@ jobs:
           python3 ci/junit_report.py test-results \
             --title "mGPU ${{ matrix.framework == 'pytorch' && 'Torch' || 'JAX' }} (${{ matrix.arch_label }})"
 
+      - name: Export Python coverage
+        if: always()
+        run: |
+          # JUnit XML is pass/fail of test cases. coverage.json is which
+          # installed transformer_engine Python functions this job executed.
+          # A missing artifact means export failed; do not fail the job.
+          docker exec \
+            -e TE_PATH=/workspace \
+            -e CI_ARCH=${{ matrix.arch_label }} \
+            -e CI_SUITE=mgpu-${{ matrix.framework }} \
+            -e COVERAGE_FILE=/workspace/python-coverage/.coverage \
+            te-runner bash /workspace/ci/export_python_coverage.sh /workspace/python-coverage \
+            || true
+
       - name: Upload logs
         if: always()
         uses: actions/upload-artifact@v4
@@ -378,6 +424,17 @@ jobs:
           if-no-files-found: ignore
           retention-days: 5
 
+      - name: Upload Python coverage
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: python-coverage-mgpu-${{ matrix.arch_label }}-${{ matrix.framework }}
+          path: |
+            python-coverage/coverage.json
+            python-coverage/coverage-meta.txt
+          if-no-files-found: ignore
+          retention-days: 14
+
       - name: Cleanup container
         if: always()
         run: docker rm -f te-runner || true

ci/_utils.sh

@@ -371,8 +371,22 @@ pytest_run() {
         _sink_plugin="-p te_ci_result_sink"
         _pytest_pythonpath="${TE_PATH}ci${PYTHONPATH:+:$PYTHONPATH}"
     fi
+    # Optional Python execution coverage (TE_COVERAGE=1). Off by default so
+    # a local ci/pytorch.sh run is unchanged. This wraps the same pytest
+    # invocation; a coverage failure to start falls back to plain pytest.
+    # Parallel pytest (sGPU) writes coverage.py shards; CI combines them
+    # after the job. This is not JUnit: it records which installed
+    # transformer_engine functions ran, not which test cases passed.
+    _pytest_launcher="python3 -m"
+    if [ "${TE_COVERAGE:-}" = 1 ]; then
+        if command -v coverage >/dev/null 2>&1; then
+            _pytest_launcher="coverage run --branch --parallel-mode --source=transformer_engine -m"
+        else
+            echo "TE_COVERAGE=1 but coverage is not installed; running pytest without coverage" >&2
+        fi
+    fi
     TE_RESULT_SINK="$_result_sink" PYTHONPATH="$_pytest_pythonpath" \
-        python3 -m pytest -v -rfEs \
+        $_pytest_launcher pytest -v -rfEs \
         --timeout=$PYTEST_TIMEOUT --timeout-method=$PYTEST_TIMEOUT_METHOD \
         $_sink_plugin $_junitxml_arg $TEST_PYTEST_ARGS "$TEST_DIR/$@"
     _pytest_rc=$?

ci/export_python_coverage.sh

@@ -0,0 +1,80 @@
+#!/bin/sh
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+#
+# Combine coverage.py shards from pytest_run and write coverage.json.
+#
+# This is execution coverage of the installed transformer_engine package
+# (which Python functions/lines ran). It is not JUnit: pass/fail of test
+# cases stays in the XML artifacts.
+#
+# Never fails the CI job. Never writes an empty stub — a missing artifact
+# is the honest signal that coverage was not produced.
+
+OUT_DIR="${1:-/workspace/python-coverage}"
+mkdir -p "$OUT_DIR" || exit 0
+cd "$OUT_DIR" || exit 0
+
+if ! command -v coverage >/dev/null 2>&1; then
+    echo "coverage is not installed; skipping Python coverage export"
+    exit 0
+fi
+
+export COVERAGE_FILE="${COVERAGE_FILE:-$OUT_DIR/.coverage}"
+JSON="$OUT_DIR/coverage.json"
+META="$OUT_DIR/coverage-meta.txt"
+rm -f "$JSON" "$META"
+
+# Shards from parallel pytest_run invocations live next to COVERAGE_FILE.
+coverage combine || true
+if ! coverage json -o "$JSON"; then
+    echo "coverage json failed; not uploading a stub"
+    rm -f "$JSON"
+    exit 0
+fi
+
+if [ ! -s "$JSON" ]; then
+    echo "coverage.json is empty; not uploading a stub"
+    rm -f "$JSON"
+    exit 0
+fi
+
+if ! python3 - "$JSON" <<'PY'
+import json
+import sys
+from pathlib import Path
+
+path = Path(sys.argv[1])
+try:
+    document = json.loads(path.read_text(encoding="utf-8"))
+except (OSError, json.JSONDecodeError, UnicodeError):
+    raise SystemExit(1)
+if not document.get("files"):
+    raise SystemExit(1)
+PY
+then
+    echo "coverage.json has no measured files; not uploading a stub"
+    rm -f "$JSON"
+    exit 0
+fi
+
+COMMIT="unknown"
+if command -v git >/dev/null 2>&1; then
+    COMMIT=$(git -C "${TE_PATH:-/workspace}" rev-parse HEAD 2>/dev/null || echo unknown)
+fi
+
+{
+    echo "commit=$COMMIT"
+    [ -n "${CI_ARCH:-}" ] && echo "arch=$CI_ARCH"
+    [ -n "${CI_SUITE:-}" ] && echo "suite=$CI_SUITE"
+    echo "format=coverage.py JSON"
+    echo "source=transformer_engine (installed package measured by CI pytest)"
+    echo "not=JUnit pass/fail of test cases"
+    echo "not=C++/HIP kernels (.cu/.hip); those need llvm-cov"
+    echo "not=torchrun/mpirun child processes"
+    echo "not=direct python script invocations (examples, checkpoint, benchmarks)"
+} > "$META"
+
+echo "Wrote $JSON (commit=$COMMIT)"
+exit 0