diff --git a/CHANGELOG.md b/CHANGELOG.md index e2ed8f4b7f2d134dc7df9dbedee0fce100b4d8bc..266690a82cb01076f24d8780629b3e49eea3beb1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,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 b9153d3bc3f3805589d461568e8e757b26fd48cb..a3048069f9febe9ec7ea9b5112b3b7b6fab537f9 100644 --- a/core/src/scheduler/SchedInst.h +++ b/core/src/scheduler/SchedInst.h @@ -23,10 +23,14 @@ #include "Scheduler.h" #include "optimizer/HybridPass.h" #include "optimizer/LargeSQ8HPass.h" +#include "optimizer/OnlyCPUPass.h" +#include "optimizer/OnlyGPUPass.h" #include "optimizer/Optimizer.h" +#include "server/Config.h" #include #include +#include #include namespace milvus { @@ -93,9 +97,21 @@ class OptimizerInst { if (instance == nullptr) { std::lock_guard lock(mutex_); if (instance == nullptr) { + server::Config& config = server::Config::GetInstance(); + std::vector search_resources; + bool has_cpu = false; + config.GetResourceConfigSearchResources(search_resources); + for (auto& resource : search_resources) { + if (resource == "cpu") { + has_cpu = true; + } + } + 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(has_cpu)); 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 0000000000000000000000000000000000000000..238a91a82c77c05fe74576e2fedf33123214dbfc --- /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/Utils.h" +#include "scheduler/task/SearchTask.h" +#include "scheduler/tasklabel/SpecResLabel.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 (not gpu_id.empty()) + return false; + + ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); + auto label = std::make_shared(std::weak_ptr(res_ptr)); + task->label() = label; + return true; +} + +} // 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 0000000000000000000000000000000000000000..76b42e376609977b6e6b4cf852e32f93fbb8f648 --- /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 0000000000000000000000000000000000000000..3fcda0e8a347bc6ca62890c78b394314f426c9e1 --- /dev/null +++ b/core/src/scheduler/optimizer/OnlyGPUPass.cpp @@ -0,0 +1,54 @@ +// 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/Utils.h" +#include "scheduler/task/SearchTask.h" +#include "scheduler/tasklabel/SpecResLabel.h" + +namespace milvus { +namespace scheduler { + +OnlyGPUPass::OnlyGPUPass(bool has_cpu) : has_cpu_(has_cpu) { +} + +bool +OnlyGPUPass::Run(const TaskPtr& task) { + if (task->Type() != TaskType::SearchTask || has_cpu_) + 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()) + return false; + + 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; + + specified_gpu_id_ = specified_gpu_id_++ % gpu_id.size(); + 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 0000000000000000000000000000000000000000..10d909d30e1888d5025df8a80ee899bc348eafc9 --- /dev/null +++ b/core/src/scheduler/optimizer/OnlyGPUPass.h @@ -0,0 +1,51 @@ +// 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: + explicit OnlyGPUPass(bool has_cpu); + + public: + bool + Run(const TaskPtr& task) override; + + private: + uint64_t specified_gpu_id_ = 0; + bool has_cpu_ = false; +}; + +using OnlyGPUPassPtr = std::shared_ptr; + +} // namespace scheduler +} // namespace milvus