fix(ci): restore torch E2E coverage in unit-vs-E2E summary

.github/workflows/ci.yaml

@@ -424,7 +424,7 @@ jobs:
           TIER_ARGS=()
           [[ "${PRIMUS_CI_FULL:-0}" == "1" ]] && TIER_ARGS=(--run-slow)
           pytest --maxfail=1 -s ./tests/unit_tests/ "${TIER_ARGS[@]}" \
-            --cov=primus --cov-report=term-missing:skip-covered \
+            --cov=primus --cov-branch --cov-report=term-missing:skip-covered \
             --junitxml="${GITHUB_WORKSPACE}/test-reports/core-unit.xml" \
             --deselect=tests/unit_tests/megatron/cco/test_tp_overlap.py::TPOverlapTestCase::test_fp8_te_linear \
             --deselect=tests/unit_tests/megatron/cco/test_tp_overlap.py::TPOverlapTestCase::test_te_linear \
@@ -439,10 +439,9 @@ jobs:
         run: |
           docker exec -i -e GITHUB_WORKSPACE="${CWS}" -w "${CWS}" "${UT_CONTAINER}" bash -s <<'IN'
           set +e
-          # Keep the unit data for the final unit-vs-E2E comparison table
-          # (.coverage.unit is fed to `coverage combine` after the E2E steps).
+          # .coverage_unit (not .coverage.unit) avoids coverage combine treating it as a .coverage.* shard.
           python -m coverage json -o coverage_unit.json 2>/dev/null || true
-          cp .coverage "${GITHUB_WORKSPACE}/.coverage.unit" 2>/dev/null || true
+          cp .coverage "${GITHUB_WORKSPACE}/.coverage_unit" 2>/dev/null || true
           IN
       - name: Setup E2E training coverage
         if: always()
@@ -458,6 +457,9 @@ jobs:
           set -e
           SP=$(python -c "import site; print(site.getsitepackages()[0])")
           echo "import coverage; coverage.process_startup()" > "$SP/primus_e2e_coverage.pth"
+          # No branch=true: it crashed all 8 ranks (SIGABRT) in
+          # test_memory_benchmark under real GPU training. E2E stays
+          # line-only; unit gets --cov-branch below instead.
           printf '[run]\nparallel = true\nsource = primus\nsigterm = true\n' > "${GITHUB_WORKSPACE}/.coveragerc_e2e"
           IN
           echo "COVERAGE_PROCESS_START=${CWS}/.coveragerc_e2e" >> "$GITHUB_ENV"
@@ -608,15 +610,36 @@ jobs:
         continue-on-error: true
         run: |
           docker exec -i -e GITHUB_WORKSPACE="${CWS}" -w "${CWS}" "${UT_CONTAINER}" bash -s <<'IN'
-          set +e
-          # Don't instrument coverage's own helper processes here.
+          set -euo pipefail
+          cd "${GITHUB_WORKSPACE}"
           unset COVERAGE_PROCESS_START
