milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

selector code refactoring (#3617) 

* selector code refactoring

Signed-off-by: cqy <yaya645@126.com>

* modify test

Signed-off-by: cqy <yaya645@126.com>
pull/3660/head
cqy123456 2020-09-08 13:58:50 +08:00 committed by GitHub
parent 18c630a4be
commit 8a52fc70d9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
13 changed files with 38 additions and 469 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
core/src/scheduler/SchedInst.h
View File

@ -19,10 +19,7 @@
#include "Utils.h"
#include "selector/BuildIndexPass.h"
#include "selector/FaissFlatPass.h"
#include "selector/FaissIVFFlatPass.h"
#include "selector/FaissIVFPQPass.h"
#include "selector/FaissIVFSQ8HPass.h"
#include "selector/FaissIVFSQ8Pass.h"
#include "selector/FaissIVFPass.h"
#include "selector/FallbackPass.h"
#include "selector/Optimizer.h"
@ -122,10 +119,7 @@ class OptimizerInst {
pass_list.push_back(std::make_shared<BuildIndexPass>());
pass_list.push_back(std::make_shared<FaissFlatPass>());
pass_list.push_back(std::make_shared<FaissIVFFlatPass>());
pass_list.push_back(std::make_shared<FaissIVFSQ8Pass>());
pass_list.push_back(std::make_shared<FaissIVFSQ8HPass>());
pass_list.push_back(std::make_shared<FaissIVFPQPass>());
pass_list.push_back(std::make_shared<FaissIVFPass>());
}
#endif
pass_list.push_back(std::make_shared<FallbackPass>());

51
core/src/scheduler/selector/FaissIVFFlatPass.h
View File

@ -1,51 +0,0 @@
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.
#ifdef MILVUS_GPU_VERSION
#pragma once
#include <condition_variable>
#include <deque>
#include <limits>
#include <list>
#include <memory>
#include <mutex>
#include <queue>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
#include "config/handler/GpuResourceConfigHandler.h"
#include "scheduler/selector/Pass.h"
namespace milvus {
namespace scheduler {
class FaissIVFFlatPass : public Pass, public server::GpuResourceConfigHandler {
public:
FaissIVFFlatPass() = default;
public:
void
Init() override;
bool
Run(const TaskPtr& task) override;
private:
int64_t idx_ = 0;
};
using FaissIVFFlatPassPtr = std::shared_ptr<FaissIVFFlatPass>;
} // namespace scheduler
} // namespace milvus
#endif

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

@ -1,86 +0,0 @@
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.
#ifdef MILVUS_GPU_VERSION
#include "scheduler/selector/FaissIVFPQPass.h"
#include "cache/GpuCacheMgr.h"
#include "config/Config.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "scheduler/SchedInst.h"
#include "scheduler/Utils.h"
#include "scheduler/task/SearchTask.h"
#include "scheduler/tasklabel/SpecResLabel.h"
#include "utils/Log.h"
#include "utils/ValidationUtil.h"
#include <fiu-local.h>
namespace milvus {
namespace scheduler {
void
FaissIVFPQPass::Init() {
#ifdef MILVUS_GPU_VERSION
server::Config& config = server::Config::GetInstance();
Status s = config.GetGpuResourceConfigGpuSearchThreshold(threshold_);
if (!s.ok()) {
threshold_ = std::numeric_limits<int32_t>::max();
}
s = config.GetGpuResourceConfigSearchResources(search_gpus_);
if (!s.ok()) {
throw std::exception();
}
SetIdentity("FaissIVFPQPass");
AddGpuEnableListener();
AddGpuSearchThresholdListener();
AddGpuSearchResourcesListener();
#endif
}
bool
FaissIVFPQPass::Run(const TaskPtr& task) {
if (task->Type() != TaskType::SearchTask) {
return false;
}
auto search_task = std::static_pointer_cast<XSearchTask>(task);
if (search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_PQ) {
return false;
}
auto search_job = std::static_pointer_cast<SearchJob>(search_task->job_.lock());
ResourcePtr res_ptr;
if (!gpu_enable_) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFPQPass: gpu disable, specify cpu to search!", "search", 0);
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else if (search_job->nq() < (uint64_t)threshold_) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFPQPass: nq < gpu_search_threshold, specify cpu to search!",
"search", 0);
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else if (search_job->extra_params()[knowhere::IndexParams::nprobe].get<int64_t>() <
milvus::server::QUERY_MAX_TOPK) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFFlatPass: nprobe > gpu_nprobe_threshold, specify cpu to search!",
"search", 0);
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFPQPass: nq >= gpu_search_threshold, specify gpu %d to search!",
"search", 0, search_gpus_[idx_]);
res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, search_gpus_[idx_]);
idx_ = (idx_ + 1) % search_gpus_.size();
}
auto label = std::make_shared<SpecResLabel>(res_ptr);
task->label() = label;
return true;
}
} // namespace scheduler
} // namespace milvus
#endif

