diff --git a/cpp/src/scheduler/SchedInst.cpp b/cpp/src/scheduler/SchedInst.cpp index 7db58404a4..df71e69af3 100644 --- a/cpp/src/scheduler/SchedInst.cpp +++ b/cpp/src/scheduler/SchedInst.cpp @@ -60,18 +60,7 @@ load_simple_config() { for (auto &gpu_id : gpu_ids) { ResMgrInst::GetInstance()->Add(ResourceFactory::Create(std::to_string(gpu_id), "GPU", gpu_id, true, true)); ResMgrInst::GetInstance()->Connect("cpu", std::to_string(gpu_id), io); - auto pinned_memory = 300; - auto temp_memory = 300; - auto resource_num = 2; - pinned_memory = 1024 * 1024 * pinned_memory; - temp_memory = 1024 * 1024 * temp_memory; - knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(gpu_id, - pinned_memory, - temp_memory, - resource_num); } - - knowhere::FaissGpuResourceMgr::GetInstance().InitResource(); } } diff --git a/cpp/src/server/Server.cpp b/cpp/src/server/Server.cpp index 4a456ff726..b9e6494eb8 100644 --- a/cpp/src/server/Server.cpp +++ b/cpp/src/server/Server.cpp @@ -20,7 +20,7 @@ #include #include #include -#include "knowhere/index/vector_index/gpu_ivf.h" +#include "wrapper/knowhere/KnowhereResource.h" #include "metrics/Metrics.h" #include "DBWrapper.h" @@ -250,6 +250,7 @@ Server::LoadConfig() { void Server::StartService() { + engine::KnowhereResource::Initialize(); engine::StartSchedulerService(); DBWrapper::GetInstance().StartService(); grpc::GrpcMilvusServer::StartService(); @@ -260,7 +261,7 @@ Server::StopService() { grpc::GrpcMilvusServer::StopService(); DBWrapper::GetInstance().StopService(); engine::StopSchedulerService(); - knowhere::FaissGpuResourceMgr::GetInstance().Free(); // free gpu resource. + engine::KnowhereResource::Finalize(); } } diff --git a/cpp/src/wrapper/knowhere/KnowhereResource.cpp b/cpp/src/wrapper/knowhere/KnowhereResource.cpp new file mode 100644 index 0000000000..d5a9c2f6dc --- /dev/null +++ b/cpp/src/wrapper/knowhere/KnowhereResource.cpp @@ -0,0 +1,65 @@ +//////////////////////////////////////////////////////////////////////////////// +// Copyright 上海赜睿信息科技有限公司(Zilliz) - All Rights Reserved +// Unauthorized copying of this file, via any medium is strictly prohibited. +// Proprietary and confidential. +//////////////////////////////////////////////////////////////////////////////// + +#include "KnowhereResource.h" +#include "server/ServerConfig.h" + +#include + +namespace zilliz { +namespace milvus { +namespace engine { + +constexpr int64_t M_BYTE = 1024 * 1024; + +ErrorCode KnowhereResource::Initialize() { + struct GpuResourceSetting { + int64_t pinned_memory = 300*M_BYTE; + int64_t temp_memory = 300*M_BYTE; + int64_t resource_num = 2; + }; + using GpuResourcesArray = std::map; + GpuResourcesArray gpu_resources; + + //get build index gpu resource + server::ServerConfig& root_config = server::ServerConfig::GetInstance(); + server::ConfigNode& db_config = root_config.GetConfig(server::CONFIG_DB); + + int32_t build_index_gpu = db_config.GetInt32Value(server::CONFIG_DB_BUILD_INDEX_GPU, 0); + gpu_resources.insert(std::make_pair(build_index_gpu, GpuResourceSetting())); + + //get search gpu resource + server::ConfigNode& res_config = root_config.GetConfig(server::CONFIG_RESOURCE); + auto resources = res_config.GetSequence("resources"); + std::set gpu_ids; + for (auto &resource : resources) { + if (resource.length() < 4 || resource.substr(0, 3) != "gpu") { + // invalid + continue; + } + auto gpu_id = std::stoi(resource.substr(3)); + gpu_resources.insert(std::make_pair(gpu_id, GpuResourceSetting())); + } + + //init gpu resources + for(auto iter = gpu_resources.begin(); iter != gpu_resources.end(); ++iter) { + knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(iter->first, + iter->second.pinned_memory, + iter->second.temp_memory, + iter->second.resource_num); + } + + return KNOWHERE_SUCCESS; +} + +ErrorCode KnowhereResource::Finalize() { + knowhere::FaissGpuResourceMgr::GetInstance().Free(); // free gpu resource. + return KNOWHERE_SUCCESS; +} + +} +} +} \ No newline at end of file diff --git a/cpp/src/wrapper/knowhere/KnowhereResource.h b/cpp/src/wrapper/knowhere/KnowhereResource.h new file mode 100644 index 0000000000..b7b320d9cc --- /dev/null +++ b/cpp/src/wrapper/knowhere/KnowhereResource.h @@ -0,0 +1,25 @@ +//////////////////////////////////////////////////////////////////////////////// +// Copyright 上海赜睿信息科技有限公司(Zilliz) - All Rights Reserved +// Unauthorized copying of this file, via any medium is strictly prohibited. +// Proprietary and confidential. +//////////////////////////////////////////////////////////////////////////////// + +#pragma once + +#include "utils/Error.h" +#include "knowhere/index/vector_index/gpu_ivf.h" + +namespace zilliz { +namespace milvus { +namespace engine { + +class KnowhereResource { +public: + static ErrorCode Initialize(); + static ErrorCode Finalize(); +}; + + +} +} +} diff --git a/cpp/unittest/knowhere/CMakeLists.txt b/cpp/unittest/knowhere/CMakeLists.txt index 089ff992fa..ddb8eb2d9e 100644 --- a/cpp/unittest/knowhere/CMakeLists.txt +++ b/cpp/unittest/knowhere/CMakeLists.txt @@ -1,7 +1,10 @@ include_directories("${CUDA_TOOLKIT_ROOT_DIR}/include") link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64") -aux_source_directory(${MILVUS_ENGINE_SRC}/wrapper/knowhere knowhere_src) +set(knowhere_src + ${MILVUS_ENGINE_SRC}/wrapper/knowhere/data_transfer.cpp + ${MILVUS_ENGINE_SRC}/wrapper/knowhere/vec_impl.cpp + ${MILVUS_ENGINE_SRC}/wrapper/knowhere/vec_index.cpp) set(helper utils.cpp)