Merge branch 'deepmodeling:devel' into devel

.github/labeler.yml

@@ -5,8 +5,8 @@ Python:
 Docs: doc/**/*
 Examples: examples/**/*
 Core: source/lib/**/*
-CUDA: source/lib/src/cuda/**/*
-ROCM: source/lib/src/rocm/**/*
+CUDA: source/lib/src/gpu/**/*
+ROCM: source/lib/src/gpu/**/*
 OP: source/op/**/*
 C++: source/api_cc/**/*
 C: source/api_c/**/*

.github/workflows/build_cc.yml

@@ -21,8 +21,6 @@ jobs:
           dp_variant: clang
     steps:
     - uses: actions/checkout@v4
-      with:
-        submodules: true
     - uses: actions/setup-python@v4
       with:
         python-version: '3.11'

.github/workflows/build_wheel.yml

@@ -5,8 +5,24 @@ on:
   pull_request:
 
 jobs:
+  determine-arm64-runner:
+    runs-on: ubuntu-latest
+    permissions: read-all
+    outputs:
+      runner: ${{ steps.set-runner.outputs.runner }}
+    steps:
+      - name: Determine which runner to use for ARM64 build
+        id: set-runner
+        run: |
+          if [ "${{ github.repository_owner }}" == "deepmodeling" ]; then
+            echo "runner=[\"Linux\",\"ARM64\"]" >> $GITHUB_OUTPUT
+          else
+            echo "runner=\"ubuntu-latest\"" >> $GITHUB_OUTPUT
+          fi
+
   build_wheels:
     name: Build wheels for cp${{ matrix.python }}-${{ matrix.platform_id }}
+    needs: determine-arm64-runner
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -33,21 +49,20 @@ jobs:
             platform_id: win_amd64
             dp_variant: cpu
           # linux-aarch64
-          - os: ubuntu-latest
+          - os: ${{ fromJson(needs.determine-arm64-runner.outputs.runner) }}
             python: 310
             platform_id: manylinux_aarch64
             dp_variant: cpu
     steps:
       - uses: actions/checkout@v4
         with:
-          submodules: true
           # https://github.com/pypa/setuptools_scm/issues/480
           fetch-depth: 0
-      - uses: docker/setup-qemu-action@v2
+      - uses: docker/setup-qemu-action@v3
         name: Setup QEMU
-        if: matrix.platform_id == 'manylinux_aarch64'
+        if: matrix.platform_id == 'manylinux_aarch64' && matrix.os == 'ubuntu-latest'
       - name: Build wheels
-        uses: pypa/cibuildwheel@v2.15
+        uses: pypa/cibuildwheel@v2.16
         env:
           CIBW_BUILD_VERBOSITY: 1
           CIBW_ARCHS: all
@@ -61,8 +76,6 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-        with:
-          submodules: true
       - uses: actions/setup-python@v4
         name: Install Python
         with:
@@ -101,20 +114,20 @@ jobs:
           name: artifact
           path: source/install/docker/dist
       - name: Log in to the Container registry
-        uses: docker/login-action@465a07811f14bebb1938fbed4728c6a1ff8901fc
+        uses: docker/login-action@343f7c4344506bcbf9b4de18042ae17996df046d
         with:
           registry: ghcr.io
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Extract metadata (tags, labels) for Docker
         id: meta
-        uses: docker/metadata-action@818d4b7b91585d195f67373fd9cb0332e31a7175
+        uses: docker/metadata-action@96383f45573cb7f253c731d3b3ab81c87ef81934
         with:
           images: ghcr.io/deepmodeling/deepmd-kit
 
       - name: Build and push Docker image
-        uses: docker/build-push-action@0a97817b6ade9f46837855d676c4cca3a2471fc9
+        uses: docker/build-push-action@0565240e2d4ab88bba5387d719585280857ece09
         with:
           context: source/install/docker
           push: ${{ github.repository_owner == 'deepmodeling' && github.event_name == 'push' }}

.github/workflows/remove_test_cuda_label.yml

@@ -0,0 +1,18 @@
+on:
+  pull_request_target:
+    types:
+      - "labeled"
+name: Test CUDA
+jobs:
+  remove_label:
+    permissions:
+      contents: read
+      pull-requests: write
+    # so one can re-trigger the workflow without manually removing the label
+    runs-on: ubuntu-latest
+    if: github.repository_owner == 'deepmodeling' && github.event.label.name == 'Test CUDA'
+    steps:
+    - uses: actions-ecosystem/action-remove-labels@v1
+      with:
+        labels: Test CUDA
+        number: ${{ github.event.pull_request.number }}

.github/workflows/test_cuda.yml