32
core/src/scheduler/selector/FaissIVFFlatPass.cpp → core/src/scheduler/selector/FaissIVFPass.cpp
View File

@ -9,7 +9,7 @@
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.
#ifdef MILVUS_GPU_VERSION
#include "scheduler/selector/FaissIVFFlatPass.h"
#include "scheduler/selector/FaissIVFPass.h"
#include "cache/GpuCacheMgr.h"
#include "config/Config.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
@ -24,8 +24,7 @@ namespace milvus {
namespace scheduler {
void
FaissIVFFlatPass::Init() {
#ifdef MILVUS_GPU_VERSION
FaissIVFPass::Init() {
server::Config& config = server::Config::GetInstance();
Status s = config.GetGpuResourceConfigGpuSearchThreshold(threshold_);
if (!s.ok()) {
@ -36,40 +35,45 @@ FaissIVFFlatPass::Init() {
throw std::exception();
}
SetIdentity("FaissIVFFlatPass");
SetIdentity("FaissIVFPass");
AddGpuEnableListener();
AddGpuSearchThresholdListener();
AddGpuSearchResourcesListener();
#endif
}
bool
FaissIVFFlatPass::Run(const TaskPtr& task) {
FaissIVFPass::Run(const TaskPtr& task) {
if (task->Type() != TaskType::SearchTask) {
return false;
}
auto search_task = std::static_pointer_cast<XSearchTask>(task);
if (search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFFLAT) {
return false;
switch (static_cast<engine::EngineType>(search_task->file_->engine_type_)) {
case engine::EngineType::FAISS_IVFFLAT:
case engine::EngineType::FAISS_PQ:
case engine::EngineType::FAISS_IVFSQ8:
case engine::EngineType::FAISS_IVFSQ8H:
break;
default:
return false;
}
auto search_job = std::static_pointer_cast<SearchJob>(search_task->job_.lock());
ResourcePtr res_ptr;
if (!gpu_enable_) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFFlatPass: gpu disable, specify cpu to search!", "search", 0);
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFPass: gpu disable, specify cpu to search!", "search", 0);
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else if (search_job->nq() < (uint64_t)threshold_) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFFlatPass: nq < gpu_search_threshold, specify cpu to search!",
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFPass: nq < gpu_search_threshold, specify cpu to search!",
"search", 0);
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else if (search_job->extra_params()[knowhere::IndexParams::nprobe].get<int64_t>() <
milvus::server::QUERY_MAX_TOPK) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFFlatPass: nprobe > gpu_nprobe_threshold, specify cpu to search!",
} else if (search_job->extra_params()[knowhere::IndexParams::nprobe].get<int64_t>() >
milvus::server::GPU_QUERY_MAX_NPROBE) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFPass: nprobe > gpu_nprobe_threshold, specify cpu to search!",
"search", 0);
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFFlatPass: nq >= gpu_search_threshold, specify gpu %d to search!",
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFPass: nq >= gpu_search_threshold, specify gpu %d to search!",
"search", 0, search_gpus_[idx_]);
res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, search_gpus_[idx_]);
idx_ = (idx_ + 1) % search_gpus_.size();

7
core/src/scheduler/selector/FaissIVFPQPass.h → core/src/scheduler/selector/FaissIVFPass.h
View File

@ -29,14 +29,15 @@
namespace milvus {
namespace scheduler {
class FaissIVFPQPass : public Pass, public server::GpuResourceConfigHandler {
class FaissIVFPass : public Pass, public server::GpuResourceConfigHandler {
public:
FaissIVFPQPass() = default;
FaissIVFPass() = default;
public:
void
Init() override;
public:
bool
Run(const TaskPtr& task) override;
@ -44,7 +45,7 @@ class FaissIVFPQPass : public Pass, public server::GpuResourceConfigHandler {
int64_t idx_ = 0;
};
using FaissIVFPQPassPtr = std::shared_ptr<FaissIVFPQPass>;
using FaissIVFPassPtr = std::shared_ptr<FaissIVFPass>;
} // namespace scheduler
} // namespace milvus

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

@ -1,84 +0,0 @@
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.
#ifdef MILVUS_GPU_VERSION
#include "scheduler/selector/FaissIVFSQ8HPass.h"
#include "cache/GpuCacheMgr.h"
#include "config/Config.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "scheduler/SchedInst.h"
#include "scheduler/Utils.h"
#include "scheduler/task/SearchTask.h"
#include "scheduler/tasklabel/SpecResLabel.h"
#include "utils/Log.h"
#include "utils/ValidationUtil.h"
namespace milvus {
namespace scheduler {
void
FaissIVFSQ8HPass::Init() {
server::Config& config = server::Config::GetInstance();
Status s = config.GetGpuResourceConfigGpuSearchThreshold(threshold_);
if (!s.ok()) {
threshold_ = std::numeric_limits<int64_t>::max();
}
s = config.GetGpuResourceConfigSearchResources(search_gpus_);
if (!s.ok()) {
throw std::exception();
}
SetIdentity("FaissIVFSQ8HPass");
AddGpuEnableListener();
AddGpuSearchThresholdListener();
AddGpuSearchResourcesListener();
}
bool
FaissIVFSQ8HPass::Run(const TaskPtr& task) {
if (task->Type() != TaskType::SearchTask) {
return false;
}
auto search_task = std::static_pointer_cast<XSearchTask>(task);
if (search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFSQ8H) {
return false;
}
auto search_job = std::static_pointer_cast<SearchJob>(search_task->job_.lock());
ResourcePtr res_ptr;
if (!gpu_enable_) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFSQ8HPass: gpu disable, specify cpu to search!", "search", 0);
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
}
if (search_job->nq() < (uint64_t)threshold_) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFSQ8HPass: nq < gpu_search_threshold, specify cpu to search!",
"search", 0);
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else if (search_job->extra_params()[knowhere::IndexParams::nprobe].get<int64_t>() <
milvus::server::QUERY_MAX_TOPK) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFFlatPass: nprobe > gpu_nprobe_threshold, specify cpu to search!",
"search", 0);
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFSQ8HPass: nq >= gpu_search_threshold, specify gpu %d to search!",
"search", 0, search_gpus_[idx_]);
res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, search_gpus_[idx_]);
idx_ = (idx_ + 1) % search_gpus_.size();
}
auto label = std::make_shared<SpecResLabel>(res_ptr);
task->label() = label;
return true;
}
} // namespace scheduler
} // namespace milvus
#endif

