Merge pull request #151 from fishpenguin/branch-0.5.1-yk

Add only GPU and only CPU version for IVF_SQ8 and IVF_FLAT

Former-commit-id: bdca3006bc7b2c4456b1874d8c5eb23e67053b16

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

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 <memory>
 #include <mutex>
+#include <string>
 #include <vector>
 
 namespace milvus {
@@ -93,9 +97,21 @@ class OptimizerInst {
         if (instance == nullptr) {
             std::lock_guard<std::mutex> lock(mutex_);
             if (instance == nullptr) {
+                server::Config& config = server::Config::GetInstance();
+                std::vector<std::string> search_resources;
+                bool has_cpu = false;
+                config.GetResourceConfigSearchResources(search_resources);
+                for (auto& resource : search_resources) {
+                    if (resource == "cpu") {
+                        has_cpu = true;
+                    }
+                }
+
                 std::vector<PassPtr> pass_list;
                 pass_list.push_back(std::make_shared<LargeSQ8HPass>());
                 pass_list.push_back(std::make_shared<HybridPass>());
+                pass_list.push_back(std::make_shared<OnlyCPUPass>());
+                pass_list.push_back(std::make_shared<OnlyGPUPass>(has_cpu));
                 instance = std::make_shared<Optimizer>(pass_list);
             }
         }

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<XSearchTask>(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<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));
+    task->label() = label;
+    return true;
+}
+
+}  // namespace scheduler
+}  // namespace milvus

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 <condition_variable>
+#include <deque>
+#include <list>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <string>
+#include <thread>
+#include <unordered_map>
+#include <vector>
+
+#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<OnlyCPUPass>;
+
+}  // namespace scheduler
+}  // namespace milvus

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<XSearchTask>(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<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));
+    task->label() = label;
+
+    specified_gpu_id_ = specified_gpu_id_++ % gpu_id.size();
+    return true;
+}
+
+}  // namespace scheduler
+}  // namespace milvus

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 <condition_variable>
+#include <deque>
+#include <list>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <string>
+#include <thread>
+#include <unordered_map>
+#include <vector>
+
+#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<OnlyGPUPass>;
+
+}  // namespace scheduler
+}  // namespace milvus