@@ -0,0 +1,60 @@
+on:
+  # manually trigger
+  workflow_dispatch:
+  pull_request:
+    types:
+      - "labeled"
+name: Test CUDA
+jobs:
+  test_cuda:
+    name: Test Python and C++ on CUDA
+    runs-on: nvidia
+    if: github.repository_owner == 'deepmodeling' && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch'
+    steps:
+    - uses: actions/checkout@v4
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.11'
+        cache: 'pip'
+    - name: Setup MPI
+      uses: mpi4py/setup-mpi@v1
+      with:
+        mpi: mpich
+    - uses: lukka/get-cmake@latest
+    - run: |
+         wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
+         && sudo dpkg -i cuda-keyring_1.0-1_all.deb \
+         && sudo apt-get update \
+         && sudo apt-get -y install cuda-11-8 libcudnn8=8.9.5.*-1+cuda11.8
+    - run: python -m pip install -U "pip>=21.3.1,!=23.0.0"
+    - run: pip install -v -e .[gpu,test,lmp,cu11] "ase @ https://github.com/rosswhitfield/ase/archive/edd03571aff6944b77b4a4b055239f3c3e4eeb66.zip"
+      env:
+        DP_BUILD_TESTING: 1
+        DP_VARIANT: cuda
+        CUDA_PATH: /usr/local/cuda-11.8
+    - run: dp --version
+    - run: pytest -s --cov=deepmd --cov=deepmd_cli source/tests --durations=0
+    - run: source/install/test_cc_local.sh
+      env:
+        OMP_NUM_THREADS: 1
+        TF_INTRA_OP_PARALLELISM_THREADS: 1
+        TF_INTER_OP_PARALLELISM_THREADS: 1
+        LMP_CXX11_ABI_0: 1
+        CMAKE_GENERATOR: Ninja
+        DP_VARIANT: cuda
+        DP_USE_MPICH2: 1
+        CUDA_PATH: /usr/local/cuda-11.8
+    - run: |
+        export LD_LIBRARY_PATH=${{ github.workspace }}/dp_test/lib:$CUDA_PATH/lib64:$LD_LIBRARY_PATH
+        export PATH=${{ github.workspace }}/dp_test/bin:$PATH
+        pytest -s --cov=deepmd source/lmp/tests
+        pytest -s --cov=deepmd source/ipi/tests
+      env:
+        OMP_NUM_THREADS: 1
+        TF_INTRA_OP_PARALLELISM_THREADS: 1
+        TF_INTER_OP_PARALLELISM_THREADS: 1
+        LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp
+        CUDA_PATH: /usr/local/cuda-11.8
+    - uses: codecov/codecov-action@v3
+      with:
+        gcov: true

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
       files: \.py$
 -   repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.0.288
+    rev: v0.0.291
     hooks:
     - id: ruff
       args: ["--fix"]
@@ -53,7 +53,7 @@ repos:
     rev: v16.0.6
     hooks:
     -   id: clang-format
-        exclude: ^source/3rdparty|source/lib/src/cuda/cudart/.+\.inc
+        exclude: ^source/3rdparty|source/lib/src/gpu/cudart/.+\.inc
 # CSS
 -   repo: https://github.com/pre-commit/mirrors-csslint
     rev: v1.0.5
@@ -83,7 +83,7 @@ repos:
         - --comment-style
         - //
         - --no-extra-eol
-        exclude: ^source/3rdparty|source/lib/src/cuda/cudart/.+\.inc
+        exclude: ^source/3rdparty|source/lib/src/gpu/cudart/.+\.inc
     # CSS
     -   id: insert-license
         files: \.(css|scss)$

backend/dynamic_metadata.py

@@ -55,8 +55,7 @@ def dynamic_metadata(
                 "sphinxcontrib-bibtex",
             ],
             "lmp": [
-                "lammps~=2023.8.2.0.0; platform_system=='Linux'",
-                "lammps~=2023.8.2.0.0; platform_system!='Linux'",
+                "lammps~=2023.8.2.1.0",
                 *find_libpython_requires,
             ],
             "ipi": [

backend/read_env.py

@@ -57,8 +57,8 @@ def get_argument_from_env() -> Tuple[str, list, list, dict, str]:
         if rocm_root:
             cmake_args.append(f"-DCMAKE_HIP_COMPILER_ROCM_ROOT:STRING={rocm_root}")
         hipcc_flags = os.environ.get("HIP_HIPCC_FLAGS")
-        if hipcc_flags:
-            cmake_args.append(f"-DHIP_HIPCC_FLAGS:STRING={hipcc_flags}")
+        if hipcc_flags is not None:
+            os.environ["HIPFLAGS"] = os.environ.get("HIPFLAGS", "") + " " + hipcc_flags
     else:
         raise RuntimeError("Unsupported DP_VARIANT option: %s" % dp_variant)
 

deepmd/__init__.py

@@ -32,7 +32,7 @@
 set_mkl()
 
 try:
-    from ._version import version as __version__
+    from deepmd_cli._version import version as __version__
 except ImportError:
     from .__about__ import (
         __version__,

deepmd/descriptor/__init__.py

@@ -14,6 +14,9 @@
 from .se_a_ebd import (
     DescrptSeAEbd,
 )
+from .se_a_ebd_v2 import (
+    DescrptSeAEbdV2,
+)
 from .se_a_ef import (
     DescrptSeAEf,
     DescrptSeAEfLower,
@@ -39,6 +42,7 @@
     "DescrptHybrid",
     "DescrptLocFrame",
     "DescrptSeA",
+    "DescrptSeAEbdV2",
     "DescrptSeAEbd",
     "DescrptSeAEf",
     "DescrptSeAEfLower",