51
core/src/scheduler/selector/FaissIVFSQ8HPass.h
View File

@ -1,51 +0,0 @@
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.
#ifdef MILVUS_GPU_VERSION
#pragma once
#include <condition_variable>
#include <deque>
#include <limits>
#include <list>
#include <memory>
#include <mutex>
#include <queue>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
#include "config/handler/GpuResourceConfigHandler.h"
#include "scheduler/selector/Pass.h"
namespace milvus {
namespace scheduler {
class FaissIVFSQ8HPass : public Pass, public server::GpuResourceConfigHandler {
public:
FaissIVFSQ8HPass() = default;
public:
void
Init() override;
bool
Run(const TaskPtr& task) override;
private:
int64_t idx_ = 0;
};
using FaissIVFSQ8HPassPtr = std::shared_ptr<FaissIVFSQ8HPass>;
} // namespace scheduler
} // namespace milvus
#endif

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

@ -1,84 +0,0 @@
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.
#ifdef MILVUS_GPU_VERSION
#include "scheduler/selector/FaissIVFSQ8Pass.h"
#include "cache/GpuCacheMgr.h"
#include "config/Config.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "scheduler/SchedInst.h"
#include "scheduler/Utils.h"
#include "scheduler/task/SearchTask.h"
#include "scheduler/tasklabel/SpecResLabel.h"
#include "utils/Log.h"
#include "utils/ValidationUtil.h"
namespace milvus {
namespace scheduler {
void
FaissIVFSQ8Pass::Init() {
#ifdef MILVUS_GPU_VERSION
server::Config& config = server::Config::GetInstance();
Status s = config.GetGpuResourceConfigGpuSearchThreshold(threshold_);
if (!s.ok()) {
threshold_ = std::numeric_limits<int32_t>::max();
}
s = config.GetGpuResourceConfigSearchResources(search_gpus_);
if (!s.ok()) {
throw std::exception();
}
SetIdentity("FaissIVFSQ8Pass");
AddGpuEnableListener();
AddGpuSearchThresholdListener();
AddGpuSearchResourcesListener();
#endif
}
bool
FaissIVFSQ8Pass::Run(const TaskPtr& task) {
if (task->Type() != TaskType::SearchTask) {
return false;
}
auto search_task = std::static_pointer_cast<XSearchTask>(task);
if (search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFSQ8) {
return false;
}
auto search_job = std::static_pointer_cast<SearchJob>(search_task->job_.lock());
ResourcePtr res_ptr;
if (!gpu_enable_) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFSQ8Pass: gpu disable, specify cpu to search!", "search", 0);
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else if (search_job->nq() < (uint64_t)threshold_) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFSQ8Pass: nq < gpu_search_threshold, specify cpu to search!",
"search", 0);
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else if (search_job->extra_params()[knowhere::IndexParams::nprobe].get<int64_t>() <
milvus::server::QUERY_MAX_TOPK) {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFFlatPass: nprobe > gpu_nprobe_threshold, specify cpu to search!",
"search", 0);
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else {
LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFSQ8Pass: nq >= gpu_search_threshold, specify gpu %d to search!",
"search", 0, search_gpus_[idx_]);
res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, search_gpus_[idx_]);
idx_ = (idx_ + 1) % search_gpus_.size();
}
auto label = std::make_shared<SpecResLabel>(res_ptr);
task->label() = label;
return true;
}
} // namespace scheduler
} // namespace milvus
#endif

