diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 3b120bca3..94a33ba9a 100644
--- a/.github/workflows/ci.yaml
+++ b/.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"
diff --git a/tools/ci/coverage_summary.py b/tools/ci/coverage_summary.py
index 400136a74..c15a7ab78 100644
--- a/tools/ci/coverage_summary.py
+++ b/tools/ci/coverage_summary.py
@@ -99,7 +99,7 @@ def _priority(group: str) -> int:
"core": "\U0001F511 Primus core;
imported by every run",
"backends/megatron": (
"100+ patches gated by fp8 / MoE /
zero-bubble-pp / fsdp2 flags;
"
- "CI E2E runs only 1-2 configs"
+ "CI E2E now runs 15+ model configs/PR
(weekly tier adds ~10 more)"
),
"backends/transformer_engine": "fp8 GEMM / attn-overlap kernels;
only hit when an E2E enables fp8",
"backends/diffusion": "no E2E trainer suite yet
(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 |", "|---|--:|--:|--:|---|"]