Merge upstream/master into overlap-aoti-gpu-cuda-ci

Master's deepmodeling#5896 split the CUDA workflow into separate test_python and
test_cc jobs on separate runners, superseding this PR's lane-1 cc-prep
tail. Resolution keeps deepmodeling#5896's two-job structure and re-applies this
PR's remaining pieces: the aoti_compile/gpu two-lane pytest overlap
inside test_python (hides the CPU-bound .pt2 compile tests behind the
GPU lane), the per-lane JUnit reports, and the persistent AOTInductor
compile cache -- now restored in BOTH jobs, since the test_cc fixture
builds (gen_*.py) drive the same inductor backend and reuse the cached
kernels.

Author: Han Wang

.github/workflows/test_cc.yml

@@ -14,9 +14,11 @@ name: Test C++
 jobs:
   testcc:
     name: Test C++
-    runs-on: ubuntu-22.04
+    runs-on: ${{ matrix.check_memleak == false && github.repository_owner == 'deepmodeling' && 'cpu' || 'ubuntu-22.04' }}
     strategy:
       matrix:
+        # Only regular jobs run the LAMMPS tests. Use self-hosted CPU runners
+        # for those jobs in deepmodeling, with a hosted fallback for forks.
         include:
           - check_memleak: true
             enable_tensorflow: true
@@ -36,6 +38,11 @@ jobs:
             enable_paddle: true
     steps:
       - uses: actions/checkout@v7
+      - name: Install GCC
+        run: |
+          sudo apt-get update
+          sudo apt-get install --yes build-essential
+        if: ${{ matrix.check_memleak == false && github.repository_owner == 'deepmodeling' }}
       - uses: actions/setup-python@v6
         with:
           python-version: "3.11"

.github/workflows/test_cuda.yml

@@ -14,17 +14,14 @@ concurrency:
   cancel-in-progress: true
 name: Test CUDA
 jobs:
-  test_cuda:
-    name: Test Python and C++ on CUDA
+  test_python:
+    name: Test Python on CUDA
     runs-on: gpu
-    # 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.
+    # The serial Python test suite can take several hours on the GPU runner.
     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.
+    # 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
@@ -81,47 +78,32 @@ jobs:
           DP_ENABLE_NATIVE_OPTIMIZATION: 1
           DP_ENABLE_PYTORCH: 1
       - run: dp --version
-      - 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.
+      - 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).
-          # 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=$?
             # 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 )) ) &
+            exit $rc ) &
           pid_a=$!
-          # 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=$?
@@ -130,7 +112,7 @@ jobs:
           pid_b=$!
           rc_a=0; wait $pid_a || rc_a=$?
           rc_b=0; wait $pid_b || rc_b=$?
-          echo "lane exit codes: aoti_compile+cc-prep=$rc_a gpu=$rc_b"
+          echo "lane exit codes: aoti_compile=$rc_a gpu=$rc_b"
           exit $(( rc_a != 0 || rc_b != 0 ))
         env:
           NUM_WORKERS: 0
@@ -150,10 +132,71 @@ jobs:
             junit-aoti.xml
             junit-gpu.xml
           if-no-files-found: warn
-      - name: C++ tests (ctest; build + fixtures done in lane 1 above)
-        run: |
+
+  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@v6
+        with:
+          python-version: "3.11"
+          # cache: 'pip'
+      - name: Install wget and unzip
+        run: sudo apt-get update && sudo apt-get install -y wget unzip
+      - uses: lukka/get-cmake@latest
+        with:
+          useLocalCache: true
+          useCloudCache: false
+      - run: |
+          UBUNTU_VERSION=$(lsb_release -rs | tr -d '.')
+          OS_NAME="ubuntu${UBUNTU_VERSION}"
+          echo "Current OS: ${OS_NAME}"
+
+          wget https://developer.download.nvidia.com/compute/cuda/repos/${OS_NAME}/x86_64/cuda-keyring_1.1-1_all.deb \
+          && sudo dpkg -i cuda-keyring_1.1-1_all.deb \
+          && sudo apt-get update \
+          && sudo apt-get -y install cuda-toolkit-12-9 cudnn9-cuda-12
+          echo "CUDA_PATH=/usr/local/cuda-12.9" >> $GITHUB_ENV
+          echo "/usr/local/cuda-12.9/bin" >> $GITHUB_PATH
+        # if: false # skip as we use nvidia image
+      - run: python -m pip install -U uv
+      - run: source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_gpu --group pin_pytorch_gpu --group pin_jax_gpu
+      - run: |
+          export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
+          export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
+          pip install --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cu126/paddlepaddle-gpu/" --index-url https://pypi.org/simple --trusted-host www.paddlepaddle.org.cn --trusted-host paddlepaddle.org.cn "paddlepaddle-gpu==3.4.0.dev20260310"
+          source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
+          # See https://github.com/jax-ml/jax/issues/29042
+          source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13'
+        env:
+          DP_VARIANT: cuda
+          DP_ENABLE_NATIVE_OPTIMIZATION: 1
+          DP_ENABLE_PYTORCH: 1
+      - run: dp --version
+      - name: Convert models
+        run: source/tests/infer/convert-models.sh
+      - 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
+          source/install/test_cc_local.sh
         env:
           OMP_NUM_THREADS: 1
           TF_INTRA_OP_PARALLELISM_THREADS: 1