51
core/src/scheduler/selector/FaissIVFSQ8Pass.h
View File

@ -1,51 +0,0 @@
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.
#ifdef MILVUS_GPU_VERSION
#pragma once
#include <condition_variable>
#include <deque>
#include <limits>
#include <list>
#include <memory>
#include <mutex>
#include <queue>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
#include "config/handler/GpuResourceConfigHandler.h"
#include "scheduler/selector/Pass.h"
namespace milvus {
namespace scheduler {
class FaissIVFSQ8Pass : public Pass, public server::GpuResourceConfigHandler {
public:
FaissIVFSQ8Pass() = default;
public:
void
Init() override;
bool
Run(const TaskPtr& task) override;
private:
int64_t idx_ = 0;
};
using FaissIVFSQ8PassPtr = std::shared_ptr<FaissIVFSQ8Pass>;
} // namespace scheduler
} // namespace milvus
#endif

1
core/src/utils/ValidationUtil.h
View File

@ -23,6 +23,7 @@ namespace milvus {
namespace server {
constexpr int64_t QUERY_MAX_TOPK = 2048;
constexpr int64_t GPU_QUERY_MAX_NPROBE = 2048;
class ValidationUtil {
private:

5
core/unittest/scheduler/test_action.cpp
View File

@ -13,13 +13,10 @@
#include <gtest/gtest.h>
#include <src/scheduler/task/BuildIndexTask.h>
#include <src/scheduler/task/SearchTask.h>
#include <src/scheduler/optimizer/FaissIVFFlatPass.h>
#include "scheduler/optimizer/BuildIndexPass.h"
#include "scheduler/optimizer/FaissFlatPass.h"
#include "scheduler/optimizer/FaissIVFPQPass.h"
#include "scheduler/optimizer/FaissIVFSQ8HPass.h"
#include "scheduler/optimizer/FaissIVFSQ8Pass.h"
#include "scheduler/optimizer/FaissIVFPass.h"
namespace milvus {
namespace scheduler {

39
core/unittest/scheduler/test_selector.cpp
View File

@ -18,10 +18,7 @@
#include "scheduler/resource/CpuResource.h"
#include "scheduler/selector/BuildIndexPass.h"
#include "scheduler/selector/FaissFlatPass.h"
#include "scheduler/selector/FaissIVFFlatPass.h"
#include "scheduler/selector/FaissIVFPQPass.h"
#include "scheduler/selector/FaissIVFSQ8HPass.h"
#include "scheduler/selector/FaissIVFSQ8Pass.h"
#include "scheduler/selector/FaissIVFPass.h"
#include "scheduler/selector/FallbackPass.h"
namespace milvus {
@ -47,42 +44,20 @@ TEST(OptimizerTest, TEST_OPTIMIZER) {
fiu_disable("get_gpu_config_search_resources.disable_gpu_resource_fail");
fiu_disable("check_config_gpu_search_threshold_fail");
FaissIVFFlatPass faiss_ivf_flat_pass;
FaissIVFPass faiss_ivf_pass;
fiu_enable("check_config_gpu_search_threshold_fail", 1, NULL, 0);
fiu_enable("get_gpu_config_search_resources.disable_gpu_resource_fail", 1, NULL, 0);
ASSERT_ANY_THROW(faiss_ivf_flat_pass.Init(););
fiu_disable("get_gpu_config_search_resources.disable_gpu_resource_fail");
fiu_disable("check_config_gpu_search_threshold_fail");
FaissIVFPQPass faiss_ivf_pq_pass;
fiu_enable("check_config_gpu_search_threshold_fail", 1, NULL, 0);
fiu_enable("get_gpu_config_search_resources.disable_gpu_resource_fail", 1, NULL, 0);
ASSERT_ANY_THROW(faiss_ivf_pq_pass.Init(););
ASSERT_ANY_THROW(faiss_ivf_pass.Init(););
fiu_disable("get_gpu_config_search_resources.disable_gpu_resource_fail");
fiu_disable("check_config_gpu_search_threshold_fail");
auto file = std::make_shared<SegmentSchema>();
file->engine_type_ = (int)engine::EngineType::FAISS_IVFFLAT;
file->index_params_ = "{ \"nlist\": 100 }";
file->engine_type_ = (int)engine::EngineType::FAISS_IDMAP;
file->index_params_ = "";
file->dimension_ = 64;
auto search_task = std::make_shared<XSearchTask>(nullptr, file, nullptr);
ASSERT_FALSE(faiss_ivf_pq_pass.Run(search_task));
FaissIVFSQ8HPass faiss_ivf_q8h_pass;
fiu_enable("check_config_gpu_search_threshold_fail", 1, NULL, 0);
faiss_ivf_q8h_pass.Init();
fiu_disable("check_config_gpu_search_threshold_fail");
auto search_task2 = std::make_shared<XSearchTask>(nullptr, file, nullptr);
ASSERT_FALSE(faiss_ivf_q8h_pass.Run(build_index_task));
ASSERT_FALSE(faiss_ivf_q8h_pass.Run(search_task2));
FaissIVFSQ8Pass faiss_ivf_q8_pass;
fiu_enable("check_config_gpu_search_threshold_fail", 1, NULL, 0);
fiu_enable("get_gpu_config_search_resources.disable_gpu_resource_fail", 1, NULL, 0);
ASSERT_ANY_THROW(faiss_ivf_q8_pass.Init(););
fiu_disable("get_gpu_config_search_resources.disable_gpu_resource_fail");
fiu_disable("check_config_gpu_search_threshold_fail");
ASSERT_FALSE(faiss_ivf_pass.Run(search_task));
ASSERT_FALSE(faiss_ivf_pass.Run(build_index_task));
FallbackPass fall_back_pass;
fall_back_pass.Init();

6
tests/milvus_python_test/entity/test_search.py
View File

@ -1223,7 +1223,11 @@ class TestSearchParamsInvalid(object):
search_param = {"nprobe": nprobe}
query_vecs = gen_vectors(nprobe, dim)
status, result = connect.search(collection, top_k, query_vecs, params=search_param)
assert not status.OK()
# IVFSQ8H supported later
if index["index_type"] == IndexType.IVF_SQ8H:
assert not status.OK()
else:
assert status.OK()
@pytest.fixture(
scope="function",