From e1ae5e07f7ccb658fa66a9f854591badb1da3d6b Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sun, 24 Sep 2023 20:56:06 -0400
Subject: [PATCH] cmake: use pip to install tensorflow (#2858)

When setting the `INSTALL_TENSORFLOW` variable, previously CMake calls
conda to install tensorflow. Now we use pip to install it.

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 source/CMakeLists.txt             |  7 +++----
 source/cmake/Findtensorflow.cmake | 19 ++++++++++++-------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt
index c3b5fde054..1a6b82bcf0 100644
--- a/source/CMakeLists.txt
+++ b/source/CMakeLists.txt
@@ -25,10 +25,6 @@ if((NOT BUILD_PY_IF) AND (NOT BUILD_CPP_IF))
 endif()
 
 if(BUILD_CPP_IF AND BUILD_TESTING)
-  if(NOT INSTALL_TENSORFLOW)
-    # some errors in conda packages...
-    find_package(GTest)
-  endif()
   if(NOT GTEST_LIBRARIES)
     configure_file(${CMAKE_CURRENT_SOURCE_DIR}/cmake/googletest.cmake.in
                    googletest-download/CMakeLists.txt @ONLY)
@@ -127,6 +123,9 @@ endif(USE_ROCM_TOOLKIT)
 set(DEEPMD_SOURCE_DIR ${PROJECT_SOURCE_DIR}/..)
 
 # setup tensorflow libraries by python
+if(INSTALL_TENSORFLOW)
+  set(USE_TF_PYTHON_LIBS TRUE)
+endif(INSTALL_TENSORFLOW)
 if(USE_TF_PYTHON_LIBS)
   if(NOT "$ENV{CIBUILDWHEEL}" STREQUAL "1")
     find_package(
diff --git a/source/cmake/Findtensorflow.cmake b/source/cmake/Findtensorflow.cmake
index 6ab5747805..3ebbd4ea62 100644
--- a/source/cmake/Findtensorflow.cmake
+++ b/source/cmake/Findtensorflow.cmake
@@ -19,25 +19,30 @@ if(SKBUILD)
 endif(SKBUILD)
 
 if(BUILD_CPP_IF AND INSTALL_TENSORFLOW)
-  # Here we try to install libtensorflow_cc using conda install.
+  # Here we try to install libtensorflow_cc using pip install.
 
   if(USE_CUDA_TOOLKIT)
-    set(VARIANT cuda)
+    set(VARIANT "")
   else()
-    set(VARIANT cpu)
+    set(VARIANT "-cpu")
   endif()
 
   if(NOT DEFINED TENSORFLOW_ROOT)
     set(TENSORFLOW_ROOT ${CMAKE_INSTALL_PREFIX})
   endif()
-  # execute conda install
-  execute_process(COMMAND conda create libtensorflow_cc=*=${VARIANT}* -c
-                          deepmodeling -y -p ${TENSORFLOW_ROOT})
+  # execute pip install
+  execute_process(
+    COMMAND ${Python_EXECUTABLE} -m pip install tensorflow${VARIANT} --no-deps
+            --target=${TENSORFLOW_ROOT})
+  set(TENSORFLOW_ROOT
+      ${TENSORFLOW_ROOT}/lib/python${Python_VERSION_MAJOR}.${Python_VERSION_MINOR}/site-packages/tensorflow
+  )
 endif()
 
 if(BUILD_CPP_IF
    AND USE_TF_PYTHON_LIBS
-   AND NOT SKBUILD)
+   AND NOT SKBUILD
+   AND NOT INSTALL_TENSORFLOW)
   # Here we try to install libtensorflow_cc.so as well as
   # libtensorflow_framework.so using libs within the python site-package
   # tensorflow folder.