Skip to content

ci(cuda): overlap AOTI .pt2 compiles with GPU unit tests #5882

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

Open
wanghan-iapcm wants to merge 5 commits into
base: master
Choose a base branch
from
+329 −85
Open

ci(cuda): overlap AOTI .pt2 compiles with GPU unit tests #5882

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a9dab3a
ci(cuda): overlap AOTI .pt2 compiles with GPU unit tests
Jul 20, 2026
f660993
test: make aoti_compile partition explicit + fail-closed
Jul 20, 2026
b4a7996
ci(cuda): measure and tune the overlapped pytest lanes
Jul 20, 2026
6f8d76b
fix: robustify lane timing and guard cleanup (review)
Jul 20, 2026
70b5175
ci(cuda): overlap C++ build + fixture gen with the Python GPU lane
Jul 21, 2026
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
99 changes: 94 additions & 5 deletions .github/workflows/test_cuda.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,31 @@ jobs:
# The full CUDA suite (serial pytest ~3h43m + C++ + LAMMPS) exceeds the
# default 360-min job limit; the self-hosted GPU runner has no hard cap.
timeout-minutes: 480
# Share ONE AOTInductor on-disk compile cache across every step: the two
# pytest lanes AND the C++ fixture builds (gen_*.py in test_cc_local.sh)
# all drive the same inductor backend on the same architectures, so the C++
# fixture compiles reuse the kernels the Python lane just built. 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

@coderabbitai coderabbitai Bot Jul 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔒 Security & Privacy | 🟡 Minor | ⚡ Quick win

Disable credential persistence to improve workflow security.

The actions/checkout action persists the GITHUB_TOKEN in the local git config by default. If untrusted code or tests are executed within this environment, the token could be exposed. It is a security best practice to explicitly set persist-credentials: false unless the workflow requires authenticating subsequent git commands (like pushing a commit).

🛡️ Proposed fix 
-      - uses: actions/checkout@v7
+      - uses: actions/checkout@v7
+        with:
+          persist-credentials: false
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
- uses: actions/checkout@v7
- uses: actions/checkout@v7
with:
persist-credentials: false
🧰 Tools
🪛 zizmor (1.26.1)

[warning] 36-36: credential persistence through GitHub Actions artifacts (artipacked): does not set persist-credentials: false

(artipacked)

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In @.github/workflows/test_cuda.yml at line 36, Update the actions/checkout step
in the CUDA test workflow to explicitly disable credential persistence by
setting persist-credentials to false, while leaving the existing checkout
behavior unchanged.

Source: Linters/SAST tools

- 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@v6
with:
python-version: "3.11"
Expand Down Expand Up @@ -63,19 +81,90 @@ jobs:
DP_ENABLE_NATIVE_OPTIMIZATION: 1
DP_ENABLE_PYTORCH: 1
- run: dp --version
- run: python -m pytest source/tests
- name: Build C++ and generate fixtures (overlaps the Python GPU lane)
# The C++ build + gen_*.py .pt2 fixtures are CPU-bound with the GPU idle
# (~1h of the job) and independent of the Python tests. Run them as a
# background step so they overlap the Python GPU lane's idle-CPU tail;
# the ctest step waits on this (- wait: cc-prep) before consuming the
# fixtures. Niced hard (19) so both Python lanes keep CPU priority --
# this only fills otherwise-idle cores. DP_CC_SKIP_CTEST builds + gens
# only; ctest runs later once this joins.
id: cc-prep
background: true
run: |
export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH
source/tests/infer/convert-models.sh
DP_CC_SKIP_CTEST=1 nice -n 19 source/install/test_cc_local.sh
env:
OMP_NUM_THREADS: 1
TF_INTRA_OP_PARALLELISM_THREADS: 1
TF_INTER_OP_PARALLELISM_THREADS: 1
CMAKE_GENERATOR: Ninja
DP_VARIANT: cuda
DP_USE_MPICH2: 1
- name: Run Python tests (overlap AOTI compiles with GPU unit tests)
# The tests that freeze a .pt2 (-m aoti_compile) are CPU-bound: the GPU
# sits ~idle (measured ~98% idle, <250 MiB) while inductor/g++/ptxas
# generate code. Run them concurrently with the GPU-bound tests that do
# NOT compile a .pt2 so the compile CPU-time overlaps the GPU unit tests
# instead of serializing in front of them. The compile group is niced so
# the latency-sensitive GPU dispatch of the GPU group keeps CPU priority.
# Peak GPU memory of the compile group is negligible, so the two groups
# coexist on one device without contention.
run: |
# Cap the compile lane's inductor compile threads to ~half the cores
# so the GPU lane keeps CPU for its host-side dispatch; nice it 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 want most
# when a lane fails).
( t0=$(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=$?
echo "[lane aoti_compile] wall=$(( $(date +%s) - t0 ))s exit=$rc"
exit $rc ) &
pid_a=$!
( t0=$(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) - t0 ))s exit=$rc"
exit $rc ) &
pid_b=$!
rc_a=0; wait $pid_a || rc_a=$?
rc_b=0; wait $pid_b || rc_b=$?
# pytest exit 5 == "no tests collected"; harmless for the compile group
# (e.g. if a branch has no .pt2-freezing tests), a real error otherwise.
[ "$rc_a" = 5 ] && rc_a=0
echo "lane exit codes: aoti_compile=$rc_a gpu=$rc_b"
exit $(( rc_a != 0 || rc_b != 0 ))
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: Convert models
run: source/tests/infer/convert-models.sh
- run: |
- 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@v4
with:
name: cuda-pytest-junit
path: |
junit-aoti.xml
junit-gpu.xml
if-no-files-found: warn
# Barrier: the background C++ build + fixture generation must be done
# before ctest consumes the fixtures. By now the Python step has run for
# ~2.5h, so the ~1h of C++ prep has almost always already finished during
# it -- this wait is usually instant.
- wait: cc-prep
- name: C++ tests (ctest; build + fixtures done in the background step)
run: |
export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH
source/install/test_cc_local.sh
DP_CC_SKIP_BUILD=1 source/install/test_cc_local.sh
env:
OMP_NUM_THREADS: 1
TF_INTRA_OP_PARALLELISM_THREADS: 1
Expand Down
5 changes: 4 additions & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -495,7 +495,10 @@ runtime-evaluated-base-classes = ["torch.nn.Module"]
"examples/**/*.py" = ["T20", "TID253"]

[tool.pytest.ini_options]
markers = "run"
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)",
]
timeout = 1800

[tool.coverage.run]
Expand Down
176 changes: 97 additions & 79 deletions source/install/test_cc_local.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,14 @@
#!/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 @@ -20,27 +28,37 @@ 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}
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 '

# 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.
python -c '
import shutil, sys
from pathlib import Path
from deepmd.env import SHARED_LIB_DIR
Expand All @@ -55,67 +73,67 @@ else:
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"
# 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
fi
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
# 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
fi
fi
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

if [ "${DP_CC_SKIP_CTEST}" != "1" ]; then
ctest --output-on-failure
fi
ctest --output-on-failure
Loading
Loading