Merge branch 'lmp_units' of github.com:Yi-FanLi/deepmd-kit into lmp_u…

…nits

.gitignore

@@ -22,7 +22,7 @@ _skbuild
 deepmd_kit.egg-info/
 dist
 .eggs
-_version.py
+/deepmd/_version.py
 venv*
 .vscode/**
 _build

backend/__init__.py

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

backend/dp_backend.py

@@ -4,12 +4,14 @@
     List,
 )
 
-from find_tensorflow import (
+from scikit_build_core import build as _orig
+
+from .find_tensorflow import (
     find_tensorflow,
 )
-
-# TODO: switch to scikit_build_core after it is available
-from setuptools import build_meta as _orig
+from .read_env import (
+    set_scikit_build_env,
+)
 
 __all__ = [
     "build_sdist",
@@ -24,10 +26,14 @@ def __dir__() -> List[str]:
     return __all__
 
 
+set_scikit_build_env()
+
 prepare_metadata_for_build_wheel = _orig.prepare_metadata_for_build_wheel
 build_wheel = _orig.build_wheel
 build_sdist = _orig.build_sdist
 get_requires_for_build_sdist = _orig.get_requires_for_build_sdist
+prepare_metadata_for_build_editable = _orig.prepare_metadata_for_build_editable
+build_editable = _orig.build_editable
 
 
 def get_requires_for_build_wheel(
@@ -36,7 +42,7 @@ def get_requires_for_build_wheel(
     return _orig.get_requires_for_build_wheel(config_settings) + find_tensorflow()[1]
 
 
-# TODO: export get_requires_for_build_editable, prepare_metadata_for_build_editable, build_editable
-# after scikit-build is ready
-# See https://github.com/scikit-build/scikit-build/issues/740
-# Now we use the legacy-editable mode
+def get_requires_for_build_editable(
+    config_settings: dict,
+) -> List[str]:
+    return _orig.get_requires_for_build_editable(config_settings) + find_tensorflow()[1]

backend/dynamic_metadata.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Dict,
+    List,
+    Optional,
+)
+
+from .find_tensorflow import (
+    get_tf_requirement,
+)
+from .read_env import (
+    get_argument_from_env,
+)
+
+__all__ = ["dynamic_metadata"]
+
+
+def __dir__() -> List[str]:
+    return __all__
+
+
+def dynamic_metadata(
+    field: str,
+    settings: Optional[Dict[str, object]] = None,
+) -> str:
+    assert field in ["optional-dependencies", "entry-points", "scripts"]
+    _, _, find_libpython_requires, extra_scripts, tf_version = get_argument_from_env()
+    if field == "scripts":
+        return {
+            "dp": "deepmd_cli.main:main",
+            **extra_scripts,
+        }
+    elif field == "optional-dependencies":
+        return {
+            "test": [
+                "dpdata>=0.1.9",
+                "ase",
+                "pytest",
+                "pytest-cov",
+                "pytest-sugar",
+            ],
+            "docs": [
+                "sphinx>=3.1.1",
+                "sphinx_rtd_theme>=1.0.0rc1",
+                "sphinx_markdown_tables",
+                "myst-nb",
+                "breathe",
+                "exhale",
+                "numpydoc",
+                "ase",
+                "deepmodeling-sphinx>=0.1.0",
+                "dargs>=0.3.4",
+                "sphinx-argparse",
+                "pygments-lammps",
+                "sphinxcontrib-bibtex",
+            ],
+            "lmp": [
+                "lammps~=2023.8.2.0.0; platform_system=='Linux'",
+                "lammps~=2023.8.2.0.0; platform_system!='Linux'",
+                *find_libpython_requires,
+            ],
+            "ipi": [
+                "i-PI",
+                *find_libpython_requires,
+            ],
+            **get_tf_requirement(tf_version),
+            "cu11": [
+                "nvidia-cuda-runtime-cu11",
+                "nvidia-cublas-cu11",
+                "nvidia-cufft-cu11",
+                "nvidia-curand-cu11",
+                "nvidia-cusolver-cu11",
+                "nvidia-cusparse-cu11",
+                "nvidia-cudnn-cu11",
+                "nvidia-cuda-nvcc-cu11",
+            ],
+            "cu12": [
+                "nvidia-cuda-runtime-cu12",
+                "nvidia-cublas-cu12",
+                "nvidia-cufft-cu12",
+                "nvidia-curand-cu12",
+                "nvidia-cusolver-cu12",
+                "nvidia-cusparse-cu12",
+                "nvidia-cudnn-cu12",
+                "nvidia-cuda-nvcc-cu12",
+            ],
+        }

backend/find_tensorflow.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import os
 import site
+from functools import (
+    lru_cache,
+)
 from importlib.machinery import (
     FileFinder,
 )
@@ -25,6 +28,7 @@
 )
 
 
+@lru_cache()
 def find_tensorflow() -> Tuple[Optional[str], List[str]]:
     """Find TensorFlow library.
 
@@ -89,6 +93,7 @@ def find_tensorflow() -> Tuple[Optional[str], List[str]]:
     return tf_install_dir, requires
 
 
+@lru_cache()
 def get_tf_requirement(tf_version: str = "") -> dict:
     """Get TensorFlow requirement (CPU) when TF is not installed.
 
@@ -143,6 +148,7 @@ def get_tf_requirement(tf_version: str = "") -> dict:
         }
 
 
+@lru_cache()
 def get_tf_version(tf_path: Union[str, Path]) -> str:
     """Get TF version from a TF Python library path.
 

backend/read_env.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Read environment variables to configure the build."""
+
+import os
+from functools import (
+    lru_cache,
+)
+from typing import (
+    Tuple,
+)
+
+from packaging.version import (
+    Version,
+)
+
+from .find_tensorflow import (
+    find_tensorflow,
+    get_tf_version,
+)
+
+
+@lru_cache()
+def get_argument_from_env() -> Tuple[str, list, list, dict, str]:
+    """Get the arguments from environment variables.
+
+    The environment variables are assumed to be not changed during the build.
+
+    Returns
+    -------
+    str
+        The minimum required CMake version.
+    list of str
+        The CMake arguments.
+    list of str
+        The requirements for the build.
+    dict
+        The extra scripts to be installed.
+    str
+        The TensorFlow version.
+    """
+    cmake_args = []
+    extra_scripts = {}
+    # get variant option from the environment varibles, available: cpu, cuda, rocm
+    dp_variant = os.environ.get("DP_VARIANT", "cpu").lower()
+    if dp_variant == "cpu" or dp_variant == "":
+        cmake_minimum_required_version = "3.16"
+    elif dp_variant == "cuda":
+        cmake_minimum_required_version = "3.23"
+        cmake_args.append("-DUSE_CUDA_TOOLKIT:BOOL=TRUE")
+        cuda_root = os.environ.get("CUDAToolkit_ROOT")
+        if cuda_root:
+            cmake_args.append(f"-DCUDAToolkit_ROOT:STRING={cuda_root}")
+    elif dp_variant == "rocm":
+        cmake_minimum_required_version = "3.21"
+        cmake_args.append("-DUSE_ROCM_TOOLKIT:BOOL=TRUE")
+        rocm_root = os.environ.get("ROCM_ROOT")
+        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}")
+    else:
+        raise RuntimeError("Unsupported DP_VARIANT option: %s" % dp_variant)
+
+    if os.environ.get("DP_BUILD_TESTING", "0") == "1":
+        cmake_args.append("-DBUILD_TESTING:BOOL=TRUE")
+    if os.environ.get("DP_ENABLE_NATIVE_OPTIMIZATION", "0") == "1":
+        cmake_args.append("-DENABLE_NATIVE_OPTIMIZATION:BOOL=TRUE")
+    dp_lammps_version = os.environ.get("DP_LAMMPS_VERSION", "")
+    dp_ipi = os.environ.get("DP_ENABLE_IPI", "0")
+    if dp_lammps_version != "" or dp_ipi == "1":
+        cmake_args.append("-DBUILD_CPP_IF:BOOL=TRUE")
+        cmake_args.append("-DUSE_TF_PYTHON_LIBS:BOOL=TRUE")
+    else:
+        cmake_args.append("-DBUILD_CPP_IF:BOOL=FALSE")
+
+    if dp_lammps_version != "":
+        cmake_args.append(f"-DLAMMPS_VERSION={dp_lammps_version}")
+    if dp_ipi == "1":
+        cmake_args.append("-DENABLE_IPI:BOOL=TRUE")
+        extra_scripts["dp_ipi"] = "deepmd.entrypoints.ipi:dp_ipi"
+
+    tf_install_dir, _ = find_tensorflow()
+    tf_version = get_tf_version(tf_install_dir)
+    if tf_version == "" or Version(tf_version) >= Version("2.12"):
+        find_libpython_requires = []
+    else:
+        find_libpython_requires = ["find_libpython"]
+    cmake_args.append(f"-DTENSORFLOW_VERSION={tf_version}")
+
+    cmake_args = [
+        f"-DTENSORFLOW_ROOT:PATH={tf_install_dir}",
+        "-DBUILD_PY_IF:BOOL=TRUE",
+        *cmake_args,
+    ]
+    return (
+        cmake_minimum_required_version,
+        cmake_args,
+        find_libpython_requires,
+        extra_scripts,
+        tf_version,
+    )
+
+
+def set_scikit_build_env():
+    """Set scikit-build environment variables before executing scikit-build."""
+    cmake_minimum_required_version, cmake_args, _, _, _ = get_argument_from_env()
+    os.environ["SKBUILD_CMAKE_MINIMUM_VERSION"] = cmake_minimum_required_version
+    os.environ["SKBUILD_CMAKE_ARGS"] = ";".join(cmake_args)

deepmd/.gitignore

@@ -1,4 +1,2 @@
-op/_*.py
 pkg_config
 run_config.ini
-!op/__init__.py

deepmd/env.py

@@ -27,6 +27,8 @@
     Version,
 )
 
+import deepmd.lib
+
 if TYPE_CHECKING:
     from types import (
         ModuleType,
@@ -101,7 +103,9 @@ def dlopen_library(module: str, filename: str):
     "TF_VERSION",
 ]
 
-SHARED_LIB_MODULE = "op"
+SHARED_LIB_MODULE = "lib"
+SHARED_LIB_DIR = Path(deepmd.lib.__path__[0])
+CONFIG_FILE = SHARED_LIB_DIR / "run_config.ini"
 
 # Python library version
 try:
@@ -361,11 +365,7 @@ def get_module(module_name: str) -> "ModuleType":
         ext = ".so"
         prefix = "lib"
 
-    module_file = (
-        (Path(__file__).parent / SHARED_LIB_MODULE / (prefix + module_name))
-        .with_suffix(ext)
-        .resolve()
-    )
+    module_file = (SHARED_LIB_DIR / (prefix + module_name)).with_suffix(ext).resolve()
 
     if not module_file.is_file():
         raise FileNotFoundError(f"module {module_name} does not exist")
@@ -433,7 +433,7 @@ def get_module(module_name: str) -> "ModuleType":
 
 
 def _get_package_constants(
-    config_file: Path = Path(__file__).parent / "run_config.ini",
+    config_file: Path = CONFIG_FILE,
 ) -> Dict[str, str]:
     """Read package constants set at compile time by CMake to dictionary.
 

deepmd/lmp.py

@@ -64,7 +64,7 @@ def get_library_path(module: str) -> List[str]:
     raise RuntimeError("Unsupported platform")
 
 tf_dir = tf.sysconfig.get_lib()
-op_dir = str((Path(__file__).parent / "op").absolute())
+op_dir = str((Path(__file__).parent / "lib").absolute())
 
 
 cuda_library_paths = []

deepmd_cli/main.py

@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import argparse
-import imp
+import importlib.util
 import logging
+import os
+import sys
 import textwrap
 from typing import (
     List,
@@ -12,11 +14,15 @@
 def load_child_module(name):
     """Load a child module without loading its parent module."""
     names = name.split(".")
-    path = None
-    for name in names:
-        f, path, info = imp.find_module(name, path)
-        path = [path]
-    return imp.load_module(name, f, path[0], info)
+    parent_spec = importlib.util.find_spec(names[0])
+    paths = os.path.join(*names[1:]) + ".py"
+    spec = importlib.util.spec_from_file_location(
+        name, os.path.join(parent_spec.submodule_search_locations[0], paths)
+    )
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    return module
 
 
 __version__ = load_child_module("deepmd._version").__version__