pd: support paddle backend and water/se_e2_a

.github/workflows/test_python.yml

@@ -31,6 +31,7 @@ jobs:
         export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
         source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py
         source/install/uv_with_retry.sh pip install --system horovod --no-build-isolation
+        source/install/uv_with_retry.sh pip install --system --pre "paddlepaddle" -i https://www.paddlepaddle.org.cn/packages/nightly/cpu/
       env:
         # Please note that uv has some issues with finding
         # existing TensorFlow package. Currently, it uses

backend/find_paddle.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import importlib
+import os
+import site
+from functools import (
+    lru_cache,
+)
+from importlib.machinery import (
+    FileFinder,
+)
+from importlib.util import (
+    find_spec,
+)
+from pathlib import (
+    Path,
+)
+from sysconfig import (
+    get_path,
+)
+from typing import (
+    Optional,
+    Union,
+)
+
+
+@lru_cache
+def find_paddle() -> tuple[Optional[str], list[str]]:
+    """Find PaddlePadle library.
+
+    Tries to find PaddlePadle in the order of:
+
+    1. Environment variable `PADDLE_ROOT` if set
+    2. The current Python environment.
+    3. user site packages directory if enabled
+    4. system site packages directory (purelib)
+
+    Considering the default PaddlePadle package still uses old CXX11 ABI, we
+    cannot install it automatically.
+
+    Returns
+    -------
+    str, optional
+        PaddlePadle library path if found.
+    list of str
+        Paddle requirement if not found. Empty if found.
+    """
+    if os.environ.get("DP_ENABLE_PADDLE", "0") == "0":
+        return None, []
+    requires = []
+    pd_spec = None
+
+    if (pd_spec is None or not pd_spec) and os.environ.get("PADDLE_ROOT") is not None:
+        site_packages = Path(os.environ.get("PADDLE_ROOT")).parent.absolute()
+        pd_spec = FileFinder(str(site_packages)).find_spec("paddle")
+
+    # get paddle spec
+    # note: isolated build will not work for backend
+    if pd_spec is None or not pd_spec:
+        pd_spec = find_spec("paddle")
+
+    if not pd_spec and site.ENABLE_USER_SITE:
+        # first search TF from user site-packages before global site-packages
+        site_packages = site.getusersitepackages()
+        if site_packages:
+            pd_spec = FileFinder(site_packages).find_spec("paddle")
+
+    if not pd_spec:
+        # purelib gets site-packages path
+        site_packages = get_path("purelib")
+        if site_packages:
+            pd_spec = FileFinder(site_packages).find_spec("paddle")
+
+    # get install dir from spec
+    try:
+        pd_install_dir = pd_spec.submodule_search_locations[0]  # type: ignore
+        # AttributeError if ft_spec is None
+        # TypeError if submodule_search_locations are None
+        # IndexError if submodule_search_locations is an empty list
+    except (AttributeError, TypeError, IndexError):
+        pd_install_dir = None
+        requires.extend(get_pd_requirement()["paddle"])
+    return pd_install_dir, requires
+
+
+@lru_cache
+def get_pd_requirement(pd_version: str = "") -> dict:
+    """Get PaddlePadle requirement when Paddle is not installed.
+
+    If pd_version is not given and the environment variable `PADDLE_VERSION` is set, use it as the requirement.
+
+    Parameters
+    ----------
+    pd_version : str, optional
+        Paddle version
+
+    Returns
+    -------
+    dict
+        PaddlePadle requirement.
+    """
+    if pd_version is None:
+        return {"paddle": []}
+    if pd_version == "":
+        pd_version = os.environ.get("PADDLE_VERSION", "")
+
+    return {
+        "paddle": [
+            "paddlepaddle>=3.0.0b1" if pd_version != "" else "paddlepaddle>=3.0.0b1",
+        ],
+    }
+
+
+@lru_cache
+def get_pd_version(pd_path: Optional[Union[str, Path]]) -> str:
+    """Get Paddle version from a Paddle Python library path.
+
+    Parameters
+    ----------
+    pd_path : str or Path
+        Paddle Python library path, e.g. "/python3.10/site-packages/paddle/"
+
+    Returns
+    -------
+    str
+        version
+    """
+    if pd_path is None or pd_path == "":
+        return ""
+    version_file = Path(pd_path) / "version" / "__init__.py"
+    spec = importlib.util.spec_from_file_location("paddle.version", version_file)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.full_version

