Merge branch 'devel' into tests-mv-init-models

.github/workflows/test_cuda.yml

@@ -76,7 +76,7 @@ jobs:
     - run: |
         export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/dp_test/lib:$GITHUB_WORKSPACE/libtorch/lib:$CUDA_PATH/lib64:$LD_LIBRARY_PATH
         export PATH=$GITHUB_WORKSPACE/dp_test/bin:$PATH
-        python -m pytest source/lmp/tests
+        python -m pytest -s source/lmp/tests || (cat log.lammps && exit 1)
         python -m pytest source/ipi/tests
       env:
         OMP_NUM_THREADS: 1

.github/workflows/test_python.yml

@@ -13,14 +13,10 @@ jobs:
     name: Test Python
     runs-on: ubuntu-22.04
     strategy:
+      fail-fast: false
       matrix:
-        include:
-          - python: 3.8
-            tf:
-            torch:
-          - python: "3.11"
-            tf:
-            torch:
+        group: [1, 2, 3, 4, 5, 6]
+        python: ["3.8", "3.11"]
 
     steps:
     - uses: actions/checkout@v4
@@ -45,15 +41,56 @@ jobs:
         HOROVOD_WITHOUT_PYTORCH: 1
         HOROVOD_WITHOUT_GLOO: 1
     - run: dp --version
-    - run: pytest --cov=deepmd source/tests --durations=0
+    - name: Get durations from cache
+      uses: actions/cache@v4
+      with:
+        path: test_durations
+        # the key must never match, even when restarting workflows, as that
+        # will cause durations to get out of sync between groups, the
+        # combined durations will be loaded if available
+        key: test-durations-split-${{ github.run_id }}-${{ github.run_number}}-${{ matrix.python }}-${{ matrix.group }}
+        restore-keys: |
+          test-durations-combined-${{ matrix.python }}-${{ github.sha }}
+          test-durations-combined-${{ matrix.python }}
+    - run: pytest --cov=deepmd source/tests --durations=0 --splits 6 --group ${{ matrix.group }} --store-durations --durations-path=.test_durations_${{ matrix.group }} --splitting-algorithm least_duration
       env:
         NUM_WORKERS: 0
+    - name: Upload partial durations
+      uses: actions/upload-artifact@v4
+      with:
+        name: split-${{ matrix.python }}-${{ matrix.group }}
+        path: .test_durations_${{ matrix.group }}
     - uses: codecov/codecov-action@v4
       env:
         CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+  update_durations:
+    name: Combine and update integration test durations
+    runs-on: ubuntu-22.04
+    strategy:
+      fail-fast: false
+      matrix:
+        python: ["3.8", "3.11"]
+    needs: testpython
+    steps:
+    - name: Get durations from cache
+      uses: actions/cache@v4
+      with:
+        path: .test_durations
+        # key won't match during the first run for the given commit, but
+        # restore-key will if there's a previous stored durations file,
+        # so cache will both be loaded and stored
+        key: test-durations-combined-${{ matrix.python }}-${{ github.sha }}
+        restore-keys: test-durations-combined-${{ matrix.python }}
+    - name: Download artifacts
+      uses: actions/download-artifact@v4
+      with:
+        pattern: split-${{ matrix.python }}-*
+        merge-multiple: true
+    - name: Combine test durations
+      run: jq '. + input' .test_durations_* > .test_durations
   pass:
     name: Pass testing Python
-    needs: [testpython]
+    needs: [testpython, update_durations]
     runs-on: ubuntu-latest
     if: always()
     steps:

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.4.1
+    rev: v0.4.2
     hooks:
       - id: ruff
         args: ["--fix"]

backend/find_tensorflow.py

@@ -155,7 +155,7 @@ def get_tf_requirement(tf_version: str = "") -> dict:
                 "tensorflow-cpu; platform_machine!='aarch64' and (platform_machine!='arm64' or platform_system != 'Darwin')",
                 "tensorflow; platform_machine=='aarch64' or (platform_machine=='arm64' and platform_system == 'Darwin')",
                 # https://github.com/tensorflow/tensorflow/issues/61830
-                "tensorflow-cpu<2.15; platform_system=='Windows'",
+                "tensorflow-cpu!=2.15.*; platform_system=='Windows'",
                 *extra_requires,
             ],
             "gpu": [

backend/read_env.py

@@ -63,7 +63,7 @@ def get_argument_from_env() -> Tuple[str, list, list, dict, str]:
         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)
+        raise RuntimeError(f"Unsupported DP_VARIANT option: {dp_variant}")
 
     if os.environ.get("DP_BUILD_TESTING", "0") == "1":
         cmake_args.append("-DBUILD_TESTING:BOOL=TRUE")

data/raw/copy_raw.py

@@ -85,7 +85,7 @@ def _main():
     )
     args = parser.parse_args()
 
