Skip to content

revert(ci): undo AOTI .pt2 compile / GPU test overlap (#5882) #5916

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wanghan-iapcm merged 1 commit into from
Jul 28, 2026
+107 −351
Merged

revert(ci): undo AOTI .pt2 compile / GPU test overlap (#5882) #5916

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 1 addition & 78 deletions .github/workflows/test_cuda.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,29 +19,13 @@ jobs:
runs-on: gpu
# The serial Python test suite can take several hours on the GPU runner.
timeout-minutes: 480
# Share ONE AOTInductor on-disk compile cache across the two pytest lanes;
# actions/cache (below) persists it across runs so unchanged models never
# recompile.
env:
TORCHINDUCTOR_CACHE_DIR: ${{ github.workspace }}/.inductor-cache
# https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845
# container:
# image: nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04
# options: --gpus all
if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
steps:
- uses: actions/checkout@v7
- name: Cache AOTInductor compile artifacts
# Restored AFTER checkout so checkout's git-clean can't wipe it. The repo
# Actions cache is capped at 10 GB; if the inductor cache outgrows it the
# oldest entries are evicted -- correctness is unaffected (inductor
# re-validates each graph; a miss just recompiles).
uses: actions/cache@v6
with:
path: ${{ github.workspace }}/.inductor-cache
key: inductor-cuda-${{ github.sha }}
restore-keys: |
inductor-cuda-
- uses: actions/setup-python@v7
with:
python-version: "3.11"
Expand Down Expand Up @@ -78,84 +62,23 @@ jobs:
DP_ENABLE_NATIVE_OPTIMIZATION: 1
DP_ENABLE_PYTORCH: 1
- run: dp --version
- name: Run Python tests (overlap CPU-bound aoti compiles with the GPU lane)
# Two pytest lanes on one GPU box:
# Lane 1 (CPU stream): the .pt2-freezing tests (-m aoti_compile) --
# CPU-bound, GPU ~98% idle while inductor/g++/ptxas run.
# Lane 2 (GPU stream): the GPU-bound remainder (-m "not aoti_compile")
# -- the bulk of the suite and the wall-clock bottleneck. It does not
# freeze .pt2, so it only reads the inductor cache, never competing
# to write.
# Wall-clock becomes max(lane 1, lane 2) instead of their sum; the whole
# compile stream is hidden behind the GPU lane. (The C++ build + fixture
# generation runs in the separate test_cc job.)
run: |
# Cap the compile lane's inductor threads to ~half the cores so lane 2
# keeps CPU for its host-side GPU dispatch; nice the compile work too.
nc=$(( $(nproc) / 2 )); [ "$nc" -lt 1 ] && nc=1
# rc=0 + `|| rc=$?` so a lane's failure does not let the step's default
# `set -e` skip the wall-clock echo (the timing we most want on failure).
( ta=$(date +%s); rc=0
TORCHINDUCTOR_COMPILE_THREADS=$nc nice -n 15 \
python -m pytest source/tests -m "aoti_compile" \
-p no:cacheprovider --junit-xml=junit-aoti.xml || rc=$?
# pytest exit 5 == "no tests collected"; harmless for the compile lane.
[ "$rc" = 5 ] && rc=0
echo "[lane aoti_compile] wall=$(( $(date +%s) - ta ))s exit=$rc"
exit $rc ) &
pid_a=$!
( tb=$(date +%s); rc=0
python -m pytest source/tests -m "not aoti_compile" \
-p no:cacheprovider --junit-xml=junit-gpu.xml || rc=$?
echo "[lane gpu] wall=$(( $(date +%s) - tb ))s exit=$rc"
exit $rc ) &
pid_b=$!
rc_a=0; wait $pid_a || rc_a=$?
rc_b=0; wait $pid_b || rc_b=$?
echo "lane exit codes: aoti_compile=$rc_a gpu=$rc_b"
exit $(( rc_a != 0 || rc_b != 0 ))
- run: python -m pytest source/tests
env:
NUM_WORKERS: 0
CUDA_VISIBLE_DEVICES: 0
# See https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
XLA_PYTHON_CLIENT_PREALLOCATE: false
XLA_PYTHON_CLIENT_ALLOCATOR: platform
FLAGS_use_stride_compute_kernel: 0
- name: Upload per-lane JUnit reports
# junit-aoti.xml = compile lane (T_A), junit-gpu.xml = GPU lane (T_B);
# used to compare the two lanes against the pre-split serial baseline.
if: always()
uses: actions/upload-artifact@v7
with:
name: cuda-pytest-junit
path: |
junit-aoti.xml
junit-gpu.xml
if-no-files-found: warn

test_cc:
name: Test C++ on CUDA
runs-on: gpu
# The gen_*.py fixture builds drive the same AOTInductor backend on the same
# architectures as the Python lanes, so they reuse the kernels cached by
# test_python / previous runs.
env:
TORCHINDUCTOR_CACHE_DIR: ${{ github.workspace }}/.inductor-cache
if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
steps:
# Jobs run on separate runners, so the C++ job needs its own complete
# CUDA toolchain and Python dependency installation.
- uses: actions/checkout@v7
- name: Cache AOTInductor compile artifacts
# Restored AFTER checkout so checkout's git-clean can't wipe it. Both
# CUDA jobs share the same key; if both try to save it, the second save
# is skipped with a warning (caches are immutable), which is harmless.
uses: actions/cache@v6
with:
path: ${{ github.workspace }}/.inductor-cache
key: inductor-cuda-${{ github.sha }}
restore-keys: |
inductor-cuda-
- uses: actions/setup-python@v7
with:
python-version: "3.11"
Expand Down
5 changes: 1 addition & 4 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -495,10 +495,7 @@ runtime-evaluated-base-classes = ["torch.nn.Module"]
"examples/**/*.py" = ["T20", "TID253"]

[tool.pytest.ini_options]
markers = [
"run",
"aoti_compile: test freezes a .pt2 (AOTInductor) artifact -- a CPU-bound compile stage the CUDA CI runs concurrently with GPU-bound tests (auto-applied in source/tests/conftest.py)",
]
markers = "run"
timeout = 1800

[tool.coverage.run]
Expand Down
239 changes: 105 additions & 134 deletions source/install/test_cc_local.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,6 @@
#!/bin/bash
set -ex

# Phase gating (default: run everything, so existing callers like test_cc.yml
# are unaffected). The CUDA workflow splits the CPU-bound build+fixture-gen
# from the GPU-bound ctest so the former can overlap the Python GPU tests:
# DP_CC_SKIP_CTEST=1 -> configure + build + install + gen fixtures, no ctest
# DP_CC_SKIP_BUILD=1 -> skip build/gen, run only ctest on the built tree
DP_CC_SKIP_BUILD=${DP_CC_SKIP_BUILD:-0}
DP_CC_SKIP_CTEST=${DP_CC_SKIP_CTEST:-0}

if [ "$DP_VARIANT" = "cuda" ]; then
CUDA_ARGS="-DUSE_CUDA_TOOLKIT=TRUE"
elif [ "$DP_VARIANT" = "rocm" ]; then
Expand All @@ -28,47 +20,28 @@ BUILD_TMP_DIR=${SCRIPT_PATH}/../build_tests
PADDLE_INFERENCE_DIR=${BUILD_TMP_DIR}/paddle_inference_install_dir
mkdir -p ${BUILD_TMP_DIR}
cd ${BUILD_TMP_DIR}

# LD_LIBRARY_PATH additions needed by BOTH the gen scripts (which import
# deepmd.pt and dlopen the custom op .so that depends on libdeepmd.so in the
# install prefix) AND ctest. Set once up front so either phase works alone.
# The install prefix may not exist yet during the build; a missing dir in
# LD_LIBRARY_PATH is harmless.
export LD_LIBRARY_PATH=${INSTALL_PREFIX}/lib:${LD_LIBRARY_PATH}
if [ "${ENABLE_PADDLE:-TRUE}" == "TRUE" ]; then
export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PADDLE_INFERENCE_DIR}/third_party/install/onednn/lib:${PADDLE_INFERENCE_DIR}/third_party/install/mklml/lib
fi

if [ "${DP_CC_SKIP_BUILD}" != "1" ]; then
cmake \
-D ENABLE_TENSORFLOW=${ENABLE_TENSORFLOW:-TRUE} \
-D ENABLE_PYTORCH=${ENABLE_PYTORCH:-TRUE} \
-D ENABLE_PADDLE=${ENABLE_PADDLE:-TRUE} \
-D INSTALL_TENSORFLOW=FALSE \
-D USE_TF_PYTHON_LIBS=${ENABLE_TENSORFLOW:-TRUE} \
-D USE_PT_PYTHON_LIBS=${ENABLE_PYTORCH:-TRUE} \
-D CMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} \
-D BUILD_TESTING:BOOL=TRUE \
-D LAMMPS_VERSION=stable_22Jul2025_update2 \
${CUDA_ARGS} ..
cmake --build . -j${NPROC}
cmake --install .
# Generate PT/PT2 model files for C++ tests.
# Must run after cmake --build so that libdeepmd_op_pt.so (custom ops) is available.
if [ "${ENABLE_PYTORCH:-TRUE}" == "TRUE" ]; then
# Install the custom op .so to SHARED_LIB_DIR so that `import deepmd.pt`
# loads it via cxx_op.py.
#
# The install MUST be atomic (copy to a temp file in the same dir, then
# os.replace). shutil.copy2 overwrites the destination inode IN PLACE,
# which corrupts the mmap'd code pages of any process that already
# dlopen'd this .so -- e.g. the concurrent Python test lane in the CUDA
# CI overlap -- and SIGSEGVs it (see #5882). os.replace swaps the
# directory entry to a NEW inode instead: live mappings keep the old
# (refcounted) inode intact, and new dlopens (the gen_*.py scripts) pick
# up the new file.
python -c '
import os, shutil, sys
cmake \
-D ENABLE_TENSORFLOW=${ENABLE_TENSORFLOW:-TRUE} \
-D ENABLE_PYTORCH=${ENABLE_PYTORCH:-TRUE} \
-D ENABLE_PADDLE=${ENABLE_PADDLE:-TRUE} \
-D INSTALL_TENSORFLOW=FALSE \
-D USE_TF_PYTHON_LIBS=${ENABLE_TENSORFLOW:-TRUE} \
-D USE_PT_PYTHON_LIBS=${ENABLE_PYTORCH:-TRUE} \
-D CMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} \
-D BUILD_TESTING:BOOL=TRUE \
-D LAMMPS_VERSION=stable_22Jul2025_update2 \
${CUDA_ARGS} ..
cmake --build . -j${NPROC}
cmake --install .
# Generate PT/PT2 model files for C++ tests.
# Must run after cmake --build so that libdeepmd_op_pt.so (custom ops) is available.
if [ "${ENABLE_PYTORCH:-TRUE}" == "TRUE" ]; then
# Install the custom op .so to SHARED_LIB_DIR so that `import deepmd.pt`
# loads it via cxx_op.py. The .so depends on libdeepmd.so (compute
# kernels) in the install prefix, so add that to LD_LIBRARY_PATH too.
export LD_LIBRARY_PATH=${INSTALL_PREFIX}/lib:${LD_LIBRARY_PATH}
python -c '
import shutil, sys
from pathlib import Path
from deepmd.env import SHARED_LIB_DIR
so = Path("'"${BUILD_TMP_DIR}"'") / "op" / "pt" / "libdeepmd_op_pt.so"
Expand All @@ -79,97 +52,95 @@ elif dst.exists() and dst.resolve() == so.resolve():
print(f"Already linked: {dst} -> {so}")
else:
SHARED_LIB_DIR.mkdir(parents=True, exist_ok=True)
tmp = dst.with_name(dst.name + f".tmp{os.getpid()}")
shutil.copy2(str(so), str(tmp))
os.replace(str(tmp), str(dst))
shutil.copy2(str(so), str(dst))
print(f"Installed {so} -> {dst}")
'
# When the build uses -fsanitize=leak, the custom op .so requires the LSAN
# runtime to be preloaded (otherwise dlopen fails). We disable leak detection
# in the gen scripts to avoid false reports from torch/paddle internals.
INFER_SCRIPT_PATH=${SCRIPT_PATH}/../tests/infer
# Remove stale generated model files so they can't be accidentally reused
# if gen scripts change format or the code version changes.
rm -f ${INFER_SCRIPT_PATH}/*.pt2 ${INFER_SCRIPT_PATH}/*.pte
_GEN_ENV=""
if echo "${CXXFLAGS:-}" | grep -q fsanitize=leak; then
_LSAN_LIB=$(gcc -print-file-name=liblsan.so 2>/dev/null || true)
if [ -n "${_LSAN_LIB}" ] && [ -f "${_LSAN_LIB}" ]; then
# DP_GEN_UNDER_SANITIZER: explicit signal for gen scripts that need
# to skip sanitizer-incompatible sections (e.g. gen_dpa2.py's
# AOTInductor graph .pt2 eval, which can SEGV under the LSAN
# runtime). Sniffing LD_PRELOAD inside the gen script is NOT
# reliable: the sanitizer runtime removes its own entry from the
# process environment during startup.
_GEN_ENV="LD_PRELOAD=${_LSAN_LIB} LSAN_OPTIONS=detect_leaks=0 DP_GEN_UNDER_SANITIZER=lsan"
fi
# When the build uses -fsanitize=leak, the custom op .so requires the LSAN
# runtime to be preloaded (otherwise dlopen fails). We disable leak detection
# in the gen scripts to avoid false reports from torch/paddle internals.
INFER_SCRIPT_PATH=${SCRIPT_PATH}/../tests/infer
# Remove stale generated model files so they can't be accidentally reused
# if gen scripts change format or the code version changes.
rm -f ${INFER_SCRIPT_PATH}/*.pt2 ${INFER_SCRIPT_PATH}/*.pte
_GEN_ENV=""
if echo "${CXXFLAGS:-}" | grep -q fsanitize=leak; then
_LSAN_LIB=$(gcc -print-file-name=liblsan.so 2>/dev/null || true)
if [ -n "${_LSAN_LIB}" ] && [ -f "${_LSAN_LIB}" ]; then
# DP_GEN_UNDER_SANITIZER: explicit signal for gen scripts that need
# to skip sanitizer-incompatible sections (e.g. gen_dpa2.py's
# AOTInductor graph .pt2 eval, which can SEGV under the LSAN
# runtime). Sniffing LD_PRELOAD inside the gen script is NOT
# reliable: the sanitizer runtime removes its own entry from the
# process environment during startup.
_GEN_ENV="LD_PRELOAD=${_LSAN_LIB} LSAN_OPTIONS=detect_leaks=0 DP_GEN_UNDER_SANITIZER=lsan"
fi
# Run gen scripts in parallel for faster model generation.
# Wait on each PID separately so any failure is caught by set -e.
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_sea.py &
PID1=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa1.py &
PID2=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa2.py &
PID3=$!
wait $PID1
wait $PID2
wait $PID3
fi
# Run gen scripts in parallel for faster model generation.
# Wait on each PID separately so any failure is caught by set -e.
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_sea.py &
PID1=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa1.py &
PID2=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa2.py &
PID3=$!
wait $PID1
wait $PID2
wait $PID3

env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa3.py &
PID4=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_fparam_aparam.py &
PID5=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_model_devi.py &
PID6=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_chg_spin.py &
PID9=$!
wait $PID4
wait $PID5
wait $PID6
wait $PID9
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa3.py &
PID4=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_fparam_aparam.py &
PID5=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_model_devi.py &
PID6=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_chg_spin.py &
PID9=$!
wait $PID4
wait $PID5
wait $PID6
wait $PID9

env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4.py &
PID9=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa1_pairexcl.py &
PID10=$!
wait $PID9
wait $PID10
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4.py &
PID9=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa1_pairexcl.py &
PID10=$!
wait $PID9
wait $PID10

env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_spin.py &
PID7=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_spin_model_devi.py &
PID8=$!
wait $PID7
wait $PID8
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_spin.py &
PID7=$!
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_spin_model_devi.py &
PID8=$!
wait $PID7
wait $PID8

# Native-spin DPA4 graph archives (baseline + model-level pair
# exclusion). Without this the whole native-spin graph C++ suite
# GTEST_SKIPs on the missing fixture, which is how a dead
# applyPairExclusion seam in DeepSpinPTExpt went unnoticed.
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_spin.py &
PID11=$!
# DPA4 + analytical ZBL bridging: a linear COMPOSITION on the graph
# lower, which no other C++ fixture covers.
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_zbl.py &
PID12=$!
# Native-spin DPA4 + charge-spin FiLM: the ONLY fixture with both
# is_spin=true and dim_chg_spin>0, i.e. the only one on which the
# runtime charge_spin argument of DeepSpin::compute is not inert.
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_spin_chgspin.py &
PID13=$!
# Native-spin DPA4 + ZBL bridging COMBINED: spin reaching a linear
# composition, and the only fixture pinning that a bridged model
# exports NO with-comm artifact (single-rank only by construction).
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_spin_zbl.py &
PID14=$!
wait $PID11
wait $PID12
wait $PID13
wait $PID14
fi
# Native-spin DPA4 graph archives (baseline + model-level pair
# exclusion). Without this the whole native-spin graph C++ suite
# GTEST_SKIPs on the missing fixture, which is how a dead
# applyPairExclusion seam in DeepSpinPTExpt went unnoticed.
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_spin.py &
PID11=$!
# DPA4 + analytical ZBL bridging: a linear COMPOSITION on the graph
# lower, which no other C++ fixture covers.
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_zbl.py &
PID12=$!
# Native-spin DPA4 + charge-spin FiLM: the ONLY fixture with both
# is_spin=true and dim_chg_spin>0, i.e. the only one on which the
# runtime charge_spin argument of DeepSpin::compute is not inert.
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_spin_chgspin.py &
PID13=$!
# Native-spin DPA4 + ZBL bridging COMBINED: spin reaching a linear
# composition, and the only fixture pinning that a bridged model
# exports NO with-comm artifact (single-rank only by construction).
env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_spin_zbl.py &
PID14=$!
wait $PID11
wait $PID12
wait $PID13
wait $PID14
fi

if [ "${DP_CC_SKIP_CTEST}" != "1" ]; then
ctest --output-on-failure
if [ "${ENABLE_PADDLE:-TRUE}" == "TRUE" ]; then
PADDLE_INFERENCE_DIR=${BUILD_TMP_DIR}/paddle_inference_install_dir
export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PADDLE_INFERENCE_DIR}/third_party/install/onednn/lib:${PADDLE_INFERENCE_DIR}/third_party/install/mklml/lib
fi
ctest --output-on-failure
Loading
Loading