From a9dab3a39e2be48bc9a0f88c6ccd4613b742f708 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 20 Jul 2026 08:04:06 +0800 Subject: [PATCH 1/6] ci(cuda): overlap AOTI .pt2 compiles with GPU unit tests The CUDA job ran source/tests as one serial process, so the .pt2-freezing tests (torch.export + AOTInductor) held the step for minutes while inductor/g++/ptxas generated code on the CPU -- the GPU measured ~98% idle (<250 MiB) during those compiles. Split the pytest step into two groups that run concurrently on the one GPU: a CPU-bound compile group (-m aoti_compile, niced so the GPU group keeps CPU priority) and the GPU-bound remainder (-m "not aoti_compile"), so the compile CPU-time overlaps the GPU unit tests. The compile group's GPU footprint is negligible, so the two coexist on one device without contention; wall-clock becomes max(compile, gpu) rather than their sum. The aoti_compile marker is auto-applied in source/tests/conftest.py to any test whose module references a .pt2 freeze entry point, so the partition stays correct as tests are added without a hand-maintained file list. --- .github/workflows/test_cuda.yml | 24 ++++++++++++++++- pyproject.toml | 5 +++- source/tests/conftest.py | 46 +++++++++++++++++++++++++++++++++ 3 files changed, 73 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index 280da63686..ca1ea0b44d 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -63,7 +63,29 @@ jobs: DP_ENABLE_NATIVE_OPTIMIZATION: 1 DP_ENABLE_PYTORCH: 1 - run: dp --version - - run: python -m pytest source/tests + - 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: | + nice -n 15 python -m pytest source/tests -m "aoti_compile" \ + -p no:cacheprovider --junit-xml=junit-aoti.xml & + pid_a=$! + python -m pytest source/tests -m "not aoti_compile" \ + -p no:cacheprovider --junit-xml=junit-gpu.xml & + 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 "aoti_compile group exit=$rc_a ; gpu group exit=$rc_b" + exit $(( rc_a != 0 || rc_b != 0 )) env: NUM_WORKERS: 0 CUDA_VISIBLE_DEVICES: 0 diff --git a/pyproject.toml b/pyproject.toml index 6216438ac4..8cb293beaa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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] diff --git a/source/tests/conftest.py b/source/tests/conftest.py index eabaa25cd4..c6a743b6c0 100644 --- a/source/tests/conftest.py +++ b/source/tests/conftest.py @@ -1,8 +1,54 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import os +from pathlib import ( + Path, +) + +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 + + +# A test module is "AOTI-compiling" if it freezes a ``.pt2`` at runtime, i.e. it +# references one of the pt_expt AOTInductor freeze entry points below. Those +# tests are CPU-bound (minutes of inductor / g++ / ptxas code generation) while +# the GPU sits ~idle (measured ~98% idle, <250 MiB peak), so CI can run them +# concurrently with the GPU-bound tests that do NOT compile a ``.pt2`` -- the +# compile CPU-time then overlaps the GPU unit tests instead of serializing in +# front of them (see ``.github/workflows/test_cuda.yml``). Detecting the set by +# source scan keeps the ``aoti_compile`` partition correct as tests are added, +# without a hand-maintained file list. +_AOTI_FREEZE_SYMBOLS = ( + "deserialize_to_file", + "_trace_and_export", + "aoti_compile_and_package", +) +_aoti_module_cache: dict[str, bool] = {} + + +def _module_compiles_pt2(path: str) -> bool: + hit = _aoti_module_cache.get(path) + if hit is None: + try: + src = Path(path).read_text(encoding="utf-8") + except OSError: + src = "" + hit = _aoti_module_cache[path] = any(s in src for s in _AOTI_FREEZE_SYMBOLS) + return hit + + +def pytest_collection_modifyitems(items) -> None: + """Auto-tag every test whose module freezes a ``.pt2`` with ``aoti_compile``. + + Lets the CUDA CI split the suite into a CPU-bound compile group + (``-m aoti_compile``) and a GPU-bound group (``-m "not aoti_compile"``) + that run concurrently on the same GPU. + """ + for item in items: + path = getattr(item, "path", None) or getattr(item, "fspath", None) + if path is not None and _module_compiles_pt2(str(path)): + item.add_marker(pytest.mark.aoti_compile) From f660993b3de2df306f057c52804aea5ccf1c8744 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 20 Jul 2026 14:41:46 +0800 Subject: [PATCH 2/6] test: make aoti_compile partition explicit + fail-closed Replace the source-grep aoti_compile auto-marker with an explicit curated list of the modules that actually freeze a .pt2, plus an always-on runtime guard. The grep both missed freeze(...)-driven compiles (test_dp_freeze, test_finetune, test_multitask, test_dp_test, test_change_bias, test_models -- none name the low-level entry points) and over-tagged files that only mock the compiler or emit .pth/.pte. The producer list was derived by running the suite under a detector that wraps the real compile entry points (aoti_compile_and_package, _trace_and_export). The guard wraps those same entry points every run and fails the session, with the offending nodeids, if any test freezes a .pt2 without the aoti_compile marker -- so the list cannot silently drift, and a compile can never escape into the GPU lane unnoticed. It runs in the CPU test_python CI too, so drift is caught cheaply rather than only in the expensive CUDA job. --- source/tests/conftest.py | 152 ++++++++++++++++++++++++++++++--------- 1 file changed, 119 insertions(+), 33 deletions(-) diff --git a/source/tests/conftest.py b/source/tests/conftest.py index c6a743b6c0..fb01dc1e17 100644 --- a/source/tests/conftest.py +++ b/source/tests/conftest.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import os -from pathlib import ( - Path, -) import pytest @@ -13,42 +10,131 @@ import tensorflow # noqa: F401 -# A test module is "AOTI-compiling" if it freezes a ``.pt2`` at runtime, i.e. it -# references one of the pt_expt AOTInductor freeze entry points below. Those -# tests are CPU-bound (minutes of inductor / g++ / ptxas code generation) while -# the GPU sits ~idle (measured ~98% idle, <250 MiB peak), so CI can run them -# concurrently with the GPU-bound tests that do NOT compile a ``.pt2`` -- the -# compile CPU-time then overlaps the GPU unit tests instead of serializing in -# front of them (see ``.github/workflows/test_cuda.yml``). Detecting the set by -# source scan keeps the ``aoti_compile`` partition correct as tests are added, -# without a hand-maintained file list. -_AOTI_FREEZE_SYMBOLS = ( - "deserialize_to_file", - "_trace_and_export", - "aoti_compile_and_package", +# --- 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_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", + } ) -_aoti_module_cache: dict[str, bool] = {} -def _module_compiles_pt2(path: str) -> bool: - hit = _aoti_module_cache.get(path) - if hit is None: - try: - src = Path(path).read_text(encoding="utf-8") - except OSError: - src = "" - hit = _aoti_module_cache[path] = any(s in src for s in _AOTI_FREEZE_SYMBOLS) - return hit +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: - """Auto-tag every test whose module freezes a ``.pt2`` with ``aoti_compile``. - - Lets the CUDA CI split the suite into a CPU-bound compile group - (``-m aoti_compile``) and a GPU-bound group (``-m "not aoti_compile"``) - that run concurrently on the same GPU. - """ + """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 _module_compiles_pt2(str(path)): + 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 + yield + _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}") From b4a79960df9a5fdb3ecb19840e335323d68e1eb7 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 20 Jul 2026 14:42:08 +0800 Subject: [PATCH 3/6] ci(cuda): measure and tune the overlapped pytest lanes - echo each lane's wall-clock ([lane aoti_compile] / [lane gpu]) so T_A and T_B (vs the ~3h46m serial baseline) appear directly in the log - cap the compile lane at TORCHINDUCTOR_COMPILE_THREADS=nproc/2 (plus nice) so the GPU lane keeps CPU for its host-side dispatch - upload junit-aoti.xml / junit-gpu.xml as artifacts for post-run inspection - share one AOTInductor compile cache across all steps via a job-level TORCHINDUCTOR_CACHE_DIR, so the C++ fixture builds (gen_*.py in test_cc_local.sh) reuse the kernels the Python lane compiled, and persist it across runs with actions/cache (10 GB repo-cache limit noted inline) --- .github/workflows/test_cuda.yml | 49 +++++++++++++++++++++++++++++---- 1 file changed, 44 insertions(+), 5 deletions(-) diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index ca1ea0b44d..ad730b997a 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -20,6 +20,13 @@ 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 @@ -27,6 +34,17 @@ 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@v6 with: python-version: "3.11" @@ -73,18 +91,28 @@ jobs: # Peak GPU memory of the compile group is negligible, so the two groups # coexist on one device without contention. run: | - nice -n 15 python -m pytest source/tests -m "aoti_compile" \ - -p no:cacheprovider --junit-xml=junit-aoti.xml & + # 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 + ( t0=$(date +%s) + 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=$! - python -m pytest source/tests -m "not aoti_compile" \ - -p no:cacheprovider --junit-xml=junit-gpu.xml & + ( t0=$(date +%s) + 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 "aoti_compile group exit=$rc_a ; gpu group exit=$rc_b" + echo "lane exit codes: aoti_compile=$rc_a gpu=$rc_b" exit $(( rc_a != 0 || rc_b != 0 )) env: NUM_WORKERS: 0 @@ -93,6 +121,17 @@ 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@v4 + with: + name: cuda-pytest-junit + path: | + junit-aoti.xml + junit-gpu.xml + if-no-files-found: warn - name: Convert models run: source/tests/infer/convert-models.sh - run: | From 6f8d76bee397e2ac574b45ee1ebe3c64e7129801 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 20 Jul 2026 18:03:16 +0800 Subject: [PATCH 4/6] fix: robustify lane timing and guard cleanup (review) - CUDA lanes: init rc=0 and `|| rc=$?` so the step's default `set -e` cannot skip the per-lane wall-clock echo when a lane fails (that timing is exactly what we want on a failure). - conftest guard: wrap the hookwrapper yield in try/finally so _current_item is reset even if the wrapped hook raises, avoiding misattributed compiles in teardown hooks. --- .github/workflows/test_cuda.yml | 15 +++++++++------ source/tests/conftest.py | 6 ++++-- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index ad730b997a..7dbd3f6a9f 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -94,17 +94,20 @@ jobs: # 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 - ( t0=$(date +%s) + # 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" + -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) + ( 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" + -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=$? diff --git a/source/tests/conftest.py b/source/tests/conftest.py index fb01dc1e17..9e42da4517 100644 --- a/source/tests/conftest.py +++ b/source/tests/conftest.py @@ -118,8 +118,10 @@ def probe(*args, **kwargs): @pytest.hookimpl(hookwrapper=True) def pytest_runtest_protocol(item, nextitem): _current_item["item"] = item - yield - _current_item["item"] = None + try: + yield + finally: + _current_item["item"] = None def pytest_sessionfinish(session, exitstatus) -> None: From 70b517569dd786e10df76e72305c02c1b42cff01 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 21 Jul 2026 11:09:36 +0800 Subject: [PATCH 5/6] ci(cuda): overlap C++ build + fixture gen with the Python GPU lane The C++ portion ran serially after the Python tests: cmake --build plus the gen_*.py .pt2 fixtures are ~1h of CPU-bound work with the GPU idle (measured: 56m gen + 7.5m build + 5m convert-models), followed by ~39m of GPU tests (ctest + LAMMPS). The Python GPU lane meanwhile has a ~1h23m idle-CPU tail (compile lane ends ~1h13m, GPU lane runs to ~2h36m). Run the C++ CPU-prep as a background step (GitHub Actions parallel steps) so it overlaps the Python GPU lane, with a `- wait: cc-prep` barrier before the ctest step consumes the fixtures. The prep is niced -19 so both Python lanes keep CPU priority and it only fills idle cores. test_cc_local.sh gains DP_CC_SKIP_CTEST / DP_CC_SKIP_BUILD phase gates (default runs both, so test_cc.yml is unaffected); convert-models.sh moves into the prep step. Projected: hides ~1h14m of GPU-idle compile, whole CUDA job ~4h57m -> ~3h43m. --- .github/workflows/test_cuda.yml | 33 +++++- source/install/test_cc_local.sh | 176 ++++++++++++++++++-------------- 2 files changed, 126 insertions(+), 83 deletions(-) diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index 7dbd3f6a9f..8238ebf9dd 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -81,6 +81,27 @@ jobs: DP_ENABLE_NATIVE_OPTIMIZATION: 1 DP_ENABLE_PYTORCH: 1 - run: dp --version + - 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 @@ -135,11 +156,15 @@ jobs: junit-aoti.xml junit-gpu.xml if-no-files-found: warn - - name: Convert models - run: source/tests/infer/convert-models.sh - - run: | + # 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 diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index d3ae4d2886..7b24b49024 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -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 @@ -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 @@ -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 From ec84518e5100585469f6e79f7e1b5478750190d5 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 21 Jul 2026 23:04:38 +0800 Subject: [PATCH 6/6] ci(cuda): serialize C++ prep behind the aoti_compile lane to end 3-way contention The previous layout ran the C++ build + gen_*.py fixture compiles as a third `background:` step concurrent with BOTH Python lanes. That 3-way overlap (two pytest lanes + the C++ build, all sharing one TORCHINDUCTOR_CACHE_DIR) SIGSEGV'd both Python lanes; the identical suite passed on the prior sha that had only the two Python lanes. Collapse it back to two workstreams. Lane 1 now runs the aoti_compile pytest group and THEN, sequentially, the C++ build + fixture generation -- both are CPU-bound with the GPU idle, so serializing them keeps only one inductor compiler writing the cache at a time and lets cc-prep reuse the kernels the aoti tests just built. Lane 2 keeps the GPU-bound bulk. ctest + LAMMPS run once both lanes join. At any instant only two heavy streams run, matching the green baseline, and the whole CPU stream stays hidden under lane 2. Drops the `background:`/`- wait:` steps; test_cc_local.sh's DP_CC_SKIP_CTEST / DP_CC_SKIP_BUILD gating is unchanged. --- .github/workflows/test_cuda.yml | 89 +++++++++++++++------------------ 1 file changed, 39 insertions(+), 50 deletions(-) diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index 8238ebf9dd..6038eba503 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -81,62 +81,56 @@ jobs: DP_ENABLE_NATIVE_OPTIMIZATION: 1 DP_ENABLE_PYTORCH: 1 - run: dp --version - - 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 + - name: Run Python tests + C++ build (overlap CPU compile work with the GPU lane) + # Two lanes on one GPU box, sized so only TWO heavy workstreams ever run + # at once -- the shape that was green before the C++ prep was (wrongly) + # added as a THIRD concurrent stream (that 3-way overlap SIGFAULTed both + # Python lanes via TORCHINDUCTOR_CACHE_DIR contention; see #5882): + # Lane 1 (CPU compile stream): the .pt2-freezing tests (-m aoti_compile) + # -- CPU-bound, GPU ~98% idle while inductor/g++/ptxas run -- THEN the + # C++ build + gen_*.py .pt2 fixtures (also CPU-bound, GPU idle). They + # run SEQUENTIALLY in one lane, so the compile work never races the + # shared inductor cache and cc-prep reuses the kernels the aoti tests + # just built. + # Lane 2 (GPU stream): the GPU-bound tests (-m "not aoti_compile") -- the + # ~19k-test bulk and the wall-clock bottleneck. It does not freeze + # .pt2, so it only reads the inductor cache, never competing to write. + # The whole CPU stream is hidden under lane 2; ctest + LAMMPS run after + # BOTH lanes join. 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. + # 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 want most - # when a lane fails). - ( t0=$(date +%s); rc=0 + # 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). + # Lane 1: aoti_compile pytest, THEN the C++ build + fixture generation. + ( 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=$? - echo "[lane aoti_compile] wall=$(( $(date +%s) - t0 ))s exit=$rc" - exit $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" + tc=$(date +%s); rc_cc=0 + ( export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH + export 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 + source/tests/infer/convert-models.sh + DP_CC_SKIP_CTEST=1 nice -n 19 source/install/test_cc_local.sh ) || rc_cc=$? + echo "[lane cc-prep] wall=$(( $(date +%s) - tc ))s exit=$rc_cc" + exit $(( rc != 0 || rc_cc != 0 )) ) & pid_a=$! - ( t0=$(date +%s); rc=0 + # Lane 2: the GPU-bound unit tests (the bulk). + ( 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) - t0 ))s exit=$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=$? - # 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" + echo "lane exit codes: aoti_compile+cc-prep=$rc_a gpu=$rc_b" exit $(( rc_a != 0 || rc_b != 0 )) env: NUM_WORKERS: 0 @@ -156,12 +150,7 @@ jobs: 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) + - name: C++ tests (ctest; build + fixtures done in lane 1 above) run: | export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH DP_CC_SKIP_BUILD=1 source/install/test_cc_local.sh