@@ -177,12 +220,12 @@ jobs:
           CUDA_VISIBLE_DEVICES: 0
   pass:
     name: Pass testing on CUDA
-    needs: [test_cuda]
+    needs: [test_python, test_cc]
     runs-on: ubuntu-slim
     if: always()
     steps:
       - name: Decide whether the needed jobs succeeded or failed
         uses: re-actors/alls-green@release/v1
         with:
           jobs: ${{ toJSON(needs) }}
-          allowed-skips: test_cuda
+          allowed-skips: test_python, test_cc

.pre-commit-config.yaml

@@ -23,14 +23,14 @@ repos:
       - id: check-toml
   # Python
   - repo: https://github.com/PyCQA/isort
-    rev: 9.0.0a3
+    rev: 9.0.0b1
     hooks:
       - id: isort
         files: \.py$
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.15.18
+    rev: v0.15.22
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -62,7 +62,7 @@ repos:
           - mdformat-gfm-alerts==2.0.0
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v22.1.5
+    rev: v22.1.8
     hooks:
       - id: clang-format
         exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|source/tests/infer/.+\.json$)

source/lmp/pair_base.cpp

@@ -437,6 +437,12 @@ PairDeepBaseModel::PairDeepBaseModel(
   out_rel = 0;
   out_rel_v = 0;
   stdf_comm_buff_size = 0;
+  counts = nullptr;
+  displacements = nullptr;
+  tagsend = nullptr;
+  tagrecv = nullptr;
+  stdfsend = nullptr;
+  stdfrecv = nullptr;
   eps = 0.;
   eps_v = 0.;
   scale = NULL;
@@ -485,13 +491,40 @@ void PairDeepBaseModel::print_summary(const string pre) const {
 }
 
 PairDeepBaseModel::~PairDeepBaseModel() {
+  destroy_model_deviation_buffers();
   if (allocated) {
     memory->destroy(setflag);
     memory->destroy(cutsq);
     memory->destroy(scale);
   }
 }
 
+void PairDeepBaseModel::ensure_model_deviation_buffers() {
+  if (counts == nullptr) {
+    memory->create(counts, comm->nprocs, "deepmd:counts");
+    memory->create(displacements, comm->nprocs, "deepmd:displacements");
+  }
+
+  const int ntotal = atom->natoms;
+  if (ntotal > stdf_comm_buff_size) {
+    memory->grow(stdfsend, ntotal, "deepmd:stdfsendall");
+    memory->grow(stdfrecv, ntotal, "deepmd:stdfrecvall");
+    memory->grow(tagsend, ntotal, "deepmd:tagsendall");
+    memory->grow(tagrecv, ntotal, "deepmd:tagrecvall");
+    stdf_comm_buff_size = ntotal;
+  }
+}
+
+void PairDeepBaseModel::destroy_model_deviation_buffers() {
+  memory->destroy(counts);
+  memory->destroy(displacements);
+  memory->destroy(stdfsend);
+  memory->destroy(stdfrecv);
+  memory->destroy(tagsend);
+  memory->destroy(tagrecv);
+  stdf_comm_buff_size = 0;
+}
+
 void PairDeepBaseModel::allocate() {
   allocated = 1;
   int n = atom->ntypes;
@@ -537,17 +570,7 @@ void PairDeepBaseModel::init_style() {
   // neighbor->requests[irequest]->newton = 2;
 #endif
   if (out_each == 1) {
-    int ntotal = atom->natoms;
-    int nprocs = comm->nprocs;
-    if (ntotal > stdf_comm_buff_size) {
-      stdf_comm_buff_size = ntotal;
-    }
-    memory->create(counts, nprocs, "deepmd:counts");
-    memory->create(displacements, nprocs, "deepmd:displacements");
-    memory->create(stdfsend, ntotal, "deepmd:stdfsendall");
-    memory->create(stdfrecv, ntotal, "deepmd:stdfrecvall");
-    memory->create(tagsend, ntotal, "deepmd:tagsendall");
-    memory->create(tagrecv, ntotal, "deepmd:tagrecvall");
+    ensure_model_deviation_buffers();
   }
 }
 

source/lmp/pair_base.h

@@ -61,6 +61,11 @@ class PairDeepBaseModel : public Pair {
   deepmd_compat::DeepBaseModel& deep_base;
   deepmd_compat::DeepBaseModelDevi& deep_base_model_devi;
   virtual void allocate();
+  // The model-deviation gather arrays are owned by this pair style.  Keep
+  // their allocation in one place so repeated LAMMPS initialization cannot
+  // overwrite a live pointer and leak the previous allocation.
+  void ensure_model_deviation_buffers();
+  void destroy_model_deviation_buffers();
   double** scale;
   unsigned numb_models;
   double cutoff;

source/lmp/pair_deepmd.cpp

@@ -557,18 +557,7 @@ void PairDeepMD::compute(int eflag, int vflag) {
           // Gather std_f and tags
           tagint* tag = atom->tag;
           int nprocs = comm->nprocs;
-          // Grow arrays if necessary
-          if (atom->natoms > stdf_comm_buff_size) {
-            stdf_comm_buff_size = atom->natoms;
-            memory->destroy(stdfsend);
-            memory->destroy(stdfrecv);
-            memory->destroy(tagsend);
-            memory->destroy(tagrecv);
-            memory->create(stdfsend, stdf_comm_buff_size, "deepmd:stdfsendall");
-            memory->create(stdfrecv, stdf_comm_buff_size, "deepmd:stdfrecvall");
-            memory->create(tagsend, stdf_comm_buff_size, "deepmd:tagsendall");
-            memory->create(tagrecv, stdf_comm_buff_size, "deepmd:tagrecvall");
-          }
+          ensure_model_deviation_buffers();
           for (int ii = 0; ii < nlocal; ii++) {
             tagsend[ii] = tag[ii];
             stdfsend[ii] = std_f[ii];

source/lmp/pair_deepspin.cpp

@@ -472,18 +472,7 @@ void PairDeepSpin::compute(int eflag, int vflag) {
           // Gather std_f and tags
           tagint* tag = atom->tag;
           int nprocs = comm->nprocs;
-          // Grow arrays if necessary
-          if (atom->natoms > stdf_comm_buff_size) {
-            stdf_comm_buff_size = atom->natoms;
-            memory->destroy(stdfsend);
-            memory->destroy(stdfrecv);
-            memory->destroy(tagsend);
-            memory->destroy(tagrecv);
-            memory->create(stdfsend, stdf_comm_buff_size, "deepmd:stdfsendall");
-            memory->create(stdfrecv, stdf_comm_buff_size, "deepmd:stdfrecvall");
-            memory->create(tagsend, stdf_comm_buff_size, "deepmd:tagsendall");
-            memory->create(tagrecv, stdf_comm_buff_size, "deepmd:tagrecvall");
-          }
+          ensure_model_deviation_buffers();
           for (int ii = 0; ii < nlocal; ii++) {
             tagsend[ii] = tag[ii];
             stdfsend[ii] = std_f[ii];

source/lmp/tests/test_lammps.py

@@ -347,6 +347,21 @@ def test_pair_deepmd_model_devi(lammps) -> None:
     assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v))))
 
 
+def test_pair_deepmd_model_devi_atomic_reinit_clear_smoke(lammps) -> None:
+    """Smoke-test repeated pair initialization followed by pair destruction."""
+    lammps.pair_style(
+        f"deepmd {pb_file.resolve()} {pb_file2.resolve()} "
+        f"out_file {md_file.resolve()} out_freq 1 atomic"
+    )
+    lammps.pair_coeff("* *")
+
+    # Re-enter init_style() on one pair, then destroy it.  This guards against
+    # future double-free or use-after-free mistakes in the buffer lifecycle.
+    lammps.run(0)
+    lammps.run(0)
+    lammps.clear()
+
+
 def test_pair_deepmd_model_devi_virial(lammps) -> None:
     lammps.pair_style(
         f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic"

source/op/tf/prod_env_mat_multi_device.cc

@@ -1964,8 +1964,11 @@ class ProdEnvMatAMixOp : public OpKernel {
         deepmd::prod_env_mat_a_cpu(
             em, em_deriv, rij, nlist, coord, type, inlist, max_nbor_size, avg,
             std, nloc, frame_nall, rcut_r, rcut_r_smth, sec_a, f_type);
-        // do nlist mapping if coords were copied
-        _map_nei_info_cpu(nlist, ntype, nmask, type, &idx_mapping[0], nloc,
+        // The mapping is populated only when PBC coordinate copies are made.
+        // Other neighbor-list modes must not form a pointer into the empty
+        // vector.
+        _map_nei_info_cpu(nlist, ntype, nmask, type,
+                          b_nlist_map ? idx_mapping.data() : nullptr, nloc,
                           nnei, ntypes, b_nlist_map);
       }
     }