From 43562c3d173d10632cc53c59bd1a8aa41dc9fd3e Mon Sep 17 00:00:00 2001
From: xulongteng <xulongteng@baidu.com>
Date: Wed, 31 Jul 2019 20:19:29 +0800
Subject: [PATCH] add cube init

---
 configure/proto/server_configure.proto |  2 ++
 predictor/common/constant.cpp          |  1 +
 predictor/common/constant.h            |  3 ++
 predictor/framework/resource.cpp       | 46 ++++++++++++++++++++++++--
 predictor/framework/resource.h         |  4 ++-
 predictor/src/pdserving.cpp            |  8 +++++
 6 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/configure/proto/server_configure.proto b/configure/proto/server_configure.proto
index 3f236fcf..d1d50860 100644
--- a/configure/proto/server_configure.proto
+++ b/configure/proto/server_configure.proto
@@ -35,6 +35,8 @@ message ModelToolkitConf { repeated EngineDesc engines = 1; };
 message ResourceConf {
   required string model_toolkit_path = 1;
   required string model_toolkit_file = 2;
+  optional string cube_config_path = 3;
+  optional string cube_config_file = 4;
 };
 
 // DAG node depency info
diff --git a/predictor/common/constant.cpp b/predictor/common/constant.cpp
index f9dac852..c1cd1c8b 100644
--- a/predictor/common/constant.cpp
+++ b/predictor/common/constant.cpp
@@ -40,6 +40,7 @@ DEFINE_int32(
 DEFINE_int32(reload_interval_s, 10, "");
 DEFINE_bool(enable_model_toolkit, false, "enable model toolkit");
 DEFINE_string(enable_protocol_list, "baidu_std", "set protocol list");
+DEFINE_bool(enable_cube, false, "enable cube");
 
 const char* START_OP_NAME = "startup_op";
 }  // namespace predictor
diff --git a/predictor/common/constant.h b/predictor/common/constant.h
index 857acd35..da44103e 100644
--- a/predictor/common/constant.h
+++ b/predictor/common/constant.h
@@ -39,6 +39,9 @@ DECLARE_int32(num_threads);
 DECLARE_int32(reload_interval_s);
 DECLARE_bool(enable_model_toolkit);
 DECLARE_string(enable_protocol_list);
+DECLARE_bool(enable_cube);
+DECLARE_string(cube_config_path);
+DECLARE_string(cube_config_file);
 
 // STATIC Variables
 extern const char* START_OP_NAME;
diff --git a/predictor/framework/resource.cpp b/predictor/framework/resource.cpp
index a2ed40cb..1ce765de 100644
--- a/predictor/framework/resource.cpp
+++ b/predictor/framework/resource.cpp
@@ -22,7 +22,7 @@ namespace paddle_serving {
 namespace predictor {
 
 using configure::ResourceConf;
-
+using rec::mcube::CubeAPI;
 // __thread bool p_thread_initialized = false;
 
 static void dynamic_resource_deleter(void* d) {
@@ -91,6 +91,45 @@ int Resource::initialize(const std::string& path, const std::string& file) {
   return 0;
 }
 
+int Resource::cube_initialize(const std::string& path,
+                              const std::string& file) {
+  // cube
+  if (!FLAGS_enable_cube) {
+    return 0;
+  }
+
+  ResourceConf resource_conf;
+  if (configure::read_proto_conf(path, file, &resource_conf) != 0) {
+    LOG(ERROR) << "Failed initialize resource from: " << path << "/" << file;
+    return -1;
+  }
+
+  int err = 0;
+  std::string cube_config_path = resource_conf.cube_config_path();
+  if (err != 0) {
+    LOG(ERROR) << "reade cube_config_path failed, path[" << path << "], file["
+               << cube_config_path << "]";
+    return -1;
+  }
+  std::string cube_config_file = resource_conf.cube_config_file();
+  if (err != 0) {
+    LOG(ERROR) << "reade cube_config_file failed, path[" << path << "], file["
+               << cube_config_file << "]";
+    return -1;
+  }
+  err = CubeAPI::instance()->init(cube_config_path.c_str(),
+                                  cube_config_file.c_str());
+  if (err != 0) {
+    LOG(ERROR) << "failed initialize cube, config: " << cube_config_path << "/"
+               << cube_config_file << " error code : " << err;
+    return -1;
+  }
+
+  LOG(INFO) << "Successfully initialize cube";
+
+  return 0;
+}
+
 int Resource::thread_initialize() {
   // mempool
   if (MempoolWrapper::instance().thread_initialize() != 0) {
@@ -192,7 +231,10 @@ int Resource::finalize() {
     LOG(ERROR) << "Failed proc finalize infer manager";
     return -1;
   }
-
+  if (CubeAPI::instance()->destroy() != 0) {
+    LOG(ERROR) << "Destory cube api failed ";
+    return -1;
+  }
   THREAD_KEY_DELETE(_tls_bspec_key);
 
   return 0;
diff --git a/predictor/framework/resource.h b/predictor/framework/resource.h
index a9dde0ab..01893722 100644
--- a/predictor/framework/resource.h
+++ b/predictor/framework/resource.h
@@ -13,7 +13,9 @@
 // limitations under the License.
 
 #pragma once
+#include <memory>
 #include <string>
+#include "include/cube/cube_api.h"
 #include "kvdb/paddle_rocksdb.h"
 #include "predictor/common/inner_common.h"
 #include "predictor/framework/memory.h"
@@ -45,7 +47,7 @@ class Resource {
   }
 
   int initialize(const std::string& path, const std::string& file);
-
+  int cube_initialize(const std::string& path, const std::string& file);
   int thread_initialize();
 
   int thread_clear();
diff --git a/predictor/src/pdserving.cpp b/predictor/src/pdserving.cpp
index 46c9b468..a86b39ab 100644
--- a/predictor/src/pdserving.cpp
+++ b/predictor/src/pdserving.cpp
@@ -209,6 +209,14 @@ int main(int argc, char** argv) {
   }
   LOG(INFO) << "Succ call pthread worker start function";
 
+  if (Resource::instance().cube_initialize(FLAGS_resource_path,
+                                           FLAGS_resource_file) != 0) {
+    LOG(ERROR) << "Failed initialize cube, conf: " << FLAGS_resource_path << "/"
+               << FLAGS_resource_file;
+    return -1;
+  }
+  LOG(INFO) << "Succ initialize cube";
+
   FLAGS_logtostderr = false;
 
   if (ServerManager::instance().start_and_wait() != 0) {
-- 
GitLab