admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2026-01-07 19:31:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

GPU search turn into CPU search when nprobe >gpu_max_k (#3586)

Browse Source 
Signed-off-by: shengjun.li <shengjun.li@zilliz.com>
This commit is contained in:
cqy123456 2020-09-04 18:52:50 +08:00 committed by shengjun.li
parent af2e7d9a88
commit dddf697e90
4 changed files with 25 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
core/src/scheduler/selector/FaissIVFFlatPass.cpp
View File

@ -12,6 +12,8 @@
#include "scheduler/selector/FaissIVFFlatPass.h"
#include "cache/GpuCacheMgr.h"
#include "config/ServerConfig.h"
#include "faiss/gpu/utils/DeviceUtils.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "scheduler/SchedInst.h"
#include "scheduler/Utils.h"
#include "scheduler/task/SearchTask.h"
@ -54,7 +56,11 @@ FaissIVFFlatPass::Run(const TaskPtr& task) {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: gpu disable, specify cpu to search!");
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else if (search_task->nq() < threshold_) {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nq < gpu_search_threshold, specify cpu to search!");
LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nq < gpu_search_threshold, specify cpu to search! ");
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else if (search_task->ExtraParam()[knowhere::IndexParams::nprobe].get<int64_t>() >
faiss::gpu::getMaxKSelection()) {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nprobe > gpu_max_nprobe_threshold, specify cpu to search!");
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nq >= gpu_search_threshold, specify gpu %d to search!",

6
core/src/scheduler/selector/FaissIVFPQPass.cpp
View File

@ -12,6 +12,8 @@
#include "scheduler/selector/FaissIVFPQPass.h"
#include "cache/GpuCacheMgr.h"
#include "config/ServerConfig.h"
#include "faiss/gpu/utils/DeviceUtils.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "scheduler/SchedInst.h"
#include "scheduler/Utils.h"
#include "scheduler/task/SearchTask.h"
@ -58,6 +60,10 @@ FaissIVFPQPass::Run(const TaskPtr& task) {
} else if (search_task->nq() < threshold_) {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFPQPass: nq < gpu_search_threshold, specify cpu to search!");
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else if (search_task->ExtraParam()[knowhere::IndexParams::nprobe].get<int64_t>() >
faiss::gpu::getMaxKSelection()) {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nprobe > gpu_max_nprobe_threshold, specify cpu to search!");
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFPQPass: nq >= gpu_search_threshold, specify gpu %d to search!",
search_gpus_[idx_]);

6
core/src/scheduler/selector/FaissIVFSQ8HPass.cpp
View File

@ -13,6 +13,8 @@
#include "scheduler/selector/FaissIVFSQ8HPass.h"
#include "cache/GpuCacheMgr.h"
#include "config/ServerConfig.h"
#include "faiss/gpu/utils/DeviceUtils.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "scheduler/SchedInst.h"
#include "scheduler/Utils.h"
#include "scheduler/task/SearchTask.h"
@ -55,6 +57,10 @@ FaissIVFSQ8HPass::Run(const TaskPtr& task) {
} else if (search_task->nq() < threshold_) {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFSQ8HPass: nq < gpu_search_threshold, specify cpu to search!");
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else if (search_task->ExtraParam()[knowhere::IndexParams::nprobe].get<int64_t>() >
faiss::gpu::getMaxKSelection()) {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nprobe > gpu_max_nprobe_threshold, specify cpu to search!");
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFSQ8HPass: nq >= gpu_search_threshold, specify gpu %d to search!",
search_gpus_[idx_]);

6
core/src/scheduler/selector/FaissIVFSQ8Pass.cpp
View File

@ -12,6 +12,8 @@
#include "scheduler/selector/FaissIVFSQ8Pass.h"
#include "cache/GpuCacheMgr.h"
#include "config/ServerConfig.h"
#include "faiss/gpu/utils/DeviceUtils.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "scheduler/SchedInst.h"
#include "scheduler/Utils.h"
#include "scheduler/task/SearchTask.h"
@ -56,6 +58,10 @@ FaissIVFSQ8Pass::Run(const TaskPtr& task) {
} else if (search_task->nq() < threshold_) {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFSQ8Pass: nq < gpu_search_threshold, specify cpu to search!");
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else if (search_task->ExtraParam()[knowhere::IndexParams::nprobe].get<int64_t>() >
faiss::gpu::getMaxKSelection()) {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nprobe > gpu_max_nprobe_threshold, specify cpu to search!");
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else {
LOG_SERVER_DEBUG_ << LogOut("FaissIVFSQ8Pass: nq >= gpu_search_threshold, specify gpu %d to search!",
search_gpus_[idx_]);