diff --git a/cpp/bench/ann/src/common/ann_types.hpp b/cpp/bench/ann/src/common/ann_types.hpp index 2c1105a272..e964a81efa 100644 --- a/cpp/bench/ann/src/common/ann_types.hpp +++ b/cpp/bench/ann/src/common/ann_types.hpp @@ -120,7 +120,7 @@ class ANN : public AnnBase { // The advantage of this way is that index has smaller size // and many indices can share one dataset. // - // AlgoProperty::need_dataset_when_search of such algorithm should be true, + // SearchParam::needs_dataset() of such algorithm should be true, // and set_search_dataset() should save the passed-in pointer somewhere. // The client code should call set_search_dataset() before searching, // and should not release dataset before searching is finished. diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index 10a256bd63..3a930d288e 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -215,11 +215,13 @@ void bench_search(::benchmark::State& state, search_param->metric_objective = metric_objective; } catch (const std::exception& e) { state.SkipWithError("Failed to create an algo: " + std::string(e.what())); + return; } - algo->set_search_param(*search_param); + auto algo_property = parse_algo_property(algo->get_preference(), sp_json); current_algo_props = std::make_shared(algo_property.dataset_memory_type, algo_property.query_memory_type); + if (search_param->needs_dataset()) { try { algo->set_search_dataset(dataset->base_set(current_algo_props->dataset_memory_type), @@ -231,6 +233,14 @@ void bench_search(::benchmark::State& state, return; } } + + try { + algo->set_search_param(*search_param); + + } catch (const std::exception& ex) { + state.SkipWithError("An error occurred setting search parameters: " + std::string(ex.what())); + return; + } } const auto algo_property = *current_algo_props; diff --git a/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp b/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp index 8a0a3ffd37..97d1bbf307 100644 --- a/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp +++ b/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp @@ -49,7 +49,7 @@ void parse_build_param(const nlohmann::json& conf, typename raft::bench::ann::FaissCpuIVFPQ::BuildParam& param) { parse_base_build_param(conf, param); - param.M_ratio = conf.at("M_ratio"); + param.M = conf.at("M"); if (conf.contains("usePrecomputed")) { param.usePrecomputed = conf.at("usePrecomputed"); } else { diff --git a/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h b/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h index 626e52b086..755fe9f197 100644 --- a/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h +++ b/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h @@ -152,6 +152,7 @@ void FaissCpu::build(const T* dataset, size_t nrow, cudaStream_t stream) index_->train(nrow, dataset); // faiss::IndexFlat::train() will do nothing assert(index_->is_trained); index_->add(nrow, dataset); + index_refine_ = std::make_unique(this->index_.get(), dataset); } template @@ -163,7 +164,6 @@ void FaissCpu::set_search_param(const AnnSearchParam& param) dynamic_cast(index_.get())->nprobe = nprobe; if (search_param.refine_ratio > 1.0) { - this->index_refine_ = std::make_unique(this->index_.get()); this->index_refine_.get()->k_factor = search_param.refine_ratio; } @@ -229,7 +229,7 @@ template class FaissCpuIVFPQ : public FaissCpu { public: struct BuildParam : public FaissCpu::BuildParam { - int M_ratio; + int M; int bitsPerCode; bool usePrecomputed; }; @@ -237,12 +237,8 @@ class FaissCpuIVFPQ : public FaissCpu { FaissCpuIVFPQ(Metric metric, int dim, const BuildParam& param) : FaissCpu(metric, dim, param) { this->init_quantizer(dim); - this->index_ = std::make_unique(this->quantizer_.get(), - dim, - param.nlist, - dim / param.M_ratio, - param.bitsPerCode, - this->metric_type_); + this->index_ = std::make_unique( + this->quantizer_.get(), dim, param.nlist, param.M, param.bitsPerCode, this->metric_type_); } void save(const std::string& file) const override diff --git a/cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu b/cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu index 8de8973f16..8b04ba1980 100644 --- a/cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu +++ b/cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu @@ -50,7 +50,7 @@ void parse_build_param(const nlohmann::json& conf, typename raft::bench::ann::FaissGpuIVFPQ::BuildParam& param) { parse_base_build_param(conf, param); - param.M_ratio = conf.at("M_ratio"); + param.M = conf.at("M"); if (conf.contains("usePrecomputed")) { param.usePrecomputed = conf.at("usePrecomputed"); } else { diff --git a/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h b/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h index 8db8e29ef7..4f13ff8a49 100644 --- a/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h +++ b/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h @@ -35,6 +35,9 @@ #include #include +#include +#include + #include #include #include @@ -84,6 +87,7 @@ class FaissGpu : public ANN { struct SearchParam : public AnnSearchParam { int nprobe; float refine_ratio = 1.0; + auto needs_dataset() const -> bool override { return refine_ratio > 1.0f; } }; struct BuildParam { @@ -101,6 +105,7 @@ class FaissGpu : public ANN { RAFT_CUDA_TRY(cudaGetDevice(&device_)); RAFT_CUDA_TRY(cudaEventCreate(&sync_, cudaEventDisableTiming)); faiss_default_stream_ = gpu_resource_.getDefaultStream(device_); + raft::resource::set_cuda_stream(handle_, faiss_default_stream_); } virtual ~FaissGpu() noexcept { RAFT_CUDA_TRY_NO_THROW(cudaEventDestroy(sync_)); } @@ -109,6 +114,8 @@ class FaissGpu : public ANN { virtual void set_search_param(const FaissGpu::AnnSearchParam& param) {} + void set_search_dataset(const T* dataset, size_t nrow) override { dataset_ = dataset; } + // TODO: if the number of results is less than k, the remaining elements of 'neighbors' // will be filled with (size_t)-1 void search(const T* queries, @@ -123,7 +130,7 @@ class FaissGpu : public ANN { AlgoProperty property; // to enable building big dataset which is larger than GPU memory property.dataset_memory_type = MemoryType::Host; - property.query_memory_type = MemoryType::Device; + property.query_memory_type = MemoryType::Host; return property; } @@ -142,7 +149,7 @@ class FaissGpu : public ANN { mutable faiss::gpu::StandardGpuResources gpu_resource_; std::unique_ptr index_; - std::unique_ptr index_refine_; + std::unique_ptr index_refine_{nullptr}; faiss::MetricType metric_type_; int nlist_; int device_; @@ -150,6 +157,9 @@ class FaissGpu : public ANN { cudaStream_t faiss_default_stream_{nullptr}; double training_sample_fraction_; std::unique_ptr search_params_; + const T* dataset_; + raft::device_resources handle_; + float refine_ratio_ = 1.0; }; template @@ -194,7 +204,25 @@ void FaissGpu::search(const T* queries, { static_assert(sizeof(size_t) == sizeof(faiss::idx_t), "sizes of size_t and faiss::idx_t are different"); - index_->search(batch_size, queries, k, distances, reinterpret_cast(neighbors)); + + if (this->refine_ratio_ > 1.0) { + // TODO: FAISS changed their search APIs to accept the search parameters as a struct object + // but their refine API doesn't allow the struct to be passed in. Once this is fixed, we + // need to re-enable refinement below + // index_refine_->search(batch_size, queries, k, distances, + // reinterpret_cast(neighbors), this->search_params_.get()); Related FAISS issue: + // https://github.com/facebookresearch/faiss/issues/3118 + throw std::runtime_error( + "FAISS doesn't support refinement in their new APIs so this feature is disabled in the " + "benchmarks for the time being."); + } else { + index_->search(batch_size, + queries, + k, + distances, + reinterpret_cast(neighbors), + this->search_params_.get()); + } stream_wait(stream); } @@ -217,7 +245,13 @@ void FaissGpu::load_(const std::string& file) std::unique_ptr cpu_index(dynamic_cast(faiss::read_index(file.c_str()))); assert(cpu_index); - dynamic_cast(index_.get())->copyFrom(cpu_index.get()); + + try { + dynamic_cast(index_.get())->copyFrom(cpu_index.get()); + + } catch (const std::exception& e) { + std::cout << "Error loading index file: " << std::string(e.what()) << std::endl; + } } template @@ -242,11 +276,7 @@ class FaissGpuIVFFlat : public FaissGpu { faiss::IVFSearchParameters faiss_search_params; faiss_search_params.nprobe = nprobe; this->search_params_ = std::make_unique(faiss_search_params); - - if (search_param.refine_ratio > 1.0) { - this->index_refine_ = std::make_unique(this->index_.get()); - this->index_refine_.get()->k_factor = search_param.refine_ratio; - } + this->refine_ratio_ = search_param.refine_ratio; } void save(const std::string& file) const override @@ -263,7 +293,7 @@ template class FaissGpuIVFPQ : public FaissGpu { public: struct BuildParam : public FaissGpu::BuildParam { - int M_ratio; + int M; bool useFloat16; bool usePrecomputed; }; @@ -279,7 +309,7 @@ class FaissGpuIVFPQ : public FaissGpu { std::make_unique(&(this->gpu_resource_), dim, param.nlist, - dim / param.M_ratio, + param.M, 8, // FAISS only supports bitsPerCode=8 this->metric_type_, config); @@ -290,14 +320,15 @@ class FaissGpuIVFPQ : public FaissGpu { auto search_param = dynamic_cast::SearchParam&>(param); int nprobe = search_param.nprobe; assert(nprobe <= nlist_); - + this->refine_ratio_ = search_param.refine_ratio; faiss::IVFPQSearchParameters faiss_search_params; faiss_search_params.nprobe = nprobe; this->search_params_ = std::make_unique(faiss_search_params); if (search_param.refine_ratio > 1.0) { - this->index_refine_ = std::make_unique(this->index_.get()); + this->index_refine_ = + std::make_unique(this->index_.get(), this->dataset_); this->index_refine_.get()->k_factor = search_param.refine_ratio; } } @@ -349,9 +380,10 @@ class FaissGpuIVFSQ : public FaissGpu { faiss_search_params.nprobe = nprobe; this->search_params_ = std::make_unique(faiss_search_params); - + this->refine_ratio_ = search_param.refine_ratio; if (search_param.refine_ratio > 1.0) { - this->index_refine_ = std::make_unique(this->index_.get()); + this->index_refine_ = + std::make_unique(this->index_.get(), this->dataset_); this->index_refine_.get()->k_factor = search_param.refine_ratio; } } diff --git a/docs/source/ann_benchmarks_param_tuning.md b/docs/source/ann_benchmarks_param_tuning.md index 90a5e54e32..d787a96955 100644 --- a/docs/source/ann_benchmarks_param_tuning.md +++ b/docs/source/ann_benchmarks_param_tuning.md @@ -91,7 +91,7 @@ IVF-pq is an inverted-file index, which partitions the vectors into a series of | `M_ratio` | `build_param` | Y | Positive Integer Power of 2 [8-64] | | Ratio of numbeer of chunks or subquantizers for each vector. Computed by `dims` / `M_ratio` | | `usePrecomputed` | `build_param` | N | Boolean. Default=`false` | `false` | Use pre-computed lookup tables to speed up search at the cost of increased memory usage. | | `useFloat16` | `build_param` | N | Boolean. Default=`false` | `false` | Use half-precision floats for clustering step. | -| `numProbes` | `search_params` | Y | Positive Integer >0 | | The closest number of clusters to search for each query vector. Larger values will improve recall but will search more points in the index. | +| `nprobe` | `search_params` | Y | Positive Integer >0 | | The closest number of clusters to search for each query vector. Larger values will improve recall but will search more points in the index. | | `refine_ratio` | `search_params` | N| Positive Number >=0 | 0 | `refine_ratio * k` nearest neighbors are queried from the index initially and an additional refinement step improves recall by selecting only the best `k` neighbors. | ### `faiss_cpu_flat` @@ -118,16 +118,16 @@ Use FAISS IVF-Flat index on CPU Use FAISS IVF-PQ index on CPU -| Parameter | Type | Required | Data Type | Default | Description | -|-----------------|----------------|----------|------------------------------------|---------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `nlist` | `build_param` | Y | Positive Integer >0 | | Number of clusters to partition the vectors into. Larger values will put less points into each cluster but this will impact index build time as more clusters need to be trained. | -| `ratio` | `build_param` | N | Positive Integer >0 | 2 | `1/ratio` is the number of training points which should be used to train the clusters. | -| `M` | `build_param` | Y | Positive Integer Power of 2 [8-64] | | Number of chunks or subquantizers for each vector. | +| Parameter | Type | Required | Data Type | Default | Description | +|------------------|----------------|----------|------------------------------------|---------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `nlist` | `build_param` | Y | Positive Integer >0 | | Number of clusters to partition the vectors into. Larger values will put less points into each cluster but this will impact index build time as more clusters need to be trained. | +| `ratio` | `build_param` | N | Positive Integer >0 | 2 | `1/ratio` is the number of training points which should be used to train the clusters. | +| `M` | `build_param` | Y | Positive Integer Power of 2 [8-64] | | Number of chunks or subquantizers for each vector. | | `usePrecomputed` | `build_param` | N | Boolean. Default=`false` | `false` | Use pre-computed lookup tables to speed up search at the cost of increased memory usage. | -| `bitsPerCode` | `build_param` | N | Positive Integer [4-8] | 8 | Number of bits to use for each code. | -| `numProbes` | `search_params` | Y | Positive Integer >0 | | The closest number of clusters to search for each query vector. Larger values will improve recall but will search more points in the index. | -| `refine_ratio` | `search_params` | N| Positive Number >=0 | 0 | `refine_ratio * k` nearest neighbors are queried from the index initially and an additional refinement step improves recall by selecting only the best `k` neighbors. | -| `numThreads` | `search_params` | N | Positive Integer >0 | 1 | Number of threads to use for queries. | +| `bitsPerCode` | `build_param` | N | Positive Integer [4-8] | 8 | Number of bits to use for each code. | +| `nprobe` | `search_params` | Y | Positive Integer >0 | | The closest number of clusters to search for each query vector. Larger values will improve recall but will search more points in the index. | +| `refine_ratio` | `search_params` | N| Positive Number >=0 | 0 | `refine_ratio * k` nearest neighbors are queried from the index initially and an additional refinement step improves recall by selecting only the best `k` neighbors. | +| `numThreads` | `search_params` | N | Positive Integer >0 | 1 | Number of threads to use for queries. | ## HNSW diff --git a/docs/source/raft_ann_benchmarks.md b/docs/source/raft_ann_benchmarks.md index 305cb836e3..e0ff0970b2 100644 --- a/docs/source/raft_ann_benchmarks.md +++ b/docs/source/raft_ann_benchmarks.md @@ -425,6 +425,7 @@ A single configuration will often define a set of algorithms, with associated in base_file: sift-128-euclidean/base.fbin query_file: sift-128-euclidean/query.fbin groundtruth_neighbors_file: sift-128-euclidean/groundtruth.neighbors.ibin + dims: 128 distance: euclidean ``` diff --git a/docs/source/wiki_all_dataset.md b/docs/source/wiki_all_dataset.md index 5c7f972b3d..983cbefd5a 100644 --- a/docs/source/wiki_all_dataset.md +++ b/docs/source/wiki_all_dataset.md @@ -2,8 +2,6 @@ The `wiki-all` dataset was created to stress vector search algorithms at scale with both a large number of vectors and dimensions. The entire dataset contains 88M vectors with 768 dimensions and is meant for testing the types of vectors one would typically encounter in retrieval augmented generation (RAG) workloads. The full dataset is ~251GB in size, which is intentionally larger than the typical memory of GPUs. The massive scale is intended to promote the use of compression and efficient out-of-core methods for both indexing and search. -## Getting the dataset - The dataset is composed of all the available languages of in the [Cohere Wikipedia dataset](https://huggingface.co/datasets/Cohere/wikipedia-22-12). An [English version]( https://www.kaggle.com/datasets/jjinho/wikipedia-20230701) is also available. @@ -13,13 +11,15 @@ Cohere's English Texts are older (2022) and smaller than the Kaggle English Wiki To form the final dataset, the Wiki texts were chunked into 85 million 128-token pieces. For reference, Cohere chunks Wiki texts into 104-token pieces. Finally, the embeddings of each chunk were computed using the [paraphrase-multilingual-mpnet-base-v2](https://huggingface.co/sentence-transformers/paraphrase-multilingual-mpnet-base-v2) embedding model. The resulting dataset is an embedding matrix of size 88 million by 768. Also included with the dataset is a query file containing 10k query vectors and a groundtruth file to evaluate nearest neighbors algorithms. +## Getting the dataset + ### Full dataset A version of the dataset is made available in the binary format that can be used directly by the [raft-ann-bench](https://docs.rapids.ai/api/raft/nightly/raft_ann_benchmarks/) tool. The full 88M dataset is ~251GB and the download link below contains tarballs that have been split into multiple parts. The following will download all 10 the parts and untar them to a `wiki_all_88M` directory: ```bash -curl -s https://data.rapids.ai/raft/datasets/wiki_all/wiki_all.tar.{00..9} | tar -xf - -C /datasets/wiki_all_88M/ +curl -s https://data.rapids.ai/raft/datasets/wiki_all/wiki_all.tar.{00..9} | tar -xf - -C wiki_all_88M/ ``` The above has the unfortunate drawback that if the command should fail for any reason, all the parts need to be re-downloaded. The files can also be downloaded individually and then untarred to the directory. Each file is ~27GB and there are 10 of them. @@ -29,7 +29,7 @@ curl -s https://data.rapids.ai/raft/datasets/wiki_all/wiki_all.tar.00 ... curl -s https://data.rapids.ai/raft/datasets/wiki_all/wiki_all.tar.09 -cat wiki_all.tar.* | tar -xf - -C /datasets/wiki_all_88M/ +cat wiki_all.tar.* | tar -xf - -C wiki_all_88M/ ``` ### 1M and 10M subsets diff --git a/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py b/python/raft-ann-bench/src/raft-ann-bench/constraints/__init__.py similarity index 76% rename from python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py rename to python/raft-ann-bench/src/raft-ann-bench/constraints/__init__.py index 03bf707e6b..b61d185356 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/constraints/__init__.py @@ -16,25 +16,25 @@ DTYPE_SIZES = {"float": 4, "half": 2, "fp8": 1} -def raft_ivf_pq_build_validator(params, dims): +def raft_ivf_pq_build_constraints(params, dims): if "pq_dim" in params: return params["pq_dim"] <= dims return True -def raft_ivf_pq_search_validator(params, build_params, k, batch_size): +def raft_ivf_pq_search_constraints(params, build_params, k, batch_size): ret = True if "internalDistanceDtype" in params and "smemLutDtype" in params: ret = ( DTYPE_SIZES[params["smemLutDtype"]] - < DTYPE_SIZES[params["internalDistanceDtype"]] + <= DTYPE_SIZES[params["internalDistanceDtype"]] ) if "nlist" in build_params and "nprobe" in params: - ret = build_params["nlist"] <= params["nprobe"] + ret = ret and build_params["nlist"] >= params["nprobe"] return ret -def raft_cagra_search_validator(params, build_params, k, batch_size): +def raft_cagra_search_constraints(params, build_params, k, batch_size): if "itopk" in params: return params["itopk"] >= k diff --git a/python/raft-ann-bench/src/raft-ann-bench/data_export/__main__.py b/python/raft-ann-bench/src/raft-ann-bench/data_export/__main__.py index fe2670411e..dd338c0c45 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/data_export/__main__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/data_export/__main__.py @@ -33,6 +33,7 @@ def read_file(dataset, dataset_path, method): def convert_json_to_csv_build(dataset, dataset_path): for file, algo_name, df in read_file(dataset, dataset_path, "build"): + algo_name = algo_name.replace("_base", "") df["name"] = df["name"].str.split("/").str[0] write = pd.DataFrame( { @@ -51,6 +52,7 @@ def convert_json_to_csv_build(dataset, dataset_path): def convert_json_to_csv_search(dataset, dataset_path): for file, algo_name, df in read_file(dataset, dataset_path, "search"): + algo_name = algo_name.replace("_base", "") df["name"] = df["name"].str.split("/").str[0] write = pd.DataFrame( { diff --git a/python/raft-ann-bench/src/raft-ann-bench/plot/__main__.py b/python/raft-ann-bench/src/raft-ann-bench/plot/__main__.py index ef81768f4d..78f8aea8b8 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/plot/__main__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/plot/__main__.py @@ -372,12 +372,10 @@ def load_all_results( algo_group_files = [ result_filename for result_filename in result_files ] + for i in range(len(algo_group_files)): algo_group = algo_group_files[i].replace(".csv", "").split("_") - if len(algo_group) == 2: - algo_group_files[i] = ("_".join(algo_group), "base") - else: - algo_group_files[i] = ("_".join(algo_group[:-1]), algo_group[-1]) + algo_group_files[i] = ("_".join(algo_group[:-1]), algo_group[-1]) algo_group_files = list(zip(*algo_group_files)) if len(algorithms) > 0: diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/__main__.py b/python/raft-ann-bench/src/raft-ann-bench/run/__main__.py index 477c289666..57f03d59ec 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/__main__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/run/__main__.py @@ -50,10 +50,7 @@ def validate_algorithm(algos_conf, algo, gpu_present): def find_executable(algos_conf, algo, group, k, batch_size): executable = algos_conf[algo]["executable"] - if group != "base": - return_str = f"{algo}_{group}-{k}-{batch_size}" - else: - return_str = f"{algo}-{k}-{batch_size}" + return_str = f"{algo}_{group}-{k}-{batch_size}" build_path = os.getenv("RAFT_HOME") if build_path is not None: @@ -128,14 +125,14 @@ def run_build_and_search( ann_executable_path, "--search", "--data_prefix=" + dataset_path, - "--benchmark_counters_tabular", + "--benchmark_counters_tabular=true", "--override_kv=k:%s" % k, "--override_kv=n_queries:%s" % batch_size, "--benchmark_min_warmup_time=0.01", "--benchmark_out_format=json", + "--mode=%s" % mode, "--benchmark_out=" + f"{os.path.join(search_folder, f'{algo}.json')}", - "--mode=%s" % mode, ] if force: cmd = cmd + ["--overwrite"] @@ -338,8 +335,10 @@ def add_algo_group(group_list): algos_conf[algo["name"]]["groups"][group] = algo[ "groups" ][group] - if "validators" in algo: - algos_conf[algo["name"]]["validators"] = algo["validators"] + if "constraints" in algo: + algos_conf[algo["name"]]["constraints"] = algo[ + "constraints" + ] if insert_algo: add_algo_group(named_groups) @@ -382,15 +381,20 @@ def add_algo_group(group_list): index["build_param"][param_names[i]] = params[i] index_name += "." + f"{param_names[i]}{params[i]}" - if "validators" in algos_conf[algo]: - if "build" in algos_conf[algo]["validators"]: - importable = algos_conf[algo]["validators"]["build"] + if "constraints" in algos_conf[algo]: + if "build" in algos_conf[algo]["constraints"]: + importable = algos_conf[algo]["constraints"]["build"] importable = importable.split(".") module = ".".join(importable[:-1]) func = importable[-1] validator = import_module(module) - build_validator = getattr(validator, func) - if not build_validator( + build_constraints = getattr(validator, func) + if "dims" not in conf_file["dataset"]: + raise ValueError( + "`dims` needed for build constraints but not " + "specified in datasets.yaml" + ) + if not build_constraints( index["build_param"], conf_file["dataset"]["dims"] ): continue @@ -405,25 +409,30 @@ def add_algo_group(group_list): search_dict = dict() for i in range(len(search_params)): search_dict[search_param_names[i]] = search_params[i] - if "validators" in algos_conf[algo]: - if "search" in algos_conf[algo]["validators"]: - importable = algos_conf[algo]["validators"][ + if "constraints" in algos_conf[algo]: + if "search" in algos_conf[algo]["constraints"]: + importable = algos_conf[algo]["constraints"][ "search" ] importable = importable.split(".") module = ".".join(importable[:-1]) func = importable[-1] validator = import_module(module) - search_validator = getattr(validator, func) - if search_validator( + search_constraints = getattr(validator, func) + if search_constraints( search_dict, index["build_param"], k, batch_size, ): index["search_params"].append(search_dict) + else: + index["search_params"].append(search_dict) executables_to_run[executable]["index"].append(index) + if len(index["search_params"]) == 0: + print("No search parameters were added to configuration") + run_build_and_search( conf_file, f"{args.dataset}.json", diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_flat.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_flat.yaml index 6542bbab4c..ed237becb3 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_flat.yaml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_flat.yaml @@ -2,9 +2,9 @@ name: faiss_gpu_ivf_flat groups: base: build: - nlists: [1024, 2048, 4096, 8192, 16000, 32000] - ratio: [1, 10, 25] - useFloat16: [True, False] + nlist: [2048] + ratio: [1, 4, 10] + useFloat16: [False] search: - numProbes: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] - refine_ratio: [1, 2, 4, 10] \ No newline at end of file + nprobe: [2048] + refine_ratio: [1] diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yaml index 7e453d506e..87c3afc727 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yaml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yaml @@ -2,11 +2,11 @@ name: faiss_gpu_ivf_pq groups: base: build: - nlist: [1024, 2048, 4096, 8192, 16000, 32000, 64000, 100000] - M_ratio: [2, 4] - ratio: [1, 10, 25] - usePrecomputed: [True, False] - useFloat16: [True, False] + nlist: [1024, 2048, 4096, 8192] + M: [8, 16] + ratio: [10, 25] + usePrecomputed: [False] + useFloat16: [False] search: - numProbes: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] - refine_ratio: [1, 2, 4, 10] \ No newline at end of file + nprobe: [1, 5, 10, 50, 100, 200] + refine_ratio: [1] \ No newline at end of file diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/hnswlib.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/hnswlib.yaml new file mode 100644 index 0000000000..021ca96bcd --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/hnswlib.yaml @@ -0,0 +1,8 @@ +name: hnswlib +groups: + base: + build: + M: [12, 16, 24, 36] + efConstruction: [64, 128, 256, 512] + search: + ef: [10, 20, 40, 60, 80, 120, 200, 400, 600, 800] \ No newline at end of file diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml index 3e90494968..927b90557a 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml @@ -1,11 +1,12 @@ name: raft_cagra -validators: - search: raft-ann-bench.validators.raft_cagra_search_validator +constraints: + search: raft-ann-bench.constraints.raft_cagra_search_constraints groups: base: build: graph_degree: [32, 64] intermediate_graph_degree: [64, 96] + graph_build_algo: ["NN_DESCENT"] search: itopk: [32, 64, 128] diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yaml index 2e1912c6b0..fac383119a 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yaml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yaml @@ -1,17 +1,17 @@ name: raft_ivf_pq -validators: - build: raft-ann-bench.validators.raft_ivf_pq_build_validator - search: raft-ann-bench.validators.raft_ivf_pq_search_validator +constraints: + build: raft-ann-bench.constraints.raft_ivf_pq_build_constraints + search: raft-ann-bench.constraints.raft_ivf_pq_search_constraints groups: base: build: nlist: [1024, 2048, 4096, 8192] pq_dim: [64, 32] pq_bits: [8, 6, 5, 4] - ratio: [1, 10, 25] + ratio: [10, 25] niter: [25] search: - nprobe: [1, 5, 10, 50, 100, 200, 500] - internalDistanceDtype: ["float", "half"] + nprobe: [1, 5, 10, 50, 100, 200] + internalDistanceDtype: ["float"] smemLutDtype: ["float", "fp8", "half"] refine_ratio: [1, 2, 4] \ No newline at end of file diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml index 23476cc056..11a91f91fd 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml @@ -1,9 +1,9 @@ -- name: bigann-100M +- name: bigann-1B base_file: bigann-1B/base.1B.u8bin subset_size: 100000000 dims: 128 query_file: bigann-1B/query.public.10K.u8bin - groundtruth_neighbors_file: bigann-100M/groundtruth.neighbors.ibin + groundtruth_neighbors_file: bigann-1B/groundtruth.neighbors.ibin distance: euclidean - name: deep-1B @@ -14,10 +14,10 @@ distance: inner_product - name: bigann-100M - base_file: bigann-1B/base.1B.u8bin + base_file: bigann-100M/base.1B.u8bin subset_size: 100000000 dims: 128 - query_file: bigann-1B/query.public.10K.u8bin + query_file: bigann-100M/query.public.10K.u8bin groundtruth_neighbors_file: bigann-100M/groundtruth.neighbors.ibin distance: euclidean @@ -99,23 +99,23 @@ groundtruth_neighbors_file: sift-128-euclidean/groundtruth.neighbors.ibin distance: euclidean -- name: wiki_all_1M, - dims: 784 - base_file: wiki_all_1M/base.1MM.fbin, - query_file: wiki_all_1M/queries.fbin, - groundtruth_neighbors_file: wiki_all_1M/groundtruth.1M.neighbors.ibin, +- name: wiki_all_1M + dims: 768 + base_file: wiki_all_1M/base.1MM.fbin + query_file: wiki_all_1M/queries.fbin + groundtruth_neighbors_file: wiki_all_1M/groundtruth.1M.neighbors.ibin distance: euclidean - name: wiki_all_10M, - dims: 784 - base_file: wiki_all_10M/base.10M.fbin, - query_file: wiki_all_10M/queries.fbin, - groundtruth_neighbors_file: wiki_all_10M/groundtruth.10M.neighbors.ibin, + dims: 768 + base_file: wiki_all_10M/base.10M.fbin + query_file: wiki_all_10M/queries.fbin + groundtruth_neighbors_file: wiki_all_10M/groundtruth.10M.neighbors.ibin distance: euclidean -- name: wiki_all_88M, - dims: 784 - base_file: wiki_all_88M/base.88M.fbin, - query_file: wiki_all_88M/queries.fbin, - groundtruth_neighbors_file: wiki_all_88M/groundtruth.88M.neighbors.ibin, +- name: wiki_all_88M + dims: 768 + base_file: wiki_all_88M/base.88M.fbin + query_file: wiki_all_88M/queries.fbin + groundtruth_neighbors_file: wiki_all_88M/groundtruth.88M.neighbors.ibin distance: euclidean