diff --git a/CHANGELOG.md b/CHANGELOG.md index 5176ce2b79..f277ec2a59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ Please mark all change in change log and use the ticket from JIRA. - \#104 - test_scheduler core dump - \#115 - Using new structure for tasktable - \#139 - New config opion use_gpu_threshold +- \#146 - Add only GPU and only CPU version for IVF_SQ8 and IVF_FLAT ## Improvement - \#64 - Improvement dump function in scheduler diff --git a/core/src/scheduler/SchedInst.h b/core/src/scheduler/SchedInst.h index b9153d3bc3..e758f37851 100644 --- a/core/src/scheduler/SchedInst.h +++ b/core/src/scheduler/SchedInst.h @@ -23,6 +23,8 @@ #include "Scheduler.h" #include "optimizer/HybridPass.h" #include "optimizer/LargeSQ8HPass.h" +#include "optimizer/OnlyCPUPass.h" +#include "optimizer/OnlyGPUPass.h" #include "optimizer/Optimizer.h" #include @@ -96,6 +98,8 @@ class OptimizerInst { std::vector pass_list; pass_list.push_back(std::make_shared()); pass_list.push_back(std::make_shared()); + pass_list.push_back(std::make_shared()); + pass_list.push_back(std::make_shared()); instance = std::make_shared(pass_list); } } diff --git a/core/src/scheduler/optimizer/OnlyCPUPass.cpp b/core/src/scheduler/optimizer/OnlyCPUPass.cpp new file mode 100644 index 0000000000..d974887bb7 --- /dev/null +++ b/core/src/scheduler/optimizer/OnlyCPUPass.cpp @@ -0,0 +1,48 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +#include "scheduler/optimizer/OnlyCPUPass.h" +#include "scheduler/SchedInst.h" +#include "scheduler/task/SearchTask.h" +#include "scheduler/tasklabel/SpecResLabel.h" +#include "scheduler/Utils.h" + +namespace milvus { +namespace scheduler { + +bool +OnlyCPUPass::Run(const TaskPtr& task) { + if (task->Type() != TaskType::SearchTask) + return false; + auto search_task = std::static_pointer_cast(task); + if (search_task->file_->engine_type_ != (int) engine::EngineType::FAISS_IVFSQ8 && + search_task->file_->engine_type_ != (int) engine::EngineType::FAISS_IVFFLAT) { + return false; + } + + auto gpu_id = get_gpu_pool(); + if (gpu_id.empty()) { + ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); + auto label = std::make_shared(std::weak_ptr(res_ptr)); + task->label() = label; + return true; + } + return false; +} + +} // namespace scheduler +} // namespace milvus diff --git a/core/src/scheduler/optimizer/OnlyCPUPass.h b/core/src/scheduler/optimizer/OnlyCPUPass.h new file mode 100644 index 0000000000..edc3848ed4 --- /dev/null +++ b/core/src/scheduler/optimizer/OnlyCPUPass.h @@ -0,0 +1,47 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "Pass.h" + +namespace milvus { +namespace scheduler { + +class OnlyCPUPass : public Pass { +public: + OnlyCPUPass() = default; + +public: + bool + Run(const TaskPtr &task) override; +}; + +using OnlyCPUPassPtr = std::shared_ptr; + +} // namespace scheduler +} // namespace milvus diff --git a/core/src/scheduler/optimizer/OnlyGPUPass.cpp b/core/src/scheduler/optimizer/OnlyGPUPass.cpp new file mode 100644 index 0000000000..20127aefb8 --- /dev/null +++ b/core/src/scheduler/optimizer/OnlyGPUPass.cpp @@ -0,0 +1,66 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +#include "scheduler/optimizer/OnlyGPUPass.h" +#include "scheduler/SchedInst.h" +#include "scheduler/task/SearchTask.h" +#include "scheduler/tasklabel/SpecResLabel.h" +#include "scheduler/Utils.h" +#include "server/Config.h" + +namespace milvus { +namespace scheduler { + +bool +OnlyGPUPass::Run(const TaskPtr& task) { + if (task->Type() != TaskType::SearchTask) + return false; + + auto search_task = std::static_pointer_cast(task); + if (search_task->file_->engine_type_ != (int) engine::EngineType::FAISS_IVFSQ8 && + search_task->file_->engine_type_ != (int) engine::EngineType::FAISS_IVFFLAT) { + return false; + } + + server::Config& config = server::Config::GetInstance(); + std::vector search_resources; + config.GetResourceConfigSearchResources(search_resources); + for (auto &resource : search_resources) { + if (resource == "cpu") { + return false; + } + } + + auto gpu_id = get_gpu_pool(); + if (!gpu_id.empty()) { + ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, gpu_id[specified_gpu_id_]); + auto label = std::make_shared(std::weak_ptr(res_ptr)); + task->label() = label; + } else { + return false; + } + + if (specified_gpu_id_ < gpu_id.size() - 1) { + ++specified_gpu_id_; + } else { + specified_gpu_id_ = 0; + } + return true; +} + +} // namespace scheduler +} // namespace milvus diff --git a/core/src/scheduler/optimizer/OnlyGPUPass.h b/core/src/scheduler/optimizer/OnlyGPUPass.h new file mode 100644 index 0000000000..9220bab11d --- /dev/null +++ b/core/src/scheduler/optimizer/OnlyGPUPass.h @@ -0,0 +1,50 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "Pass.h" + +namespace milvus { +namespace scheduler { + +class OnlyGPUPass : public Pass { +public: + OnlyGPUPass() = default; + +public: + bool + Run(const TaskPtr &task) override; + +private: + uint64_t specified_gpu_id_ = 0; +}; + +using OnlyGPUPassPtr = std::shared_ptr; + +} // namespace scheduler +} // namespace milvus