From 21bfc1616c4452fe6b25bb977e56986ca1c7df0d Mon Sep 17 00:00:00 2001
From: qianhaohao <qianhaohao@baidu.com>
Date: Fri, 3 Dec 2021 13:06:00 +0800
Subject: [PATCH] add cmake for python grpc proto

---
 cmake/external/grpc.cmake                     |  1 +
 cmake/generic.cmake                           | 30 +++++++++++++++++++
 core/he/CMakeLists.txt                        |  7 +++++
 .../paddle_fl/feature_engineering/README.md   |  8 +----
 4 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/cmake/external/grpc.cmake b/cmake/external/grpc.cmake
index bdbd13a4..91b3173f 100644
--- a/cmake/external/grpc.cmake
+++ b/cmake/external/grpc.cmake
@@ -19,6 +19,7 @@ SET(GRPC_SOURCES_DIR ${THIRD_PARTY_PATH}/grpc)
 SET(GRPC_INSTALL_DIR ${THIRD_PARTY_PATH}/install/grpc)
 SET(GRPC_INCLUDE_DIR "${GRPC_INSTALL_DIR}/include/" CACHE PATH "grpc include directory." FORCE)
 SET(GRPC_CPP_PLUGIN "${GRPC_INSTALL_DIR}/bin/grpc_cpp_plugin" CACHE FILEPATH "GRPC_CPP_PLUGIN" FORCE)
+SET(GRPC_PYTHON_PLUGIN "${GRPC_INSTALL_DIR}/bin/grpc_python_plugin" CACHE FILEPATH "GRPC_PYTHON_PLUGIN" FORCE)
 
 include(ProcessorCount)
 ProcessorCount(NUM_OF_PROCESSOR)
diff --git a/cmake/generic.cmake b/cmake/generic.cmake
index c0275cac..f078b535 100644
--- a/cmake/generic.cmake
+++ b/cmake/generic.cmake
@@ -714,6 +714,36 @@ function(py_proto_compile TARGET_NAME)
   add_custom_target(${TARGET_NAME} ALL DEPENDS ${py_srcs} protobuf)
 endfunction()
 
+function(py_proto_grpc_compile TARGET_NAME)
+  set(oneValueArgs PROTO PROTO_PATH)
+  set(multiValueArgs SRCS)
+  set(options "")
+  cmake_parse_arguments(py_proto_grpc_compile "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
+
+  get_filename_component(ABS_PROTO ${py_proto_grpc_compile_PROTO} ABSOLUTE)
+  get_filename_component(PROTO_WE ${py_proto_grpc_compile_PROTO} NAME_WE)
+  #get_filename_component(PROTO_PATH ${ABS_PROTO} PATH)
+  set(PROTO_PATH ${py_proto_grpc_compile_PROTO_PATH})
+
+  set(grpc_proto_py "${CMAKE_CURRENT_BINARY_DIR}/${PROTO_WE}_pb2.py")
+  set(grpc_grpc_py "${CMAKE_CURRENT_BINARY_DIR}/${PROTO_WE}_pb2_grpc.py")
+  
+  set(py_srcs)
+  list(APPEND py_srcs "${grpc_proto_py}" "${grpc_grpc_py}")
+
+  add_custom_command(
+          OUTPUT "${grpc_proto_py}" "${grpc_grpc_py}"
+          COMMAND ${PROTOBUF_PROTOC_EXECUTABLE}
+          ARGS --grpc_python_out "${CMAKE_CURRENT_BINARY_DIR}" -I "${PROTO_PATH}"
+          --plugin=protoc-gen-grpc_python="${GRPC_PYTHON_PLUGIN}" "${ABS_PROTO}"
+          COMMAND ${PROTOBUF_PROTOC_EXECUTABLE}
+          ARGS --python_out "${CMAKE_CURRENT_BINARY_DIR}" -I "${PROTO_PATH}"
+          "${ABS_PROTO}"
+          DEPENDS "${ABS_PROTO}" ${PROTOBUF_PROTOC_EXECUTABLE} extern_grpc)
+
+  add_custom_target(${TARGET_NAME} ALL DEPENDS ${py_srcs} protobuf extern_grpc)
+endfunction()
+
 function(py_test TARGET_NAME)
   if(WITH_TESTING)
     set(options "")
diff --git a/core/he/CMakeLists.txt b/core/he/CMakeLists.txt
index f5a33972..05539a38 100644
--- a/core/he/CMakeLists.txt
+++ b/core/he/CMakeLists.txt
@@ -13,5 +13,12 @@ add_library(he_utils MODULE ${PYBIND_HE_SRCS})
 target_link_libraries(he_utils PRIVATE pybind gmp gmpxx)
 set_target_properties(he_utils PROPERTIES PREFIX "${PYTHON_MODULE_PREFIX}")
 
+set(FE_PROTO_MODULE_PATH "paddle_fl/feature_engineering/proto")
+py_proto_grpc_compile(my_target PROTO "${CMAKE_SOURCE_DIR}/python/${FE_PROTO_MODULE_PATH}/metrics.proto"
+                                PROTO_PATH "${CMAKE_SOURCE_DIR}/python")
+install(FILES "${CMAKE_CURRENT_BINARY_DIR}/${FE_PROTO_MODULE_PATH}/metrics_pb2.py"
+              "${CMAKE_CURRENT_BINARY_DIR}/${FE_PROTO_MODULE_PATH}/metrics_pb2_grpc.py"
+        DESTINATION "${CMAKE_SOURCE_DIR}/python/${FE_PROTO_MODULE_PATH}")
+
 set(FEATURE_LIB "${CMAKE_SOURCE_DIR}/python/paddle_fl/feature_engineering/libs")
 install(TARGETS he_utils LIBRARY DESTINATION ${FEATURE_LIB})
diff --git a/python/paddle_fl/feature_engineering/README.md b/python/paddle_fl/feature_engineering/README.md
index 33637911..a844833f 100644
--- a/python/paddle_fl/feature_engineering/README.md
+++ b/python/paddle_fl/feature_engineering/README.md
@@ -35,13 +35,7 @@ make -j48
 make install
 ```
 
-3.生成grpc_pb
-```
-cd /path/to/PaddleFL/python
-python3 paddle_fl/feature_engineering/proto/run_protogen.py
-```
-
-4.pip打包并安装
+3.pip打包并安装
 ```
 cd /path/to/PaddleFL/python/paddle_fl
 mkdir build && cd build