Add only GPU and only CPU version for IVF_SQ8 and IVF_FLAT

Former-commit-id: 14f5b095df2988d5f65ecd1cb2628fd3a3791a14

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

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 <memory>
@@ -96,6 +98,8 @@ class OptimizerInst {
                 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>());
                 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/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<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()) {
+        ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
+        auto label = std::make_shared<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));
+        task->label() = label;
+        return true;
+    }
+    return false;
+}
+
+}  // 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,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<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;
+    }
+
+    server::Config& config = server::Config::GetInstance();
+    std::vector<std::string> 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<SpecResLabel>(std::weak_ptr<Resource>(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

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 <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:
+    OnlyGPUPass() = default;
+
+public:
+    bool
+    Run(const TaskPtr &task) override;
+
+private:
+    uint64_t specified_gpu_id_ = 0;
+};
+
+using OnlyGPUPassPtr = std::shared_ptr<OnlyGPUPass>;
+
+}  // namespace scheduler
+}  // namespace milvus