-    print("# copy the system by %s copies" % args.ncopies)  # noqa: T201
+    print(f"# copy the system by {args.ncopies} copies")  # noqa: T201
     assert np.all(
         np.array(args.ncopies, dtype=int) >= np.array([1, 1, 1], dtype=int)
     ), "number of copies should be larger than or equal to 1"

deepmd/dpmodel/descriptor/dpa1.py

@@ -505,11 +505,11 @@ def call(
 
         # calculate gg
         gg = self.cal_g(ss, 0)
-        input_r = dmatrix.reshape(-1, nnei, 4)[:, :, 1:4] / (
+        input_r = dmatrix.reshape(-1, nnei, 4)[:, :, 1:4] / np.maximum(
             np.linalg.norm(
                 dmatrix.reshape(-1, nnei, 4)[:, :, 1:4], axis=-1, keepdims=True
-            )
-            + 1e-12
+            ),
+            1e-12,
         )
         gg = self.dpa1_attention(
             gg, nlist_mask, input_r=input_r, sw=sw
@@ -527,6 +527,7 @@ def call(
         # nf x nloc x (ng x ng1 + tebd_dim)
         if self.concat_output_tebd:
             grrg = np.concatenate([grrg, atype_embd.reshape(nf, nloc, -1)], axis=-1)
+        gr = gr.reshape(nf, nloc, *gr.shape[1:])
         return grrg, gr[..., 1:], None, None, sw
 
     def serialize(self) -> dict:

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -135,7 +135,7 @@ def deserialize(cls, data: dict) -> "BD":
             """
             if cls is BD:
                 return BD.get_class_by_type(data["type"]).deserialize(data)
-            raise NotImplementedError("Not implemented in class %s" % cls.__name__)
+            raise NotImplementedError(f"Not implemented in class {cls.__name__}")
 
         @classmethod
         @abstractmethod

deepmd/dpmodel/fitting/make_base_fitting.py

@@ -88,7 +88,7 @@ def deserialize(cls, data: dict) -> "BF":
             """
             if cls is BF:
                 return BF.get_class_by_type(data["type"]).deserialize(data)
-            raise NotImplementedError("Not implemented in class %s" % cls.__name__)
+            raise NotImplementedError(f"Not implemented in class {cls.__name__}")
 
     setattr(BF, fwd_method_name, BF.fwd)
     delattr(BF, "fwd")

deepmd/dpmodel/model/base_model.py

@@ -126,7 +126,7 @@ def deserialize(cls, data: dict) -> "BaseBaseModel":
                 if model_type == "standard":
                     model_type = data.get("fitting", {}).get("type", "ener")
                 return cls.get_class_by_type(model_type).deserialize(data)
-            raise NotImplementedError("Not implemented in class %s" % cls.__name__)
+            raise NotImplementedError(f"Not implemented in class {cls.__name__}")
 
         model_def_script: str
 

deepmd/entrypoints/doc.py

@@ -16,5 +16,5 @@ def doc_train_input(*, out_type: str = "rst", **kwargs):
     elif out_type == "json":
         doc_str = gen_json()
     else:
-        raise RuntimeError("Unsupported out type %s" % out_type)
+        raise RuntimeError(f"Unsupported out type {out_type}")
     print(doc_str)  # noqa: T201

deepmd/entrypoints/neighbor_stat.py

@@ -91,6 +91,6 @@ def neighbor_stat(
     data.get_batch()
     nei = NeighborStat(data.get_ntypes(), rcut, mixed_type=mixed_type)
     min_nbor_dist, max_nbor_size = nei.get_stat(data)
-    log.info("min_nbor_dist: %f" % min_nbor_dist)
-    log.info("max_nbor_size: %s" % str(max_nbor_size))
+    log.info(f"min_nbor_dist: {min_nbor_dist:f}")
+    log.info(f"max_nbor_size: {max_nbor_size!s}")
     return min_nbor_dist, max_nbor_size

deepmd/entrypoints/test.py

@@ -460,14 +460,14 @@ def test_ener(
         save_txt_file(
             detail_path.with_suffix(".e.out"),
             pe,
-            header="%s: data_e pred_e" % system,
+            header=f"{system}: data_e pred_e",
             append=append_detail,
         )
         pe_atom = pe / natoms
         save_txt_file(
             detail_path.with_suffix(".e_peratom.out"),
             pe_atom,
-            header="%s: data_e pred_e" % system,
+            header=f"{system}: data_e pred_e",
             append=append_detail,
         )
         if not out_put_spin:
@@ -481,7 +481,7 @@ def test_ener(
             save_txt_file(
                 detail_path.with_suffix(".f.out"),
                 pf,
-                header="%s: data_fx data_fy data_fz pred_fx pred_fy pred_fz" % system,
+                header=f"{system}: data_fx data_fy data_fz pred_fx pred_fy pred_fz",
                 append=append_detail,
             )
         else:
@@ -496,14 +496,13 @@ def test_ener(
             save_txt_file(
                 detail_path.with_suffix(".fr.out"),
                 pf_real,
-                header="%s: data_fx data_fy data_fz pred_fx pred_fy pred_fz" % system,
+                header=f"{system}: data_fx data_fy data_fz pred_fx pred_fy pred_fz",
                 append=append_detail,
             )
             save_txt_file(
                 detail_path.with_suffix(".fm.out"),
                 pf_mag,
-                header="%s: data_fmx data_fmy data_fmz pred_fmx pred_fmy pred_fmz"
-                % system,
+                header=f"{system}: data_fmx data_fmy data_fmz pred_fmx pred_fmy pred_fmz",
                 append=append_detail,
             )
         pv = np.concatenate(

deepmd/env.py

@@ -46,9 +46,9 @@
     global_float_prec = "float"
 else:
     raise RuntimeError(
-        "Unsupported float precision option: %s. Supported: high,"
+        f"Unsupported float precision option: {dp_float_prec}. Supported: high,"
         "low. Please set precision with environmental variable "
-        "DP_INTERFACE_PREC." % dp_float_prec
+        "DP_INTERFACE_PREC."
     )
 
 

deepmd/main.py

@@ -650,7 +650,7 @@ def main_parser() -> argparse.ArgumentParser:
 
     # --version
     parser.add_argument(
-        "--version", action="version", version="DeePMD-kit v%s" % __version__
+        "--version", action="version", version=f"DeePMD-kit v{__version__}"
     )
 
     # * train nvnmd script ******************************************************************

deepmd/pt/entrypoints/main.py

@@ -286,10 +286,14 @@ def train(FLAGS):
 
 def freeze(FLAGS):
     model = torch.jit.script(inference.Tester(FLAGS.model, head=FLAGS.head).model)
+    if '"type": "dpa2"' in model.model_def_script:
+        extra_files = {"type": "dpa2"}
+    else:
+        extra_files = {"type": "else"}
     torch.jit.save(
         model,
         FLAGS.output,
-        {},
+        extra_files,
     )
 
 

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -186,6 +186,7 @@ def forward_common_atomic(
         mapping: Optional[torch.Tensor] = None,
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
+        comm_dict: Optional[Dict[str, torch.Tensor]] = None,
     ) -> Dict[str, torch.Tensor]:
         """Common interface for atomic inference.
 
@@ -207,6 +208,8 @@ def forward_common_atomic(
             frame parameters, shape: nf x dim_fparam
         aparam
             atomic parameter, shape: nf x nloc x dim_aparam
+        comm_dict
+            The data needed for communication for parallel inference.
 
         Returns
         -------
@@ -234,6 +237,7 @@ def forward_common_atomic(
             mapping=mapping,
             fparam=fparam,
             aparam=aparam,
+            comm_dict=comm_dict,
         )
         ret_dict = self.apply_out_stat(ret_dict, atype)
 

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -130,6 +130,7 @@ def forward_atomic(
         mapping: Optional[torch.Tensor] = None,
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
+        comm_dict: Optional[Dict[str, torch.Tensor]] = None,
     ) -> Dict[str, torch.Tensor]:
         """Return atomic prediction.
 
@@ -163,6 +164,7 @@ def forward_atomic(
             extended_atype,
             nlist,
             mapping=mapping,
+            comm_dict=comm_dict,
         )
         assert descriptor is not None
         # energy, force

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -144,6 +144,7 @@ def forward_atomic(
         mapping: Optional[torch.Tensor] = None,
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
+        comm_dict: Optional[Dict[str, torch.Tensor]] = None,
     ) -> Dict[str, torch.Tensor]:
         """Return atomic prediction.
 

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -228,6 +228,7 @@ def forward_atomic(
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
+        comm_dict: Optional[Dict[str, torch.Tensor]] = None,
     ) -> Dict[str, torch.Tensor]:
         nframes, nloc, nnei = nlist.shape
         extended_coord = extended_coord.view(nframes, -1, 3)

deepmd/pt/model/descriptor/dpa1.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
     Callable,
+    Dict,
     List,
     Optional,
     Tuple,
@@ -453,6 +454,7 @@ def forward(
         extended_atype: torch.Tensor,
         nlist: torch.Tensor,
         mapping: Optional[torch.Tensor] = None,
+        comm_dict: Optional[Dict[str, torch.Tensor]] = None,
     ):
         """Compute the descriptor.
 
@@ -466,6 +468,8 @@ def forward(
             The neighbor list. shape: nf x nloc x nnei
         mapping
             The index mapping, not required by this descriptor.
+        comm_dict
+            The data needed for communication for parallel inference.
 
         Returns
         -------