pip install lammps module (deepmodeling#2186)

## add lammps module to the wheel

In the next version, one can install DeePMD-kit using

```sh
pip install deepmd-kit[gpu,lmp]
```

Then one can use `lmp` to run the LAMMPS program with `pair deepmd`.
This works for Linux and macOS, and requires `libpython` installed to
load TF libraries. The LAMMPS Python package was prebuilt in [my
repo](https://github.com/njzjz/lammps-wheel).

## Build a Python package with LAMMPS module from the source

```sh
export DP_LAMMPS_VERSION=stable_23Jun2022_update2
export DP_variant=cuda
pip install -v .
```
But note that the same compilation condition (CXX ABI, MPI) should be
used to build LAMMPS and its plugin.

## Other fixes
- rename `op_abi` module to `deepmd_op` to share the same name between
Python and C++
- fix macOS OP library suffix in api_cc: `libdeepmd_op.dylib` ->
`libdeepmd_op.so`
- fix Windows dlopen APIs in api_cc
- fix compilation error on windows: `an explicit specialization or
instantiation of a function template cannot have any default arguments`
- avoid linking libdeepmd_lmp with the MPI library, which has been
linked with the main LAMMPS

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

Author: njzjz

deepmd/env.py

@@ -376,7 +376,7 @@ def _get_package_constants(
 TF_VERSION = GLOBAL_CONFIG["tf_version"]
 TF_CXX11_ABI_FLAG = int(GLOBAL_CONFIG["tf_cxx11_abi_flag"])
 
-op_module = get_module("op_abi")
+op_module = get_module("deepmd_op")
 op_grads_module = get_module("op_grads")
 
 # FLOAT_PREC

deepmd/lmp.py

@@ -0,0 +1,49 @@
+"""Register entry points for lammps-wheel."""
+import os
+import platform
+from pathlib import Path
+from typing import List, Optional
+
+from find_libpython import find_libpython
+
+from deepmd.env import tf
+
+
+def get_env(paths: List[Optional[str]]) -> str:
+    """Get the environment variable from given paths"""
+    return ":".join((p for p in paths if p is not None))
+
+
+if platform.system() == "Linux":
+    lib_env = "LD_LIBRARY_PATH"
+elif platform.system() == "Darwin":
+    lib_env = "DYLD_FALLBACK_LIBRARY_PATH"
+else:
+    raise RuntimeError("Unsupported platform")
+
+tf_dir = tf.sysconfig.get_lib()
+op_dir = str((Path(__file__).parent / "op").absolute())
+# set LD_LIBRARY_PATH
+os.environ[lib_env] = get_env([
+    os.environ.get(lib_env),
+    tf_dir,
+    os.path.join(tf_dir, "python"),
+    op_dir,
+])
+
+# preload python library
+libpython = find_libpython()
+if platform.system() == "Linux":
+    preload_env = "LD_PRELOAD"
+elif platform.system() == "Darwin":
+    preload_env = "DYLD_INSERT_LIBRARIES"
+else:
+    raise RuntimeError("Unsupported platform")
+os.environ[preload_env] = get_env([
+    os.environ.get(preload_env),
+    libpython,
+])
+
+def get_op_dir() -> str:
+    """Get the directory of the deepmd-kit OP library"""
+    return op_dir

doc/install/easy-install.md

@@ -84,7 +84,7 @@ docker pull deepmodeling/dpmdkit-rocm:dp2.0.3-rocm4.5.2-tf2.6-lmp29Sep2021
 
 ## Install Python interface with pip
 
-If you only want to install the Python interface and have no existing TensorFlow installed, you can use `pip` to install the pre-built package of the Python interface with CUDA 11 supported:
+If you have no existing TensorFlow installed, you can use `pip` to install the pre-built package of the Python interface with CUDA 11 supported:
 
 ```bash
 pip install deepmd-kit[gpu]
@@ -95,6 +95,13 @@ Or install the CPU version without CUDA supported:
 pip install deepmd-kit[cpu]
 ```
 
+[LAMMPS module](../third-party/lammps-command.md) is only provided on Linux and macOS. To enable it, add `lmp` to extras:
+```bash
+pip install deepmd-kit[gpu,lmp]
+```
+MPICH is required for parallel running.
+
+It is suggested to install the package into an isolated environment.
 The supported platform includes Linux x86-64 and aarch64 with GNU C Library 2.28 or above, macOS x86-64, and Windows x86-64.
 A specific version of TensorFlow which is compatible with DeePMD-kit will be also installed.
 

pyproject.toml

@@ -49,25 +49,35 @@ repository = "https://github.com/deepmodeling/deepmd-kit"
 write_to = "deepmd/_version.py"
 
 [tool.cibuildwheel]
-test-command = "dp -h"
-test-requires = "tensorflow"
+test-command = [
+    "dp -h",
+]
+test-extras = ["cpu"]
 build = ["cp310-*"]
 skip = ["*-win32", "*-manylinux_i686", "*-musllinux*"]
 # TODO: bump to "latest" tag when CUDA supports GCC 12
 manylinux-x86_64-image = "quay.io/pypa/manylinux_2_28_x86_64:2022-11-19-1b19e81"
 manylinux-aarch64-image = "quay.io/pypa/manylinux_2_28_aarch64:2022-11-19-1b19e81"
 
 [tool.cibuildwheel.macos]
+environment = { DP_LAMMPS_VERSION="stable_23Jun2022_update2" }
+before-all = ["brew install mpich"]
 repair-wheel-command = "delocate-wheel --require-archs {delocate_archs} -w {dest_dir} -v {wheel} --ignore-missing-dependencies"
+test-extras = ["cpu", "test", "lmp"]
+test-command = [
+    "dp -h",
+    "pytest {project}/source/tests/test_lammps.py"
+]
 
 [tool.cibuildwheel.linux]
-repair-wheel-command = "auditwheel repair --exclude libtensorflow_framework.so.2 --exclude libtensorflow_framework.so.1 --exclude libtensorflow_framework.so -w {dest_dir} {wheel}"
+repair-wheel-command = "auditwheel repair --exclude libtensorflow_framework.so.2 --exclude libtensorflow_framework.so.1 --exclude libtensorflow_framework.so --exclude _pywrap_tensorflow_internal.so -w {dest_dir} {wheel}"
 environment-pass = ["CIBW_BUILD", "DP_VARIANT"]
-before-all = """
-if [ "$(uname -m)" = "x86_64" ]; then
-    yum config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && yum install -y cuda-11-8
-fi
-"""
+environment = { DP_VARIANT="cuda", DP_LAMMPS_VERSION="stable_23Jun2022_update2", MPI_HOME="/usr/lib64/mpich", PATH="/usr/lib64/mpich/bin:$PATH" }
+before-all = [
+    """{ if [ "$(uname -m)" = "x86_64" ] ; then yum config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && yum install -y cuda-11-8; fi }""",
+    "yum install -y mpich-devel",
+]
+
 
 # selectively turn of lintner warnings, always include reasoning why any warning should
 # be silenced

setup.py

@@ -34,6 +34,13 @@
     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", "")
+if dp_lammps_version != "":
+    cmake_args.append("-DBUILD_CPP_IF:BOOL=TRUE")
+    cmake_args.append("-DUSE_TF_PYTHON_LIBS:BOOL=TRUE")
+    cmake_args.append(f"-DLAMMPS_VERSION={dp_lammps_version}")
+else:
+    cmake_args.append("-DBUILD_CPP_IF:BOOL=FALSE")
 
 tf_install_dir, _ = find_tensorflow()
 tf_version = get_tf_version(tf_install_dir)
@@ -73,7 +80,6 @@ def get_tag(self):
     cmake_args=[
         f"-DTENSORFLOW_ROOT:PATH={tf_install_dir}",
         "-DBUILD_PY_IF:BOOL=TRUE",
-        "-DBUILD_CPP_IF:BOOL=FALSE",
         *cmake_args,
     ],
     cmake_source_dir="source",
@@ -95,9 +101,17 @@ def get_tag(self):
             "sphinx-argparse",
             "pygments-lammps",
             ],
+        "lmp": [
+            "lammps-manylinux-2-28~=2022.6.23.2.2; platform_system=='Linux'",
+            "lammps~=2022.6.23.2.2; platform_system!='Linux'",
+            "find_libpython",
+        ],
         **get_tf_requirement(tf_version),
     },
-    entry_points={"console_scripts": ["dp = deepmd.entrypoints.main:main"]},
+    entry_points={
+        "console_scripts": ["dp = deepmd.entrypoints.main:main"],
+        "lammps.plugins": ["deepmd = deepmd.lmp:get_op_dir"],
+    },
     cmdclass = {
         "bdist_wheel": bdist_wheel_abi3,
     },

source/CMakeLists.txt

@@ -112,7 +112,11 @@ set(DEEPMD_SOURCE_DIR ${PROJECT_SOURCE_DIR}/..)
 
 # setup tensorflow libraries by python
 if (USE_TF_PYTHON_LIBS)
-  find_package (Python COMPONENTS Interpreter Development REQUIRED)
+  if(NOT $ENV{CIBUILDWHEEL} STREQUAL "1")
+    find_package (Python COMPONENTS Interpreter Development REQUIRED)
+  else()
+    set(Python_LIBRARIES ${Python_LIBRARY})
+  endif()
 endif(USE_TF_PYTHON_LIBS)
 
 # find tensorflow, I need tf abi info
@@ -149,8 +153,8 @@ endif()
 
 # define names of libs
 set (LIB_DEEPMD		"deepmd")
+set (LIB_DEEPMD_OP		"deepmd_op")
 if (BUILD_CPP_IF)
-  set (LIB_DEEPMD_OP		"deepmd_op")
   set (LIB_DEEPMD_CC		"deepmd_cc")
   set (LIB_DEEPMD_C     "deepmd_c")
   if (USE_CUDA_TOOLKIT)
@@ -181,8 +185,10 @@ if (BUILD_CPP_IF)
   add_subdirectory (lmp/)
   if (CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 4.8)
     # add_subdirectory (md/)
+    if(NOT BUILD_PY_IF)
     add_subdirectory (ipi/)
     add_subdirectory (gmx/)
+    endif()
   endif ()
 endif (BUILD_CPP_IF)
 
@@ -205,7 +211,7 @@ add_custom_target(lammps
     COMMAND ${CMAKE_COMMAND} -P ${CMAKE_CURRENT_BINARY_DIR}/cmake_lammps.cmake)
 
 # add configure file
-if(BUILD_CPP_IF)
+if(BUILD_CPP_IF AND NOT BUILD_PY_IF)
   include(CMakePackageConfigHelpers)
   set(targets_export_name ${CMAKE_PROJECT_NAME}Targets CACHE INTERNAL "")
   set(generated_dir "${CMAKE_CURRENT_BINARY_DIR}/generated" CACHE INTERNAL "")
@@ -220,4 +226,4 @@ if(BUILD_CPP_IF)
     "${config_file}" INSTALL_DESTINATION ${cmake_files_install_dir})
   install(FILES ${version_file} ${config_file}
     DESTINATION ${cmake_files_install_dir})
