diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp index 923ca0db56..d496d8085c 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp @@ -133,9 +133,11 @@ GPUIVF::search_impl(int64_t n, const float* data, int64_t k, float* distances, i device_index->nprobe = search_cfg->nprobe; // assert(device_index->getNumProbes() == search_cfg->nprobe); - { + try { ResScope rs(res_, gpu_id_); device_index->search(n, (float*)data, k, distances, labels); + } catch (faiss::FaissException& e) { + KNOWHERE_THROW_MSG(e.what()); } } else { KNOWHERE_THROW_MSG("Not a GpuIndexIVF type."); diff --git a/core/src/scheduler/task/SearchTask.cpp b/core/src/scheduler/task/SearchTask.cpp index 08bc6525aa..34c8fc4b4b 100644 --- a/core/src/scheduler/task/SearchTask.cpp +++ b/core/src/scheduler/task/SearchTask.cpp @@ -212,7 +212,13 @@ XSearchTask::Execute() { ResMgrInst::GetInstance()->GetResource(path().Last())->type() == ResourceType::CPU) { hybrid = true; } - index_engine_->Search(nq, vectors, topk, nprobe, output_distance.data(), output_ids.data(), hybrid); + Status s = + index_engine_->Search(nq, vectors, topk, nprobe, output_distance.data(), output_ids.data(), hybrid); + if (!s.ok()) { + search_job->GetStatus() = s; + search_job->SearchDone(index_id_); + return; + } double span = rc.RecordSection(hdr + ", do search"); // search_job->AccumSearchCost(span);