throw a message if tf runtime is incompatible

fix deepmodeling#557 and deepmodeling#796.
njzjz · Jun 25, 2021 · 8d40a3b · 8d40a3b
1 parent 4555034
commit 8d40a3b
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 4 deletions.
diff --git a/deepmd/env.py b/deepmd/env.py
@@ -153,8 +153,6 @@ def get_module(module_name: str) -> "ModuleType":
         return module
 
 
-op_module = get_module("libop_abi")
-op_grads_module = get_module("libop_grads")
 
 
 def _get_package_constants(
@@ -165,7 +163,7 @@ def _get_package_constants(
     Parameters
     ----------
     config_file : str, optional
-        path to CONFIG file, by default "config/run_config.ini"
+        path to CONFIG file, by default "pkg_config/run_config.ini"
 
     Returns
     -------
@@ -178,6 +176,24 @@ def _get_package_constants(
 
 GLOBAL_CONFIG = _get_package_constants()
 MODEL_VERSION = GLOBAL_CONFIG["model_version"]
+TF_VERSION = GLOBAL_CONFIG["tf_version"]
+
+try:
+    op_module = get_module("libop_abi")
+    op_grads_module = get_module("libop_grads")
+except tf.errors.NotFoundError:
+    # different versions may cause incompatibility, see #557 and #796 for example
+    # throw a message if versions are different
+    if TF_VERSION != tf.version.VERSION:
+        raise RuntimeError("The version of TensorFlow used to compile this "
+            "deepmd-kit package is %s, but the version of TensorFlow runtime you "
+            "are using is %s. These two versions are incompatible and thus an error "
+            "is raised. You need to install TensorFlow %s, or rebuild deepmd-kit "
+            "using TensorFlow %s.\nIf you are using a wheel from pypi, you "
+            "may consider to install deepmd-kit execuating "
+            "`pip install deepmd-kit --no-binary deepmd-kit` instead." % (
+                TF_VERSION, tf.version.VERSION, TF_VERSION, tf.version.VERSION,
+            ))
 
 if GLOBAL_CONFIG["precision"] == "-DHIGH_PREC":
     GLOBAL_TF_FLOAT_PRECISION = tf.float64

diff --git a/source/cmake/Findtensorflow.cmake b/source/cmake/Findtensorflow.cmake
@@ -137,10 +137,27 @@ else (BUILD_CPP_IF)
   endif ()
 endif (BUILD_CPP_IF)
 
+# detect TensorFlow version
+try_run(
+  TENSORFLOW_VERSION_RUN_RESULT_VAR TENSORFLOW_VERSION_COMPILE_RESULT_VAR
+  ${CMAKE_CURRENT_BINARY_DIR}/tf_version
+  "${CMAKE_CURRENT_LIST_DIR}/tf_version.cpp"
+  LINK_LIBRARIES ${TensorFlowFramework_LIBRARY}
+  CMAKE_FLAGS "-DINCLUDE_DIRECTORIES:STRING=${TensorFlow_INCLUDE_DIRS}"
+  RUN_OUTPUT_VARIABLE TENSORFLOW_VERSION
+  COMPILE_OUTPUT_VARIABLE TENSORFLOW_VERSION_COMPILE_OUTPUT_VAR
+)
+if (NOT ${TENSORFLOW_VERSION_COMPILE_RESULT_VAR})
+  message(FATAL_ERROR "Failed to compile: \n ${TENSORFLOW_VERSION_COMPILE_OUTPUT_VAR}" )
+endif()
+if (NOT ${TENSORFLOW_VERSION_RUN_RESULT_VAR} EQUAL "0")
+  message(FATAL_ERROR "Failed to run, return code: ${TENSORFLOW_VERSION}" )
+endif()
+
 # print message
 if (NOT TensorFlow_FIND_QUIETLY)
   message(STATUS "Found TensorFlow: ${TensorFlow_INCLUDE_DIRS}, ${TensorFlow_LIBRARY}, ${TensorFlowFramework_LIBRARY} "
-    " in ${TensorFlow_search_PATHS}")
+    " in ${TensorFlow_search_PATHS} (found version \"${TENSORFLOW_VERSION}\")")
 endif ()
 
 unset(TensorFlow_search_PATHS)
diff --git a/source/cmake/tf_version.cpp b/source/cmake/tf_version.cpp
@@ -0,0 +1,10 @@
+#include <iostream>
+#include "tensorflow/core/public/version.h"
+
+int main(int argc, char * argv[])
+{
+  // See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/public/version.h
+  // TF_VERSION_STRING has been avaiable since TensorFlow v0.6
+  std::cout << TF_VERSION_STRING;
+  return 0;
+}
diff --git a/source/config/run_config.ini b/source/config/run_config.ini
@@ -6,5 +6,6 @@ GIT_DATE = @GIT_DATE@
 GIT_BRANCH = @GIT_BRANCH@
 TF_INCLUDE_DIR = @TensorFlow_INCLUDE_DIRS@
 TF_LIBS = @TensorFlow_LIBRARY@
+TF_VERSION = @TENSORFLOW_VERSION@
 PRECISION = @PREC_DEF@
 MODEL_VERSION=@MODEL_VERSION@