From 26fd0d18d17bafff75e592d59c282c78dae4ac44 Mon Sep 17 00:00:00 2001 From: Luis Melo Date: Fri, 2 Feb 2024 15:42:24 -0300 Subject: [PATCH 1/3] Added GPU IndexTypes to support Milvus-GPU --- src/impl/TypeUtils.cpp | 10 ++++++++++ src/impl/types/IndexDesc.cpp | 7 +++++++ src/include/milvus/types/IndexType.h | 2 ++ test/it/TestCreateIndex.cpp | 29 ++++++++++++++++++++++++++++ 4 files changed, 48 insertions(+) diff --git a/src/impl/TypeUtils.cpp b/src/impl/TypeUtils.cpp index ae213d03..b642fec4 100644 --- a/src/impl/TypeUtils.cpp +++ b/src/impl/TypeUtils.cpp @@ -393,6 +393,12 @@ IndexTypeCast(const std::string& type) { if (type == "BIN_IVF_FLAT") { return IndexType::BIN_IVF_FLAT; } + if (type == "GPU_IVF_FLAT"){ + return IndexType::GPU_IVF_FLAT; + } + if (type == "GPU_IVF_PQ"){ + return IndexType::GPU_IVF_PQ; + } return IndexType::INVALID; } @@ -897,6 +903,10 @@ to_string(milvus::IndexType index_type) { return "BIN_FLAT"; case milvus::IndexType::BIN_IVF_FLAT: return "BIN_IVF_FLAT"; + case milvus::IndexType::GPU_IVF_FLAT: + return "GPU_IVF_FLAT"; + case milvus::IndexType::GPU_IVF_PQ: + return "GPU_IVF_PQ"; default: return "INVALID"; } diff --git a/src/impl/types/IndexDesc.cpp b/src/impl/types/IndexDesc.cpp index 76939668..e0db98a8 100644 --- a/src/impl/types/IndexDesc.cpp +++ b/src/impl/types/IndexDesc.cpp @@ -57,6 +57,7 @@ Status validate_index_and_metric(const MetricType metric_type, const IndexType index_type) { if ((metric_type == milvus::MetricType::IP || metric_type == milvus::MetricType::L2) && (index_type == milvus::IndexType::FLAT || index_type == milvus::IndexType::IVF_FLAT || + index_type == milvus::IndexType::GPU_IVF_FLAT || index_type == milvus::IndexType::GPU_IVF_PQ || index_type == milvus::IndexType::IVF_SQ8 || index_type == milvus::IndexType::IVF_PQ || index_type == milvus::IndexType::HNSW || index_type == milvus::IndexType::IVF_HNSW || index_type == milvus::IndexType::RHNSW_FLAT || index_type == milvus::IndexType::RHNSW_SQ || @@ -90,6 +91,12 @@ validate_params(const IndexDesc& data, const std::unordered_mapCreateIndex(collection_name, index_desc, progress_monitor); EXPECT_FALSE(status.IsOk()); +} + +TEST_F(MilvusMockedTest, TestCreateGPUIndexInstantly){ + milvus::ConnectParam connect_param{"127.0.0.1", server_.ListenPort()}; + client_->Connect(connect_param); + + std::string collection_name = "test_collection"; + std::string field_name = "test_field"; + std::string index_name = "test_gpu_index"; + auto index_type = milvus::IndexType::GPU_IVF_FLAT; + auto metric_type = milvus::MetricType::IP; + int64_t index_id = 0; + + milvus::IndexDesc index_desc(field_name, "", index_type, metric_type, index_id); + index_desc.AddExtraParam("nlist", 1024); + const auto progress_monitor = ::milvus::ProgressMonitor::NoWait(); + + EXPECT_CALL(service_, Flush(_, AllOf(Property(&FlushRequest::collection_names, ElementsAre(collection_name))), _)) + .WillOnce([&](::grpc::ServerContext*, const FlushRequest*, FlushResponse*) { return ::grpc::Status{}; }); + + EXPECT_CALL(service_, CreateIndex(_, + AllOf(Property(&CreateIndexRequest::collection_name, collection_name), + Property(&CreateIndexRequest::field_name, field_name)), + _)) + .WillOnce([](::grpc::ServerContext*, const CreateIndexRequest*, ::milvus::proto::common::Status*) { + return ::grpc::Status{}; + }); + auto status = client_->CreateIndex(collection_name, index_desc, progress_monitor); + EXPECT_TRUE(status.IsOk()); } \ No newline at end of file From 2708af60ec056dc985a1f8596f14864b068382de Mon Sep 17 00:00:00 2001 From: Luis Melo Date: Mon, 5 Feb 2024 11:01:58 -0300 Subject: [PATCH 2/3] Fix Linting errors --- src/impl/TypeUtils.cpp | 4 ++-- src/impl/types/IndexDesc.cpp | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/impl/TypeUtils.cpp b/src/impl/TypeUtils.cpp index b642fec4..9c85ead5 100644 --- a/src/impl/TypeUtils.cpp +++ b/src/impl/TypeUtils.cpp @@ -393,10 +393,10 @@ IndexTypeCast(const std::string& type) { if (type == "BIN_IVF_FLAT") { return IndexType::BIN_IVF_FLAT; } - if (type == "GPU_IVF_FLAT"){ + if (type == "GPU_IVF_FLAT") { return IndexType::GPU_IVF_FLAT; } - if (type == "GPU_IVF_PQ"){ + if (type == "GPU_IVF_PQ") { return IndexType::GPU_IVF_PQ; } return IndexType::INVALID; diff --git a/src/impl/types/IndexDesc.cpp b/src/impl/types/IndexDesc.cpp index e0db98a8..d95fce45 100644 --- a/src/impl/types/IndexDesc.cpp +++ b/src/impl/types/IndexDesc.cpp @@ -92,7 +92,7 @@ validate_params(const IndexDesc& data, const std::unordered_map Date: Tue, 20 Feb 2024 08:41:47 -0300 Subject: [PATCH 3/3] Clang format from build directory --- test/it/TestCreateIndex.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/it/TestCreateIndex.cpp b/test/it/TestCreateIndex.cpp index 18c0b30d..ec8a2e80 100644 --- a/test/it/TestCreateIndex.cpp +++ b/test/it/TestCreateIndex.cpp @@ -189,8 +189,8 @@ TEST_F(MilvusMockedTest, TestCreateIndexFailed) { EXPECT_FALSE(status.IsOk()); } -TEST_F(MilvusMockedTest, TestCreateGPUIndexInstantly){ - milvus::ConnectParam connect_param{"127.0.0.1", server_.ListenPort()}; +TEST_F(MilvusMockedTest, TestCreateGPUIndexInstantly) { + milvus::ConnectParam connect_param{"127.0.0.1", server_.ListenPort()}; client_->Connect(connect_param); std::string collection_name = "test_collection";