-endif(BUILD_CPP_IF)
+endif(BUILD_CPP_IF AND NOT BUILD_PY_IF)

source/api_c/CMakeLists.txt

@@ -23,6 +23,12 @@ if (CMAKE_TESTING_ENABLED)
   target_link_libraries(${libname} PRIVATE coverage_config)
 endif()
 
+if(BUILD_PY_IF)
+  install(
+    TARGETS ${libname}
+    DESTINATION deepmd/op/
+  )
+else(BUILD_PY_IF)
 install(
   TARGETS ${libname}
   EXPORT ${CMAKE_PROJECT_NAME}Targets
@@ -33,6 +39,7 @@ install(
   FILES		${INC_SRC}
   DESTINATION	include/deepmd
 )
+endif(BUILD_PY_IF)
 
 if (PACKAGE_C)
   MESSAGE(STATUS "Packaging C API library")

source/api_cc/CMakeLists.txt

@@ -42,6 +42,12 @@ if (CMAKE_TESTING_ENABLED)
 endif()
 target_compile_features(${libname} PUBLIC cxx_std_11)
 
+if(BUILD_PY_IF)
+  install(
+    TARGETS ${libname}
+    DESTINATION deepmd/op/
+  )
+else(BUILD_PY_IF)
 install(
   TARGETS ${libname}
   EXPORT ${CMAKE_PROJECT_NAME}Targets
@@ -64,3 +70,4 @@ ${CMAKE_INSTALL_PREFIX}/lib/${CMAKE_SHARED_LIBRARY_PREFIX}${libname}${LOW_PREC_V
 if (CMAKE_TESTING_ENABLED)
   add_subdirectory(tests)
 endif()
+endif(BUILD_PY_IF)

source/api_cc/src/DeepPot.cc

@@ -87,7 +87,7 @@ run_model <double, double> (ENERGYTYPE &			dener,
 	   Session *			session, 
 	   const std::vector<std::pair<std::string, Tensor>> & input_tensors,
 	   const AtomMap&	atommap, 
-	   const int			nghost = 0);
+	   const int			nghost);
 
 template
 void
@@ -97,7 +97,7 @@ run_model <double, float> (ENERGYTYPE &			dener,
      Session *			session, 
      const std::vector<std::pair<std::string, Tensor>> & input_tensors,
      const AtomMap&	atommap, 
-     const int			nghost = 0);
+     const int			nghost);
 
 template
 void
@@ -107,7 +107,7 @@ run_model <float, double> (ENERGYTYPE &			dener,
 	   Session *			session, 
 	   const std::vector<std::pair<std::string, Tensor>> & input_tensors,
 	   const AtomMap&	atommap, 
-	   const int			nghost = 0);
+	   const int			nghost);
 
 template
 void
