提交 59878b1d 编写于 作者: Z zhenshan.cao 提交者: yefu.chen

Add lazy connect logic for indexbuilder

Signed-off-by: Nzhenshan.cao <zhenshan.cao@zilliz.com>
上级 6cd8c252
......@@ -67,65 +67,75 @@ IndexWrapper::parse() {
config_[key] = value;
}
auto stoi_closure = [](const std::string& s) -> int { return std::stoi(s); };
/***************************** meta *******************************/
check_parameter<int>(milvus::knowhere::meta::DIM, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::meta::TOPK, stoi_closure, std::nullopt);
/***************************** IVF Params *******************************/
check_parameter<int>(milvus::knowhere::IndexParams::nprobe, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::IndexParams::nlist, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::IndexParams::m, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::IndexParams::nbits, stoi_closure, std::nullopt);
if (!config_.contains(milvus::knowhere::meta::DIM)) {
// should raise exception here?
PanicInfo("dim must be specific in type params or index params!");
} else {
auto dim = config_[milvus::knowhere::meta::DIM].get<std::string>();
config_[milvus::knowhere::meta::DIM] = std::stoi(dim);
}
/************************** NSG Parameter **************************/
check_parameter<int>(milvus::knowhere::IndexParams::knng, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::IndexParams::search_length, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::IndexParams::out_degree, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::IndexParams::candidate, stoi_closure, std::nullopt);
if (!config_.contains(milvus::knowhere::meta::TOPK)) {
} else {
auto topk = config_[milvus::knowhere::meta::TOPK].get<std::string>();
config_[milvus::knowhere::meta::TOPK] = std::stoi(topk);
}
/************************** HNSW Params *****************************/
check_parameter<int>(milvus::knowhere::IndexParams::efConstruction, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::IndexParams::M, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::IndexParams::ef, stoi_closure, std::nullopt);
if (!config_.contains(milvus::knowhere::IndexParams::nlist)) {
} else {
auto nlist = config_[milvus::knowhere::IndexParams::nlist].get<std::string>();
config_[milvus::knowhere::IndexParams::nlist] = std::stoi(nlist);
}
/************************** Annoy Params *****************************/
check_parameter<int>(milvus::knowhere::IndexParams::n_trees, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::IndexParams::search_k, stoi_closure, std::nullopt);
if (!config_.contains(milvus::knowhere::IndexParams::nprobe)) {
} else {
auto nprobe = config_[milvus::knowhere::IndexParams::nprobe].get<std::string>();
config_[milvus::knowhere::IndexParams::nprobe] = std::stoi(nprobe);
}
/************************** PQ Params *****************************/
check_parameter<int>(milvus::knowhere::IndexParams::PQM, stoi_closure, std::nullopt);
if (!config_.contains(milvus::knowhere::IndexParams::nbits)) {
} else {
auto nbits = config_[milvus::knowhere::IndexParams::nbits].get<std::string>();
config_[milvus::knowhere::IndexParams::nbits] = std::stoi(nbits);
}
/************************** NGT Params *****************************/
check_parameter<int>(milvus::knowhere::IndexParams::edge_size, stoi_closure, std::nullopt);
if (!config_.contains(milvus::knowhere::IndexParams::m)) {
} else {
auto m = config_[milvus::knowhere::IndexParams::m].get<std::string>();
config_[milvus::knowhere::IndexParams::m] = std::stoi(m);
}
/************************** NGT Search Params *****************************/
check_parameter<int>(milvus::knowhere::IndexParams::epsilon, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::IndexParams::max_search_edges, stoi_closure, std::nullopt);
/************************** NSG Parameter **************************/
if (!config_.contains(milvus::knowhere::IndexParams::knng)) {
} else {
auto knng = config_[milvus::knowhere::IndexParams::knng].get<std::string>();
config_[milvus::knowhere::IndexParams::knng] = std::stoi(knng);
}
/************************** NGT_PANNG Params *****************************/
check_parameter<int>(milvus::knowhere::IndexParams::forcedly_pruned_edge_size, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::IndexParams::selectively_pruned_edge_size, stoi_closure, std::nullopt);
if (!config_.contains(milvus::knowhere::IndexParams::search_length)) {
} else {
auto search_length = config_[milvus::knowhere::IndexParams::search_length].get<std::string>();
config_[milvus::knowhere::IndexParams::search_length] = std::stoi(search_length);
}
/************************** NGT_ONNG Params *****************************/
check_parameter<int>(milvus::knowhere::IndexParams::outgoing_edge_size, stoi_closure, std::nullopt);
check_parameter<int>(milvus::knowhere::IndexParams::incoming_edge_size, stoi_closure, std::nullopt);
if (!config_.contains(milvus::knowhere::IndexParams::out_degree)) {
} else {
auto out_degree = config_[milvus::knowhere::IndexParams::out_degree].get<std::string>();
config_[milvus::knowhere::IndexParams::out_degree] = std::stoi(out_degree);
}
/************************** Serialize Params *******************************/
check_parameter<int>(milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE, stoi_closure, std::optional{4});
}
if (!config_.contains(milvus::knowhere::IndexParams::candidate)) {
} else {
auto candidate = config_[milvus::knowhere::IndexParams::candidate].get<std::string>();
config_[milvus::knowhere::IndexParams::candidate] = std::stoi(candidate);
}
template <typename T>
void
IndexWrapper::check_parameter(const std::string& key, std::function<T(std::string)> fn, std::optional<T> default_v) {
if (!config_.contains(key)) {
if (default_v.has_value()) {
config_[key] = default_v.value();
}
/************************** Serialize *******************************/
if (!config_.contains(milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE)) {
config_[milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE] = 4;
} else {
auto value = config_[key];
config_[key] = fn(value);
auto slice_size = config_[milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE].get<std::string>();
config_[milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE] = std::stoi(slice_size);
}
}
......
......@@ -52,12 +52,6 @@ class IndexWrapper {
void
StoreRawData(const knowhere::DatasetPtr& dataset);
template <typename T>
void
check_parameter(const std::string& key,
std::function<T(std::string)> fn,
std::optional<T> default_v = std::nullopt);
public:
void
BuildWithIds(const knowhere::DatasetPtr& dataset);
......
......@@ -25,17 +25,14 @@ NM_List() {
static std::vector<std::string> ret{
milvus::knowhere::IndexEnum::INDEX_FAISS_IVFFLAT,
milvus::knowhere::IndexEnum::INDEX_NSG,
milvus::knowhere::IndexEnum::INDEX_RHNSWFlat,
};
return ret;
}
std::vector<std::string>
BIN_List() {
static std::vector<std::string> ret{
milvus::knowhere::IndexEnum::INDEX_FAISS_BIN_IDMAP,
milvus::knowhere::IndexEnum::INDEX_FAISS_BIN_IVFFLAT,
};
static std::vector<std::string> ret{milvus::knowhere::IndexEnum::INDEX_FAISS_BIN_IDMAP,
milvus::knowhere::IndexEnum::INDEX_FAISS_BIN_IVFFLAT};
return ret;
}
......@@ -43,7 +40,7 @@ std::vector<std::string>
Need_ID_List() {
static std::vector<std::string> ret{
// milvus::knowhere::IndexEnum::INDEX_FAISS_BIN_IVFFLAT,
// milvus::knowhere::IndexEnum::INDEX_NSG,
// milvus::knowhere::IndexEnum::INDEX_NSG
};
return ret;
......@@ -51,9 +48,7 @@ Need_ID_List() {
std::vector<std::string>
Need_BuildAll_list() {
static std::vector<std::string> ret{
milvus::knowhere::IndexEnum::INDEX_NSG,
};
static std::vector<std::string> ret{milvus::knowhere::IndexEnum::INDEX_NSG};
return ret;
}
......
......@@ -99,105 +99,14 @@ generate_conf(const milvus::knowhere::IndexType& index_type, const milvus::knowh
{milvus::knowhere::Metric::TYPE, metric_type},
};
} else if (index_type == milvus::knowhere::IndexEnum::INDEX_NSG) {
return milvus::knowhere::Config{
{milvus::knowhere::meta::DIM, DIM},
{milvus::knowhere::IndexParams::nlist, 163},
{milvus::knowhere::IndexParams::nprobe, 8},
{milvus::knowhere::IndexParams::knng, 20},
{milvus::knowhere::IndexParams::search_length, 40},
{milvus::knowhere::IndexParams::out_degree, 30},
{milvus::knowhere::IndexParams::candidate, 100},
{milvus::knowhere::Metric::TYPE, metric_type},
};
#ifdef MILVUS_SUPPORT_SPTAG
} else if (index_type == milvus::knowhere::IndexEnum::INDEX_SPTAG_KDT_RNT) {
return milvus::knowhere::Config{
{milvus::knowhere::meta::DIM, DIM},
// {milvus::knowhere::meta::TOPK, 10},
{milvus::knowhere::Metric::TYPE, metric_type},
{milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE, 4},
};
} else if (index_type == milvus::knowhere::IndexEnum::INDEX_SPTAG_BKT_RNT) {
return milvus::knowhere::Config{
{milvus::knowhere::meta::DIM, DIM},
// {milvus::knowhere::meta::TOPK, 10},
{milvus::knowhere::Metric::TYPE, metric_type},
{milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE, 4},
};
#endif
} else if (index_type == milvus::knowhere::IndexEnum::INDEX_HNSW) {
return milvus::knowhere::Config{
{milvus::knowhere::meta::DIM, DIM},
// {milvus::knowhere::meta::TOPK, 10},
{milvus::knowhere::IndexParams::M, 16},
{milvus::knowhere::IndexParams::efConstruction, 200},
{milvus::knowhere::IndexParams::ef, 200},
{milvus::knowhere::Metric::TYPE, metric_type},
};
} else if (index_type == milvus::knowhere::IndexEnum::INDEX_ANNOY) {
return milvus::knowhere::Config{
{milvus::knowhere::meta::DIM, DIM},
// {milvus::knowhere::meta::TOPK, 10},
{milvus::knowhere::IndexParams::n_trees, 4},
{milvus::knowhere::IndexParams::search_k, 100},
{milvus::knowhere::Metric::TYPE, metric_type},
{milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE, 4},
};
} else if (index_type == milvus::knowhere::IndexEnum::INDEX_RHNSWFlat) {
return milvus::knowhere::Config{
{milvus::knowhere::meta::DIM, DIM},
// {milvus::knowhere::meta::TOPK, 10},
{milvus::knowhere::IndexParams::M, 16},
{milvus::knowhere::IndexParams::efConstruction, 200},
{milvus::knowhere::IndexParams::ef, 200},
{milvus::knowhere::Metric::TYPE, metric_type},
{milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE, 4},
};
} else if (index_type == milvus::knowhere::IndexEnum::INDEX_RHNSWPQ) {
return milvus::knowhere::Config{
{milvus::knowhere::meta::DIM, DIM},
// {milvus::knowhere::meta::TOPK, 10},
{milvus::knowhere::IndexParams::M, 16},
{milvus::knowhere::IndexParams::efConstruction, 200},
{milvus::knowhere::IndexParams::ef, 200},
{milvus::knowhere::Metric::TYPE, metric_type},
{milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE, 4},
{milvus::knowhere::IndexParams::PQM, 8},
};
} else if (index_type == milvus::knowhere::IndexEnum::INDEX_RHNSWSQ) {
return milvus::knowhere::Config{
{milvus::knowhere::meta::DIM, DIM},
// {milvus::knowhere::meta::TOPK, 10},
{milvus::knowhere::IndexParams::M, 16},
{milvus::knowhere::IndexParams::efConstruction, 200},
{milvus::knowhere::IndexParams::ef, 200},
{milvus::knowhere::Metric::TYPE, metric_type},
{milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE, 4},
};
} else if (index_type == milvus::knowhere::IndexEnum::INDEX_NGTPANNG) {
return milvus::knowhere::Config{
{milvus::knowhere::meta::DIM, DIM},
// {milvus::knowhere::meta::TOPK, 10},
{milvus::knowhere::Metric::TYPE, metric_type},
{milvus::knowhere::IndexParams::edge_size, 10},
{milvus::knowhere::IndexParams::epsilon, 0.1},
{milvus::knowhere::IndexParams::max_search_edges, 50},
{milvus::knowhere::IndexParams::forcedly_pruned_edge_size, 60},
{milvus::knowhere::IndexParams::selectively_pruned_edge_size, 30},
{milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE, 4},
};
} else if (index_type == milvus::knowhere::IndexEnum::INDEX_NGTONNG) {
return milvus::knowhere::Config{
{milvus::knowhere::meta::DIM, DIM},
// {milvus::knowhere::meta::TOPK, 10},
{milvus::knowhere::Metric::TYPE, metric_type},
{milvus::knowhere::IndexParams::edge_size, 20},
{milvus::knowhere::IndexParams::epsilon, 0.1},
{milvus::knowhere::IndexParams::max_search_edges, 50},
{milvus::knowhere::IndexParams::outgoing_edge_size, 5},
{milvus::knowhere::IndexParams::incoming_edge_size, 40},
{milvus::knowhere::INDEX_FILE_SLICE_SIZE_IN_MEGABYTE, 4},
};
return milvus::knowhere::Config{{milvus::knowhere::meta::DIM, DIM},
{milvus::knowhere::IndexParams::nlist, 163},
{milvus::knowhere::IndexParams::nprobe, 8},
{milvus::knowhere::IndexParams::knng, 20},
{milvus::knowhere::IndexParams::search_length, 40},
{milvus::knowhere::IndexParams::out_degree, 30},
{milvus::knowhere::IndexParams::candidate, 100},
{milvus::knowhere::Metric::TYPE, metric_type}};
}
return milvus::knowhere::Config();
}
......@@ -457,17 +366,6 @@ INSTANTIATE_TEST_CASE_P(
std::pair(milvus::knowhere::IndexEnum::INDEX_FAISS_BIN_IVFFLAT,
milvus::knowhere::Metric::JACCARD),
std::pair(milvus::knowhere::IndexEnum::INDEX_FAISS_BIN_IDMAP, milvus::knowhere::Metric::JACCARD),
#ifdef MILVUS_SUPPORT_SPTAG
std::pair(milvus::knowhere::IndexEnum::INDEX_SPTAG_KDT_RNT, milvus::knowhere::Metric::L2),
std::pair(milvus::knowhere::IndexEnum::INDEX_SPTAG_BKT_RNT, milvus::knowhere::Metric::L2),
#endif
std::pair(milvus::knowhere::IndexEnum::INDEX_HNSW, milvus::knowhere::Metric::L2),
std::pair(milvus::knowhere::IndexEnum::INDEX_ANNOY, milvus::knowhere::Metric::L2),
std::pair(milvus::knowhere::IndexEnum::INDEX_RHNSWFlat, milvus::knowhere::Metric::L2),
std::pair(milvus::knowhere::IndexEnum::INDEX_RHNSWPQ, milvus::knowhere::Metric::L2),
std::pair(milvus::knowhere::IndexEnum::INDEX_RHNSWSQ, milvus::knowhere::Metric::L2),
std::pair(milvus::knowhere::IndexEnum::INDEX_NGTPANNG, milvus::knowhere::Metric::L2),
std::pair(milvus::knowhere::IndexEnum::INDEX_NGTONNG, milvus::knowhere::Metric::L2),
std::pair(milvus::knowhere::IndexEnum::INDEX_NSG, milvus::knowhere::Metric::L2)));
TEST_P(IndexWrapperTest, Constructor) {
......
......@@ -14,7 +14,9 @@ import (
type UniqueID = typeutil.UniqueID
type Client struct {
client indexbuilderpb.IndexBuildServiceClient
client indexbuilderpb.IndexBuildServiceClient
address string
ctx context.Context
}
type IndexDescription struct {
......@@ -26,12 +28,9 @@ type IndexDescription struct {
}
func NewBuildIndexClient(ctx context.Context, address string) (*Client, error) {
conn, err := grpc.DialContext(ctx, address, grpc.WithInsecure(), grpc.WithBlock())
if err != nil {
return nil, err
}
return &Client{
client: indexbuilderpb.NewIndexBuildServiceClient(conn),
address: address,
ctx: ctx,
}, nil
}
......@@ -39,7 +38,22 @@ func parseTS(t int64) time.Time {
return time.Unix(0, t)
}
func (c *Client) tryConnect() error {
if c.client != nil {
return nil
}
conn, err := grpc.DialContext(c.ctx, c.address, grpc.WithInsecure(), grpc.WithBlock())
if err != nil {
return err
}
c.client = indexbuilderpb.NewIndexBuildServiceClient(conn)
return nil
}
func (c *Client) BuildIndexWithoutID(columnDataPaths []string, typeParams map[string]string, indexParams map[string]string) (UniqueID, error) {
if c.tryConnect() != nil {
panic("BuildIndexWithoutID: failed to connect index builder")
}
var typeParamsKV []*commonpb.KeyValuePair
for typeParam := range typeParams {
typeParamsKV = append(typeParamsKV, &commonpb.KeyValuePair{
......@@ -72,6 +86,9 @@ func (c *Client) BuildIndexWithoutID(columnDataPaths []string, typeParams map[st
}
func (c *Client) DescribeIndex(indexID UniqueID) (*IndexDescription, error) {
if c.tryConnect() != nil {
panic("DescribeIndex: failed to connect index builder")
}
ctx := context.TODO()
request := &indexbuilderpb.DescribleIndexRequest{
IndexID: indexID,
......@@ -92,6 +109,9 @@ func (c *Client) DescribeIndex(indexID UniqueID) (*IndexDescription, error) {
}
func (c *Client) GetIndexFilePaths(indexID UniqueID) ([]string, error) {
if c.tryConnect() != nil {
panic("GetIndexFilePaths: failed to connect index builder")
}
ctx := context.TODO()
request := &indexbuilderpb.GetIndexFilePathsRequest{
IndexID: indexID,
......
......@@ -181,8 +181,6 @@ func TestCIndex_Codec(t *testing.T) {
err = index.Delete()
assert.Equal(t, err, nil)
err = copyIndex.Delete()
assert.Equal(t, err, nil)
}
}
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册