diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index dfb15db6ab..f978f83f91 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -19,11 +19,6 @@ 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 @@ -31,17 +26,6 @@ jobs: 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" @@ -78,42 +62,7 @@ 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 @@ -121,41 +70,15 @@ jobs: 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" diff --git a/pyproject.toml b/pyproject.toml index b205b9c160..4307833ca7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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] diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index d31a883607..da763bd215 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -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 @@ -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" @@ -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 diff --git a/source/tests/conftest.py b/source/tests/conftest.py index 2db5d625c1..eabaa25cd4 100644 --- a/source/tests/conftest.py +++ b/source/tests/conftest.py @@ -2,142 +2,7 @@ import os -import pytest - if os.environ.get("DP_CI_IMPORT_PADDLE_BEFORE_TF", "0") == "1": # Paddle must be loaded before TensorFlow in the CI test process. import paddle # noqa: F401 import tensorflow # noqa: F401 - - -# --- aoti_compile partition: explicit producer list + runtime guard ---------- -# -# Tests that freeze a ``.pt2`` (torch.export / make_fx trace + AOTInductor -# compile) are CPU-bound: the GPU sits ~idle (measured ~98% idle, <250 MiB) -# while inductor / g++ / ptxas generate code. The CUDA CI marks them -# ``aoti_compile`` and runs them in a CPU lane concurrent with the GPU-bound -# tests (see ``.github/workflows/test_cuda.yml``). -# -# The producer set is listed EXPLICITLY below rather than inferred from test -# source: a source grep both MISSES ``freeze(...)``-driven compiles (e.g. -# ``test_dp_freeze`` / ``test_finetune`` never name the low-level entry points) -# and OVER-tags files that only mock the compiler or produce ``.pth`` / ``.pte``. -# The list was derived by running the suite under a detector that wraps the real -# compile entry points; the always-on ``pytest_sessionfinish`` guard below -# re-checks it every run and fails if any test freezes a ``.pt2`` without the -# marker, so the list cannot silently drift out of date. -_AOTI_COMPILE_MODULES = frozenset( - { - # direct AOTInductor / _trace_and_export API - "pt_expt/descriptor/test_dpa1_cuda.py", - "pt_expt/model/test_model_compression.py", - "pt_expt/model/test_dpa4_export.py", - "pt_expt/model/test_export_pipeline.py", - "pt_expt/model/test_export_with_comm.py", - "pt_expt/model/test_dpa1_graph_lower.py", - "pt_expt/model/test_graph_export.py", - "pt_expt/model/test_zbl_bridging.py", - "pt_expt/model/test_graph_export_with_comm.py", - "pt_expt/utils/test_graph_pt2_metadata.py", - "pt_expt/infer/test_deep_eval_metadata_only.py", - "pt_expt/infer/test_deep_eval_pt_checkpoint.py", - "pt_expt/infer/test_deep_eval_spin.py", - "pt_expt/infer/test_deep_eval.py", - "pt_expt/infer/test_deep_eval_property.py", - "pt_expt/infer/test_dpa4_deep_eval.py", - "pt_expt/infer/test_graph_deepeval.py", - # freeze(...) / finetune / multitask driven (a grep on the low-level - # entry points misses these; confirmed by the runtime detector) - "pt_expt/test_dp_freeze.py", - "pt_expt/test_finetune.py", - "pt_expt/test_multitask.py", - "pt_expt/test_dp_test.py", - "pt_expt/test_change_bias.py", - "infer/test_models.py", - # pt / jax producers - "pt/model/test_sezm_export.py", - "pt/model/test_sezm_spin_model.py", - "jax/test_deep_dos.py", - "jax/test_training.py", - } -) - - -def _tests_relpath(path: object) -> str: - """Normalize an item path to ``/.../test_x.py`` under ``source/tests``.""" - p = str(path).replace(os.sep, "/") - marker = "source/tests/" - i = p.rfind(marker) - return p[i + len(marker) :] if i >= 0 else p - - -def pytest_collection_modifyitems(items) -> None: - """Tag tests in the explicit producer modules with ``aoti_compile``.""" - for item in items: - path = getattr(item, "path", None) or getattr(item, "fspath", None) - if path is not None and _tests_relpath(path) in _AOTI_COMPILE_MODULES: - item.add_marker(pytest.mark.aoti_compile) - - -# --- guard: fail-closed if a .pt2 compile escapes the aoti_compile marker ----- -_UNMARKED_COMPILERS: set[str] = set() -_current_item: dict = {"item": None} - - -def pytest_configure(config) -> None: - """Wrap the real freeze/compile entry points to record any test that - triggers one without the ``aoti_compile`` marker (drift detection). - """ - try: - import torch._inductor as ti - - import deepmd.pt_expt.utils.serialization as ser - except Exception: - return # pt_expt / inductor unavailable -> nothing can compile a .pt2 - - def wrap(mod, name): - orig = getattr(mod, name, None) - if orig is None or getattr(orig, "_dp_aoti_guard", False): - return - - def probe(*args, **kwargs): - item = _current_item["item"] - if item is not None and item.get_closest_marker("aoti_compile") is None: - _UNMARKED_COMPILERS.add(item.nodeid) - return orig(*args, **kwargs) - - probe._dp_aoti_guard = True - setattr(mod, name, probe) - - # aoti_compile_and_package covers every real .pt2 compile (incl. freeze / - # finetune driven, re-imported from torch._inductor at call time); - # _trace_and_export covers the trace-only export tests. - wrap(ti, "aoti_compile_and_package") - wrap(ser, "_trace_and_export") - - -@pytest.hookimpl(hookwrapper=True) -def pytest_runtest_protocol(item, nextitem): - _current_item["item"] = item - try: - yield - finally: - _current_item["item"] = None - - -def pytest_sessionfinish(session, exitstatus) -> None: - if _UNMARKED_COMPILERS and session.exitstatus == 0: - session.exitstatus = 1 - - -def pytest_terminal_summary(terminalreporter, exitstatus, config) -> None: - if not _UNMARKED_COMPILERS: - return - terminalreporter.section("aoti_compile marker drift") - terminalreporter.write_line( - "These tests froze a .pt2 but lack the 'aoti_compile' marker, so the " - "CUDA CI would run their compile in the GPU lane. Add each one's module " - "to _AOTI_COMPILE_MODULES in source/tests/conftest.py:" - ) - for nodeid in sorted(_UNMARKED_COMPILERS): - terminalreporter.write_line(f" {nodeid}")