diff --git a/knowhere/index/vector_index/IndexDiskANNConfig.cpp b/knowhere/index/vector_index/IndexDiskANNConfig.cpp index 1d2624355..2a934adf8 100644 --- a/knowhere/index/vector_index/IndexDiskANNConfig.cpp +++ b/knowhere/index/vector_index/IndexDiskANNConfig.cpp @@ -11,6 +11,7 @@ #include "knowhere/index/vector_index/IndexDiskANNConfig.h" +#include #include #include #include @@ -59,6 +60,7 @@ static constexpr uint32_t kBuildNumThreadsMinValue = 1; static constexpr uint32_t kBuildNumThreadsMaxValue = 128; static constexpr uint32_t kDiskPqBytesMinValue = 0; static constexpr std::optional kDiskPqBytesMaxValue = std::nullopt; +static constexpr uint32_t kSearchListSizeMaxValue = 200; static constexpr uint32_t kBeamwidthMinValue = 1; static constexpr uint32_t kBeamwidthMaxValue = 128; static constexpr uint64_t kKMinValue = 1; @@ -200,8 +202,8 @@ from_json(const Config& config, DiskANNPrepareConfig& prep_conf) { auto num_thread_max_value = kLinuxAioMaxnrLimit / prep_conf.aio_maxnr; CheckNumericParamAndSet(config, kNumThreads, kSearchNumThreadsMinValue, num_thread_max_value, prep_conf.num_threads); - CheckNumericParamAndSet(config, kCacheDramBudgetGb, kCacheDramBudgetGbMinValue, - kCacheDramBudgetGbMaxValue, prep_conf.search_cache_budget_gb); + CheckNumericParamAndSet(config, kCacheDramBudgetGb, kCacheDramBudgetGbMinValue, kCacheDramBudgetGbMaxValue, + prep_conf.search_cache_budget_gb); CheckNonNumbericParamAndSet(config, kWarmUp, prep_conf.warm_up); CheckNonNumbericParamAndSet(config, kUseBfsCache, prep_conf.use_bfs_cache); } @@ -216,7 +218,8 @@ void from_json(const Config& config, DiskANNQueryConfig& query_conf) { CheckNumericParamAndSet(config, kK, kKMinValue, kKMaxValue, query_conf.k); // The search_list_size should be no less than the k. - CheckNumericParamAndSet(config, kSearchListSize, query_conf.k, 10 * query_conf.k, + CheckNumericParamAndSet(config, kSearchListSize, query_conf.k, + std::max(kSearchListSizeMaxValue, static_cast(10 * query_conf.k)), query_conf.search_list_size); CheckNumericParamAndSet(config, kBeamwidth, kBeamwidthMinValue, kBeamwidthMaxValue, query_conf.beamwidth); } diff --git a/thirdparty/DiskANN/src/aux_utils.cpp b/thirdparty/DiskANN/src/aux_utils.cpp index c2f02c25c..5a7778aba 100644 --- a/thirdparty/DiskANN/src/aux_utils.cpp +++ b/thirdparty/DiskANN/src/aux_utils.cpp @@ -983,7 +983,7 @@ namespace diskann { num_pq_chunks = num_pq_chunks <= 0 ? 1 : num_pq_chunks; num_pq_chunks = num_pq_chunks > dim ? dim : num_pq_chunks; - LOG(DEBUG) << "Compressing " << dim << "-dimensional data into " + LOG(INFO) << "Compressing " << dim << "-dimensional data into " << num_pq_chunks << " bytes per vector."; size_t train_size, train_dim; @@ -1021,6 +1021,7 @@ namespace diskann { if (config.compare_metric == diskann::Metric::INNER_PRODUCT) make_zero_mean = false; + auto pq_s = std::chrono::high_resolution_clock::now(); generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, (uint32_t) num_pq_chunks, NUM_KMEANS_REPS, pq_pivots_path, make_zero_mean); @@ -1028,7 +1029,9 @@ namespace diskann { generate_pq_data_from_pivots(data_file_to_use.c_str(), 256, (uint32_t) num_pq_chunks, pq_pivots_path, pq_compressed_vectors_path); - + auto pq_e = std::chrono::high_resolution_clock::now(); + std::chrono::duration pq_diff = pq_e - pq_s; + LOG(INFO) << "Training PQ codes cost: " << pq_diff.count() << "s"; delete[] train_data; train_data = nullptr; @@ -1037,12 +1040,14 @@ namespace diskann { #if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD) MallocExtension::instance()->ReleaseFreeMemory(); #endif - + auto graph_s = std::chrono::high_resolution_clock::now(); diskann::build_merged_vamana_index( data_file_to_use.c_str(), diskann::Metric::L2, L, R, config.accelerate_build, p_val, indexing_ram_budget, mem_index_path, medoids_path, centroids_path); - + auto graph_e = std::chrono::high_resolution_clock::now(); + std::chrono::duration graph_diff = graph_e - graph_s; + LOG(INFO) << "Training graph cost: " << graph_diff.count() << "s"; if (!use_disk_pq) { diskann::create_disk_layout(data_file_to_use.c_str(), mem_index_path, disk_index_path);