@@ -117,7 +117,7 @@ run_model <float, float> (ENERGYTYPE &			dener,
      Session *			session, 
      const std::vector<std::pair<std::string, Tensor>> & input_tensors,
      const AtomMap&	atommap, 
-     const int			nghost = 0);
+     const int			nghost);
 
 template <typename MODELTYPE, typename VALUETYPE>
 static void run_model (ENERGYTYPE   &		dener,
@@ -210,7 +210,7 @@ void run_model <double, double> (ENERGYTYPE   &		dener,
     Session*			session, 
     const std::vector<std::pair<std::string, Tensor>> & input_tensors,
     const deepmd::AtomMap &   atommap, 
-    const int&		nghost = 0);
+    const int&		nghost);
 
 template
 void run_model <double, float> (ENERGYTYPE   &		dener,
@@ -221,7 +221,7 @@ void run_model <double, float> (ENERGYTYPE   &		dener,
     Session*			session, 
     const std::vector<std::pair<std::string, Tensor>> & input_tensors,
     const deepmd::AtomMap &   atommap, 
-    const int&		nghost = 0);
+    const int&		nghost);
 
 template
 void run_model <float, double> (ENERGYTYPE   &		dener,
@@ -232,7 +232,7 @@ void run_model <float, double> (ENERGYTYPE   &		dener,
     Session*			session, 
     const std::vector<std::pair<std::string, Tensor>> & input_tensors,
     const deepmd::AtomMap &   atommap, 
-    const int&		nghost = 0);
+    const int&		nghost);
 
 template
 void run_model <float, float> (ENERGYTYPE   &		dener,
@@ -243,7 +243,7 @@ void run_model <float, float> (ENERGYTYPE   &		dener,
     Session*			session, 
     const std::vector<std::pair<std::string, Tensor>> & input_tensors,
     const deepmd::AtomMap &   atommap, 
-    const int&		nghost = 0);
+    const int&		nghost);
 
 DeepPot::
 DeepPot ()

