diff --git a/docs/Processor.md b/docs/Processor.md index 8ba5f886e37..36e0e9aef22 100644 --- a/docs/Processor.md +++ b/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] } ``` #### 模型路径配置 diff --git a/serving/processor/framework/graph_optimizer.cc b/serving/processor/framework/graph_optimizer.cc index 4f63340ce37..43f2bc230d6 100644 --- a/serving/processor/framework/graph_optimizer.cc +++ b/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& 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(); } diff --git a/serving/processor/framework/graph_optimizer.h b/serving/processor/framework/graph_optimizer.h index 3a99d657f52..89d3ac9f765 100644 --- a/serving/processor/framework/graph_optimizer.h +++ b/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 size; }; struct SrcInfo { @@ -263,7 +264,8 @@ class SavedModelOptimizer : public GraphOptimizer { std::unordered_map>& origin_import_nodes, std::vector& new_kv_import_nodes); - Status RewriteEmbeddingVariableAttr(embedding::StorageType st, std::string path); + Status RewriteEmbeddingVariableAttr(embedding::StorageType st, const std::string& path, + const std::vector& size); Node* storage_pointer_node_ = nullptr;// storage placeholder node Node* version_node_ = nullptr; // version placeholder node diff --git a/serving/processor/serving/model_config.cc b/serving/processor/serving/model_config.cc index 28a9a064280..96ad1d86907 100644 --- a/serving/processor/serving/model_config.cc +++ b/serving/processor/serving/model_config.cc @@ -379,6 +379,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, diff --git a/serving/processor/serving/model_config.h b/serving/processor/serving/model_config.h index 41e342e1545..86648156972 100644 --- a/serving/processor/serving/model_config.h +++ b/serving/processor/serving/model_config.h @@ -59,6 +59,7 @@ struct ModelConfig { // EmbeddingVariable Config embedding::StorageType storage_type = embedding::StorageType::INVALID; std::string storage_path; + std::vector storage_size; }; class ModelConfigFactory { diff --git a/serving/processor/serving/model_config_test.cc b/serving/processor/serving/model_config_test.cc index e7bc6c11408..2a7790918c3 100644 --- a/serving/processor/serving/model_config_test.cc +++ b/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\" \ }"; diff --git a/serving/processor/serving/model_instance.cc b/serving/processor/serving/model_instance.cc index 66280baebd3..d39a6be7e81 100644 --- a/serving/processor/serving/model_instance.cc +++ b/serving/processor/serving/model_instance.cc @@ -269,6 +269,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);