-          # Merge per-rank E2E data, then line-merge unit + E2E into one dataset.
-          # The JSON is consumed by the coverage-summary job (no per-job summary).
-          python -m coverage combine 2>/dev/null || true
-          COVERAGE_FILE="${GITHUB_WORKSPACE}/.coverage_all" python -m coverage combine --keep \
-            "${GITHUB_WORKSPACE}/.coverage.unit" "${GITHUB_WORKSPACE}/.coverage_e2e" 2>/dev/null || true
-          COVERAGE_FILE="${GITHUB_WORKSPACE}/.coverage_all" python -m coverage json -o coverage_combined.json 2>/dev/null || true
+          shopt -s nullglob
+          # COVERAGE_FILE is set explicitly below: docker exec doesn't
+          # inherit the GITHUB_ENV value set by the earlier step.
+          e2e_shards=(.coverage_e2e.*)
+          if [ ${#e2e_shards[@]} -gt 0 ]; then
+            COVERAGE_FILE=.coverage_e2e python -m coverage combine
+          fi
+          if [ ! -s .coverage_unit ] && [ ! -s .coverage_e2e ]; then
+            echo "::warning::no torch coverage data to combine (neither unit nor E2E wrote a snapshot)"
+            exit 1
+          fi
+          if [ ! -s .coverage_e2e ]; then
+            echo "::warning::no torch E2E coverage data; skipping coverage_combined.json (Unit+E2E summary will fall back to unit-only)"
+            exit 0
+          fi
+          parts=(.coverage_e2e)
+          if [ -s .coverage_unit ]; then
+            # coverage combine can't mix branch data with statement data, so
+            # project .coverage_unit to line-only first; coverage_unit.json
+            # (exported earlier) still keeps the real branch data.
+            python -c "import coverage; src=coverage.CoverageData(basename='.coverage_unit'); src.read(); dst=coverage.CoverageData(basename='.coverage_unit_lines'); [dst.add_lines({f: src.lines(f)}) for f in src.measured_files()]; dst.write()"
+            parts=(.coverage_unit_lines .coverage_e2e)
+          fi
+          echo "Combining: ${parts[*]}"
+          COVERAGE_FILE=.coverage_all python -m coverage combine --keep "${parts[@]}"
+          COVERAGE_FILE=.coverage_all python -m coverage json -o coverage_combined.json
+          test -s coverage_combined.json
           IN
       - name: Upload torch coverage json
         if: always()
@@ -645,7 +668,7 @@ jobs:
           set +e
           cd "${GITHUB_WORKSPACE}" || exit 0
           rm -rf logs test-reports ut_out .pytest_cache .hypothesis \
-                 .coverage .coverage.* .coverage_e2e* .coverage.unit .coverage_all \
+                 .coverage .coverage.* .coverage_e2e* .coverage_unit .coverage_unit_lines* .coverage_all \
                  .coveragerc_e2e .e2e_scope coverage_unit.json coverage_combined.json
           # Bytecode caches written during the run (root-owned; block checkout clean).
           find . -type d -name __pycache__ -prune -exec rm -rf {} + 2>/dev/null
@@ -757,6 +780,7 @@ jobs:
           pip install coverage >/dev/null 2>&1 || true
           SP=$(python -c "import site; print(site.getsitepackages()[0])")
           echo "import coverage; coverage.process_startup()" > "$SP/primus_e2e_coverage.pth"
+          # branch=true stays off here too -- see the torch E2E step for why.
           printf '[run]\nparallel = true\nsource = primus\nsigterm = true\n' > "$GITHUB_WORKSPACE/.coveragerc_e2e"
           echo "COVERAGE_PROCESS_START=$GITHUB_WORKSPACE/.coveragerc_e2e" >> "$GITHUB_ENV"
           echo "COVERAGE_FILE=$GITHUB_WORKSPACE/.coverage_e2e" >> "$GITHUB_ENV"

tools/ci/coverage_summary.py

@@ -99,7 +99,7 @@ def _priority(group: str) -> int:
     "core": "\U0001F511 Primus core;<br>imported by every run",
     "backends/megatron": (
         "100+ patches gated by fp8 / MoE /<br>zero-bubble-pp / fsdp2 flags;<br>"
-        "CI E2E runs only 1-2 configs"
+        "CI E2E now runs 15+ model configs/PR<br>(weekly tier adds ~10 more)"
     ),
     "backends/transformer_engine": "fp8 GEMM / attn-overlap kernels;<br>only hit when an E2E enables fp8",
     "backends/diffusion": "no E2E trainer suite yet<br>(unit-tested only)",
@@ -190,7 +190,8 @@ def group_executed(group):  # single-report mode: hide groups a partial run neve
             % (_pct(tc, tn), _pct(te, tn), format(tn, ","), excl)
         )
         out.append(
-            "_Including all modules (nothing excluded): Unit %.1f%% -> Unit+E2E %.1f%%._\n" % (p_all, s_all)
+            "_Including all modules (nothing excluded; Unit also counts "
+            "branches, E2E is lines only): Unit %.1f%% -> Unit+E2E %.1f%%._\n" % (p_all, s_all)
         )
         out.append(_LEGEND)
         out += ["| Module | Stmts | Unit | Unit+E2E | Notes |", "|---|--:|--:|--:|---|"]