[Serving] Support user defined storage size of multi-tier Embedding in serving. (#417)

Author: lixy9474

docs/Processor.md

@@ -375,7 +375,10 @@ void* initialize(const char* model_entry, const char* model_config, int* state);
 "ev_storage_type": 12,
 
 # 多级存储路径, 如果设置了多级存储
-"ev_storage_path": "/ssd/1/"
+"ev_storage_path": "/ssd/1/",
+
+# 多级存储中每级存储的大小
+"ev_storage_size": [1024, 1024]
 }
 ```
 #### 模型路径配置

serving/processor/framework/graph_optimizer.cc

@@ -607,7 +607,8 @@ Status SavedModelOptimizer::RunNativeTFGraphPass() {
   }
 
   if (option_.st) {
-    TF_RETURN_IF_ERROR(RewriteEmbeddingVariableAttr(option_.st, option_.path));
+    TF_RETURN_IF_ERROR(RewriteEmbeddingVariableAttr(option_.st,
+                        option_.path, option_.size));
   }
 
   // Add other passes here
@@ -2360,12 +2361,14 @@ Node* SavedModelOptimizer::UpdateRestoreShardNodeInputs(
 }
 
 Status SavedModelOptimizer::RewriteEmbeddingVariableAttr(
-    embedding::StorageType st, std::string path) {
+    embedding::StorageType st, const std::string& path,
+    const std::vector<int64>& storage_size) {
   for (Node* node : graph_.nodes()) {
     if (node->op_def().name() == "KvResourceImportV2" ||
         node->op_def().name() == "InitializeKvVariableOp") {
       node->AddAttr("storage_type", st);
       node->AddAttr("storage_path", path);
+      node->AddAttr("storage_size", storage_size);
 
       LOG(INFO) << "debug op: " << node->DebugString();
     }

serving/processor/framework/graph_optimizer.h

@@ -149,6 +149,7 @@ struct GraphOptimizerOption {
   // multi tiered embedding
   embedding::StorageType st = embedding::StorageType::INVALID;
   std::string path;
+  std::vector<int64> size;
 };
 
 struct SrcInfo {
@@ -263,7 +264,8 @@ class SavedModelOptimizer : public GraphOptimizer {
       std::unordered_map<std::string, std::vector<Node*>>& origin_import_nodes,
       std::vector<Node*>& new_kv_import_nodes);
 
-  Status RewriteEmbeddingVariableAttr(embedding::StorageType st, std::string path);
+  Status RewriteEmbeddingVariableAttr(embedding::StorageType st, const std::string& path,
+                                      const std::vector<int64>& size);
 
   Node* storage_pointer_node_ = nullptr;// storage placeholder node
   Node* version_node_ = nullptr; // version placeholder node

serving/processor/serving/model_config.cc

@@ -385,6 +385,12 @@ Status ModelConfigFactory::Create(const char* model_config, ModelConfig** config
       case embedding::StorageType::DRAM_SSDHASH:
         (*config)->storage_type = embedding::StorageType::DRAM_SSDHASH;
         (*config)->storage_path = json_config["ev_storage_path"].asString();
+        for (int i = 0; i < json_config["ev_storage_size"].size(); i++)
+          (*config)->storage_size.emplace_back(json_config["ev_storage_size"][i].asInt64());
+        if (json_config["ev_storage_size"].size() < 4) {
+          for (int i =  json_config["ev_storage_size"].size(); i < 4; i++)
+            (*config)->storage_size.emplace_back(1024*1024*1024);
+        }
         break;
       default:
         return Status(error::Code::INVALID_ARGUMENT,

serving/processor/serving/model_config.h

@@ -60,6 +60,7 @@ struct ModelConfig {
   // EmbeddingVariable Config
   embedding::StorageType storage_type = embedding::StorageType::INVALID;
   std::string storage_path;
+  std::vector<int64> storage_size;
 };
 
 class ModelConfigFactory {

serving/processor/serving/model_config_test.cc

@@ -236,6 +236,7 @@ const std::string oss_config = " \
     \"oss_access_id\" : \"test_id\", \
     \"ev_storage_type\" : 12, \
     \"ev_storage_path\" : \"123\", \
+    \"ev_storage_size\" : [1024, 1024], \
     \"oss_access_key\" : \"test_key\" \
   }";
 
@@ -244,6 +245,12 @@ const std::string oss_config = " \
       ModelConfigFactory::Create(oss_config.c_str(), &config).ok());
   EXPECT_EQ(12, config->storage_type);
   EXPECT_EQ("123", config->storage_path);
+  for (int i = 0; i < 2; i++) {
+    EXPECT_EQ(1024, config->storage_size[i]);
+  }
+  for (int i = 2; i < 4; i++) {
+    EXPECT_EQ(1024*1024*1024, config->storage_size[i]);
+  }
 }
 
 TEST_F(ModelConfigTest, ShouldFailureWhenConfigEmbeddingConfig) {
@@ -266,6 +273,7 @@ const std::string oss_config = " \
     \"oss_access_id\" : \"test_id\", \
     \"ev_storage_type\" : 14, \
     \"ev_storage_path\" : \"123\", \
+    \"ev_storage_size\" : [1024, 1024, 1024, 1024], \
     \"oss_access_key\" : \"test_key\" \
   }";
 

serving/processor/serving/model_instance.cc

@@ -277,6 +277,7 @@ Status LocalSessionInstance::Init(ModelConfig* config,
 
   option.st = config->storage_type;
   option.path = config->storage_path;
+  option.size = config->storage_size;
 
   optimizer_ = new SavedModelOptimizer(config->signature_name,
       &meta_graph_def_, option);