Disable query cache while tuning (#903)

Signed-off-by: chasingegg <[email protected]>
milvus-io · Jun 16, 2023 · 1822281 · 1822281
1 parent da16880
commit 1822281
Show file tree

Hide file tree

Showing 8 changed files with 15 additions and 7 deletions.
diff --git a/include/knowhere/config.h b/include/knowhere/config.h
@@ -436,6 +436,7 @@ class BaseConfig : public Config {
     CFG_FLOAT range_filter;
     CFG_BOOL trace_visit;
     CFG_BOOL enable_mmap;
+    CFG_BOOL for_tuning;
     KNOHWERE_DECLARE_CONFIG(BaseConfig) {
         KNOWHERE_CONFIG_DECLARE_FIELD(metric_type).set_default("L2").description("metric type").for_train_and_search();
         KNOWHERE_CONFIG_DECLARE_FIELD(k)
@@ -464,6 +465,7 @@ class BaseConfig : public Config {
             .set_default(false)
             .description("enable mmap for load index")
             .for_deserialize_from_file();
+        KNOWHERE_CONFIG_DECLARE_FIELD(for_tuning).set_default(false).description("for tuning").for_search();
     }
 
     int

diff --git a/src/common/config.cc b/src/common/config.cc
@@ -46,7 +46,8 @@ static const std::unordered_set<std::string> ext_legal_json_keys = {"metric_type
                                                                     "gpu_id",
                                                                     "num_threads",
                                                                     "round_decimal",
-                                                                    "offset"};
+                                                                    "offset",
+                                                                    "for_tuning"};
 
 Status
 Config::FormatAndCheck(const Config& cfg, Json& json) {

diff --git a/src/index/diskann/diskann.cc b/src/index/diskann/diskann.cc
@@ -534,6 +534,7 @@ DiskANNIndexNode<T>::Search(const DataSet& dataset, const Config& cfg, const Bit
     auto lsearch = static_cast<uint64_t>(search_conf.search_list_size);
     auto beamwidth = static_cast<uint64_t>(search_conf.beamwidth);
     auto filter_ratio = static_cast<float>(search_conf.filter_threshold);
+    auto for_tuning = static_cast<bool>(search_conf.for_tuning);
 
     auto nq = dataset.GetRows();
     auto dim = dataset.GetDim();
@@ -558,7 +559,7 @@ DiskANNIndexNode<T>::Search(const DataSet& dataset, const Config& cfg, const Bit
         futures.push_back(pool_->push([&, index = row]() {
             pq_flash_index_->cached_beam_search(xq + (index * dim), k, lsearch, p_id + (index * k),
                                                 p_dist + (index * k), beamwidth, false, nullptr, feder_result, bitset,
-                                                filter_ratio);
+                                                filter_ratio, for_tuning);
         }));
     }
     for (auto& future : futures) {

diff --git a/src/index/hnsw/hnsw.cc b/src/index/hnsw/hnsw.cc
@@ -130,7 +130,7 @@ class HnswIndexNode : public IndexNode {
         auto p_id = new int64_t[k * nq];
         auto p_dist = new float[k * nq];
 
-        hnswlib::SearchParam param{(size_t)hnsw_cfg.ef};
+        hnswlib::SearchParam param{(size_t)hnsw_cfg.ef, hnsw_cfg.for_tuning};
         bool transform =
             (index_->metric_type_ == hnswlib::Metric::INNER_PRODUCT || index_->metric_type_ == hnswlib::Metric::COSINE);
 

diff --git a/thirdparty/DiskANN/include/diskann/pq_flash_index.h b/thirdparty/DiskANN/include/diskann/pq_flash_index.h
@@ -101,7 +101,8 @@ namespace diskann {
         const bool use_reorder_data = false, QueryStats *stats = nullptr,
         const knowhere::feder::diskann::FederResultUniq &feder = nullptr,
         knowhere::BitsetView                             bitset_view = nullptr,
-        const float                                      filter_ratio = -1.0f);
+        const float                                      filter_ratio = -1.0f,
+        const bool                                       for_tuning = false);
 
     DISKANN_DLLEXPORT _u32 range_search(
         const T *query1, const double range, const _u64 min_l_search,

diff --git a/thirdparty/DiskANN/src/pq_flash_index.cpp b/thirdparty/DiskANN/src/pq_flash_index.cpp
@@ -976,7 +976,7 @@ namespace diskann {
       const T *query1, const _u64 k_search, const _u64 l_search, _s64 *indices,
       float *distances, const _u64 beam_width, const bool use_reorder_data,
       QueryStats *stats, const knowhere::feder::diskann::FederResultUniq &feder,
-      knowhere::BitsetView bitset_view, const float filter_ratio_in) {
+      knowhere::BitsetView bitset_view, const float filter_ratio_in, const bool for_tuning) {
     if (beam_width > MAX_N_SECTOR_READS)
       throw ANNException("Beamwidth can not be higher than MAX_N_SECTOR_READS",
                          -1, __FUNCSIG__, __FILE__, __LINE__);
@@ -1068,7 +1068,8 @@ namespace diskann {
     full_retset.reserve(4096);
     auto vec_hash = knowhere::hash_vec(query_float, data_dim);
     _u32 best_medoid = 0;
-    if (!lru_cache.try_get(vec_hash, best_medoid)) {
+    // for tuning, do not use cache
+    if (for_tuning || !lru_cache.try_get(vec_hash, best_medoid)) {
       float best_dist = (std::numeric_limits<float>::max)();
       std::vector<SimpleNeighbor> medoid_dists;
       for (_u64 cur_m = 0; cur_m < num_medoids; cur_m++) {

diff --git a/thirdparty/hnswlib/hnswlib/hnswalg.h b/thirdparty/hnswlib/hnswlib/hnswalg.h
@@ -1144,7 +1144,8 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
 
         tableint currObj = enterpoint_node_;
         auto vec_hash = knowhere::hash_vec((const float*)query_data, *(size_t*)dist_func_param_);
-        if (!lru_cache.try_get(vec_hash, currObj)) {
+        // for tuning, do not use cache
+        if (param->for_tuning || !lru_cache.try_get(vec_hash, currObj)) {
             dist_t curdist = calcDistance(query_data, enterpoint_node_);
 
             for (int level = maxlevel_; level > 0; level--) {

diff --git a/thirdparty/hnswlib/hnswlib/hnswlib.h b/thirdparty/hnswlib/hnswlib/hnswlib.h
@@ -192,6 +192,7 @@ class SpaceInterface {
 
 struct SearchParam {
     size_t ef_;
+    bool for_tuning;
 };
 
 template <typename dist_t>