admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-30 07:25:37 +08:00
Code Issues Packages Projects Releases Wiki Activity
milvus/core/unittest/server/test_rpc.cpp
shengjun.li 290ff2e439 Support IVF_PQ on GPU with using metric_type IP (#2593) 
* Support IVF_PQ on GPU with using metric_type IP

Signed-off-by: shengjun.li <shengjun.li@zilliz.com>

* remove useless code

Signed-off-by: shengjun.li <shengjun.li@zilliz.com>
2020-06-19 10:42:23 +08:00

1584 lines
66 KiB
C++
Raw Blame History

// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.
#include <gtest/gtest.h>
#include <opentracing/mocktracer/tracer.h>
#include <boost/filesystem.hpp>
#include <thread>
#include "config/Config.h"
#include "server/Server.h"
#include "server/delivery/RequestHandler.h"
#include "server/delivery/RequestScheduler.h"
#include "server/delivery/request/BaseRequest.h"
#include "server/grpc_impl/GrpcRequestHandler.h"
#include "src/version.h"
#include "grpc/gen-milvus/milvus.grpc.pb.h"
#include "grpc/gen-status/status.pb.h"
#include "scheduler/ResourceFactory.h"
#include "scheduler/SchedInst.h"
#include "server/DBWrapper.h"
#include "server/grpc_impl/GrpcServer.h"
#include "utils/CommonUtil.h"
#include <fiu-control.h>
#include <fiu-local.h>
namespace {
static const char* COLLECTION_NAME = "test_grpc";
static constexpr int64_t COLLECTION_DIM = 256;
static constexpr int64_t INDEX_FILE_SIZE = 1024;
static constexpr int64_t VECTOR_COUNT = 1000;
static constexpr int64_t INSERT_LOOP = 10;
constexpr int64_t SECONDS_EACH_HOUR = 3600;
void
CopyRowRecord(::milvus::grpc::RowRecord* target, const std::vector<float>& src) {
auto vector_data = target->mutable_float_data();
vector_data->Resize(static_cast<int>(src.size()), 0.0);
memcpy(vector_data->mutable_data(), src.data(), src.size() * sizeof(float));
}
void
CopyBinRowRecord(::milvus::grpc::RowRecord* target, const std::vector<uint8_t>& src) {
auto vector_data = target->mutable_binary_data();
vector_data->resize(static_cast<int>(src.size()));
memcpy(vector_data->data(), src.data(), src.size());
}
void
SearchFunc(std::shared_ptr<milvus::server::grpc::GrpcRequestHandler> handler, ::grpc::ServerContext* context,
std::shared_ptr<::milvus::grpc::SearchParam> request,
std::shared_ptr<::milvus::grpc::TopKQueryResult> result) {
handler->Search(context, request.get(), result.get());
}
class RpcHandlerTest : public testing::Test {
protected:
void
SetUp() override {
auto res_mgr = milvus::scheduler::ResMgrInst::GetInstance();
res_mgr->Clear();
res_mgr->Add(milvus::scheduler::ResourceFactory::Create("disk", "DISK", 0, false));
res_mgr->Add(milvus::scheduler::ResourceFactory::Create("cpu", "CPU", 0));
res_mgr->Add(milvus::scheduler::ResourceFactory::Create("gtx1660", "GPU", 0));
auto default_conn = milvus::scheduler::Connection("IO", 500.0);
auto PCIE = milvus::scheduler::Connection("IO", 11000.0);
res_mgr->Connect("disk", "cpu", default_conn);
res_mgr->Connect("cpu", "gtx1660", PCIE);
res_mgr->Start();
milvus::scheduler::SchedInst::GetInstance()->Start();
milvus::scheduler::JobMgrInst::GetInstance()->Start();
milvus::engine::DBOptions opt;
milvus::server::Config::GetInstance().SetGeneralConfigMetaURI("sqlite://:@:/");
milvus::server::Config::GetInstance().SetDBConfigArchiveDiskThreshold("");
milvus::server::Config::GetInstance().SetDBConfigArchiveDaysThreshold("");
milvus::server::Config::GetInstance().SetStorageConfigPath("/tmp/milvus_test");
milvus::server::Config::GetInstance().SetCacheConfigCacheInsertData("");
milvus::server::Config::GetInstance().SetEngineConfigOmpThreadNum("");
milvus::server::Config::GetInstance().SetNetworkConfigBindPort("19531");
// serverConfig.SetValue(server::CONFIG_CLUSTER_MODE, "cluster");
// DBWrapper::GetInstance().GetInstance().StartService();
// DBWrapper::GetInstance().GetInstance().StopService();
//
// serverConfig.SetValue(server::CONFIG_CLUSTER_MODE, "read_only");
// DBWrapper::GetInstance().GetInstance().StartService();
// DBWrapper::GetInstance().GetInstance().StopService();
milvus::server::DBWrapper::GetInstance().StartService();
// initialize handler, create collection
handler = std::make_shared<milvus::server::grpc::GrpcRequestHandler>(opentracing::Tracer::Global());
dummy_context = std::make_shared<milvus::server::Context>("dummy_request_id");
opentracing::mocktracer::MockTracerOptions tracer_options;
auto mock_tracer =
std::shared_ptr<opentracing::Tracer>{new opentracing::mocktracer::MockTracer{std::move(tracer_options)}};
auto mock_span = mock_tracer->StartSpan("mock_span");
auto trace_context = std::make_shared<milvus::tracing::TraceContext>(mock_span);
dummy_context->SetTraceContext(trace_context);
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
::milvus::grpc::CollectionSchema request;
::milvus::grpc::Status status;
request.set_collection_name(COLLECTION_NAME);
request.set_dimension(COLLECTION_DIM);
request.set_index_file_size(INDEX_FILE_SIZE);
request.set_metric_type(1);
handler->SetContext(&context, dummy_context);
handler->random_id();
::grpc::Status grpc_status = handler->CreateCollection(&context, &request, &status);
}
void
TearDown() override {
milvus::server::DBWrapper::GetInstance().StopService();
milvus::scheduler::JobMgrInst::GetInstance()->Stop();
milvus::scheduler::ResMgrInst::GetInstance()->Stop();
milvus::scheduler::SchedInst::GetInstance()->Stop();
boost::filesystem::remove_all("/tmp/milvus_test");
}
protected:
std::shared_ptr<milvus::server::grpc::GrpcRequestHandler> handler;
std::shared_ptr<milvus::server::Context> dummy_context;
};
void
BuildVectors(int64_t from, int64_t to, std::vector<std::vector<float>>& vector_record_array) {
if (to <= from) {
return;
}
vector_record_array.clear();
for (int64_t k = from; k < to; k++) {
std::vector<float> record;
record.resize(COLLECTION_DIM);
for (int64_t i = 0; i < COLLECTION_DIM; i++) {
record[i] = (float)(i + k);
}
vector_record_array.emplace_back(record);
}
}
void
BuildBinVectors(int64_t from, int64_t to, std::vector<std::vector<uint8_t>>& vector_record_array) {
if (to <= from) {
return;
}
vector_record_array.clear();
for (int64_t k = from; k < to; k++) {
std::vector<uint8_t> record;
record.resize(COLLECTION_DIM / 8);
for (int64_t i = 0; i < COLLECTION_DIM / 8; i++) {
record[i] = (i + k) % 256;
}
vector_record_array.emplace_back(record);
}
}
std::string
CurrentTmDate(int64_t offset_day = 0) {
time_t tt;
time(&tt);
tt = tt + 8 * SECONDS_EACH_HOUR;
tt = tt + 24 * SECONDS_EACH_HOUR * offset_day;
tm t;
gmtime_r(&tt, &t);
std::string str =
std::to_string(t.tm_year + 1900) + "-" + std::to_string(t.tm_mon + 1) + "-" + std::to_string(t.tm_mday);
return str;
}
} // namespace
TEST_F(RpcHandlerTest, HAS_COLLECTION_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
::milvus::grpc::CollectionName request;
::milvus::grpc::BoolReply reply;
::grpc::Status status = handler->HasCollection(&context, &request, &reply);
request.set_collection_name(COLLECTION_NAME);
status = handler->HasCollection(&context, &request, &reply);
ASSERT_TRUE(status.error_code() == ::grpc::Status::OK.error_code());
int error_code = reply.status().error_code();
ASSERT_EQ(error_code, ::milvus::grpc::ErrorCode::SUCCESS);
fiu_init(0);
fiu_enable("HasCollectionRequest.OnExecute.throw_std_exception", 1, NULL, 0);
handler->HasCollection(&context, &request, &reply);
ASSERT_NE(reply.status().error_code(), ::milvus::grpc::ErrorCode::SUCCESS);
fiu_disable("HasCollectionRequest.OnExecute.throw_std_exception");
}
TEST_F(RpcHandlerTest, INDEX_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
::milvus::grpc::IndexParam request;
::milvus::grpc::Status response;
::grpc::Status grpc_status = handler->CreateIndex(&context, &request, &response);
request.set_collection_name("test1");
handler->CreateIndex(&context, &request, &response);
request.set_collection_name(COLLECTION_NAME);
handler->CreateIndex(&context, &request, &response);
request.set_index_type(1);
handler->CreateIndex(&context, &request, &response);
::milvus::grpc::KeyValuePair* kv = request.add_extra_params();
kv->set_key("param");
kv->set_value("{ \"nlist\": 16384 }");
grpc_status = handler->CreateIndex(&context, &request, &response);
ASSERT_EQ(grpc_status.error_code(), ::grpc::Status::OK.error_code());
int error_code = response.error_code();
// ASSERT_EQ(error_code, ::milvus::grpc::ErrorCode::SUCCESS);
fiu_init(0);
fiu_enable("CreateIndexRequest.OnExecute.not_has_collection", 1, NULL, 0);
grpc_status = handler->CreateIndex(&context, &request, &response);
ASSERT_TRUE(grpc_status.ok());
fiu_disable("CreateIndexRequest.OnExecute.not_has_collection");
fiu_enable("CreateIndexRequest.OnExecute.throw_std.exception", 1, NULL, 0);
grpc_status = handler->CreateIndex(&context, &request, &response);
ASSERT_TRUE(grpc_status.ok());
fiu_disable("CreateIndexRequest.OnExecute.throw_std.exception");
fiu_enable("CreateIndexRequest.OnExecute.create_index_fail", 1, NULL, 0);
grpc_status = handler->CreateIndex(&context, &request, &response);
ASSERT_TRUE(grpc_status.ok());
fiu_disable("CreateIndexRequest.OnExecute.create_index_fail");
::milvus::grpc::CollectionName collection_name;
::milvus::grpc::IndexParam index_param;
handler->DescribeIndex(&context, &collection_name, &index_param);
collection_name.set_collection_name("test4");
handler->DescribeIndex(&context, &collection_name, &index_param);
collection_name.set_collection_name(COLLECTION_NAME);
handler->DescribeIndex(&context, &collection_name, &index_param);
fiu_init(0);
fiu_enable("DescribeIndexRequest.OnExecute.throw_std_exception", 1, NULL, 0);
handler->DescribeIndex(&context, &collection_name, &index_param);
fiu_disable("DescribeIndexRequest.OnExecute.throw_std_exception");
::milvus::grpc::Status status;
collection_name.Clear();
handler->DropIndex(&context, &collection_name, &status);
collection_name.set_collection_name("test5");
handler->DropIndex(&context, &collection_name, &status);
collection_name.set_collection_name(COLLECTION_NAME);
fiu_init(0);
fiu_enable("DropIndexRequest.OnExecute.collection_not_exist", 1, NULL, 0);
handler->DropIndex(&context, &collection_name, &status);
fiu_disable("DropIndexRequest.OnExecute.collection_not_exist");
fiu_enable("DropIndexRequest.OnExecute.drop_index_fail", 1, NULL, 0);
handler->DropIndex(&context, &collection_name, &status);
fiu_disable("DropIndexRequest.OnExecute.drop_index_fail");
fiu_enable("DropIndexRequest.OnExecute.throw_std_exception", 1, NULL, 0);
handler->DropIndex(&context, &collection_name, &status);
fiu_disable("DropIndexRequest.OnExecute.throw_std_exception");
handler->DropIndex(&context, &collection_name, &status);
}
TEST_F(RpcHandlerTest, INSERT_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
::milvus::grpc::InsertParam request;
::milvus::grpc::Status response;
request.set_collection_name(COLLECTION_NAME);
std::vector<std::vector<float>> record_array;
BuildVectors(0, VECTOR_COUNT, record_array);
::milvus::grpc::VectorIds vector_ids;
for (auto& record : record_array) {
::milvus::grpc::RowRecord* grpc_record = request.add_row_record_array();
CopyRowRecord(grpc_record, record);
}
handler->Insert(&context, &request, &vector_ids);
ASSERT_EQ(vector_ids.vector_id_array_size(), VECTOR_COUNT);
fiu_init(0);
fiu_enable("InsertRequest.OnExecute.id_array_error", 1, NULL, 0);
handler->Insert(&context, &request, &vector_ids);
ASSERT_NE(vector_ids.vector_id_array_size(), VECTOR_COUNT);
fiu_disable("InsertRequest.OnExecute.id_array_error");
fiu_enable("InsertRequest.OnExecute.db_not_found", 1, NULL, 0);
handler->Insert(&context, &request, &vector_ids);
ASSERT_NE(vector_ids.vector_id_array_size(), VECTOR_COUNT);
fiu_disable("InsertRequest.OnExecute.db_not_found");
fiu_enable("InsertRequest.OnExecute.describe_collection_fail", 1, NULL, 0);
handler->Insert(&context, &request, &vector_ids);
ASSERT_NE(vector_ids.vector_id_array_size(), VECTOR_COUNT);
fiu_disable("InsertRequest.OnExecute.describe_collection_fail");
fiu_enable("InsertRequest.OnExecute.illegal_vector_id", 1, NULL, 0);
handler->Insert(&context, &request, &vector_ids);
ASSERT_NE(vector_ids.vector_id_array_size(), VECTOR_COUNT);
fiu_disable("InsertRequest.OnExecute.illegal_vector_id");
fiu_enable("InsertRequest.OnExecute.illegal_vector_id2", 1, NULL, 0);
handler->Insert(&context, &request, &vector_ids);
ASSERT_NE(vector_ids.vector_id_array_size(), VECTOR_COUNT);
fiu_disable("InsertRequest.OnExecute.illegal_vector_id2");
fiu_enable("InsertRequest.OnExecute.throw_std_exception", 1, NULL, 0);
handler->Insert(&context, &request, &vector_ids);
ASSERT_NE(vector_ids.vector_id_array_size(), VECTOR_COUNT);
fiu_disable("InsertRequest.OnExecute.throw_std_exception");
fiu_enable("InsertRequest.OnExecute.insert_fail", 1, NULL, 0);
handler->Insert(&context, &request, &vector_ids);
fiu_disable("InsertRequest.OnExecute.insert_fail");
fiu_enable("InsertRequest.OnExecute.invalid_ids_size", 1, NULL, 0);
handler->Insert(&context, &request, &vector_ids);
fiu_disable("InsertRequest.OnExecute.invalid_ids_size");
// insert vectors with wrong dim
std::vector<float> record_wrong_dim(COLLECTION_DIM - 1, 0.5f);
::milvus::grpc::RowRecord* grpc_record = request.add_row_record_array();
CopyRowRecord(grpc_record, record_wrong_dim);
handler->Insert(&context, &request, &vector_ids);
ASSERT_EQ(vector_ids.status().error_code(), ::milvus::grpc::ILLEGAL_ROWRECORD);
}
TEST_F(RpcHandlerTest, SEARCH_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
::milvus::grpc::SearchParam request;
::milvus::grpc::TopKQueryResult response;
// test null input
handler->Search(&context, nullptr, &response);
// test invalid collection name
handler->Search(&context, &request, &response);
// test collection not exist
request.set_collection_name("test3");
handler->Search(&context, &request, &response);
// test invalid topk
request.set_collection_name(COLLECTION_NAME);
handler->Search(&context, &request, &response);
// test invalid nprobe
request.set_topk(10);
handler->Search(&context, &request, &response);
// test empty query record array
milvus::grpc::KeyValuePair* kv = request.add_extra_params();
kv->set_key(milvus::server::grpc::EXTRA_PARAM_KEY);
kv->set_value("{ \"nprobe\": 32 }");
handler->Search(&context, &request, &response);
ASSERT_EQ(response.ids_size(), 0UL);
std::vector<std::vector<float>> record_array;
BuildVectors(0, VECTOR_COUNT, record_array);
::milvus::grpc::InsertParam insert_param;
for (auto& record : record_array) {
::milvus::grpc::RowRecord* grpc_record = insert_param.add_row_record_array();
CopyRowRecord(grpc_record, record);
}
// insert vectors
insert_param.set_collection_name(COLLECTION_NAME);
::milvus::grpc::VectorIds vector_ids;
handler->Insert(&context, &insert_param, &vector_ids);
// flush
::milvus::grpc::Status grpc_status;
::milvus::grpc::FlushParam flush_param;
flush_param.add_collection_name_array(COLLECTION_NAME);
handler->Flush(&context, &flush_param, &grpc_status);
// search
BuildVectors(0, 10, record_array);
for (auto& record : record_array) {
::milvus::grpc::RowRecord* row_record = request.add_query_record_array();
CopyRowRecord(row_record, record);
}
handler->Search(&context, &request, &response);
ASSERT_NE(response.ids_size(), 0UL);
// wrong file id
::milvus::grpc::SearchInFilesParam search_in_files_param;
std::string* file_id = search_in_files_param.add_file_id_array();
*file_id = "test_tbl";
handler->SearchInFiles(&context, &search_in_files_param, &response);
ASSERT_EQ(response.ids_size(), 0UL);
}
TEST_F(RpcHandlerTest, COMBINE_SEARCH_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
// create collection
std::string collection_name = "search_combines";
::milvus::grpc::CollectionSchema collection_schema;
collection_schema.set_collection_name(collection_name);
collection_schema.set_dimension(COLLECTION_DIM);
collection_schema.set_index_file_size(INDEX_FILE_SIZE);
collection_schema.set_metric_type(1); // L2 metric
::milvus::grpc::Status status;
handler->CreateCollection(&context, &collection_schema, &status);
ASSERT_EQ(status.error_code(), 0) << status.reason();
// insert vectors
std::vector<std::vector<float>> record_array;
BuildVectors(0, VECTOR_COUNT, record_array);
::milvus::grpc::InsertParam insert_param;
int64_t vec_id = 0;
for (auto& record : record_array) {
::milvus::grpc::RowRecord* grpc_record = insert_param.add_row_record_array();
CopyRowRecord(grpc_record, record);
insert_param.add_row_id_array(++vec_id);
}
insert_param.set_collection_name(collection_name);
::milvus::grpc::VectorIds vector_ids;
handler->Insert(&context, &insert_param, &vector_ids);
// flush
::milvus::grpc::Status grpc_status;
::milvus::grpc::FlushParam flush_param;
flush_param.add_collection_name_array(collection_name);
handler->Flush(&context, &flush_param, &grpc_status);
// multi thread search requests will be combined
int QUERY_COUNT = 10;
int64_t NQ = 2;
int64_t TOPK = 5;
using RequestPtr = std::shared_ptr<::milvus::grpc::SearchParam>;
std::vector<RequestPtr> request_array;
for (int i = 0; i < QUERY_COUNT; i++) {
RequestPtr request = std::make_shared<::milvus::grpc::SearchParam>();
request->set_collection_name(collection_name);
request->set_topk(TOPK);
milvus::grpc::KeyValuePair* kv = request->add_extra_params();
kv->set_key(milvus::server::grpc::EXTRA_PARAM_KEY);
kv->set_value("{}");
BuildVectors(i * NQ, (i + 1) * NQ, record_array);
for (auto& record : record_array) {
::milvus::grpc::RowRecord* row_record = request->add_query_record_array();
CopyRowRecord(row_record, record);
}
request_array.emplace_back(request);
}
using ResultPtr = std::shared_ptr<::milvus::grpc::TopKQueryResult>;
std::vector<ResultPtr> result_array;
using ThreadPtr = std::shared_ptr<std::thread>;
std::vector<ThreadPtr> thread_list;
for (int i = 0; i < QUERY_COUNT; i++) {
ResultPtr result_ptr = std::make_shared<::milvus::grpc::TopKQueryResult>();
result_array.push_back(result_ptr);
ThreadPtr thread = std::make_shared<std::thread>(SearchFunc, handler, &context, request_array[i], result_ptr);
thread_list.emplace_back(thread);
std::this_thread::sleep_for(std::chrono::milliseconds(5));
}
// wait search finish
for (auto& iter : thread_list) {
iter->join();
}
// check result
int64_t index = 0;
for (auto& result_ptr : result_array) {
ASSERT_NE(result_ptr->ids_size(), 0);
ASSERT_NE(result_ptr->row_num(), 0);
std::string msg = "Result no." + std::to_string(index) + ": \n";
for (int64_t i = 0; i < NQ; i++) {
for (int64_t k = 0; k < TOPK; k++) {
msg += "[";
msg += std::to_string(result_ptr->ids(i * TOPK + k));
msg += ", ";
msg += std::to_string(result_ptr->distances(i * TOPK + k));
msg += "]";
msg += ", ";
}
msg += "\n";
ASSERT_NE(result_ptr->ids(i * TOPK), 0);
ASSERT_LT(result_ptr->distances(i * TOPK), 0.00001);
}
std::cout << msg << std::endl;
index++;
}
}
TEST_F(RpcHandlerTest, COMBINE_SEARCH_BINARY_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
// create collection
std::string collection_name = "combine_bin";
::milvus::grpc::CollectionSchema collection_schema;
collection_schema.set_collection_name(collection_name);
collection_schema.set_dimension(COLLECTION_DIM);
collection_schema.set_index_file_size(INDEX_FILE_SIZE);
collection_schema.set_metric_type(5); // tanimoto metric
::milvus::grpc::Status status;
handler->CreateCollection(&context, &collection_schema, &status);
ASSERT_EQ(status.error_code(), 0);
// insert vectors
std::vector<std::vector<uint8_t>> record_array;
BuildBinVectors(0, VECTOR_COUNT, record_array);
::milvus::grpc::InsertParam insert_param;
int64_t vec_id = 0;
for (auto& record : record_array) {
::milvus::grpc::RowRecord* grpc_record = insert_param.add_row_record_array();
CopyBinRowRecord(grpc_record, record);
insert_param.add_row_id_array(++vec_id);
}
insert_param.set_collection_name(collection_name);
::milvus::grpc::VectorIds vector_ids;
handler->Insert(&context, &insert_param, &vector_ids);
// flush
::milvus::grpc::Status grpc_status;
::milvus::grpc::FlushParam flush_param;
flush_param.add_collection_name_array(collection_name);
handler->Flush(&context, &flush_param, &grpc_status);
// multi thread search requests will be combined
int QUERY_COUNT = 10;
int64_t NQ = 2;
int64_t TOPK = 5;
using RequestPtr = std::shared_ptr<::milvus::grpc::SearchParam>;
std::vector<RequestPtr> request_array;
for (int i = 0; i < QUERY_COUNT; i++) {
RequestPtr request = std::make_shared<::milvus::grpc::SearchParam>();
request->set_collection_name(collection_name);
request->set_topk(TOPK);
milvus::grpc::KeyValuePair* kv = request->add_extra_params();
kv->set_key(milvus::server::grpc::EXTRA_PARAM_KEY);
kv->set_value("{}");
BuildBinVectors(i * NQ, (i + 1) * NQ, record_array);
for (auto& record : record_array) {
::milvus::grpc::RowRecord* row_record = request->add_query_record_array();
CopyBinRowRecord(row_record, record);
}
request_array.emplace_back(request);
}
using ResultPtr = std::shared_ptr<::milvus::grpc::TopKQueryResult>;
std::vector<ResultPtr> result_array;
using ThreadPtr = std::shared_ptr<std::thread>;
std::vector<ThreadPtr> thread_list;
for (int i = 0; i < QUERY_COUNT; i++) {
ResultPtr result_ptr = std::make_shared<::milvus::grpc::TopKQueryResult>();
result_array.push_back(result_ptr);
ThreadPtr thread = std::make_shared<std::thread>(SearchFunc, handler, &context, request_array[i], result_ptr);
thread_list.emplace_back(thread);
std::this_thread::sleep_for(std::chrono::milliseconds(5));
}
// wait search finish
for (auto& iter : thread_list) {
iter->join();
}
// check result
int64_t index = 0;
for (auto& result_ptr : result_array) {
ASSERT_NE(result_ptr->ids_size(), 0);
ASSERT_NE(result_ptr->row_num(), 0);
std::string msg = "Result no." + std::to_string(++index) + ": \n";
for (int64_t i = 0; i < NQ; i++) {
for (int64_t k = 0; k < TOPK; k++) {
msg += "[";
msg += std::to_string(result_ptr->ids(i * TOPK + k));
msg += ", ";
msg += std::to_string(result_ptr->distances(i * TOPK + k));
msg += "]";
msg += ", ";
}
msg += "\n";
ASSERT_NE(result_ptr->ids(i * TOPK), 0);
ASSERT_LT(result_ptr->distances(i * TOPK), 0.00001);
}
std::cout << msg << std::endl;
}
}
TEST_F(RpcHandlerTest, TABLES_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
::milvus::grpc::Status response;
std::string collection_name = "tbl";
{
::milvus::grpc::CollectionSchema collection_schema;
// create collection test
// test null input
handler->CreateCollection(&context, nullptr, &response);
// test invalid collection name
handler->CreateCollection(&context, &collection_schema, &response);
// test invalid collection dimension
collection_schema.set_collection_name(collection_name);
handler->CreateCollection(&context, &collection_schema, &response);
// test invalid index file size
collection_schema.set_dimension(COLLECTION_DIM);
// handler->CreateCollection(&context, &collection_schema, &response);
// test invalid index metric type
collection_schema.set_index_file_size(INDEX_FILE_SIZE);
handler->CreateCollection(&context, &collection_schema, &response);
// test collection already exist
collection_schema.set_metric_type(1);
handler->CreateCollection(&context, &collection_schema, &response);
}
// describe collection test
// test invalid collection name
::milvus::grpc::CollectionName grpc_collection_name;
{
::milvus::grpc::CollectionSchema collection_schema;
handler->DescribeCollection(&context, &grpc_collection_name, &collection_schema);
grpc_collection_name.set_collection_name(COLLECTION_NAME);
::grpc::Status status = handler->DescribeCollection(&context, &grpc_collection_name, &collection_schema);
ASSERT_EQ(status.error_code(), ::grpc::Status::OK.error_code());
fiu_init(0);
fiu_enable("DescribeCollectionRequest.OnExecute.describe_collection_fail", 1, NULL, 0);
handler->DescribeCollection(&context, &grpc_collection_name, &collection_schema);
fiu_disable("DescribeCollectionRequest.OnExecute.describe_collection_fail");
fiu_enable("DescribeCollectionRequest.OnExecute.throw_std_exception", 1, NULL, 0);
handler->DescribeCollection(&context, &grpc_collection_name, &collection_schema);
fiu_disable("DescribeCollectionRequest.OnExecute.throw_std_exception");
}
::milvus::grpc::InsertParam request;
std::vector<std::vector<float>> record_array;
BuildVectors(0, VECTOR_COUNT, record_array);
::milvus::grpc::VectorIds vector_ids;
for (int64_t i = 0; i < VECTOR_COUNT; i++) {
vector_ids.add_vector_id_array(i);
}
// Insert vectors
// test invalid collection name
handler->Insert(&context, &request, &vector_ids);
request.set_collection_name(collection_name);
// test empty row record
handler->Insert(&context, &request, &vector_ids);
for (auto& record : record_array) {
::milvus::grpc::RowRecord* grpc_record = request.add_row_record_array();
CopyRowRecord(grpc_record, record);
}
// test vector_id size not equal to row record size
vector_ids.clear_vector_id_array();
vector_ids.add_vector_id_array(1);
handler->Insert(&context, &request, &vector_ids);
// normally test
vector_ids.clear_vector_id_array();
handler->Insert(&context, &request, &vector_ids);
request.clear_row_record_array();
vector_ids.clear_vector_id_array();
for (uint64_t i = 0; i < 10; ++i) {
::milvus::grpc::RowRecord* grpc_record = request.add_row_record_array();
CopyRowRecord(grpc_record, record_array[i]);
}
handler->Insert(&context, &request, &vector_ids);
// show collections
{
::milvus::grpc::Command cmd;
::milvus::grpc::CollectionNameList collection_name_list;
auto status = handler->ShowCollections(&context, &cmd, &collection_name_list);
ASSERT_EQ(status.error_code(), ::grpc::Status::OK.error_code());
// show collection info
::milvus::grpc::CollectionInfo collection_info;
status = handler->ShowCollectionInfo(&context, &grpc_collection_name, &collection_info);
ASSERT_EQ(status.error_code(), ::grpc::Status::OK.error_code());
fiu_init(0);
fiu_enable("ShowCollectionsRequest.OnExecute.show_collections_fail", 1, NULL, 0);
handler->ShowCollections(&context, &cmd, &collection_name_list);
fiu_disable("ShowCollectionsRequest.OnExecute.show_collections_fail");
}
// Count Collection
{
::milvus::grpc::CollectionRowCount count;
grpc_collection_name.Clear();
auto status = handler->CountCollection(&context, &grpc_collection_name, &count);
grpc_collection_name.set_collection_name(collection_name);
status = handler->CountCollection(&context, &grpc_collection_name, &count);
ASSERT_EQ(status.error_code(), ::grpc::Status::OK.error_code());
// ASSERT_EQ(count.collection_row_count(), vector_ids.vector_id_array_size());
fiu_init(0);
fiu_enable("CountCollectionRequest.OnExecute.db_not_found", 1, NULL, 0);
status = handler->CountCollection(&context, &grpc_collection_name, &count);
fiu_disable("CountCollectionRequest.OnExecute.db_not_found");
fiu_enable("CountCollectionRequest.OnExecute.status_error", 1, NULL, 0);
status = handler->CountCollection(&context, &grpc_collection_name, &count);
fiu_disable("CountCollectionRequest.OnExecute.status_error");
fiu_enable("CountCollectionRequest.OnExecute.throw_std_exception", 1, NULL, 0);
status = handler->CountCollection(&context, &grpc_collection_name, &count);
fiu_disable("CountCollectionRequest.OnExecute.throw_std_exception");
}
// Preload Collection
{
::milvus::grpc::CollectionSchema collection_schema;
grpc_collection_name.Clear();
auto status = handler->PreloadCollection(&context, &grpc_collection_name, &response);
grpc_collection_name.set_collection_name(COLLECTION_NAME);
status = handler->PreloadCollection(&context, &grpc_collection_name, &response);
ASSERT_EQ(status.error_code(), ::grpc::Status::OK.error_code());
fiu_enable("PreloadCollectionRequest.OnExecute.preload_collection_fail", 1, NULL, 0);
handler->PreloadCollection(&context, &grpc_collection_name, &response);
fiu_disable("PreloadCollectionRequest.OnExecute.preload_collection_fail");
fiu_enable("PreloadCollectionRequest.OnExecute.throw_std_exception", 1, NULL, 0);
handler->PreloadCollection(&context, &grpc_collection_name, &response);
fiu_disable("PreloadCollectionRequest.OnExecute.throw_std_exception");
fiu_init(0);
fiu_enable("CreateCollectionRequest.OnExecute.invalid_index_file_size", 1, NULL, 0);
collection_schema.set_collection_name(collection_name);
handler->CreateCollection(&context, &collection_schema, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("CreateCollectionRequest.OnExecute.invalid_index_file_size");
fiu_enable("CreateCollectionRequest.OnExecute.db_already_exist", 1, NULL, 0);
collection_schema.set_collection_name(collection_name);
handler->CreateCollection(&context, &collection_schema, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("CreateCollectionRequest.OnExecute.db_already_exist");
fiu_enable("CreateCollectionRequest.OnExecute.create_collection_fail", 1, NULL, 0);
collection_schema.set_collection_name(collection_name);
handler->CreateCollection(&context, &collection_schema, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("CreateCollectionRequest.OnExecute.create_collection_fail");
fiu_enable("CreateCollectionRequest.OnExecute.throw_std_exception", 1, NULL, 0);
collection_schema.set_collection_name(collection_name);
handler->CreateCollection(&context, &collection_schema, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("CreateCollectionRequest.OnExecute.throw_std_exception");
}
// Drop collection
{
::milvus::grpc::CollectionSchema collection_schema;
collection_schema.set_dimension(128);
collection_schema.set_index_file_size(1024);
collection_schema.set_metric_type(1);
grpc_collection_name.set_collection_name("");
// test invalid collection name
::grpc::Status grpc_status = handler->DropCollection(&context, &grpc_collection_name, &response);
grpc_collection_name.set_collection_name(collection_name);
fiu_enable("DropCollectionRequest.OnExecute.db_not_found", 1, NULL, 0);
handler->DropCollection(&context, &grpc_collection_name, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("DropCollectionRequest.OnExecute.db_not_found");
fiu_enable("DropCollectionRequest.OnExecute.describe_collection_fail", 1, NULL, 0);
handler->DropCollection(&context, &grpc_collection_name, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("DropCollectionRequest.OnExecute.describe_collection_fail");
fiu_enable("DropCollectionRequest.OnExecute.throw_std_exception", 1, NULL, 0);
handler->DropCollection(&context, &grpc_collection_name, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("DropCollectionRequest.OnExecute.throw_std_exception");
grpc_status = handler->DropCollection(&context, &grpc_collection_name, &response);
ASSERT_EQ(grpc_status.error_code(), ::grpc::Status::OK.error_code());
int error_code = response.error_code();
ASSERT_EQ(error_code, ::milvus::grpc::ErrorCode::SUCCESS);
collection_schema.set_collection_name(grpc_collection_name.collection_name());
handler->DropCollection(&context, &grpc_collection_name, &response);
sleep(1);
handler->CreateCollection(&context, &collection_schema, &response);
ASSERT_EQ(response.error_code(), ::grpc::Status::OK.error_code());
fiu_enable("DropCollectionRequest.OnExecute.drop_collection_fail", 1, NULL, 0);
handler->DropCollection(&context, &grpc_collection_name, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("DropCollectionRequest.OnExecute.drop_collection_fail");
handler->DropCollection(&context, &grpc_collection_name, &response);
}
}
TEST_F(RpcHandlerTest, GET_VECTOR_BY_ID_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
::milvus::grpc::Status response;
::milvus::grpc::InsertParam request;
request.set_collection_name(COLLECTION_NAME);
std::vector<std::vector<float>> record_array;
BuildVectors(0, VECTOR_COUNT, record_array);
::milvus::grpc::VectorIds vector_ids;
for (auto& record : record_array) {
::milvus::grpc::RowRecord* grpc_record = request.add_row_record_array();
CopyRowRecord(grpc_record, record);
}
handler->Insert(&context, &request, &vector_ids);
ASSERT_EQ(vector_ids.vector_id_array_size(), VECTOR_COUNT);
::milvus::grpc::FlushParam flush_param;
flush_param.add_collection_name_array(COLLECTION_NAME);
handler->Flush(&context, &flush_param, &response);
::milvus::grpc::VectorsIdentity vectors_identity;
vectors_identity.set_collection_name(COLLECTION_NAME);
for (size_t i = 0; i < 10; i++) {
vectors_identity.mutable_id_array()->Add(vector_ids.vector_id_array(i));
}
::milvus::grpc::VectorsData vectors_data;
handler->GetVectorsByID(&context, &vectors_identity, &vectors_data);
ASSERT_EQ(10, vectors_data.vectors_data_size());
}
TEST_F(RpcHandlerTest, GET_BIN_VECTORS_BY_IDS_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
::milvus::grpc::CollectionSchema collection_schema;
::milvus::grpc::Status response;
std::string str_collection_name = "tbl_get_bin_vector_by_ids";
collection_schema.set_collection_name(str_collection_name);
collection_schema.set_dimension(COLLECTION_DIM);
collection_schema.set_index_file_size(INDEX_FILE_SIZE);
collection_schema.set_metric_type(3);
handler->CreateCollection(&context, &collection_schema, &response);
::milvus::grpc::InsertParam request;
request.set_collection_name(str_collection_name);
std::vector<std::vector<uint8_t>> bin_record_array;
BuildBinVectors(0, VECTOR_COUNT, bin_record_array);
::milvus::grpc::VectorIds vector_ids;
for (auto& record : bin_record_array) {
::milvus::grpc::RowRecord* grpc_record = request.add_row_record_array();
CopyBinRowRecord(grpc_record, record);
}
handler->Insert(&context, &request, &vector_ids);
ASSERT_EQ(vector_ids.vector_id_array_size(), VECTOR_COUNT);
::milvus::grpc::FlushParam flush_param;
flush_param.add_collection_name_array(str_collection_name);
handler->Flush(&context, &flush_param, &response);
::milvus::grpc::VectorsIdentity vectors_identity;
vectors_identity.set_collection_name(str_collection_name);
for (size_t i = 0; i < 10; i++) {
vectors_identity.mutable_id_array()->Add(vector_ids.vector_id_array(i));
}
::milvus::grpc::VectorsData vectors_data;
handler->GetVectorsByID(&context, &vectors_identity, &vectors_data);
ASSERT_EQ(10, vectors_data.vectors_data_size());
}
TEST_F(RpcHandlerTest, GET_VECTOR_IDS_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
::milvus::grpc::CollectionSchema collection_schema;
::milvus::grpc::Status response;
std::string str_collection_name = "tbl_get_vector_ids";
collection_schema.set_collection_name(str_collection_name);
collection_schema.set_dimension(COLLECTION_DIM);
collection_schema.set_index_file_size(INDEX_FILE_SIZE);
collection_schema.set_metric_type(1);
handler->CreateCollection(&context, &collection_schema, &response);
::milvus::grpc::InsertParam request;
request.set_collection_name(str_collection_name);
std::vector<std::vector<float>> record_array;
BuildVectors(0, VECTOR_COUNT, record_array);
::milvus::grpc::VectorIds vector_ids;
for (auto& record : record_array) {
::milvus::grpc::RowRecord* grpc_record = request.add_row_record_array();
CopyRowRecord(grpc_record, record);
}
handler->Insert(&context, &request, &vector_ids);
ASSERT_EQ(vector_ids.vector_id_array_size(), VECTOR_COUNT);
::milvus::grpc::FlushParam flush_param;
flush_param.add_collection_name_array(str_collection_name);
handler->Flush(&context, &flush_param, &response);
::milvus::grpc::CollectionName collection_name;
collection_name.set_collection_name(str_collection_name);
::milvus::grpc::CollectionInfo collection_info;
handler->ShowCollectionInfo(&context, &collection_name, &collection_info);
std::string json_info = collection_info.json_info();
auto info_json = nlohmann::json::parse(json_info);
std::string segment0 = info_json["partitions"][0]["segments"][0]["name"];
::milvus::grpc::GetVectorIDsParam vector_ids_param;
vector_ids_param.set_collection_name(str_collection_name);
vector_ids_param.set_segment_name(segment0);
::milvus::grpc::VectorIds vectors_ids;
handler->GetVectorIDs(&context, &vector_ids_param, &vectors_ids);
ASSERT_EQ(0, vectors_ids.status().error_code());
}
TEST_F(RpcHandlerTest, PARTITION_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
::milvus::grpc::CollectionSchema collection_schema;
::milvus::grpc::Status response;
std::string str_collection_name = "tbl_partition";
collection_schema.set_collection_name(str_collection_name);
collection_schema.set_dimension(COLLECTION_DIM);
collection_schema.set_index_file_size(INDEX_FILE_SIZE);
collection_schema.set_metric_type(1);
handler->CreateCollection(&context, &collection_schema, &response);
::milvus::grpc::PartitionParam partition_param;
partition_param.set_collection_name(str_collection_name);
std::string partition_tag = "0";
partition_param.set_tag(partition_tag);
handler->CreatePartition(&context, &partition_param, &response);
ASSERT_EQ(response.error_code(), ::grpc::Status::OK.error_code());
::milvus::grpc::CollectionName collection_name;
collection_name.set_collection_name(str_collection_name);
::milvus::grpc::PartitionList partition_list;
handler->ShowPartitions(&context, &collection_name, &partition_list);
ASSERT_EQ(response.error_code(), ::grpc::Status::OK.error_code());
ASSERT_EQ(partition_list.partition_tag_array_size(), 2);
fiu_init(0);
fiu_enable("ShowPartitionsRequest.OnExecute.invalid_collection_name", 1, NULL, 0);
handler->ShowPartitions(&context, &collection_name, &partition_list);
fiu_disable("ShowPartitionsRequest.OnExecute.invalid_collection_name");
fiu_enable("ShowPartitionsRequest.OnExecute.show_partition_fail", 1, NULL, 0);
handler->ShowPartitions(&context, &collection_name, &partition_list);
fiu_disable("ShowPartitionsRequest.OnExecute.show_partition_fail");
fiu_init(0);
fiu_enable("CreatePartitionRequest.OnExecute.invalid_collection_name", 1, NULL, 0);
handler->CreatePartition(&context, &partition_param, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("CreatePartitionRequest.OnExecute.invalid_collection_name");
fiu_enable("CreatePartitionRequest.OnExecute.invalid_partition_name", 1, NULL, 0);
handler->CreatePartition(&context, &partition_param, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("CreatePartitionRequest.OnExecute.invalid_partition_name");
fiu_enable("CreatePartitionRequest.OnExecute.invalid_partition_tags", 1, NULL, 0);
handler->CreatePartition(&context, &partition_param, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("CreatePartitionRequest.OnExecute.invalid_partition_tags");
fiu_enable("CreatePartitionRequest.OnExecute.db_already_exist", 1, NULL, 0);
handler->CreatePartition(&context, &partition_param, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("CreatePartitionRequest.OnExecute.db_already_exist");
fiu_enable("CreatePartitionRequest.OnExecute.create_partition_fail", 1, NULL, 0);
handler->CreatePartition(&context, &partition_param, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("CreatePartitionRequest.OnExecute.create_partition_fail");
fiu_enable("CreatePartitionRequest.OnExecute.throw_std_exception", 1, NULL, 0);
handler->CreatePartition(&context, &partition_param, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("CreatePartitionRequest.OnExecute.throw_std_exception");
// test has partition
::milvus::grpc::PartitionParam has_partition_param;
has_partition_param.set_collection_name(str_collection_name);
has_partition_param.set_tag(partition_tag);
::milvus::grpc::BoolReply reply;
handler->HasPartition(&context, &has_partition_param, &reply);
ASSERT_EQ(0, reply.status().error_code());
ASSERT_EQ(true, reply.bool_reply());
::milvus::grpc::PartitionParam partition_parm;
partition_parm.set_collection_name(str_collection_name);
partition_parm.set_tag(partition_tag);
fiu_enable("DropPartitionRequest.OnExecute.invalid_collection_name", 1, NULL, 0);
handler->DropPartition(&context, &partition_parm, &response);
ASSERT_NE(response.error_code(), ::grpc::Status::OK.error_code());
fiu_disable("DropPartitionRequest.OnExecute.invalid_collection_name");
handler->DropPartition(&context, &partition_parm, &response);
ASSERT_EQ(response.error_code(), ::grpc::Status::OK.error_code());
}
TEST_F(RpcHandlerTest, RELOAD_SEGMENTS_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
::milvus::grpc::CollectionSchema collection_schema;
::milvus::grpc::Status response;
std::string str_collection_name = "tbl_reload_segments";
collection_schema.set_collection_name(str_collection_name);
collection_schema.set_dimension(COLLECTION_DIM);
collection_schema.set_index_file_size(INDEX_FILE_SIZE);
collection_schema.set_metric_type(1);
handler->CreateCollection(&context, &collection_schema, &response);
::milvus::grpc::InsertParam request;
request.set_collection_name(str_collection_name);
std::vector<std::vector<float>> record_array;
BuildVectors(0, VECTOR_COUNT, record_array);
::milvus::grpc::VectorIds vector_ids;
for (auto& record : record_array) {
::milvus::grpc::RowRecord* grpc_record = request.add_row_record_array();
CopyRowRecord(grpc_record, record);
}
handler->Insert(&context, &request, &vector_ids);
ASSERT_EQ(vector_ids.vector_id_array_size(), VECTOR_COUNT);
::milvus::grpc::FlushParam flush_param;
flush_param.add_collection_name_array(str_collection_name);
handler->Flush(&context, &flush_param, &response);
::milvus::grpc::DeleteByIDParam delete_param;
delete_param.set_collection_name(str_collection_name);
for (size_t i = 0; i < 10; i++) {
delete_param.mutable_id_array()->Add(vector_ids.vector_id_array(i));
}
handler->DeleteByID(&context, &delete_param, &response);
ASSERT_EQ(0, response.error_code()) << response.reason();
// ::milvus::grpc::FlushParam flush_param;
flush_param.clear_collection_name_array();
flush_param.add_collection_name_array(str_collection_name);
handler->Flush(&context, &flush_param, &response);
fiu_enable("ReLoadSegmentsRequest.OnExecute.readonly", 1, NULL, 0);
::milvus::grpc::ReLoadSegmentsParam reload_request;
reload_request.set_collection_name(str_collection_name);
::milvus::grpc::Status reload_response;
bool found = false;
for (size_t i = 0; i < 5000; i++) {
reload_request.clear_segment_id_array();
reload_request.add_segment_id_array(std::to_string(i));
handler->ReloadSegments(&context, &reload_request, &reload_response);
if (reload_response.error_code() == 0) {
found = true;
break;
}
}
fiu_disable("ReLoadSegmentsRequest.OnExecute.readonly");
ASSERT_TRUE(found) << reload_response.reason();
}
TEST_F(RpcHandlerTest, COMPACT) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
::milvus::grpc::CollectionSchema collection_schema;
::milvus::grpc::Status response;
std::string str_collection_name = "tbl_rpc_compact";
collection_schema.set_collection_name(str_collection_name);
collection_schema.set_dimension(COLLECTION_DIM);
collection_schema.set_index_file_size(INDEX_FILE_SIZE);
collection_schema.set_metric_type(1);
handler->CreateCollection(&context, &collection_schema, &response);
::milvus::grpc::InsertParam request;
request.set_collection_name(str_collection_name);
std::vector<std::vector<float>> bin_record_array;
BuildVectors(0, VECTOR_COUNT, bin_record_array);
::milvus::grpc::VectorIds vector_ids;
for (auto& record : bin_record_array) {
::milvus::grpc::RowRecord* grpc_record = request.add_row_record_array();
CopyRowRecord(grpc_record, record);
}
handler->Insert(&context, &request, &vector_ids);
ASSERT_EQ(vector_ids.vector_id_array_size(), VECTOR_COUNT);
::milvus::grpc::DeleteByIDParam delete_param;
delete_param.set_collection_name(str_collection_name);
for (size_t i = 0; i < 10; i++) {
delete_param.mutable_id_array()->Add(vector_ids.vector_id_array(i));
}
handler->DeleteByID(&context, &delete_param, &response);
ASSERT_EQ(0, response.error_code()) << response.reason();
::milvus::grpc::CollectionName collection_name;
collection_name.set_collection_name(str_collection_name);
handler->Compact(&context, &collection_name, &response);
ASSERT_EQ(0, response.error_code());
}
TEST_F(RpcHandlerTest, CMD_TEST) {
::grpc::ServerContext context;
handler->SetContext(&context, dummy_context);
handler->RegisterRequestHandler(milvus::server::RequestHandler());
::milvus::grpc::Command command;
command.set_cmd("version");
::milvus::grpc::StringReply reply;
handler->Cmd(&context, &command, &reply);
ASSERT_EQ(reply.string_reply(), MILVUS_VERSION);
command.set_cmd("requests");
handler->Cmd(&context, &command, &reply);
ASSERT_EQ(reply.status().error_code(), ::grpc::Status::OK.error_code());
command.set_cmd("tasktable");
handler->Cmd(&context, &command, &reply);
ASSERT_EQ(reply.status().error_code(), ::grpc::Status::OK.error_code());
command.set_cmd("test");
handler->Cmd(&context, &command, &reply);
ASSERT_EQ(reply.status().error_code(), ::grpc::Status::OK.error_code());
command.set_cmd("status");
handler->Cmd(&context, &command, &reply);
ASSERT_EQ(reply.status().error_code(), ::grpc::Status::OK.error_code());
command.set_cmd("mode");
handler->Cmd(&context, &command, &reply);
ASSERT_EQ(reply.status().error_code(), ::grpc::Status::OK.error_code());
command.set_cmd("build_commit_id");
handler->Cmd(&context, &command, &reply);
ASSERT_EQ(reply.status().error_code(), ::grpc::Status::OK.error_code());
command.set_cmd("set_config");
handler->Cmd(&context, &command, &reply);
command.set_cmd("get_config");
handler->Cmd(&context, &command, &reply);
}
TEST_F(RpcHandlerTest, HYBRID_TEST) {
::grpc::ServerContext context;
milvus::grpc::Mapping mapping;
milvus::grpc::Status response;
uint64_t row_num = 1000;
uint64_t dimension = 128;
// Create Hybrid Collection
mapping.set_collection_name("test_hybrid");
auto field_0 = mapping.add_fields();
field_0->set_name("field_0");
field_0->mutable_type()->set_data_type(::milvus::grpc::DataType::INT64);
auto field_1 = mapping.add_fields();
field_1->mutable_type()->mutable_vector_param()->set_dimension(128);
field_1->set_name("field_1");
milvus::json json_param = {{"metric_type", 1}, {"engine_type", 1}};
auto extra_param = field_1->add_extra_params();
extra_param->set_key("params");
extra_param->set_value(json_param.dump());
handler->CreateHybridCollection(&context, &mapping, &response);
// Insert Entities
milvus::grpc::HInsertParam insert_param;
milvus::grpc::HEntityIDs entity_ids;
insert_param.set_collection_name("test_hybrid");
auto entity = insert_param.mutable_entities();
auto field_name_0 = entity->add_field_names();
*field_name_0 = "field_0";
auto field_name_1 = entity->add_field_names();
*field_name_1 = "field_1";
entity->set_row_num(row_num);
std::vector<int64_t> field_value(row_num, 0);
for (uint64_t i = 0; i < row_num; i++) {
field_value[i] = i;
}
entity->set_attr_records(field_value.data(), row_num * sizeof(int64_t));
std::vector<std::vector<float>> vector_field;
vector_field.resize(row_num);
for (uint64_t i = 0; i < row_num; ++i) {
vector_field[i].resize(dimension);
for (uint64_t j = 0; j < dimension; ++j) {
vector_field[i][j] = (float)((i + 10) / (j + 20));
}
}
auto vector_record = entity->add_result_values();
for (uint64_t i = 0; i < row_num; ++i) {
auto record = vector_record->mutable_vector_value()->add_value();
auto vector_data = record->mutable_float_data();
vector_data->Resize(static_cast<int>(vector_field[i].size()), 0.0);
memcpy(vector_data->mutable_data(), vector_field[i].data(), vector_field[i].size() * sizeof(float));
}
handler->InsertEntity(&context, &insert_param, &entity_ids);
ASSERT_EQ(entity_ids.entity_id_array_size(), row_num);
uint64_t nq = 10;
uint64_t topk = 10;
milvus::grpc::HSearchParam search_param;
auto general_query = search_param.mutable_general_query();
auto boolean_query_1 = general_query->mutable_boolean_query();
boolean_query_1->set_occur(milvus::grpc::Occur::MUST);
auto general_query_1 = boolean_query_1->add_general_query();
auto boolean_query_2 = general_query_1->mutable_boolean_query();
auto term_query = boolean_query_2->add_general_query()->mutable_term_query();
term_query->set_field_name("field_0");
std::vector<int64_t> term_value(nq, 0);
for (uint64_t i = 0; i < nq; ++i) {
term_value[i] = i + nq;
}
term_query->set_value_num(nq);
term_query->set_values(term_value.data(), nq * sizeof(int64_t));
auto range_query = boolean_query_2->add_general_query()->mutable_range_query();
range_query->set_field_name("field_0");
auto comp1 = range_query->add_operand();
comp1->set_operator_(::milvus::grpc::CompareOperator::GTE);
comp1->set_operand("0");
auto comp2 = range_query->add_operand();
comp2->set_operator_(::milvus::grpc::CompareOperator::LTE);
comp2->set_operand("100000");
auto vector_query = boolean_query_2->add_general_query()->mutable_vector_query();
vector_query->set_field_name("field_1");
vector_query->set_topk(topk);
vector_query->set_query_boost(2);
std::vector<std::vector<float>> query_vector;
query_vector.resize(nq);
for (uint64_t i = 0; i < nq; ++i) {
query_vector[i].resize(dimension);
for (uint64_t j = 0; j < dimension; ++j) {
query_vector[i][j] = (float)((j + 1) / (i + dimension));
}
}
for (auto record : query_vector) {
auto row_record = vector_query->add_records();
CopyRowRecord(row_record, record);
}
auto extra_param_1 = vector_query->add_extra_params();
extra_param_1->set_key("params");
milvus::json param = {{"nprobe", 16}};
extra_param_1->set_value(param.dump());
search_param.set_collection_name("test_hybrid");
auto search_extra_param = search_param.add_extra_params();
search_extra_param->set_key("params");
search_extra_param->set_value("");
milvus::grpc::TopKQueryResult topk_query_result;
handler->HybridSearch(&context, &search_param, &topk_query_result);
}
TEST_F(RpcHandlerTest, HYBRID_INVALID_TEST) {
fiu_init(0);
::grpc::ServerContext context;
milvus::grpc::Mapping mapping;
milvus::grpc::Status response;
uint64_t row_num = 1000;
uint64_t dimension = 128;
// Create Hybrid Collection
mapping.set_collection_name("test_hybrid");
auto field_0 = mapping.add_fields();
field_0->set_name("field_0");
field_0->mutable_type()->set_data_type(::milvus::grpc::DataType::INT64);
auto field_1 = mapping.add_fields();
field_1->mutable_type()->mutable_vector_param()->set_dimension(128);
field_1->set_name("field_1");
milvus::json json_param = {{"metric_type", 1}, {"engine_type", 1}};
auto extra_param = field_1->add_extra_params();
extra_param->set_key("params");
extra_param->set_value(json_param.dump());
fiu_enable("CreateHybridCollectionRequest.OnExecute.invalid_collection_name", 1, NULL, 0);
handler->CreateHybridCollection(&context, &mapping, &response);
fiu_disable("CreateHybridCollectionRequest.OnExecute.invalid_collection_name");
fiu_enable("CreateHybridCollectionRequest.OnExecute.invalid_db_execute", 1, NULL, 0);
handler->CreateHybridCollection(&context, &mapping, &response);
fiu_disable("CreateHybridCollectionRequest.OnExecute.invalid_db_execute");
handler->CreateHybridCollection(&context, &mapping, &response);
milvus::grpc::CollectionName grpc_collection_name;
grpc_collection_name.set_collection_name("test_hybrid");
fiu_enable("DescribeHybridCollectionRequest.OnExecute.invalid_db_execute", 1, NULL, 0);
handler->DescribeHybridCollection(&context, &grpc_collection_name, &mapping);
fiu_disable("DescribeHybridCollectionRequest.OnExecute.invalid_db_execute");
handler->DescribeHybridCollection(&context, &grpc_collection_name, &mapping);
// Insert Entities
milvus::grpc::HInsertParam insert_param;
milvus::grpc::HEntityIDs entity_ids;
insert_param.set_collection_name("test_hybrid");
auto entity = insert_param.mutable_entities();
auto field_name_0 = entity->add_field_names();
*field_name_0 = "field_0";
auto field_name_1 = entity->add_field_names();
*field_name_1 = "field_1";
entity->set_row_num(row_num);
std::vector<int64_t> field_value(row_num, 0);
for (uint64_t i = 0; i < row_num; i++) {
field_value[i] = i;
}
entity->set_attr_records(field_value.data(), row_num * sizeof(int64_t));
std::vector<std::vector<float>> vector_field;
vector_field.resize(row_num);
for (uint64_t i = 0; i < row_num; ++i) {
vector_field[i].resize(dimension);
for (uint64_t j = 0; j < dimension; ++j) {
vector_field[i][j] = (float)((i + 10) / (j + 20));
}
}
auto vector_record = entity->add_result_values();
for (uint64_t i = 0; i < row_num; ++i) {
auto record = vector_record->mutable_vector_value()->add_value();
auto vector_data = record->mutable_float_data();
vector_data->Resize(static_cast<int>(vector_field[i].size()), 0.0);
memcpy(vector_data->mutable_data(), vector_field[i].data(), vector_field[i].size() * sizeof(float));
}
fiu_enable("InsertEntityRequest.OnExecute.throw_std_exception", 1, NULL, 0);
handler->InsertEntity(&context, &insert_param, &entity_ids);
fiu_disable("InsertEntityRequest.OnExecute.throw_std_exception");
handler->InsertEntity(&context, &insert_param, &entity_ids);
uint64_t nq = 10;
uint64_t topk = 10;
milvus::grpc::HSearchParam search_param;
auto general_query = search_param.mutable_general_query();
auto boolean_query_1 = general_query->mutable_boolean_query();
boolean_query_1->set_occur(milvus::grpc::Occur::MUST);
auto general_query_1 = boolean_query_1->add_general_query();
auto boolean_query_2 = general_query_1->mutable_boolean_query();
auto term_query = boolean_query_2->add_general_query()->mutable_term_query();
term_query->set_field_name("field_0");
std::vector<int64_t> term_value(nq, 0);
for (uint64_t i = 0; i < nq; ++i) {
term_value[i] = i + nq;
}
term_query->set_value_num(nq);
term_query->set_values(term_value.data(), nq * sizeof(int64_t));
auto range_query = boolean_query_2->add_general_query()->mutable_range_query();
range_query->set_field_name("field_0");
auto comp1 = range_query->add_operand();
comp1->set_operator_(::milvus::grpc::CompareOperator::GTE);
comp1->set_operand("0");
auto comp2 = range_query->add_operand();
comp2->set_operator_(::milvus::grpc::CompareOperator::LTE);
comp2->set_operand("100000");
auto vector_query = boolean_query_2->add_general_query()->mutable_vector_query();
vector_query->set_field_name("field_1");
vector_query->set_topk(topk);
vector_query->set_query_boost(2);
std::vector<std::vector<float>> query_vector;
query_vector.resize(nq);
for (uint64_t i = 0; i < nq; ++i) {
query_vector[i].resize(dimension);
for (uint64_t j = 0; j < dimension; ++j) {
query_vector[i][j] = (float)((j + 1) / (i + dimension));
}
}
for (auto record : query_vector) {
auto row_record = vector_query->add_records();
CopyRowRecord(row_record, record);
}
auto extra_param_1 = vector_query->add_extra_params();
extra_param_1->set_key("params");
milvus::json param = {{"nprobe", 16}};
extra_param_1->set_value(param.dump());
search_param.set_collection_name("test_hybrid");
auto search_extra_param = search_param.add_extra_params();
search_extra_param->set_key("params");
search_extra_param->set_value("");
milvus::grpc::TopKQueryResult topk_query_result;
fiu_enable("SearchRequest.OnExecute.throw_std_exception", 1, NULL, 0);
handler->HybridSearch(&context, &search_param, &topk_query_result);
fiu_disable("SearchRequest.OnExecute.throw_std_exception");
}
//////////////////////////////////////////////////////////////////////
namespace {
class DummyRequest : public milvus::server::BaseRequest {
public:
milvus::Status
OnExecute() override {
return milvus::Status::OK();
}
static milvus::server::BaseRequestPtr
Create() {
return std::shared_ptr<milvus::server::BaseRequest>(new DummyRequest());
}
public:
DummyRequest()
: BaseRequest(std::make_shared<milvus::server::Context>("dummy_request_id"),
milvus::server::BaseRequest::kCmd) {
}
};
class RpcSchedulerTest : public testing::Test {
protected:
void
SetUp() override {
request_ptr = std::make_shared<DummyRequest>();
}
std::shared_ptr<DummyRequest> request_ptr;
};
class AsyncDummyRequest : public milvus::server::BaseRequest {
public:
milvus::Status
OnExecute() override {
return milvus::Status::OK();
}
static milvus::server::BaseRequestPtr
Create() {
return std::shared_ptr<milvus::server::BaseRequest>(new DummyRequest());
}
public:
AsyncDummyRequest()
: BaseRequest(std::make_shared<milvus::server::Context>("dummy_request_id2"), milvus::server::BaseRequest::kCmd,
true) {
}
};
} // namespace
TEST_F(RpcSchedulerTest, BASE_TASK_TEST) {
auto status = request_ptr->Execute();
ASSERT_TRUE(status.ok());
milvus::server::RequestScheduler::GetInstance().Start();
// milvus::server::RequestScheduler::GetInstance().Stop();
// milvus::server::RequestScheduler::GetInstance().Start();
milvus::server::BaseRequestPtr base_task_ptr = DummyRequest::Create();
milvus::server::RequestScheduler::ExecRequest(base_task_ptr);
milvus::server::RequestScheduler::GetInstance().ExecuteRequest(request_ptr);
fiu_init(0);
fiu_enable("RequestScheduler.ExecuteRequest.push_queue_fail", 1, NULL, 0);
milvus::server::RequestScheduler::GetInstance().ExecuteRequest(request_ptr);
fiu_disable("RequestScheduler.ExecuteRequest.push_queue_fail");
// std::string dummy2 = "dql2";
// milvus::server::BaseRequestPtr base_task_ptr2 = DummyRequest::Create(dummy2);
// fiu_enable("RequestScheduler.PutToQueue.null_queue", 1, NULL, 0);
// milvus::server::RequestScheduler::GetInstance().ExecuteRequest(base_task_ptr2);
// fiu_disable("RequestScheduler.PutToQueue.null_queue");
milvus::server::BaseRequestPtr base_task_ptr3 = DummyRequest::Create();
fiu_enable("RequestScheduler.TakeToExecute.throw_std_exception", 1, NULL, 0);
milvus::server::RequestScheduler::GetInstance().ExecuteRequest(base_task_ptr3);
fiu_disable("RequestScheduler.TakeToExecute.throw_std_exception");
milvus::server::BaseRequestPtr base_task_ptr4 = DummyRequest::Create();
fiu_enable("RequestScheduler.TakeToExecute.execute_fail", 1, NULL, 0);
milvus::server::RequestScheduler::GetInstance().ExecuteRequest(base_task_ptr4);
fiu_disable("RequestScheduler.TakeToExecute.execute_fail");
milvus::server::BaseRequestPtr base_task_ptr5 = DummyRequest::Create();
fiu_enable("RequestScheduler.PutToQueue.push_null_thread", 1, NULL, 0);
milvus::server::RequestScheduler::GetInstance().ExecuteRequest(base_task_ptr5);
fiu_disable("RequestScheduler.PutToQueue.push_null_thread");
request_ptr = nullptr;
milvus::server::RequestScheduler::GetInstance().ExecuteRequest(request_ptr);
milvus::server::BaseRequestPtr null_ptr = nullptr;
milvus::server::RequestScheduler::ExecRequest(null_ptr);
auto async_ptr = std::make_shared<AsyncDummyRequest>();
auto base_ptr = std::static_pointer_cast<milvus::server::BaseRequest>(async_ptr);
milvus::server::RequestScheduler::ExecRequest(base_ptr);
async_ptr->set_status(milvus::Status(milvus::SERVER_INVALID_ARGUMENT, ""));
milvus::server::RequestScheduler::GetInstance().Stop();
milvus::server::RequestScheduler::GetInstance().Start();
milvus::server::RequestScheduler::GetInstance().Stop();
}
TEST(RpcTest, RPC_SERVER_TEST) {
using GrpcServer = milvus::server::grpc::GrpcServer;
GrpcServer& server = GrpcServer::GetInstance();
fiu_init(0);
fiu_enable("check_config_address_fail", 1, NULL, 0);
server.Start();
sleep(2);
fiu_disable("check_config_address_fail");
server.Stop();
fiu_enable("check_config_port_fail", 1, NULL, 0);
server.Start();
sleep(2);
fiu_disable("check_config_port_fail");
server.Stop();
server.Start();
sleep(2);
server.Stop();
}
TEST(RpcTest, InterceptorHookHandlerTest) {
auto handler = std::make_shared<milvus::server::grpc::GrpcInterceptorHookHandler>();
handler->OnPostRecvInitialMetaData(nullptr, nullptr);
handler->OnPreSendMessage(nullptr, nullptr);
}
Reference in New Issue View Git Blame Copy Permalink