deepmd/backend/paddle.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from importlib.util import (
+    find_spec,
+)
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    ClassVar,
+)
+
+from deepmd.backend.backend import (
+    Backend,
+)
+
+if TYPE_CHECKING:
+    from argparse import (
+        Namespace,
+    )
+
+    from deepmd.infer.deep_eval import (
+        DeepEvalBackend,
+    )
+    from deepmd.utils.neighbor_stat import (
+        NeighborStat,
+    )
+
+
+@Backend.register("pd")
+@Backend.register("paddle")
+class PaddleBackend(Backend):
+    """Paddle backend."""
+
+    name = "Paddle"
+    """The formal name of the backend."""
+    features: ClassVar[Backend.Feature] = (
+        Backend.Feature.ENTRY_POINT
+        | Backend.Feature.DEEP_EVAL
+        | Backend.Feature.NEIGHBOR_STAT
+        | Backend.Feature.IO
+    )
+    """The features of the backend."""
+    suffixes: ClassVar[list[str]] = [".json", ".pd"]
+    """The suffixes of the backend."""
+
+    def is_available(self) -> bool:
+        """Check if the backend is available.
+
+        Returns
+        -------
+        bool
+            Whether the backend is available.
+        """
+        return find_spec("paddle") is not None
+
+    @property
+    def entry_point_hook(self) -> Callable[["Namespace"], None]:
+        """The entry point hook of the backend.
+
+        Returns
+        -------
+        Callable[[Namespace], None]
+            The entry point hook of the backend.
+        """
+        from deepmd.pd.entrypoints.main import main as deepmd_main
+
+        return deepmd_main
+
+    @property
+    def deep_eval(self) -> type["DeepEvalBackend"]:
+        """The Deep Eval backend of the backend.
+
+        Returns
+        -------
+        type[DeepEvalBackend]
+            The Deep Eval backend of the backend.
+        """
+        from deepmd.pd.infer.deep_eval import DeepEval as DeepEvalPD
+
+        return DeepEvalPD
+
+    @property
+    def neighbor_stat(self) -> type["NeighborStat"]:
+        """The neighbor statistics of the backend.
+
+        Returns
+        -------
+        type[NeighborStat]
+            The neighbor statistics of the backend.
+        """
+        from deepmd.pd.utils.neighbor_stat import (
+            NeighborStat,
+        )
+
+        return NeighborStat
+
+    @property
+    def serialize_hook(self) -> Callable[[str], dict]:
+        """The serialize hook to convert the model file to a dictionary.
+
+        Returns
+        -------
+        Callable[[str], dict]
+            The serialize hook of the backend.
+        """
+        from deepmd.pd.utils.serialization import (
+            serialize_from_file,
+        )
+
+        return serialize_from_file
+
+    @property
+    def deserialize_hook(self) -> Callable[[str, dict], None]:
+        """The deserialize hook to convert the dictionary to a model file.
+
+        Returns
+        -------
+        Callable[[str, dict], None]
+            The deserialize hook of the backend.
+        """
+        from deepmd.pd.utils.serialization import (
+            deserialize_to_file,
+        )
+
+        return deserialize_to_file

deepmd/main.py

@@ -99,9 +99,10 @@ def main_parser() -> argparse.ArgumentParser:
         formatter_class=RawTextArgumentDefaultsHelpFormatter,
         epilog=textwrap.dedent(
             """\
-        Use --tf or --pt to choose the backend:
+        Use --tf, --pt or --pd to choose the backend:
             dp --tf train input.json
             dp --pt train input.json
+            dp --pd train input.json
         """
         ),
     )

deepmd/pd/cxx_op.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+
+def load_library(module_name: str) -> bool:
+    """Load OP library.
+
+    Parameters
+    ----------
+    module_name : str
+        Name of the module
+
+    Returns
+    -------
+    bool
+        Whether the library is loaded successfully
+    """
+    # NOTE: Paddle do not support loading library from .so file yet.
+    return False
+
+
+ENABLE_CUSTOMIZED_OP = load_library("deepmd_op_pd")
+
+__all__ = [
+    "ENABLE_CUSTOMIZED_OP",
+]

deepmd/pd/entrypoints/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later