source/api_cc/src/common.cc

@@ -1,8 +1,22 @@
 #include "common.h"
 #include "AtomMap.h"
 #include "device.h"
-#include <dlfcn.h>
 #include <fcntl.h>
+#if defined(_WIN32)
+#if defined(_WIN32_WINNT)
+#undef _WIN32_WINNT
+#endif
+
+// target Windows version is windows 7 and later
+#define _WIN32_WINNT _WIN32_WINNT_WIN7
+#define PSAPI_VERSION 2
+#include <windows.h>
+#include <io.h>
+#define O_RDONLY _O_RDONLY
+#else
+// not windows
+#include <dlfcn.h>
+#endif
 #include "google/protobuf/text_format.h"
 #include "google/protobuf/io/zero_copy_stream_impl.h"
 
@@ -299,10 +313,11 @@ deepmd::
 load_op_library()
 {
   tensorflow::Env* env = tensorflow::Env::Default();
-  std::string dso_path = env->FormatLibraryFileName("deepmd_op", "");
 #if defined(_WIN32)
+  std::string dso_path = "deepmd_op.dll";
   void* dso_handle = LoadLibrary(dso_path.c_str());
 #else
+  std::string dso_path = "libdeepmd_op.so";
   void* dso_handle = dlopen(dso_path.c_str(), RTLD_NOW | RTLD_LOCAL);
 #endif
   if (!dso_handle) {