When the topk of GPU search larger than 2048,turning into the CPU search (#3647)

* topk of GPU searching

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

* change the changelog

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

Co-authored-by: shengjun.li <shengjun.li@zilliz.com>

CHANGELOG.md

@@ -11,6 +11,7 @@ Please mark all change in change log and use the issue from GitHub
 -   \#3213 Allow users to specify a distance type at runtime for Flat index
 -   \#3254 Allow more choices for the parameter ‘m’ of CPU IVF_PQ
 -   \#3606 Searching auto turn into CPU when the parameter 'nprobe' of IVF is larger than 2048
+-   \#3639 Searching auto turn into CPU when the parameter 'topk' of index is larger than 2048 
 
 ## Improvement
 

core/src/scheduler/selector/FaissFlatPass.cpp

@@ -17,6 +17,7 @@
 #include "scheduler/task/SearchTask.h"
 #include "scheduler/tasklabel/SpecResLabel.h"
 #include "utils/Log.h"
+#include "utils/ValidationUtil.h"
 
 #include <fiu-local.h>
 namespace milvus {
@@ -60,6 +61,10 @@ FaissFlatPass::Run(const TaskPtr& task) {
         LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissFlatPass: nq < gpu_search_threshold, specify cpu to search!",
                                     "search", 0);
         res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
+    } else if (search_job->topk() > milvus::server::GPU_QUERY_MAX_TOPK) {
+        LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissFlatPass: topk > gpu_topk_threshold, specify cpu to search!",
+                                    "search", 0);
+        res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
     } else {
         LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissFlatPass: nq >= gpu_search_threshold, specify gpu %d to search!",
                                     "search", 0, search_gpus_[idx_]);

core/src/scheduler/selector/FaissIVFPass.cpp

@@ -67,6 +67,10 @@ FaissIVFPass::Run(const TaskPtr& task) {
         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->topk() > milvus::server::GPU_QUERY_MAX_TOPK) {
+        LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFPass: topk > gpu_topk_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::GPU_QUERY_MAX_NPROBE) {
         LOG_SERVER_DEBUG_ << LogOut("[%s][%d] FaissIVFPass: nprobe > gpu_nprobe_threshold, specify cpu to search!",

core/src/utils/ValidationUtil.cpp

@@ -407,7 +407,7 @@ Status
 ValidationUtil::ValidateSearchTopk(int64_t top_k) {
     if (top_k <= 0 || top_k > QUERY_MAX_TOPK) {
         std::string msg =
-            "Invalid topk: " + std::to_string(top_k) + ". " + "The topk must be within the range of 1 ~ 2048.";
+            "Invalid topk: " + std::to_string(top_k) + ". " + "The topk must be within the range of 1 ~ 16384.";
         LOG_SERVER_ERROR_ << msg;
         return Status(SERVER_INVALID_TOPK, msg);
     }

core/src/utils/ValidationUtil.h

@@ -22,7 +22,8 @@
 namespace milvus {
 namespace server {
 
-constexpr int64_t QUERY_MAX_TOPK = 2048;
+constexpr int64_t QUERY_MAX_TOPK = 16384;
+constexpr int64_t GPU_QUERY_MAX_TOPK = 2048;
 constexpr int64_t GPU_QUERY_MAX_NPROBE = 2048;
 
 class ValidationUtil {

sdk/examples/qps/src/ClientTest.cpp

@@ -109,7 +109,7 @@ ClientTest::CheckParameters(const TestParameters& parameters) {
         return false;
     }
 
-    if (parameters.topk_ <= 0 || parameters.topk_ > 2048) {
+    if (parameters.topk_ <= 0 || parameters.topk_ > 16384) {
         std::cout << "Invalid query topk: " << parameters.topk_ << std::endl;
         return false;
     }

tests/milvus_python_test/entity/test_search.py

@@ -147,7 +147,7 @@ class TestSearchBase:
         query_vec = [vectors[0]]
         top_k = get_top_k
         status, result = connect.search(collection, top_k, query_vec)
-        if top_k <= 2048:
+        if top_k <= 16384:
             assert status.OK()
             assert len(result[0]) == min(len(vectors), top_k)
             assert result[0][0].distance <= epsilon