From f53f128196521923c7ea482d098c44dc7e6eadd3 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Thu, 19 Oct 2023 22:39:36 -0400 Subject: [PATCH 01/49] Adding some initial yaml param conf files --- .../run/conf/algos/faiss_gpu_ivf_pq.yml | 4 ++++ .../run/conf/algos/raft_cagra.yml | 12 ++++++++++++ .../run/conf/algos/raft_ivf_flat.yml | 8 ++++++++ .../run/conf/algos/raft_ivf_pq.yml | 17 +++++++++++++++++ 4 files changed, 41 insertions(+) create mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yml create mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yml create mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yml create mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yml new file mode 100644 index 0000000000..136f9ffcab --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yml @@ -0,0 +1,4 @@ +name: faiss_gpu_ivf_pq + - base: + build: + search: \ No newline at end of file diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yml new file mode 100644 index 0000000000..4e7e6bc392 --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yml @@ -0,0 +1,12 @@ +name: raft_cagra + - base: + build: + graph_degree: [32, 64] + search: + itopk: [32, 64, 128] + + - large: + build: + graph_degree: [32, 64] + search: + itopk: [32, 64, 128] \ No newline at end of file diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yml new file mode 100644 index 0000000000..4af425896a --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yml @@ -0,0 +1,8 @@ +name: raft_ivf_flat + - base: + build: + nlist: [1024, 16384] + ratio: [1, 2] + niter: [20, 25] + search: + nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] \ No newline at end of file diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml new file mode 100644 index 0000000000..9cfd92691b --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml @@ -0,0 +1,17 @@ +name: raft_ivf_pq + - base: + build: + nlist: [1024] + pq_dim: [128, 64] + pq_bits: [8, 6] + ratio: [1] + niter: [25] + search: + nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] + internalDistanceDtype: ["float", "half"] + smemLutDtype: ["float", "fp8", "half"] + + # Custom validator with a `validate()` function that accepts a parameter combination + # and returns false if it's not valid (for example, + # internalDistanceDtype should never be lower than smemLutDtype) + validators: raft_ivf_pq_validator \ No newline at end of file From 0385543b1c23225aa9bf5284deef13a503ef8a86 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Thu, 19 Oct 2023 22:41:35 -0400 Subject: [PATCH 02/49] Adding name --- .../src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml index 9cfd92691b..3ef90bfb0d 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml @@ -14,4 +14,4 @@ name: raft_ivf_pq # Custom validator with a `validate()` function that accepts a parameter combination # and returns false if it's not valid (for example, # internalDistanceDtype should never be lower than smemLutDtype) - validators: raft_ivf_pq_validator \ No newline at end of file + validators: raft-ann-bench.validators.raft_ivf_pq_validator \ No newline at end of file From 1bf6fa7ccfeda3d245e68ede1c59df5b22418e4f Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Thu, 19 Oct 2023 23:02:37 -0400 Subject: [PATCH 03/49] Adding datasets.yml --- .../src/raft-ann-bench/run/conf/datasets.yml | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yml diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yml new file mode 100644 index 0000000000..bffa632815 --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yml @@ -0,0 +1,85 @@ +- "name": "bigann-100M", + "base_file": "bigann-1B/base.1B.u8bin", + "subset_size": 100000000, + "query_file": "bigann-1B/query.public.10K.u8bin", + "groundtruth_neighbors_file": "bigann-100M/groundtruth.neighbors.ibin", + "distance": "euclidean" + +- "name": "deep-1B", + "base_file": "deep-1B/base.1B.fbin", + "query_file": "deep-1B/query.public.10K.fbin", + "groundtruth_neighbors_file": "deep-1B/groundtruth.neighbors.ibin", + "distance": "inner_product" + +- "name": "bigann-100M", + "base_file": "bigann-1B/base.1B.u8bin", + "subset_size": 100000000, + "query_file": "bigann-1B/query.public.10K.u8bin", + "groundtruth_neighbors_file": "bigann-100M/groundtruth.neighbors.ibin", + "distance": "euclidean" + +- "name": "deep-image-96-inner", + "base_file": "deep-image-96-inner/base.fbin", + "query_file": "deep-image-96-inner/query.fbin", + "groundtruth_neighbors_file": "deep-image-96-inner/groundtruth.neighbors.ibin", + "distance": "euclidean" + +- "name": "fashion-mnist-784-euclidean", + "base_file": "fashion-mnist-784-euclidean/base.fbin", + "query_file": "fashion-mnist-784-euclidean/query.fbin", + "groundtruth_neighbors_file": "fashion-mnist-784-euclidean/groundtruth.neighbors.ibin", + "distance": "euclidean" + +- "name": "gist-960-euclidean", + "base_file": "gist-960-euclidean/base.fbin", + "query_file": "gist-960-euclidean/query.fbin", + "distance": "euclidean" + +- "name": "glove-50-angular", + "base_file": "glove-50-angular/base.fbin", + "query_file": "glove-50-angular/query.fbin", + "distance": "euclidean" + +- "name": "glove-50-inner", + "base_file": "glove-50-inner/base.fbin", + "query_file": "glove-50-inner/query.fbin", + "distance": "euclidean" + +- "name": "glove-100-angular", + "base_file": "glove-100-angular/base.fbin", + "query_file": "glove-100-angular/query.fbin", + "distance": "euclidean" + +- "name": "glove-100-inner", + "base_file": "glove-100-inner/base.fbin", + "query_file": "glove-100-inner/query.fbin", + "distance": "euclidean" + +- "name": "lastfm-65-angular", + "base_file": "lastfm-65-angular/base.fbin", + "query_file": "lastfm-65-angular/query.fbin", + "distance": "euclidean" + +- "name": "mnist-784-euclidean", + "base_file": "mnist-784-euclidean/base.fbin", + "query_file": "mnist-784-euclidean/query.fbin", + "groundtruth_neighbors_file": "mnist-784-euclidean/groundtruth.neighbors.ibin", + "distance": "euclidean" + +- "name": "nytimes-256-angular", + "base_file": "nytimes-256-angular/base.fbin", + "query_file": "nytimes-256-angular/query.fbin", + "groundtruth_neighbors_file": "nytimes-256-angular/groundtruth.neighbors.ibin", + "distance": "euclidean" + +- "name": "nytimes-256-inner", + "base_file": "nytimes-256-inner/base.fbin", + "query_file": "nytimes-256-inner/query.fbin", + "groundtruth_neighbors_file": "nytimes-256-inner/groundtruth.neighbors.ibin", + "distance": "euclidean" + +- "name": "sift-128-euclidean", + "base_file": "sift-128-euclidean/base.fbin", + "query_file": "sift-128-euclidean/query.fbin", + "groundtruth_neighbors_file": "sift-128-euclidean/groundtruth.neighbors.ibin", + "distance": "euclidean" \ No newline at end of file From 9e52c1115b7815e8309d8c885321dc45a854af59 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Thu, 19 Oct 2023 23:05:11 -0400 Subject: [PATCH 04/49] More cleanup --- .../src/raft-ann-bench/run/conf/datasets.yml | 138 +++++++++--------- 1 file changed, 69 insertions(+), 69 deletions(-) diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yml index bffa632815..d64bb56f3b 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yml @@ -1,85 +1,85 @@ -- "name": "bigann-100M", - "base_file": "bigann-1B/base.1B.u8bin", +- name: "bigann-100M", + base_file: "bigann-1B/base.1B.u8bin", "subset_size": 100000000, - "query_file": "bigann-1B/query.public.10K.u8bin", - "groundtruth_neighbors_file": "bigann-100M/groundtruth.neighbors.ibin", - "distance": "euclidean" + query_file: "bigann-1B/query.public.10K.u8bin", + groundtruth_neighbors_file: "bigann-100M/groundtruth.neighbors.ibin", + distance: "euclidean" -- "name": "deep-1B", - "base_file": "deep-1B/base.1B.fbin", - "query_file": "deep-1B/query.public.10K.fbin", - "groundtruth_neighbors_file": "deep-1B/groundtruth.neighbors.ibin", - "distance": "inner_product" +- name: "deep-1B", + base_file: "deep-1B/base.1B.fbin", + query_file: "deep-1B/query.public.10K.fbin", + groundtruth_neighbors_file: "deep-1B/groundtruth.neighbors.ibin", + distance: "inner_product" -- "name": "bigann-100M", - "base_file": "bigann-1B/base.1B.u8bin", +- name: "bigann-100M", + base_file: "bigann-1B/base.1B.u8bin", "subset_size": 100000000, - "query_file": "bigann-1B/query.public.10K.u8bin", - "groundtruth_neighbors_file": "bigann-100M/groundtruth.neighbors.ibin", - "distance": "euclidean" + query_file: "bigann-1B/query.public.10K.u8bin", + groundtruth_neighbors_file: "bigann-100M/groundtruth.neighbors.ibin", + distance: "euclidean" -- "name": "deep-image-96-inner", - "base_file": "deep-image-96-inner/base.fbin", - "query_file": "deep-image-96-inner/query.fbin", - "groundtruth_neighbors_file": "deep-image-96-inner/groundtruth.neighbors.ibin", - "distance": "euclidean" +- name: "deep-image-96-inner", + base_file: "deep-image-96-inner/base.fbin", + query_file: "deep-image-96-inner/query.fbin", + groundtruth_neighbors_file: "deep-image-96-inner/groundtruth.neighbors.ibin", + distance: "euclidean" -- "name": "fashion-mnist-784-euclidean", - "base_file": "fashion-mnist-784-euclidean/base.fbin", - "query_file": "fashion-mnist-784-euclidean/query.fbin", - "groundtruth_neighbors_file": "fashion-mnist-784-euclidean/groundtruth.neighbors.ibin", - "distance": "euclidean" +- name: "fashion-mnist-784-euclidean", + base_file: "fashion-mnist-784-euclidean/base.fbin", + query_file: "fashion-mnist-784-euclidean/query.fbin", + groundtruth_neighbors_file: "fashion-mnist-784-euclidean/groundtruth.neighbors.ibin", + distance: "euclidean" -- "name": "gist-960-euclidean", - "base_file": "gist-960-euclidean/base.fbin", - "query_file": "gist-960-euclidean/query.fbin", - "distance": "euclidean" +- name: "gist-960-euclidean", + base_file: "gist-960-euclidean/base.fbin", + query_file: "gist-960-euclidean/query.fbin", + distance: "euclidean" -- "name": "glove-50-angular", - "base_file": "glove-50-angular/base.fbin", - "query_file": "glove-50-angular/query.fbin", - "distance": "euclidean" +- name: "glove-50-angular", + base_file: "glove-50-angular/base.fbin", + query_file: "glove-50-angular/query.fbin", + distance: "euclidean" -- "name": "glove-50-inner", - "base_file": "glove-50-inner/base.fbin", - "query_file": "glove-50-inner/query.fbin", - "distance": "euclidean" +- name: "glove-50-inner", + base_file: "glove-50-inner/base.fbin", + query_file: "glove-50-inner/query.fbin", + distance: "euclidean" -- "name": "glove-100-angular", - "base_file": "glove-100-angular/base.fbin", - "query_file": "glove-100-angular/query.fbin", - "distance": "euclidean" +- name: "glove-100-angular", + base_file: "glove-100-angular/base.fbin", + query_file: "glove-100-angular/query.fbin", + distance: "euclidean" -- "name": "glove-100-inner", - "base_file": "glove-100-inner/base.fbin", - "query_file": "glove-100-inner/query.fbin", - "distance": "euclidean" +- name: "glove-100-inner", + base_file: "glove-100-inner/base.fbin", + query_file: "glove-100-inner/query.fbin", + distance: "euclidean" -- "name": "lastfm-65-angular", - "base_file": "lastfm-65-angular/base.fbin", - "query_file": "lastfm-65-angular/query.fbin", - "distance": "euclidean" +- name: "lastfm-65-angular", + base_file: "lastfm-65-angular/base.fbin", + query_file: "lastfm-65-angular/query.fbin", + distance: "euclidean" -- "name": "mnist-784-euclidean", - "base_file": "mnist-784-euclidean/base.fbin", - "query_file": "mnist-784-euclidean/query.fbin", - "groundtruth_neighbors_file": "mnist-784-euclidean/groundtruth.neighbors.ibin", - "distance": "euclidean" +- name: "mnist-784-euclidean", + base_file: "mnist-784-euclidean/base.fbin", + query_file: "mnist-784-euclidean/query.fbin", + groundtruth_neighbors_file: "mnist-784-euclidean/groundtruth.neighbors.ibin", + distance: "euclidean" -- "name": "nytimes-256-angular", - "base_file": "nytimes-256-angular/base.fbin", - "query_file": "nytimes-256-angular/query.fbin", - "groundtruth_neighbors_file": "nytimes-256-angular/groundtruth.neighbors.ibin", - "distance": "euclidean" +- name: "nytimes-256-angular", + base_file: "nytimes-256-angular/base.fbin", + query_file: "nytimes-256-angular/query.fbin", + groundtruth_neighbors_file: "nytimes-256-angular/groundtruth.neighbors.ibin", + distance: "euclidean" -- "name": "nytimes-256-inner", - "base_file": "nytimes-256-inner/base.fbin", - "query_file": "nytimes-256-inner/query.fbin", - "groundtruth_neighbors_file": "nytimes-256-inner/groundtruth.neighbors.ibin", - "distance": "euclidean" +- name: "nytimes-256-inner", + base_file: "nytimes-256-inner/base.fbin", + query_file: "nytimes-256-inner/query.fbin", + groundtruth_neighbors_file: "nytimes-256-inner/groundtruth.neighbors.ibin", + distance: "euclidean" -- "name": "sift-128-euclidean", - "base_file": "sift-128-euclidean/base.fbin", - "query_file": "sift-128-euclidean/query.fbin", - "groundtruth_neighbors_file": "sift-128-euclidean/groundtruth.neighbors.ibin", - "distance": "euclidean" \ No newline at end of file +- name: "sift-128-euclidean", + base_file: "sift-128-euclidean/base.fbin", + query_file: "sift-128-euclidean/query.fbin", + groundtruth_neighbors_file: "sift-128-euclidean/groundtruth.neighbors.ibin", + distance: "euclidean" \ No newline at end of file From 14e2c5dce3f2e93adc09f5ccd18754d58b9c9079 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 24 Oct 2023 13:34:56 -0700 Subject: [PATCH 05/49] add cross products for yaml->json configs --- .../raft-ann-bench/data_export/__main__.py | 4 +- .../src/raft-ann-bench/plot/__main__.py | 27 +- .../src/raft-ann-bench/run/__main__.py | 278 ++++++++++++------ .../run/conf/algos/raft_cagra.yaml | 13 + .../run/conf/algos/raft_cagra.yml | 12 - .../src/raft-ann-bench/run/conf/datasets.yaml | 85 ++++++ .../src/raft-ann-bench/run/conf/datasets.yml | 85 ------ 7 files changed, 310 insertions(+), 194 deletions(-) create mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml delete mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yml create mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml delete mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yml 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 fd9d00f43c..1d48e9950d 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 @@ -41,7 +41,9 @@ def convert_json_to_csv_build(dataset, dataset_path): "time": df["real_time"], } ) - write.to_csv(file.replace(".json", ".csv"), index=False) + filepath = os.path.normpath(file).split(os.sep) + filename = filepath[-1].split("-")[0] + ".csv" + write.to_csv(os.path.join(f"{os.sep}".join(filepath[:-1]), filename), index=False) def convert_json_to_csv_search(dataset, dataset_path): 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 233607c281..ad3a2490db 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 @@ -305,7 +305,7 @@ def mean_y(algo): fig = ax.get_figure() print(f"writing build output to {fn_out}") plt.title("Build Time for Highest QPS") - plt.suptitle(f"{dataset} k={k} batch_size={batch_size}") + plt.suptitle(f"{dataset}") plt.ylabel("Build Time (s)") fig.savefig(fn_out) @@ -349,17 +349,26 @@ def load_all_results( ): results_path = os.path.join(dataset_path, "result", method) result_files = os.listdir(results_path) - result_files = [ - result_filename - for result_filename in result_files - if f"{k}-{batch_size}" in result_filename - ] - if len(algorithms) > 0: + print(result_files) + if method == "search": result_files = [ result_filename for result_filename in result_files - if result_filename.split("-")[0] in algorithms + if f"{k}-{batch_size}" in result_filename ] + if len(algorithms) > 0: + result_files = [ + result_filename + for result_filename in result_files + if result_filename.split("-")[0] in algorithms + ] + elif method == "build": + if len(algorithms) > 0: + result_files = [ + result_filename + for result_filename in result_files + if result_filename.split("-")[0] in algorithms + ] results = load_lines(results_path, result_files, method, index_key) @@ -450,7 +459,7 @@ def main(): ) build_output_filepath = os.path.join( args.output_filepath, - f"build-{args.dataset}-k{k}-batch_size{batch_size}.png", + f"build-{args.dataset}.png", ) search_results = load_all_results( 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 a0d4fabb77..01f4adbac8 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 @@ -14,10 +14,10 @@ # limitations under the License. import argparse +import itertools import json import os import subprocess - import yaml @@ -132,7 +132,7 @@ def run_build_and_search( cmd = cmd + [temp_conf_filepath] subprocess.run(cmd, check=True) - os.remove(temp_conf_filepath) + # os.remove(temp_conf_filepath) def main(): @@ -148,7 +148,7 @@ def main(): gpu_present = False with open(f"{scripts_path}/algos.yaml", "r") as f: - algos_conf = yaml.safe_load(f) + algos_yaml = yaml.safe_load(f) if "RAPIDS_DATASET_ROOT_DIR" in os.environ: default_dataset_path = os.getenv("RAPIDS_DATASET_ROOT_DIR") @@ -159,6 +159,11 @@ def main(): formatter_class=argparse.ArgumentDefaultsHelpFormatter ) + parser.add_argument( + "--subset-size", + type=positive_int, + help="the number of subset rows of the dataset to build the index", + ) parser.add_argument( "-k", "--count", @@ -173,13 +178,17 @@ def main(): type=positive_int, help="number of query vectors to use in each query trial", ) + parser.add_argument( + "--dataset-configuration", + help="path to configuration file for datasets", + ) parser.add_argument( "--configuration", - help="path to configuration file for a dataset", + help="path to configuration file or directory for algorithms", ) parser.add_argument( "--dataset", - help="dataset whose configuration file will be used", + help="name of dataset", default="glove-100-inner", ) parser.add_argument( @@ -197,11 +206,16 @@ def main(): algorithms", default=None, ) + # parser.add_argument( + # "--indices", + # help="run only comma separated list of named indices. \ + # parameter `algorithms` is ignored", + # default=None, + # ) parser.add_argument( - "--indices", - help="run only comma separated list of named indices. \ - parameter `algorithms` is ignored", - default=None, + "--groups", + help="comma separated groups of parameters to run the benchmarks for", + default="base" ) parser.add_argument( "-f", @@ -225,92 +239,182 @@ def main(): k = args.count batch_size = args.batch_size - # Read configuration file associated to dataset - if args.configuration: - conf_filepath = args.configuration - elif args.dataset: - conf_filepath = os.path.join( - scripts_path, "conf", f"{args.dataset}.json" - ) + # Read configuration file associated to datasets + if args.dataset_configuration: + dataset_conf_f = args.dataset_configuration else: - raise ValueError( - "One of parameters `configuration` or \ - `dataset` need to be provided" + dataset_conf_f = os.path.join( + scripts_path, "conf", "datasets.yaml" ) - conf_filename = conf_filepath.split("/")[-1] - conf_filedir = "/".join(conf_filepath.split("/")[:-1]) - dataset_path = args.dataset_path - if not os.path.exists(conf_filepath): - raise FileNotFoundError(conf_filename) - - with open(conf_filepath, "r") as f: - conf_file = json.load(f) - - dataset_name = conf_file["dataset"]["name"] + with open(dataset_conf_f, "r") as f: + dataset_conf_all = yaml.safe_load(f) + + for dataset in dataset_conf_all: + if args.dataset == dataset["name"]: + dataset_conf = dataset + + conf_file = dict() + conf_file["dataset"] = dataset_conf + if args.subset_size: + conf_file["dataset"]["subset_size"] = args.subset_size + + conf_file["search_basic_param"] = {} + conf_file["search_basic_param"]["k"] = k + conf_file["search_basic_param"]["batch_size"] = batch_size + + algos_conf_fs = os.listdir(os.path.join(scripts_path, "conf", "algos")) + algos_conf_fs = [os.path.join(scripts_path, "conf", "algos", f) \ + for f in algos_conf_fs] + conf_filedir = os.path.join(scripts_path, "conf", "algos") + if args.configuration: + if os.path.isdir(args.configuration): + conf_filedir = args.configuration + algos_conf_fs = algos_conf_fs + \ + [os.path.join(args.configuration, f) \ + for f in os.listdir(args.configuration)] + elif os.path.isfile(args.configuration): + conf_filedir = os.path.normpath(args.configuration).split(os.sep) + algos_conf_fs = algos_conf_fs + [args.configuration] + + filter_algos = True if args.algorithms else False + if filter_algos: + allowed_algos = args.algorithms.split(",") + named_groups = args.groups.split(",") + algos_conf = dict() + for algo_f in algos_conf_fs: + with open(algo_f, "r") as f: + if algo_f.split("/")[-1] == "raft_cagra.yaml": + algo = yaml.safe_load(f) + insert_algo = True + if filter_algos: + if algo["name"] not in allowed_algos: + insert_algo = False + if insert_algo: + if algo["name"] not in algos_conf: + algos_conf[algo["name"]] = dict() + for group in algo.keys(): + if group != "name": + if group in named_groups: + algos_conf[algo["name"]][group] = algo[group] executables_to_run = dict() - # At least one named index should exist in config file - if args.indices: - indices = set(args.indices.split(",")) - # algo associated with index should still be present in algos.yaml - # and enabled - for index in conf_file["index"]: - curr_algo = index["algo"] - if index["name"] in indices and validate_algorithm( - algos_conf, curr_algo, gpu_present - ): - executable_path = find_executable( - algos_conf, curr_algo, k, batch_size - ) - if executable_path not in executables_to_run: - executables_to_run[executable_path] = {"index": []} - executables_to_run[executable_path]["index"].append(index) - - # switch to named algorithms if indices parameter is not supplied - elif args.algorithms: - algorithms = set(args.algorithms.split(",")) - # pick out algorithms from conf file that exist - # and are enabled in algos.yaml - for index in conf_file["index"]: - curr_algo = index["algo"] - if curr_algo in algorithms and validate_algorithm( - algos_conf, curr_algo, gpu_present - ): - executable_path = find_executable( - algos_conf, curr_algo, k, batch_size - ) - if executable_path not in executables_to_run: - executables_to_run[executable_path] = {"index": []} - executables_to_run[executable_path]["index"].append(index) - - # default, try to run all available algorithms - else: - for index in conf_file["index"]: - curr_algo = index["algo"] - if validate_algorithm(algos_conf, curr_algo, gpu_present): - executable_path = find_executable( - algos_conf, curr_algo, k, batch_size - ) - if executable_path not in executables_to_run: - executables_to_run[executable_path] = {"index": []} - executables_to_run[executable_path]["index"].append(index) - - # Replace index to dataset path - for executable_path in executables_to_run: - for pos, index in enumerate( - executables_to_run[executable_path]["index"] - ): - index["file"] = os.path.join( - dataset_path, dataset_name, "index", index["name"] - ) - executables_to_run[executable_path]["index"][pos] = index - + for algo in algos_conf.keys(): + validate_algorithm(algos_yaml, algo, gpu_present) + executable = find_executable(algos_yaml, algo, k, batch_size) + if executable not in executables_to_run: + executables_to_run[executable] = {"index": []} + for group in algos_conf[algo].keys(): + build_params = algos_conf[algo][group]["build"] + search_params = algos_conf[algo][group]["search"] + + param_names = [] + param_lists = [] + for param in build_params.keys(): + param_names.append(param) + param_lists.append(build_params[param]) + + all_build_params = itertools.product(*param_lists) + + search_param_names = [] + search_param_lists = [] + for search_param in search_params.keys(): + search_param_names.append(search_param) + search_param_lists.append(search_params[search_param]) + + for params in all_build_params: + index = {"algo": algo, "build_param": {}} + index_name = f"{algo}" + for i in range(len(params)): + index["build_param"][param_names[i]] = params[i] + index_name += "." + f"{param_names[i]}{params[i]}" + index["name"] = index_name + index["file"] = os.path.join(args.dataset_path, args.dataset, + "index", index_name) + index["search_params"] = [] + all_search_params = itertools.product(*search_param_lists) + for search_params in all_search_params: + search_dict = dict() + for i in range(len(search_params)): + search_dict[search_param_names[i]] = search_params[i] + index["search_params"].append(search_dict) + executables_to_run[executable]["index"].append(index) + + print(executables_to_run) + # conf_filename = conf_filepath.split("/")[-1] + # conf_filedir = "/".join(conf_filepath.split("/")[:-1]) + # dataset_path = args.dataset_path + # if not os.path.exists(conf_filepath): + # raise FileNotFoundError(conf_filename) + + # with open(conf_filepath, "r") as f: + # conf_file = json.load(f) + + # dataset_name = conf_file["dataset"]["name"] + + # executables_to_run = dict() + # # At least one named index should exist in config file + # if args.indices: + # indices = set(args.indices.split(",")) + # # algo associated with index should still be present in algos.yaml + # # and enabled + # for index in conf_file["index"]: + # curr_algo = index["algo"] + # if index["name"] in indices and validate_algorithm( + # algos_conf, curr_algo, gpu_present + # ): + # executable_path = find_executable( + # algos_conf, curr_algo, k, batch_size + # ) + # if executable_path not in executables_to_run: + # executables_to_run[executable_path] = {"index": []} + # executables_to_run[executable_path]["index"].append(index) + + # # switch to named algorithms if indices parameter is not supplied + # elif args.algorithms: + # algorithms = set(args.algorithms.split(",")) + # # pick out algorithms from conf file that exist + # # and are enabled in algos.yaml + # for index in conf_file["index"]: + # curr_algo = index["algo"] + # if curr_algo in algorithms and validate_algorithm( + # algos_conf, curr_algo, gpu_present + # ): + # executable_path = find_executable( + # algos_conf, curr_algo, k, batch_size + # ) + # if executable_path not in executables_to_run: + # executables_to_run[executable_path] = {"index": []} + # executables_to_run[executable_path]["index"].append(index) + + # # default, try to run all available algorithms + # else: + # for index in conf_file["index"]: + # curr_algo = index["algo"] + # if validate_algorithm(algos_conf, curr_algo, gpu_present): + # executable_path = find_executable( + # algos_conf, curr_algo, k, batch_size + # ) + # if executable_path not in executables_to_run: + # executables_to_run[executable_path] = {"index": []} + # executables_to_run[executable_path]["index"].append(index) + + # # Replace index to dataset path + # for executable_path in executables_to_run: + # for pos, index in enumerate( + # executables_to_run[executable_path]["index"] + # ): + # index["file"] = os.path.join( + # dataset_path, dataset_name, "index", index["name"] + # ) + # executables_to_run[executable_path]["index"][pos] = index + + print(conf_filedir) run_build_and_search( conf_file, - conf_filename, + f"{args.dataset}.json", conf_filedir, executables_to_run, - dataset_path, + args.dataset_path, args.force, build, search, 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 new file mode 100644 index 0000000000..095f77e66a --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml @@ -0,0 +1,13 @@ +name: raft_cagra +base: + build: + graph_degree: [32, 64] + intermediate_graph_degree: [64, 96] + search: + itopk: [32, 64, 128] + +large: + build: + graph_degree: [32, 64] + search: + itopk: [32, 64, 128] diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yml deleted file mode 100644 index 4e7e6bc392..0000000000 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yml +++ /dev/null @@ -1,12 +0,0 @@ -name: raft_cagra - - base: - build: - graph_degree: [32, 64] - search: - itopk: [32, 64, 128] - - - large: - build: - graph_degree: [32, 64] - search: - itopk: [32, 64, 128] \ 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 new file mode 100644 index 0000000000..45d648f88e --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml @@ -0,0 +1,85 @@ +- name: bigann-100M + base_file: bigann-1B/base.1B.u8bin + subset_size: 100000000 + query_file: bigann-1B/query.public.10K.u8bin + groundtruth_neighbors_file: bigann-100M/groundtruth.neighbors.ibin + distance: euclidean + +- name: deep-1B + base_file: deep-1B/base.1B.fbin + query_file: deep-1B/query.public.10K.fbin + groundtruth_neighbors_file: deep-1B/groundtruth.neighbors.ibin + distance: inner_product + +- name: bigann-100M + base_file: bigann-1B/base.1B.u8bin + subset_size: 100000000 + query_file: bigann-1B/query.public.10K.u8bin + groundtruth_neighbors_file: bigann-100M/groundtruth.neighbors.ibin + distance: euclidean + +- name: deep-image-96-inner + base_file: deep-image-96-inner/base.fbin + query_file: deep-image-96-inner/query.fbin + groundtruth_neighbors_file: deep-image-96-inner/groundtruth.neighbors.ibin + distance: euclidean + +- name: fashion-mnist-784-euclidean + base_file: fashion-mnist-784-euclidean/base.fbin + query_file: fashion-mnist-784-euclidean/query.fbin + groundtruth_neighbors_file: fashion-mnist-784-euclidean/groundtruth.neighbors.ibin + distance: euclidean + +- name: gist-960-euclidean + base_file: gist-960-euclidean/base.fbin + query_file: gist-960-euclidean/query.fbin + distance: euclidean + +- name: glove-50-angular + base_file: glove-50-angular/base.fbin + query_file: glove-50-angular/query.fbin + distance: euclidean + +- name: glove-50-inner + base_file: glove-50-inner/base.fbin + query_file: glove-50-inner/query.fbin + distance: euclidean + +- name: glove-100-angular + base_file: glove-100-angular/base.fbin + query_file: glove-100-angular/query.fbin + distance: euclidean + +- name: glove-100-inner + base_file: glove-100-inner/base.fbin + query_file: glove-100-inner/query.fbin + distance: euclidean + +- name: lastfm-65-angular + base_file: lastfm-65-angular/base.fbin + query_file: lastfm-65-angular/query.fbin + distance: euclidean + +- name: mnist-784-euclidean + base_file: mnist-784-euclidean/base.fbin + query_file: mnist-784-euclidean/query.fbin + groundtruth_neighbors_file: mnist-784-euclidean/groundtruth.neighbors.ibin + distance: euclidean + +- name: nytimes-256-angular + base_file: nytimes-256-angular/base.fbin + query_file: nytimes-256-angular/query.fbin + groundtruth_neighbors_file: nytimes-256-angular/groundtruth.neighbors.ibin + distance: euclidean + +- name: nytimes-256-inner + base_file: nytimes-256-inner/base.fbin + query_file: nytimes-256-inner/query.fbin + groundtruth_neighbors_file: nytimes-256-inner/groundtruth.neighbors.ibin + distance: euclidean + +- name: sift-128-euclidean + base_file: sift-128-euclidean/base.fbin + query_file: sift-128-euclidean/query.fbin + groundtruth_neighbors_file: sift-128-euclidean/groundtruth.neighbors.ibin + distance: euclidean diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yml deleted file mode 100644 index d64bb56f3b..0000000000 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yml +++ /dev/null @@ -1,85 +0,0 @@ -- name: "bigann-100M", - base_file: "bigann-1B/base.1B.u8bin", - "subset_size": 100000000, - query_file: "bigann-1B/query.public.10K.u8bin", - groundtruth_neighbors_file: "bigann-100M/groundtruth.neighbors.ibin", - distance: "euclidean" - -- name: "deep-1B", - base_file: "deep-1B/base.1B.fbin", - query_file: "deep-1B/query.public.10K.fbin", - groundtruth_neighbors_file: "deep-1B/groundtruth.neighbors.ibin", - distance: "inner_product" - -- name: "bigann-100M", - base_file: "bigann-1B/base.1B.u8bin", - "subset_size": 100000000, - query_file: "bigann-1B/query.public.10K.u8bin", - groundtruth_neighbors_file: "bigann-100M/groundtruth.neighbors.ibin", - distance: "euclidean" - -- name: "deep-image-96-inner", - base_file: "deep-image-96-inner/base.fbin", - query_file: "deep-image-96-inner/query.fbin", - groundtruth_neighbors_file: "deep-image-96-inner/groundtruth.neighbors.ibin", - distance: "euclidean" - -- name: "fashion-mnist-784-euclidean", - base_file: "fashion-mnist-784-euclidean/base.fbin", - query_file: "fashion-mnist-784-euclidean/query.fbin", - groundtruth_neighbors_file: "fashion-mnist-784-euclidean/groundtruth.neighbors.ibin", - distance: "euclidean" - -- name: "gist-960-euclidean", - base_file: "gist-960-euclidean/base.fbin", - query_file: "gist-960-euclidean/query.fbin", - distance: "euclidean" - -- name: "glove-50-angular", - base_file: "glove-50-angular/base.fbin", - query_file: "glove-50-angular/query.fbin", - distance: "euclidean" - -- name: "glove-50-inner", - base_file: "glove-50-inner/base.fbin", - query_file: "glove-50-inner/query.fbin", - distance: "euclidean" - -- name: "glove-100-angular", - base_file: "glove-100-angular/base.fbin", - query_file: "glove-100-angular/query.fbin", - distance: "euclidean" - -- name: "glove-100-inner", - base_file: "glove-100-inner/base.fbin", - query_file: "glove-100-inner/query.fbin", - distance: "euclidean" - -- name: "lastfm-65-angular", - base_file: "lastfm-65-angular/base.fbin", - query_file: "lastfm-65-angular/query.fbin", - distance: "euclidean" - -- name: "mnist-784-euclidean", - base_file: "mnist-784-euclidean/base.fbin", - query_file: "mnist-784-euclidean/query.fbin", - groundtruth_neighbors_file: "mnist-784-euclidean/groundtruth.neighbors.ibin", - distance: "euclidean" - -- name: "nytimes-256-angular", - base_file: "nytimes-256-angular/base.fbin", - query_file: "nytimes-256-angular/query.fbin", - groundtruth_neighbors_file: "nytimes-256-angular/groundtruth.neighbors.ibin", - distance: "euclidean" - -- name: "nytimes-256-inner", - base_file: "nytimes-256-inner/base.fbin", - query_file: "nytimes-256-inner/query.fbin", - groundtruth_neighbors_file: "nytimes-256-inner/groundtruth.neighbors.ibin", - distance: "euclidean" - -- name: "sift-128-euclidean", - base_file: "sift-128-euclidean/base.fbin", - query_file: "sift-128-euclidean/query.fbin", - groundtruth_neighbors_file: "sift-128-euclidean/groundtruth.neighbors.ibin", - distance: "euclidean" \ No newline at end of file From b00942d673ac97cd4fe8767d3c6b45b7ba4ff927 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 25 Oct 2023 09:33:15 -0700 Subject: [PATCH 06/49] add algo-groups --- .../src/raft-ann-bench/plot/__main__.py | 94 +++++++++++++++---- .../src/raft-ann-bench/run/__main__.py | 47 +++++++--- 2 files changed, 113 insertions(+), 28 deletions(-) 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 ad3a2490db..0c6e554e96 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 @@ -345,32 +345,74 @@ def load_lines(results_path, result_files, method, index_key): def load_all_results( - dataset_path, algorithms, k, batch_size, method, index_key + dataset_path, algorithms, groups, algo_groups, k, batch_size, method, + index_key ): results_path = os.path.join(dataset_path, "result", method) result_files = os.listdir(results_path) - print(result_files) + result_files = [result_file for result_file in result_files \ + if ".csv" in result_file] + # print(result_files) if method == "search": result_files = [ result_filename for result_filename in result_files if f"{k}-{batch_size}" in result_filename ] - if len(algorithms) > 0: - result_files = [ - result_filename - for result_filename in result_files - if result_filename.split("-")[0] in algorithms - ] - elif method == "build": - if len(algorithms) > 0: - result_files = [ - result_filename - for result_filename in result_files - if result_filename.split("-")[0] in algorithms - ] - - results = load_lines(results_path, result_files, method, index_key) + algo_group_files = [ + result_filename.split("-")[0] + for result_filename in result_files + ] + else: + 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 = list(zip(*algo_group_files)) + # final_groups = [result_files[i] for i in range(len(result_files)) if \ + # algo_group_files[i][1] in groups] + # if len(algorithms) > 0: + # final_algos = [final_groups[i] for i in range(len(result_files)) if \ + # ("_".join(result_files[i].split("_")[:-1]) in algorithms)] + # final_results = [] + if len(algorithms) > 0: + final_results = [result_files[i] for i in range(len(result_files)) if \ + (algo_group_files[0][i] in algorithms) and \ + (algo_group_files[1][i] in groups)] + else: + final_results = [result_files[i] for i in range(len(result_files)) if \ + (algo_group_files[1][i] in groups)] + + if len(algo_groups) > 0: + split_algo_groups = [algo_group.split(".") for algo_group in algo_groups] + split_algo_groups = list(zip(*split_algo_groups)) + final_algo_groups = [result_files[i] for i in range(len(result_files)) if \ + (algo_group_files[0][i] in split_algo_groups[0]) and \ + (algo_group_files[1][i] in split_algo_groups[1])] + final_results = final_results + final_algo_groups + final_results = set(final_results) + + # if len(algorithms) > 0: + # result_files = [ + # result_filename + # for result_filename in result_files + # if result_filename.split("-")[0] in algorithms + # ] + # elif method == "build": + # if len(algorithms) > 0: + # result_files = [ + # result_filename + # for result_filename in result_files + # if result_filename.split("-")[0] in algorithms + # ] + + results = load_lines(results_path, final_results, method, index_key) return results @@ -404,6 +446,15 @@ def main(): algorithms", default=None, ) + parser.add_argument( + "--groups", + help="plot only comma separated groups of parameters", + default="base" + ) + parser.add_argument( + "--algo-groups", + help="add comma separated algorithm+groups to the plot", + ) parser.add_argument( "-k", "--count", @@ -444,6 +495,11 @@ def main(): algorithms = args.algorithms.split(",") else: algorithms = [] + groups = args.groups.split(",") + if args.algo_groups: + algo_groups = args.algo_groups.split(",") + else: + algo_groups = [] k = args.count batch_size = args.batch_size if not args.build and not args.search: @@ -465,6 +521,8 @@ def main(): search_results = load_all_results( os.path.join(args.dataset_path, args.dataset), algorithms, + groups, + algo_groups, k, batch_size, "search", @@ -487,6 +545,8 @@ def main(): build_results = load_all_results( os.path.join(args.dataset_path, args.dataset), algorithms, + groups, + algo_groups, k, batch_size, "build", 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 01f4adbac8..2cff55e68e 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 @@ -45,15 +45,20 @@ def validate_algorithm(algos_conf, algo, gpu_present): ) -def find_executable(algos_conf, algo, k, batch_size): +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}" + build_path = os.getenv("RAFT_HOME") if build_path is not None: build_path = os.path.join(build_path, "cpp", "build", executable) if os.path.exists(build_path): print(f"-- Using RAFT bench from repository in {build_path}. ") - return (executable, build_path, f"{algo}-{k}-{batch_size}") + return (executable, build_path, return_str) # if there is no build folder present, we look in the conda environment conda_path = os.getenv("CONDA_PREFIX") @@ -61,7 +66,7 @@ def find_executable(algos_conf, algo, k, batch_size): conda_path = os.path.join(conda_path, "bin", "ann", executable) if os.path.exists(conda_path): print("-- Using RAFT bench found in conda environment. ") - return (executable, conda_path, f"{algo}-{k}-{batch_size}") + return (executable, conda_path, return_str) else: raise FileNotFoundError(executable) @@ -214,9 +219,13 @@ def main(): # ) parser.add_argument( "--groups", - help="comma separated groups of parameters to run the benchmarks for", + help="run only comma separated groups of parameters", default="base" ) + parser.add_argument( + "--algo-groups", + help="add comma separated algorithm+groups to run", + ) parser.add_argument( "-f", "--force", @@ -280,30 +289,44 @@ def main(): if filter_algos: allowed_algos = args.algorithms.split(",") named_groups = args.groups.split(",") + filter_algo_groups = True if args.algo_groups else False + allowed_algo_groups = None + if filter_algo_groups: + allowed_algo_groups = [algo_group.split(".") for algo_group in args.algo_groups.split(",")] + allowed_algo_groups = list(zip(*allowed_algo_groups)) algos_conf = dict() for algo_f in algos_conf_fs: with open(algo_f, "r") as f: if algo_f.split("/")[-1] == "raft_cagra.yaml": algo = yaml.safe_load(f) insert_algo = True + insert_algo_group = False if filter_algos: if algo["name"] not in allowed_algos: insert_algo = False - if insert_algo: + if filter_algo_groups: + if algo["name"] in allowed_algo_groups[0]: + insert_algo_group = True + def add_algo_group(group_list): if algo["name"] not in algos_conf: algos_conf[algo["name"]] = dict() for group in algo.keys(): if group != "name": - if group in named_groups: + if group in group_list: algos_conf[algo["name"]][group] = algo[group] + if insert_algo: + add_algo_group(named_groups) + if insert_algo_group: + add_algo_group(allowed_algo_groups[1]) + print(algos_conf) executables_to_run = dict() for algo in algos_conf.keys(): validate_algorithm(algos_yaml, algo, gpu_present) - executable = find_executable(algos_yaml, algo, k, batch_size) - if executable not in executables_to_run: - executables_to_run[executable] = {"index": []} for group in algos_conf[algo].keys(): + executable = find_executable(algos_yaml, algo, group, k, batch_size) + if executable not in executables_to_run: + executables_to_run[executable] = {"index": []} build_params = algos_conf[algo][group]["build"] search_params = algos_conf[algo][group]["search"] @@ -323,7 +346,10 @@ def main(): for params in all_build_params: index = {"algo": algo, "build_param": {}} - index_name = f"{algo}" + if group != "base": + index_name = f"{algo}_{group}" + else: + index_name = f"{algo}" for i in range(len(params)): index["build_param"][param_names[i]] = params[i] index_name += "." + f"{param_names[i]}{params[i]}" @@ -408,7 +434,6 @@ def main(): # ) # executables_to_run[executable_path]["index"][pos] = index - print(conf_filedir) run_build_and_search( conf_file, f"{args.dataset}.json", From 1f81d455e676a55513a2b5a0aac8522c9b065d2e Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Wed, 25 Oct 2023 14:51:29 -0400 Subject: [PATCH 07/49] Updated validatoes --- .../run/conf/algos/raft_cagra.yaml | 2 ++ .../run/conf/algos/raft_ivf_pq.yaml | 14 +++++++++ .../src/raft-ann-bench/validators/__init__.py | 29 +++++++++++++++++++ 3 files changed, 45 insertions(+) create mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yaml create mode 100644 python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py 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 095f77e66a..32ed400b1a 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,4 +1,6 @@ name: raft_cagra +validators: + search: raft-ann-bench.validators.raft_cagra_build_validator base: build: graph_degree: [32, 64] 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 new file mode 100644 index 0000000000..0a1853f582 --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yaml @@ -0,0 +1,14 @@ +name: raft_ivf_pq +validators: + search: raft-ann-bench.validators.raft_cagra_build_validator +base: + build: + nlist: [1024] + pq_dim: [128, 64] + pq_bits: [8, 6] + ratio: [1] + niter: [25] + search: + nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] + internalDistanceDtype: ["float", "half"] + smemLutDtype: ["float", "fp8", "half"] \ No newline at end of file diff --git a/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py b/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py new file mode 100644 index 0000000000..04448425ee --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py @@ -0,0 +1,29 @@ +# +# Copyright (c) 2023, NVIDIA CORPORATION. +# +# Licensed 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. + +DTYPE_SIZES = {"float": 4, "half": 2, "fp8": 1} + + +def raft_ivf_pq_search_validator(params, k, batch_size): + if "internalDistanceDtype" in params and "smemLutDtype" in params: + return ( + DTYPE_SIZES[params["smemLutDtype"]] + >= DTYPE_SIZES[params["internalDistanceDtype"]] + ) + + +def raft_cagra_search_validator(params, k, batch_size): + if "itopk" in params: + return params["itopk"] >= k From bcf92c56d0580567ecd26d524d878889b4e817bf Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 26 Oct 2023 09:15:03 -0700 Subject: [PATCH 08/49] remove only cagra yaml load --- .../src/raft-ann-bench/run/__main__.py | 41 +++++++++---------- 1 file changed, 20 insertions(+), 21 deletions(-) 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 2cff55e68e..b7506b2f42 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 @@ -297,27 +297,26 @@ def main(): algos_conf = dict() for algo_f in algos_conf_fs: with open(algo_f, "r") as f: - if algo_f.split("/")[-1] == "raft_cagra.yaml": - algo = yaml.safe_load(f) - insert_algo = True - insert_algo_group = False - if filter_algos: - if algo["name"] not in allowed_algos: - insert_algo = False - if filter_algo_groups: - if algo["name"] in allowed_algo_groups[0]: - insert_algo_group = True - def add_algo_group(group_list): - if algo["name"] not in algos_conf: - algos_conf[algo["name"]] = dict() - for group in algo.keys(): - if group != "name": - if group in group_list: - algos_conf[algo["name"]][group] = algo[group] - if insert_algo: - add_algo_group(named_groups) - if insert_algo_group: - add_algo_group(allowed_algo_groups[1]) + algo = yaml.safe_load(f) + insert_algo = True + insert_algo_group = False + if filter_algos: + if algo["name"] not in allowed_algos: + insert_algo = False + if filter_algo_groups: + if algo["name"] in allowed_algo_groups[0]: + insert_algo_group = True + def add_algo_group(group_list): + if algo["name"] not in algos_conf: + algos_conf[algo["name"]] = dict() + for group in algo.keys(): + if group != "name": + if group in group_list: + algos_conf[algo["name"]][group] = algo[group] + if insert_algo: + add_algo_group(named_groups) + if insert_algo_group: + add_algo_group(allowed_algo_groups[1]) print(algos_conf) executables_to_run = dict() From 933732cea15fd9ddbbf705a1d7460126ce188911 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 26 Oct 2023 09:20:20 -0700 Subject: [PATCH 09/49] fix configs --- .../src/raft-ann-bench/run/__main__.py | 2 +- .../run/conf/algos/faiss_gpu_ivf_pq.yml | 4 ---- .../run/conf/algos/raft_ivf_flat.yaml | 8 ++++++++ .../run/conf/algos/raft_ivf_flat.yml | 8 -------- .../run/conf/algos/raft_ivf_pq.yml | 17 ----------------- 5 files changed, 9 insertions(+), 30 deletions(-) delete mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yml create mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yaml delete mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yml delete mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml 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 b7506b2f42..64f88b5ebd 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 @@ -310,7 +310,7 @@ def add_algo_group(group_list): if algo["name"] not in algos_conf: algos_conf[algo["name"]] = dict() for group in algo.keys(): - if group != "name": + if group != "name" or group != "validators": if group in group_list: algos_conf[algo["name"]][group] = algo[group] if insert_algo: diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yml deleted file mode 100644 index 136f9ffcab..0000000000 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yml +++ /dev/null @@ -1,4 +0,0 @@ -name: faiss_gpu_ivf_pq - - base: - build: - search: \ No newline at end of file diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yaml new file mode 100644 index 0000000000..2a41fb035d --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yaml @@ -0,0 +1,8 @@ +name: raft_ivf_flat +base: + build: + nlist: [1024, 16384] + ratio: [1, 2] + niter: [20, 25] + search: + nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] \ No newline at end of file diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yml deleted file mode 100644 index 4af425896a..0000000000 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yml +++ /dev/null @@ -1,8 +0,0 @@ -name: raft_ivf_flat - - base: - build: - nlist: [1024, 16384] - ratio: [1, 2] - niter: [20, 25] - search: - nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] \ No newline at end of file diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml deleted file mode 100644 index 3ef90bfb0d..0000000000 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yml +++ /dev/null @@ -1,17 +0,0 @@ -name: raft_ivf_pq - - base: - build: - nlist: [1024] - pq_dim: [128, 64] - pq_bits: [8, 6] - ratio: [1] - niter: [25] - search: - nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] - internalDistanceDtype: ["float", "half"] - smemLutDtype: ["float", "fp8", "half"] - - # Custom validator with a `validate()` function that accepts a parameter combination - # and returns false if it's not valid (for example, - # internalDistanceDtype should never be lower than smemLutDtype) - validators: raft-ann-bench.validators.raft_ivf_pq_validator \ No newline at end of file From 32430c15ba353b2afebb3a24d3b0f50f980cb289 Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 27 Oct 2023 16:47:53 -0700 Subject: [PATCH 10/49] working yaml param sweeps --- .../src/raft-ann-bench/plot/__main__.py | 21 +-- .../src/raft-ann-bench/run/__main__.py | 125 ++++++------------ .../run/conf/algos/raft_cagra.yaml | 25 ++-- .../run/conf/algos/raft_ivf_flat.yaml | 15 ++- .../run/conf/algos/raft_ivf_pq.yaml | 25 ++-- 5 files changed, 74 insertions(+), 137 deletions(-) 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 0c6e554e96..884936dac5 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 @@ -375,12 +375,7 @@ def load_all_results( else: algo_group_files[i] = ("_".join(algo_group[:-1]), algo_group[-1]) algo_group_files = list(zip(*algo_group_files)) - # final_groups = [result_files[i] for i in range(len(result_files)) if \ - # algo_group_files[i][1] in groups] - # if len(algorithms) > 0: - # final_algos = [final_groups[i] for i in range(len(result_files)) if \ - # ("_".join(result_files[i].split("_")[:-1]) in algorithms)] - # final_results = [] + if len(algorithms) > 0: final_results = [result_files[i] for i in range(len(result_files)) if \ (algo_group_files[0][i] in algorithms) and \ @@ -398,20 +393,6 @@ def load_all_results( final_results = final_results + final_algo_groups final_results = set(final_results) - # if len(algorithms) > 0: - # result_files = [ - # result_filename - # for result_filename in result_files - # if result_filename.split("-")[0] in algorithms - # ] - # elif method == "build": - # if len(algorithms) > 0: - # result_files = [ - # result_filename - # for result_filename in result_files - # if result_filename.split("-")[0] in algorithms - # ] - results = load_lines(results_path, final_results, method, index_key) return results 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 64f88b5ebd..b49dc01351 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 @@ -14,13 +14,13 @@ # limitations under the License. import argparse +from importlib import import_module import itertools import json import os import subprocess import yaml - def positive_int(input_str: str) -> int: try: i = int(input_str) @@ -137,7 +137,7 @@ def run_build_and_search( cmd = cmd + [temp_conf_filepath] subprocess.run(cmd, check=True) - # os.remove(temp_conf_filepath) + os.remove(temp_conf_filepath) def main(): @@ -211,12 +211,6 @@ def main(): algorithms", default=None, ) - # parser.add_argument( - # "--indices", - # help="run only comma separated list of named indices. \ - # parameter `algorithms` is ignored", - # default=None, - # ) parser.add_argument( "--groups", help="run only comma separated groups of parameters", @@ -258,9 +252,13 @@ def main(): with open(dataset_conf_f, "r") as f: dataset_conf_all = yaml.safe_load(f) + dataset_conf = None for dataset in dataset_conf_all: if args.dataset == dataset["name"]: dataset_conf = dataset + break + if not dataset_conf: + raise ValueError("Could not find a dataset configuration") conf_file = dict() conf_file["dataset"] = dataset_conf @@ -308,11 +306,14 @@ def main(): insert_algo_group = True def add_algo_group(group_list): if algo["name"] not in algos_conf: - algos_conf[algo["name"]] = dict() - for group in algo.keys(): - if group != "name" or group != "validators": - if group in group_list: - algos_conf[algo["name"]][group] = algo[group] + algos_conf[algo["name"]] = {"groups" : {}} + for group in algo["groups"].keys(): + if group in group_list: + algos_conf[algo["name"]]["groups"][group] = \ + algo["groups"][group] + if "validators" in algo: + algos_conf[algo["name"]]["validators"] = \ + algo["validators"] if insert_algo: add_algo_group(named_groups) if insert_algo_group: @@ -322,12 +323,12 @@ def add_algo_group(group_list): executables_to_run = dict() for algo in algos_conf.keys(): validate_algorithm(algos_yaml, algo, gpu_present) - for group in algos_conf[algo].keys(): + for group in algos_conf[algo]["groups"].keys(): executable = find_executable(algos_yaml, algo, group, k, batch_size) if executable not in executables_to_run: executables_to_run[executable] = {"index": []} - build_params = algos_conf[algo][group]["build"] - search_params = algos_conf[algo][group]["search"] + build_params = algos_conf[algo]["groups"][group]["build"] + search_params = algos_conf[algo]["groups"][group]["search"] param_names = [] param_lists = [] @@ -352,6 +353,18 @@ def add_algo_group(group_list): for i in range(len(params)): 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"].split(".") + module = ".".join(importable[:-1]) + func = importable[-1] + validator = import_module(module) + build_validator = getattr(validator, func) + if not build_validator(index["build_param"]): + continue + index["name"] = index_name index["file"] = os.path.join(args.dataset_path, args.dataset, "index", index_name) @@ -361,78 +374,18 @@ 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] - index["search_params"].append(search_dict) + if "validators" in algos_conf[algo]: + if "search" in algos_conf[algo]["validators"]: + importable = \ + algos_conf[algo]["validators"]["search"].split(".") + module = ".".join(importable[:-1]) + func = importable[-1] + validator = import_module(module) + search_validator = getattr(validator, func) + if search_validator(search_dict, k, batch_size): + index["search_params"].append(search_dict) executables_to_run[executable]["index"].append(index) - print(executables_to_run) - # conf_filename = conf_filepath.split("/")[-1] - # conf_filedir = "/".join(conf_filepath.split("/")[:-1]) - # dataset_path = args.dataset_path - # if not os.path.exists(conf_filepath): - # raise FileNotFoundError(conf_filename) - - # with open(conf_filepath, "r") as f: - # conf_file = json.load(f) - - # dataset_name = conf_file["dataset"]["name"] - - # executables_to_run = dict() - # # At least one named index should exist in config file - # if args.indices: - # indices = set(args.indices.split(",")) - # # algo associated with index should still be present in algos.yaml - # # and enabled - # for index in conf_file["index"]: - # curr_algo = index["algo"] - # if index["name"] in indices and validate_algorithm( - # algos_conf, curr_algo, gpu_present - # ): - # executable_path = find_executable( - # algos_conf, curr_algo, k, batch_size - # ) - # if executable_path not in executables_to_run: - # executables_to_run[executable_path] = {"index": []} - # executables_to_run[executable_path]["index"].append(index) - - # # switch to named algorithms if indices parameter is not supplied - # elif args.algorithms: - # algorithms = set(args.algorithms.split(",")) - # # pick out algorithms from conf file that exist - # # and are enabled in algos.yaml - # for index in conf_file["index"]: - # curr_algo = index["algo"] - # if curr_algo in algorithms and validate_algorithm( - # algos_conf, curr_algo, gpu_present - # ): - # executable_path = find_executable( - # algos_conf, curr_algo, k, batch_size - # ) - # if executable_path not in executables_to_run: - # executables_to_run[executable_path] = {"index": []} - # executables_to_run[executable_path]["index"].append(index) - - # # default, try to run all available algorithms - # else: - # for index in conf_file["index"]: - # curr_algo = index["algo"] - # if validate_algorithm(algos_conf, curr_algo, gpu_present): - # executable_path = find_executable( - # algos_conf, curr_algo, k, batch_size - # ) - # if executable_path not in executables_to_run: - # executables_to_run[executable_path] = {"index": []} - # executables_to_run[executable_path]["index"].append(index) - - # # Replace index to dataset path - # for executable_path in executables_to_run: - # for pos, index in enumerate( - # executables_to_run[executable_path]["index"] - # ): - # index["file"] = os.path.join( - # dataset_path, dataset_name, "index", index["name"] - # ) - # executables_to_run[executable_path]["index"][pos] = index - run_build_and_search( conf_file, f"{args.dataset}.json", 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 32ed400b1a..3e90494968 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,15 +1,16 @@ name: raft_cagra validators: - search: raft-ann-bench.validators.raft_cagra_build_validator -base: - build: - graph_degree: [32, 64] - intermediate_graph_degree: [64, 96] - search: - itopk: [32, 64, 128] + search: raft-ann-bench.validators.raft_cagra_search_validator +groups: + base: + build: + graph_degree: [32, 64] + intermediate_graph_degree: [64, 96] + search: + itopk: [32, 64, 128] -large: - build: - graph_degree: [32, 64] - search: - itopk: [32, 64, 128] + large: + build: + graph_degree: [32, 64] + search: + itopk: [32, 64, 128] diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yaml index 2a41fb035d..6c476af39e 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yaml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yaml @@ -1,8 +1,9 @@ name: raft_ivf_flat -base: - build: - nlist: [1024, 16384] - ratio: [1, 2] - niter: [20, 25] - search: - nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] \ No newline at end of file +groups: + base: + build: + nlist: [1024, 16384] + ratio: [1, 2] + niter: [20, 25] + search: + nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] \ No newline at end of file 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 0a1853f582..107e0ac640 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,14 +1,15 @@ name: raft_ivf_pq validators: - search: raft-ann-bench.validators.raft_cagra_build_validator -base: - build: - nlist: [1024] - pq_dim: [128, 64] - pq_bits: [8, 6] - ratio: [1] - niter: [25] - search: - nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] - internalDistanceDtype: ["float", "half"] - smemLutDtype: ["float", "fp8", "half"] \ No newline at end of file + search: raft-ann-bench.validators.raft_ivf_pq_search_validator +groups: + base: + build: + nlist: [1024] + pq_dim: [128, 64] + pq_bits: [8, 6] + ratio: [1] + niter: [25] + search: + nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] + internalDistanceDtype: ["float", "half"] + smemLutDtype: ["float", "fp8", "half"] \ No newline at end of file From b65e6d3a06542259f5b0fb20ca187ef91fa7b537 Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 27 Oct 2023 17:28:23 -0700 Subject: [PATCH 11/49] add docs --- docs/source/raft_ann_benchmarks.md | 152 +++++++++++------- .../src/raft-ann-bench/plot/__main__.py | 2 +- 2 files changed, 93 insertions(+), 61 deletions(-) diff --git a/docs/source/raft_ann_benchmarks.md b/docs/source/raft_ann_benchmarks.md index 25fdf3f0f6..a0ef8806a2 100644 --- a/docs/source/raft_ann_benchmarks.md +++ b/docs/source/raft_ann_benchmarks.md @@ -4,20 +4,27 @@ This project provides a benchmark program for various ANN search implementations ## Table of Contents -- [Installing the benchmarks](#installing-the-benchmarks) - - [Conda](#conda) - - [Docker](#docker) -- [How to run the benchmarks](#how-to-run-the-benchmarks) - - [Step 1: prepare dataset](#step-1-prepare-dataset) - - [Step 2: build and search index](#step-2-build-and-search-index) - - [Step 3: data export](#step-3-data-export) - - [Step 4: plot results](#step-4-plot-results) -- [Running the benchmarks](#running-the-benchmarks) - - [End to end: small-scale (<1M to 10M)](#end-to-end-small-scale-benchmarks-1m-to-10m) - - [End to end: large-scale (>10M)](#end-to-end-large-scale-benchmarks-10m-vectors) - - [Running with Docker containers](#running-with-docker-containers) -- [Creating and customizing dataset configurations](#creating-and-customizing-dataset-configurations) -- [Adding a new ANN algorithm](#adding-a-new-ann-algorithm) +- [RAFT ANN Benchmarks](#raft-ann-benchmarks) + - [Table of Contents](#table-of-contents) + - [Installing the benchmarks](#installing-the-benchmarks) + - [Conda](#conda) + - [Docker](#docker) + - [How to run the benchmarks](#how-to-run-the-benchmarks) + - [Step 1: Prepare Dataset](#step-1-prepare-dataset) + - [Step 2: Build and Search Index](#step-2-build-and-search-index) + - [Step 3: Data Export](#step-3-data-export) + - [Step 4: Plot Results](#step-4-plot-results) + - [Running the benchmarks](#running-the-benchmarks) + - [End to end: small-scale benchmarks (\<1M to 10M)](#end-to-end-small-scale-benchmarks-1m-to-10m) + - [End to end: large-scale benchmarks (\>10M vectors)](#end-to-end-large-scale-benchmarks-10m-vectors) + - [Running with Docker containers](#running-with-docker-containers) + - [End-to-end run on GPU](#end-to-end-run-on-gpu) + - [End-to-end run on CPU](#end-to-end-run-on-cpu) + - [Manually run the scripts inside the container](#manually-run-the-scripts-inside-the-container) + - [Creating and customizing dataset configurations](#creating-and-customizing-dataset-configurations) + - [Adding a new ANN algorithm](#adding-a-new-ann-algorithm) + - [Implementation and Configuration](#implementation-and-configuration) + - [Adding a CMake Target](#adding-a-cmake-target) ## Installing the benchmarks @@ -119,32 +126,47 @@ specified configuration. The usage of the script `raft-ann-bench.run` is: ```bash -usage: run.py [-h] [-k COUNT] [-bs BATCH_SIZE] [--configuration CONFIGURATION] [--dataset DATASET] [--dataset-path DATASET_PATH] [--build] [--search] [--algorithms ALGORITHMS] [--indices INDICES] - [-f] +usage: [-h] [--subset-size SUBSET_SIZE] [-k COUNT] [-bs BATCH_SIZE] [--dataset-configuration DATASET_CONFIGURATION] [--configuration CONFIGURATION] [--dataset DATASET] + [--dataset-path DATASET_PATH] [--build] [--search] [--algorithms ALGORITHMS] [--groups GROUPS] [--algo-groups ALGO_GROUPS] [-f] options: -h, --help show this help message and exit + --subset-size SUBSET_SIZE + the number of subset rows of the dataset to build the index (default: None) -k COUNT, --count COUNT the number of nearest neighbors to search for (default: 10) -bs BATCH_SIZE, --batch-size BATCH_SIZE number of query vectors to use in each query trial (default: 10000) + --dataset-configuration DATASET_CONFIGURATION + path to configuration file for datasets (default: None) --configuration CONFIGURATION - path to configuration file for a dataset (default: None) - --dataset DATASET dataset whose configuration file will be used (default: glove-100-inner) + path to configuration file or directory for algorithms (default: None) + --dataset DATASET name of dataset (default: glove-100-inner) --dataset-path DATASET_PATH - path to dataset folder (default: ${RAPIDS_DATASET_ROOT_DIR}) + path to dataset folder, by default will look in RAPIDS_DATASET_ROOT_DIR if defined, otherwise a datasets subdirectory from the calling directory (default: + /raid/dgala/raft/datasets/) --build --search --algorithms ALGORITHMS run only comma separated list of named algorithms (default: None) - --indices INDICES run only comma separated list of named indices. parameter `algorithms` is ignored (default: None) + --groups GROUPS run only comma separated groups of parameters (default: base) + --algo-groups ALGO_GROUPS + add comma separated algorithm+groups to run (default: None) -f, --force re-run algorithms even if their results already exist (default: False) ``` -`configuration` and `dataset` : `configuration` is a path to a configuration file for a given dataset. -The configuration file should be name as `.json`. It is optional if the name of the dataset is -provided with the `dataset` argument, in which case -a configuration file will be searched for as `python/raft-ann-bench/src/raft-ann-bench/run/conf/.json`. +`dataset`: name of the dataset to be searched in [datasets.yaml](#yaml-dataset-config) + +`dataset-configuration`: optional filepath to custom dataset YAML config which has an entry for arg `dataset` + +`configuration`: optional filepath to YAML configuration for an algorithm or to directory that contains YAML configurations for several algorithms. [Here's how to configure an algorithm.](#yaml-algo-config) + +`algorithms`: runs all algorithms that it can find in YAML configs found by `configuration` + +`groups`: run only specific groups of parameters configurations for an algorithm. Groups are defined in YAML configs (see `configuration`), and by default run `base` group + +`algo-groups`: this parameter is helpful to append any specific algorithm+group combination to run the benchmark for in addition to all the arguments from `algorithms` and `groups`. It is of the format `.`, or for example, `raft_cagra.large` + For every algorithm run by this script, it outputs an index build statistics JSON file in `/result/build/` and an index search statistics JSON file in `/result/search/`. @@ -181,18 +203,21 @@ CSV file in `/result/search/<-k{k}-batch_size{batch_size} The usage of this script is: ```bash -usage: plot.py [-h] [--dataset DATASET] [--dataset-path DATASET_PATH] [--output-filepath OUTPUT_FILEPATH] [--algorithms ALGORITHMS] [-k COUNT] [-bs BATCH_SIZE] [--build] [--search] - [--x-scale X_SCALE] [--y-scale {linear,log,symlog,logit}] [--raw] +usage: [-h] [--dataset DATASET] [--dataset-path DATASET_PATH] [--output-filepath OUTPUT_FILEPATH] [--algorithms ALGORITHMS] [--groups GROUPS] [--algo-groups ALGO_GROUPS] [-k COUNT] + [-bs BATCH_SIZE] [--build] [--search] [--x-scale X_SCALE] [--y-scale {linear,log,symlog,logit}] [--raw] -options: +optional arguments: -h, --help show this help message and exit - --dataset DATASET dataset to download (default: glove-100-inner) + --dataset DATASET dataset to plot (default: glove-100-inner) --dataset-path DATASET_PATH - path to dataset folder (default: ${RAPIDS_DATASET_ROOT_DIR}) + path to dataset folder (default: /raid/dgala/raft/datasets/) --output-filepath OUTPUT_FILEPATH - directory for PNG to be saved (default: os.getcwd()) + directory for PNG to be saved (default: /raid/dgala/raft) --algorithms ALGORITHMS plot only comma separated list of named algorithms (default: None) + --groups GROUPS plot only comma separated groups of parameters (default: base) + --algo-groups ALGO_GROUPS + add comma separated algorithm+groups to the plot (default: None) -k COUNT, --count COUNT the number of nearest neighbors to search for (default: 10) -bs BATCH_SIZE, --batch-size BATCH_SIZE @@ -204,6 +229,11 @@ options: Scale to use when drawing the Y-axis (default: linear) --raw Show raw results (not just Pareto frontier) in faded colours (default: False) ``` +`algorithms`: plots all algorithms that it can find results for the specified `dataset` + +`groups`: plot only specific groups of parameters configurations for an algorithm. Groups are defined in YAML configs (see `configuration`), and by default run `base` group + +`algo-groups`: this parameter is helpful to append any specific algorithm+group combination to plot results for in addition to all the arguments from `algorithms` and `groups`. It is of the format `.`, or for example, `raft_cagra.large` The figure below is the resulting plot of running our benchmarks as of August 2023 for a batch size of 10, on an NVIDIA H100 GPU and an Intel Xeon Platinum 8480CL CPU. It presents the throughput (in Queries-Per-Second) performance for every level of recall. @@ -352,40 +382,44 @@ Additionally, the containers can be run in detached mode without any issue. ## Creating and customizing dataset configurations -A single configuration file will often define a set of algorithms, with associated index and search parameters, for a specific dataset. A configuration file uses json format with 4 major parts: -1. Dataset information -2. Algorithm information -3. Index parameters -4. Search parameters +A single configuration will often define a set of algorithms, with associated index and search parameters, that can be generalize across datasets. We use YAML to define dataset specific and algorithm specific configurations. -Below is a simple example configuration file for the 1M-scale `sift-128-euclidean` dataset: +A default `datasets.yaml` is provided by RAFT in `${RAFT_HOME}/python/raft-ann-bench/src/raft-ann-bench/run/conf` with configurations available for several datasets. Here's a simple example entry for the `sift-128-euclidean` dataset: -```json -{ - "dataset": { - "name": "sift-128-euclidean", - "base_file": "sift-128-euclidean/base.fbin", - "query_file": "sift-128-euclidean/query.fbin", - "subset_size": 1000000, - "groundtruth_neighbors_file": "sift-128-euclidean/groundtruth.neighbors.ibin", - "distance": "euclidean" - }, - "index": [] -} +```yaml +- name: sift-128-euclidean + base_file: sift-128-euclidean/base.fbin + query_file: sift-128-euclidean/query.fbin + groundtruth_neighbors_file: sift-128-euclidean/groundtruth.neighbors.ibin + distance: euclidean ``` -The `index` section will contain a list of index objects, each of which will have the following form: -```json -{ - "name": "algo_name.unique_index_name", - "algo": "algo_name", - "file": "sift-128-euclidean/algo_name/param1_val1-param2_val2", - "build_param": { "param1": "val1", "param2": "val2" }, - "search_params": [{ "search_param1": "search_val1" }] -} +Configuration files for ANN algorithms supported by `raft-ann-bench` are provided in `${RAFT_HOME}/python/raft-ann-bench/src/raft-ann-bench/run/conf`. `raft_cagra` algorithm configuration looks like: +```yaml +name: raft_cagra +validators: + search: raft-ann-bench.validators.raft_cagra_search_validator +groups: + base: + build: + graph_degree: [32, 64] + intermediate_graph_degree: [64, 96] + search: + itopk: [32, 64, 128] + + large: + build: + graph_degree: [32, 64] + search: + itopk: [32, 64, 128] ``` +It is mandatory that every algorithm have a `base` group. -The table below contains the possible settings for the `algo` field. Each unique algorithm will have its own set of `build_param` and `search_params` settings. The [ANN Algorithm Parameter Tuning Guide](ann_benchmarks_param_tuning.md) contains detailed instructions on choosing build and search parameters for each supported algorithm. +There's several things going on here: +1. `groups` - define a run group which has a particular set of parameters. Each group has `build` parameters wherein a cross-product of all parameters will be used to build a unique index. Each group also has `search` parameters, and a cross-product of all search parameters will be used to search all unique indices that we create. +2. `validators` - This has two entries `build` and `search`. It's a string denoting a Python importable function of a module that will be used to verify whether there are any bogus parameters in any `groups` that do not make sense. If such parameters exist, they will be silently ignored. + +The table below contains all algorithms supported by RAFT. Each unique algorithm will have its own set of `build` and `search` settings. The [ANN Algorithm Parameter Tuning Guide](ann_benchmarks_param_tuning.md) contains detailed instructions on choosing build and search parameters for each supported algorithm. | Library | Algorithms | |-----------|------------------------------------------------------------------| @@ -395,8 +429,6 @@ The table below contains the possible settings for the `algo` field. Each unique | HNSWlib | `hnswlib` | | RAFT | `raft_brute_force`, `raft_cagra`, `raft_ivf_flat`, `raft_ivf_pq` | -By default, the index will be placed in `bench/ann/data//index/`. Using `sift-128-euclidean` for the dataset with the `algo` example above, the indexes would be placed in `bench/ann/data/sift-128-euclidean/index/algo_name/param1_val1-param2_val2`. - ## Adding a new ANN algorithm ### Implementation and Configuration 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 884936dac5..a4a8379eb4 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 @@ -409,7 +409,7 @@ def main(): formatter_class=argparse.ArgumentDefaultsHelpFormatter ) parser.add_argument( - "--dataset", help="dataset to download", default="glove-100-inner" + "--dataset", help="dataset to plots", default="glove-100-inner" ) parser.add_argument( "--dataset-path", From 1a810fd485019f6ac86469485183804587669d55 Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 27 Oct 2023 17:28:36 -0700 Subject: [PATCH 12/49] add bench-ann cuda12 envs --- .../bench_ann_cuda-120_arch-aarch64.yaml | 39 +++++++++++++++++++ .../bench_ann_cuda-120_arch-x86_64.yaml | 39 +++++++++++++++++++ dependencies.yaml | 2 +- 3 files changed, 79 insertions(+), 1 deletion(-) create mode 100644 conda/environments/bench_ann_cuda-120_arch-aarch64.yaml create mode 100644 conda/environments/bench_ann_cuda-120_arch-x86_64.yaml diff --git a/conda/environments/bench_ann_cuda-120_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-120_arch-aarch64.yaml new file mode 100644 index 0000000000..99ec02859d --- /dev/null +++ b/conda/environments/bench_ann_cuda-120_arch-aarch64.yaml @@ -0,0 +1,39 @@ +# This file is generated by `rapids-dependency-file-generator`. +# To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +channels: +- rapidsai +- rapidsai-nightly +- dask/label/dev +- conda-forge +- nvidia +dependencies: +- benchmark>=1.8.2 +- c-compiler +- clang-tools=16.0.6 +- clang==16.0.6 +- cmake>=3.26.4 +- cuda-cudart-dev +- cuda-nvcc +- cuda-profiler-api +- cuda-version=12.0 +- cxx-compiler +- cython>=3.0.0 +- gcc_linux-aarch64=11.* +- glog>=0.6.0 +- h5py>=3.8.0 +- hnswlib=0.7.0 +- libcublas-dev +- libcurand-dev +- libcusolver-dev +- libcusparse-dev +- matplotlib +- nccl>=2.9.9 +- ninja +- nlohmann_json>=3.11.2 +- openblas +- pandas +- pyyaml +- rmm==23.12.* +- scikit-build>=0.13.1 +- sysroot_linux-aarch64==2.17 +name: bench_ann_cuda-120_arch-aarch64 diff --git a/conda/environments/bench_ann_cuda-120_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-120_arch-x86_64.yaml new file mode 100644 index 0000000000..1fbefc5c50 --- /dev/null +++ b/conda/environments/bench_ann_cuda-120_arch-x86_64.yaml @@ -0,0 +1,39 @@ +# This file is generated by `rapids-dependency-file-generator`. +# To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +channels: +- rapidsai +- rapidsai-nightly +- dask/label/dev +- conda-forge +- nvidia +dependencies: +- benchmark>=1.8.2 +- c-compiler +- clang-tools=16.0.6 +- clang==16.0.6 +- cmake>=3.26.4 +- cuda-cudart-dev +- cuda-nvcc +- cuda-profiler-api +- cuda-version=12.0 +- cxx-compiler +- cython>=3.0.0 +- gcc_linux-64=11.* +- glog>=0.6.0 +- h5py>=3.8.0 +- hnswlib=0.7.0 +- libcublas-dev +- libcurand-dev +- libcusolver-dev +- libcusparse-dev +- matplotlib +- nccl>=2.9.9 +- ninja +- nlohmann_json>=3.11.2 +- openblas +- pandas +- pyyaml +- rmm==23.12.* +- scikit-build>=0.13.1 +- sysroot_linux-64==2.17 +name: bench_ann_cuda-120_arch-x86_64 diff --git a/dependencies.yaml b/dependencies.yaml index 0c6a35f9f2..42da15abff 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -22,7 +22,7 @@ files: bench_ann: output: conda matrix: - cuda: ["11.8"] + cuda: ["11.8", "12.0"] arch: [x86_64, aarch64] includes: - build From a2e9b98700761aece5399dc45cd9a17c3784334e Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 27 Oct 2023 17:44:26 -0700 Subject: [PATCH 13/49] style fixes --- .../raft-ann-bench/data_export/__main__.py | 3 ++- .../src/raft-ann-bench/plot/__main__.py | 25 ++++++++++-------- .../src/raft-ann-bench/run/__main__.py | 26 ++++++++++++------- 3 files changed, 32 insertions(+), 22 deletions(-) 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 1d48e9950d..e5c0b54900 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 @@ -43,7 +43,8 @@ def convert_json_to_csv_build(dataset, dataset_path): ) filepath = os.path.normpath(file).split(os.sep) filename = filepath[-1].split("-")[0] + ".csv" - write.to_csv(os.path.join(f"{os.sep}".join(filepath[:-1]), filename), index=False) + write.to_csv(os.path.join(f"{os.sep}".join(filepath[:-1]), filename), + index=False) def convert_json_to_csv_search(dataset, dataset_path): 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 a4a8379eb4..186ca9a218 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 @@ -345,12 +345,12 @@ def load_lines(results_path, result_files, method, index_key): def load_all_results( - dataset_path, algorithms, groups, algo_groups, k, batch_size, method, + dataset_path, algorithms, groups, algo_groups, k, batch_size, method, index_key ): results_path = os.path.join(dataset_path, "result", method) result_files = os.listdir(results_path) - result_files = [result_file for result_file in result_files \ + result_files = [result_file for result_file in result_files if ".csv" in result_file] # print(result_files) if method == "search": @@ -377,19 +377,22 @@ def load_all_results( algo_group_files = list(zip(*algo_group_files)) if len(algorithms) > 0: - final_results = [result_files[i] for i in range(len(result_files)) if \ - (algo_group_files[0][i] in algorithms) and \ - (algo_group_files[1][i] in groups)] + final_results = [result_files[i] for i in range(len(result_files)) if + (algo_group_files[0][i] in algorithms) + and (algo_group_files[1][i] in groups)] else: - final_results = [result_files[i] for i in range(len(result_files)) if \ + final_results = [result_files[i] for i in range(len(result_files)) if (algo_group_files[1][i] in groups)] - + if len(algo_groups) > 0: - split_algo_groups = [algo_group.split(".") for algo_group in algo_groups] + split_algo_groups = \ + [algo_group.split(".") for algo_group in algo_groups] split_algo_groups = list(zip(*split_algo_groups)) - final_algo_groups = [result_files[i] for i in range(len(result_files)) if \ - (algo_group_files[0][i] in split_algo_groups[0]) and \ - (algo_group_files[1][i] in split_algo_groups[1])] + final_algo_groups = [result_files[i] for i in range(len(result_files)) + if (algo_group_files[0][i] + in split_algo_groups[0]) + and (algo_group_files[1][i] + in split_algo_groups[1])] final_results = final_results + final_algo_groups final_results = set(final_results) 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 b49dc01351..a79b2780ae 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 @@ -21,6 +21,7 @@ import subprocess import yaml + def positive_int(input_str: str) -> int: try: i = int(input_str) @@ -259,7 +260,7 @@ def main(): break if not dataset_conf: raise ValueError("Could not find a dataset configuration") - + conf_file = dict() conf_file["dataset"] = dataset_conf if args.subset_size: @@ -270,15 +271,15 @@ def main(): conf_file["search_basic_param"]["batch_size"] = batch_size algos_conf_fs = os.listdir(os.path.join(scripts_path, "conf", "algos")) - algos_conf_fs = [os.path.join(scripts_path, "conf", "algos", f) \ - for f in algos_conf_fs] + algos_conf_fs = [os.path.join(scripts_path, "conf", "algos", f) + for f in algos_conf_fs] conf_filedir = os.path.join(scripts_path, "conf", "algos") if args.configuration: if os.path.isdir(args.configuration): conf_filedir = args.configuration algos_conf_fs = algos_conf_fs + \ - [os.path.join(args.configuration, f) \ - for f in os.listdir(args.configuration)] + [os.path.join(args.configuration, f) + for f in os.listdir(args.configuration)] elif os.path.isfile(args.configuration): conf_filedir = os.path.normpath(args.configuration).split(os.sep) algos_conf_fs = algos_conf_fs + [args.configuration] @@ -290,7 +291,8 @@ def main(): filter_algo_groups = True if args.algo_groups else False allowed_algo_groups = None if filter_algo_groups: - allowed_algo_groups = [algo_group.split(".") for algo_group in args.algo_groups.split(",")] + allowed_algo_groups = [algo_group.split(".") for algo_group + in args.algo_groups.split(",")] allowed_algo_groups = list(zip(*allowed_algo_groups)) algos_conf = dict() for algo_f in algos_conf_fs: @@ -304,6 +306,7 @@ def main(): if filter_algo_groups: if algo["name"] in allowed_algo_groups[0]: insert_algo_group = True + def add_algo_group(group_list): if algo["name"] not in algos_conf: algos_conf[algo["name"]] = {"groups" : {}} @@ -314,6 +317,7 @@ def add_algo_group(group_list): if "validators" in algo: algos_conf[algo["name"]]["validators"] = \ algo["validators"] + if insert_algo: add_algo_group(named_groups) if insert_algo_group: @@ -324,7 +328,8 @@ def add_algo_group(group_list): for algo in algos_conf.keys(): validate_algorithm(algos_yaml, algo, gpu_present) for group in algos_conf[algo]["groups"].keys(): - executable = find_executable(algos_yaml, algo, group, k, batch_size) + executable = find_executable(algos_yaml, algo, group, k, + batch_size) if executable not in executables_to_run: executables_to_run[executable] = {"index": []} build_params = algos_conf[algo]["groups"][group]["build"] @@ -343,7 +348,7 @@ def add_algo_group(group_list): for search_param in search_params.keys(): search_param_names.append(search_param) search_param_lists.append(search_params[search_param]) - + for params in all_build_params: index = {"algo": algo, "build_param": {}} if group != "base": @@ -366,7 +371,7 @@ def add_algo_group(group_list): continue index["name"] = index_name - index["file"] = os.path.join(args.dataset_path, args.dataset, + index["file"] = os.path.join(args.dataset_path, args.dataset, "index", index_name) index["search_params"] = [] all_search_params = itertools.product(*search_param_lists) @@ -377,7 +382,8 @@ def add_algo_group(group_list): if "validators" in algos_conf[algo]: if "search" in algos_conf[algo]["validators"]: importable = \ - algos_conf[algo]["validators"]["search"].split(".") + algos_conf[algo]["validators"]["search"] + importable = importable.split(".") module = ".".join(importable[:-1]) func = importable[-1] validator = import_module(module) From fa193b8b5c6397a70e0d7a800b2bacf97b929672 Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 27 Oct 2023 17:46:45 -0700 Subject: [PATCH 14/49] fix filename --- docs/source/raft_ann_benchmarks.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/raft_ann_benchmarks.md b/docs/source/raft_ann_benchmarks.md index f21885e968..69590eee96 100644 --- a/docs/source/raft_ann_benchmarks.md +++ b/docs/source/raft_ann_benchmarks.md @@ -167,8 +167,8 @@ options: `algo-groups`: this parameter is helpful to append any specific algorithm+group combination to run the benchmark for in addition to all the arguments from `algorithms` and `groups`. It is of the format `.`, or for example, `raft_cagra.large` -For every algorithm run by this script, it outputs an index build statistics JSON file in `/result/build/` -and an index search statistics JSON file in `/result/search/`. +For every algorithm run by this script, it outputs an index build statistics JSON file in `/result/build/` +and an index search statistics JSON file in `/result/search/`. `dataset-path` : 1. data is read from `/` From 3b0a95608cd11337abd6d3a9b789d7e8a20c13bd Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 27 Oct 2023 17:48:26 -0700 Subject: [PATCH 15/49] correct filename again --- docs/source/raft_ann_benchmarks.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/raft_ann_benchmarks.md b/docs/source/raft_ann_benchmarks.md index 69590eee96..641eb6fb7c 100644 --- a/docs/source/raft_ann_benchmarks.md +++ b/docs/source/raft_ann_benchmarks.md @@ -167,8 +167,8 @@ options: `algo-groups`: this parameter is helpful to append any specific algorithm+group combination to run the benchmark for in addition to all the arguments from `algorithms` and `groups`. It is of the format `.`, or for example, `raft_cagra.large` -For every algorithm run by this script, it outputs an index build statistics JSON file in `/result/build/` -and an index search statistics JSON file in `/result/search/`. +For every algorithm run by this script, it outputs an index build statistics JSON file in `/result/build/` +and an index search statistics JSON file in `/result/search/`. NOTE: The filenams will not have "_{group}" if `group = "base"`. `dataset-path` : 1. data is read from `/` From 4345124a9b45a91ded24489a3f950ec3e325cc60 Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 27 Oct 2023 17:52:11 -0700 Subject: [PATCH 16/49] remove debug print --- python/raft-ann-bench/src/raft-ann-bench/run/__main__.py | 1 - 1 file changed, 1 deletion(-) 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 a79b2780ae..1eaa9d0b79 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 @@ -323,7 +323,6 @@ def add_algo_group(group_list): if insert_algo_group: add_algo_group(allowed_algo_groups[1]) - print(algos_conf) executables_to_run = dict() for algo in algos_conf.keys(): validate_algorithm(algos_yaml, algo, gpu_present) From 834f567c7cc44782394d06043120af2887f73f2a Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 27 Oct 2023 18:00:58 -0700 Subject: [PATCH 17/49] try to fix bad merge --- docs/source/raft_ann_benchmarks.md | 37 +++++++++++++----------------- 1 file changed, 16 insertions(+), 21 deletions(-) diff --git a/docs/source/raft_ann_benchmarks.md b/docs/source/raft_ann_benchmarks.md index 641eb6fb7c..577893c3a7 100644 --- a/docs/source/raft_ann_benchmarks.md +++ b/docs/source/raft_ann_benchmarks.md @@ -4,27 +4,22 @@ This project provides a benchmark program for various ANN search implementations ## Table of Contents -- [RAFT ANN Benchmarks](#raft-ann-benchmarks) - - [Table of Contents](#table-of-contents) - - [Installing the benchmarks](#installing-the-benchmarks) - - [Conda](#conda) - - [Docker](#docker) - - [How to run the benchmarks](#how-to-run-the-benchmarks) - - [Step 1: Prepare Dataset](#step-1-prepare-dataset) - - [Step 2: Build and Search Index](#step-2-build-and-search-index) - - [Step 3: Data Export](#step-3-data-export) - - [Step 4: Plot Results](#step-4-plot-results) - - [Running the benchmarks](#running-the-benchmarks) - - [End to end: small-scale benchmarks (\<1M to 10M)](#end-to-end-small-scale-benchmarks-1m-to-10m) - - [End to end: large-scale benchmarks (\>10M vectors)](#end-to-end-large-scale-benchmarks-10m-vectors) - - [Running with Docker containers](#running-with-docker-containers) - - [End-to-end run on GPU](#end-to-end-run-on-gpu) - - [End-to-end run on CPU](#end-to-end-run-on-cpu) - - [Manually run the scripts inside the container](#manually-run-the-scripts-inside-the-container) - - [Creating and customizing dataset configurations](#creating-and-customizing-dataset-configurations) - - [Adding a new ANN algorithm](#adding-a-new-ann-algorithm) - - [Implementation and Configuration](#implementation-and-configuration) - - [Adding a CMake Target](#adding-a-cmake-target) +- [Installing the benchmarks](#installing-the-benchmarks) + - [Conda](#conda) + - [Docker](#docker) +- [How to run the benchmarks](#how-to-run-the-benchmarks) + - [Step 1: prepare dataset](#step-1-prepare-dataset) + - [Step 2: build and search index](#step-2-build-and-search-index) + - [Step 3: data export](#step-3-data-export) + - [Step 4: plot results](#step-4-plot-results) +- [Running the benchmarks](#running-the-benchmarks) + - [End to end: small-scale (<1M to 10M)](#end-to-end-small-scale-benchmarks-1m-to-10m) + - [End to end: large-scale (>10M)](#end-to-end-large-scale-benchmarks-10m-vectors) + - [Running with Docker containers](#running-with-docker-containers) +- [Creating and customizing dataset configurations](#creating-and-customizing-dataset-configurations) +- [Adding a new ANN algorithm](#adding-a-new-ann-algorithm) +- [Parameter tuning guide](https://docs.rapids.ai/api/raft/nightly/ann_benchmarks_param_tuning/) +- [Wiki-all RAG/LLM Dataset](https://docs.rapids.ai/api/raft/nightly/wiki_all_dataset/) ## Installing the benchmarks From 4d764f05ffb2c8e9ca268b412531d341ff7d0a24 Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 27 Oct 2023 18:03:28 -0700 Subject: [PATCH 18/49] remove comment --- python/raft-ann-bench/src/raft-ann-bench/plot/__main__.py | 1 - 1 file changed, 1 deletion(-) 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 186ca9a218..9ad8af7dbc 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 @@ -352,7 +352,6 @@ def load_all_results( result_files = os.listdir(results_path) result_files = [result_file for result_file in result_files if ".csv" in result_file] - # print(result_files) if method == "search": result_files = [ result_filename From f644dfe607cd11a533543388f3e35a9f85b5ce6e Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 30 Oct 2023 11:11:27 -0700 Subject: [PATCH 19/49] address review comments --- docs/source/raft_ann_benchmarks.md | 90 ++++++++++--------- .../src/raft-ann-bench/plot/__main__.py | 9 +- .../src/raft-ann-bench/run/__main__.py | 13 ++- 3 files changed, 61 insertions(+), 51 deletions(-) diff --git a/docs/source/raft_ann_benchmarks.md b/docs/source/raft_ann_benchmarks.md index d610a4be84..32be9f15c7 100644 --- a/docs/source/raft_ann_benchmarks.md +++ b/docs/source/raft_ann_benchmarks.md @@ -4,23 +4,28 @@ This project provides a benchmark program for various ANN search implementations ## Table of Contents -- [Installing the benchmarks](#installing-the-benchmarks) - - [Conda](#conda) - - [Docker](#docker) -- [How to run the benchmarks](#how-to-run-the-benchmarks) - - [Step 1: prepare dataset](#step-1-prepare-dataset) - - [Step 2: build and search index](#step-2-build-and-search-index) - - [Step 3: data export](#step-3-data-export) - - [Step 4: plot results](#step-4-plot-results) -- [Running the benchmarks](#running-the-benchmarks) - - [End to end: small-scale (<1M to 10M)](#end-to-end-small-scale-benchmarks-1m-to-10m) - - [End to end: large-scale (>10M)](#end-to-end-large-scale-benchmarks-10m-vectors) - - [Running with Docker containers](#running-with-docker-containers) - - [Evaluating the results](#evaluating-the-results) -- [Creating and customizing dataset configurations](#creating-and-customizing-dataset-configurations) -- [Adding a new ANN algorithm](#adding-a-new-ann-algorithm) -- [Parameter tuning guide](https://docs.rapids.ai/api/raft/nightly/ann_benchmarks_param_tuning/) -- [Wiki-all RAG/LLM Dataset](https://docs.rapids.ai/api/raft/nightly/wiki_all_dataset/) +- [RAFT ANN Benchmarks](#raft-ann-benchmarks) + - [Table of Contents](#table-of-contents) + - [Installing the benchmarks](#installing-the-benchmarks) + - [Conda](#conda) + - [Docker](#docker) + - [How to run the benchmarks](#how-to-run-the-benchmarks) + - [Step 1: Prepare Dataset](#step-1-prepare-dataset) + - [Step 2: Build and Search Index](#step-2-build-and-search-index) + - [Step 3: Data Export](#step-3-data-export) + - [Step 4: Plot Results](#step-4-plot-results) + - [Running the benchmarks](#running-the-benchmarks) + - [End to end: small-scale benchmarks (\<1M to 10M)](#end-to-end-small-scale-benchmarks-1m-to-10m) + - [End to end: large-scale benchmarks (\>10M vectors)](#end-to-end-large-scale-benchmarks-10m-vectors) + - [Running with Docker containers](#running-with-docker-containers) + - [End-to-end run on GPU](#end-to-end-run-on-gpu) + - [End-to-end run on CPU](#end-to-end-run-on-cpu) + - [Manually run the scripts inside the container](#manually-run-the-scripts-inside-the-container) + - [Evaluating the results](#evaluating-the-results) + - [Creating and customizing dataset configurations](#creating-and-customizing-dataset-configurations) + - [Adding a new ANN algorithm](#adding-a-new-ann-algorithm) + - [Implementation and Configuration](#implementation-and-configuration) + - [Adding a CMake Target](#adding-a-cmake-target) ## Installing the benchmarks @@ -122,8 +127,8 @@ specified configuration. The usage of the script `raft-ann-bench.run` is: ```bash -usage: [-h] [--subset-size SUBSET_SIZE] [-k COUNT] [-bs BATCH_SIZE] [--dataset-configuration DATASET_CONFIGURATION] [--configuration CONFIGURATION] [--dataset DATASET] - [--dataset-path DATASET_PATH] [--build] [--search] [--algorithms ALGORITHMS] [--groups GROUPS] [--algo-groups ALGO_GROUPS] [-f] +usage: __main__.py [-h] [--subset-size SUBSET_SIZE] [-k COUNT] [-bs BATCH_SIZE] [--dataset-configuration DATASET_CONFIGURATION] [--configuration CONFIGURATION] [--dataset DATASET] + [--dataset-path DATASET_PATH] [--build] [--search] [--algorithms ALGORITHMS] [--groups GROUPS] [--algo-groups ALGO_GROUPS] [-f] [-m SEARCH_MODE] options: -h, --help show this help message and exit @@ -134,25 +139,24 @@ options: -bs BATCH_SIZE, --batch-size BATCH_SIZE number of query vectors to use in each query trial (default: 10000) --dataset-configuration DATASET_CONFIGURATION - path to configuration file for datasets (default: None) + path to YAML configuration file for datasets (default: None) --configuration CONFIGURATION - path to configuration file or directory for algorithms (default: None) + path to YAML configuration file or directory for algorithms Any run groups found in the specified file/directory will automatically override groups of the same name + present in the default configurations, including `base` (default: None) --dataset DATASET name of dataset (default: glove-100-inner) --dataset-path DATASET_PATH path to dataset folder, by default will look in RAPIDS_DATASET_ROOT_DIR if defined, otherwise a datasets subdirectory from the calling directory (default: - /raid/dgala/raft/datasets/) + os.getcwd()/datasets/) --build --search --algorithms ALGORITHMS - run only comma separated list of named algorithms (default: None) + run only comma separated list of named algorithms. If parameters `groups` and `algo-groups are both undefined, then group `base` is run by default (default: None) --groups GROUPS run only comma separated groups of parameters (default: base) --algo-groups ALGO_GROUPS - add comma separated algorithm+groups to run (default: None) + add comma separated . to run. Example usage: "--algo-groups=raft_cagra.large,hnswlib.large" (default: None) -f, --force re-run algorithms even if their results already exist (default: False) - -m MODE, --search-mode MODE - run search in 'latency' (measure individual batches) or - 'throughput' (pipeline batches and measure end-to-end) mode. - (default: 'latency') + -m SEARCH_MODE, --search-mode SEARCH_MODE + run search in 'latency' (measure individual batches) or 'throughput' (pipeline batches and measure end-to-end) mode (default: throughput) ``` `dataset`: name of the dataset to be searched in [datasets.yaml](#yaml-dataset-config) @@ -161,7 +165,7 @@ options: `configuration`: optional filepath to YAML configuration for an algorithm or to directory that contains YAML configurations for several algorithms. [Here's how to configure an algorithm.](#yaml-algo-config) -`algorithms`: runs all algorithms that it can find in YAML configs found by `configuration` +`algorithms`: runs all algorithms that it can find in YAML configs found by `configuration`. By default, only `base` group will be run. `groups`: run only specific groups of parameters configurations for an algorithm. Groups are defined in YAML configs (see `configuration`), and by default run `base` group @@ -203,21 +207,21 @@ CSV file in `/result/search/<-k{k}-batch_size{batch_size} The usage of this script is: ```bash -usage: [-h] [--dataset DATASET] [--dataset-path DATASET_PATH] [--output-filepath OUTPUT_FILEPATH] [--algorithms ALGORITHMS] [--groups GROUPS] [--algo-groups ALGO_GROUPS] [-k COUNT] - [-bs BATCH_SIZE] [--build] [--search] [--x-scale X_SCALE] [--y-scale {linear,log,symlog,logit}] [--raw] +usage: __main__.py [-h] [--dataset DATASET] [--dataset-path DATASET_PATH] [--output-filepath OUTPUT_FILEPATH] [--algorithms ALGORITHMS] [--groups GROUPS] [--algo-groups ALGO_GROUPS] [-k COUNT] + [-bs BATCH_SIZE] [--build] [--search] [--x-scale X_SCALE] [--y-scale {linear,log,symlog,logit}] [--raw] -optional arguments: +options: -h, --help show this help message and exit --dataset DATASET dataset to plot (default: glove-100-inner) --dataset-path DATASET_PATH - path to dataset folder (default: /raid/dgala/raft/datasets/) + path to dataset folder (default: os.getcwd()/datasets/) --output-filepath OUTPUT_FILEPATH - directory for PNG to be saved (default: /raid/dgala/raft) + directory for PNG to be saved (default: os.getcwd()) --algorithms ALGORITHMS - plot only comma separated list of named algorithms (default: None) + plot only comma separated list of named algorithms. If parameters `groups` and `algo-groups are both undefined, then group `base` is plot by default (default: None) --groups GROUPS plot only comma separated groups of parameters (default: base) - --algo-groups ALGO_GROUPS - add comma separated algorithm+groups to the plot (default: None) + --algo-groups ALGO_GROUPS, --algo-groups ALGO_GROUPS + add comma separated . to plot. Example usage: "--algo-groups=raft_cagra.large,hnswlib.large" (default: None) -k COUNT, --count COUNT the number of nearest neighbors to search for (default: 10) -bs BATCH_SIZE, --batch-size BATCH_SIZE @@ -229,7 +233,7 @@ optional arguments: Scale to use when drawing the Y-axis (default: linear) --raw Show raw results (not just Pareto frontier) in faded colours (default: False) ``` -`algorithms`: plots all algorithms that it can find results for the specified `dataset` +`algorithms`: plots all algorithms that it can find results for the specified `dataset`. By default, only `base` group will be plotted. `groups`: plot only specific groups of parameters configurations for an algorithm. Groups are defined in YAML configs (see `configuration`), and by default run `base` group @@ -432,8 +436,6 @@ A single configuration will often define a set of algorithms, with associated in Configuration files for ANN algorithms supported by `raft-ann-bench` are provided in `${RAFT_HOME}/python/raft-ann-bench/src/raft-ann-bench/run/conf`. `raft_cagra` algorithm configuration looks like: ```yaml name: raft_cagra -validators: - search: raft-ann-bench.validators.raft_cagra_search_validator groups: base: build: @@ -448,11 +450,11 @@ groups: search: itopk: [32, 64, 128] ``` -It is mandatory that every algorithm have a `base` group. +The default parameters for which the benchmarks are run can be overridden by creating a custom YAML file for algorithms with a `base` group. -There's several things going on here: -1. `groups` - define a run group which has a particular set of parameters. Each group has `build` parameters wherein a cross-product of all parameters will be used to build a unique index. Each group also has `search` parameters, and a cross-product of all search parameters will be used to search all unique indices that we create. -2. `validators` - This has two entries `build` and `search`. It's a string denoting a Python importable function of a module that will be used to verify whether there are any bogus parameters in any `groups` that do not make sense. If such parameters exist, they will be silently ignored. +There config above has 2 fields: +1. `name` - define the name of the algorithm for which the parameters are being specified. +2. `groups` - define a run group which has a particular set of parameters. Each group helps create a cross-product of all hyper-parameter fields for `build` and `search`. The table below contains all algorithms supported by RAFT. Each unique algorithm will have its own set of `build` and `search` settings. The [ANN Algorithm Parameter Tuning Guide](ann_benchmarks_param_tuning.md) contains detailed instructions on choosing build and search parameters for each supported algorithm. 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 9ad8af7dbc..94436eb1cc 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 @@ -411,7 +411,7 @@ def main(): formatter_class=argparse.ArgumentDefaultsHelpFormatter ) parser.add_argument( - "--dataset", help="dataset to plots", default="glove-100-inner" + "--dataset", help="dataset to plot", default="glove-100-inner" ) parser.add_argument( "--dataset-path", @@ -426,7 +426,8 @@ def main(): parser.add_argument( "--algorithms", help="plot only comma separated list of named \ - algorithms", + algorithms. If parameters `groups` and `algo-groups \ + are both undefined, then group `base` is plot by default", default=None, ) parser.add_argument( @@ -436,7 +437,9 @@ def main(): ) parser.add_argument( "--algo-groups", - help="add comma separated algorithm+groups to the plot", + "--algo-groups", + help="add comma separated . to plot. \ + Example usage: \"--algo-groups=raft_cagra.large,hnswlib.large\"", ) parser.add_argument( "-k", 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 9ea98f5db1..9a14c378a7 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 @@ -189,11 +189,14 @@ def main(): ) parser.add_argument( "--dataset-configuration", - help="path to configuration file for datasets", + help="path to YAML configuration file for datasets", ) parser.add_argument( "--configuration", - help="path to configuration file or directory for algorithms", + help="path to YAML configuration file or directory for algorithms\ + Any run groups found in the specified file/directory will \ + automatically override groups of the same name present in the \ + default configurations, including `base`", ) parser.add_argument( "--dataset", @@ -212,7 +215,8 @@ def main(): parser.add_argument( "--algorithms", help="run only comma separated list of named \ - algorithms", + algorithms. If parameters `groups` and `algo-groups \ + are both undefined, then group `base` is run by default", default=None, ) parser.add_argument( @@ -222,7 +226,8 @@ def main(): ) parser.add_argument( "--algo-groups", - help="add comma separated algorithm+groups to run", + help="add comma separated . to run. \ + Example usage: \"--algo-groups=raft_cagra.large,hnswlib.large\"", ) parser.add_argument( "-f", From f470d86260426d0a7c0fba776f322718579eb2d6 Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 30 Oct 2023 11:14:30 -0700 Subject: [PATCH 20/49] add wiki datasets to datasets.yaml --- .../src/raft-ann-bench/run/conf/datasets.yaml | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) 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 45d648f88e..33f07b9f22 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 @@ -83,3 +83,21 @@ query_file: sift-128-euclidean/query.fbin groundtruth_neighbors_file: sift-128-euclidean/groundtruth.neighbors.ibin distance: euclidean + +- name: wiki_all_1M, + base_file: wiki_all_1M/base.88M.fbin, + query_file: wiki_all_1M/queries.fbin, + groundtruth_neighbors_file: wiki_all_1M/groundtruth.88M.neighbors.ibin, + distance: euclidean + +- name: wiki_all_10M, + base_file: wiki_all_10M/base.88M.fbin, + query_file: wiki_all_10M/queries.fbin, + groundtruth_neighbors_file: wiki_all_10M/groundtruth.88M.neighbors.ibin, + distance: euclidean + +- name: wiki_all_88M, + 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 From 2b21ad0705553641563779c8144d4eec2358578e Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 30 Oct 2023 13:25:19 -0700 Subject: [PATCH 21/49] fix style --- .../raft-ann-bench/data_export/__main__.py | 6 +- .../src/raft-ann-bench/plot/__main__.py | 57 ++++++++++------- .../src/raft-ann-bench/run/__main__.py | 61 +++++++++++-------- 3 files changed, 73 insertions(+), 51 deletions(-) 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 e5c0b54900..97572e30fd 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 @@ -43,8 +43,10 @@ def convert_json_to_csv_build(dataset, dataset_path): ) filepath = os.path.normpath(file).split(os.sep) filename = filepath[-1].split("-")[0] + ".csv" - write.to_csv(os.path.join(f"{os.sep}".join(filepath[:-1]), filename), - index=False) + write.to_csv( + os.path.join(f"{os.sep}".join(filepath[:-1]), filename), + index=False + ) def convert_json_to_csv_search(dataset, dataset_path): 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 94436eb1cc..84d11f3f14 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 @@ -345,13 +345,20 @@ def load_lines(results_path, result_files, method, index_key): def load_all_results( - dataset_path, algorithms, groups, algo_groups, k, batch_size, method, + dataset_path, + algorithms, + groups, + algo_groups, + k, + batch_size, + method, index_key ): results_path = os.path.join(dataset_path, "result", method) result_files = os.listdir(results_path) - result_files = [result_file for result_file in result_files - if ".csv" in result_file] + result_files = [ + result_file for result_file in result_files if ".csv" in result_file + ] if method == "search": result_files = [ result_filename @@ -359,13 +366,11 @@ def load_all_results( if f"{k}-{batch_size}" in result_filename ] algo_group_files = [ - result_filename.split("-")[0] - for result_filename in result_files + result_filename.split("-")[0] for result_filename in result_files ] else: algo_group_files = [ - result_filename - for result_filename in result_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("_") @@ -376,22 +381,30 @@ def load_all_results( algo_group_files = list(zip(*algo_group_files)) if len(algorithms) > 0: - final_results = [result_files[i] for i in range(len(result_files)) if - (algo_group_files[0][i] in algorithms) - and (algo_group_files[1][i] in groups)] + final_results = [ + result_files[i] + for i in range(len(result_files)) + if (algo_group_files[0][i] in algorithms) + and (algo_group_files[1][i] in groups) + ] else: - final_results = [result_files[i] for i in range(len(result_files)) if - (algo_group_files[1][i] in groups)] + final_results = [ + result_files[i] + for i in range(len(result_files)) + if (algo_group_files[1][i] in groups) + ] if len(algo_groups) > 0: - split_algo_groups = \ - [algo_group.split(".") for algo_group in algo_groups] + split_algo_groups = [ + algo_group.split(".") for algo_group in algo_groups + ] split_algo_groups = list(zip(*split_algo_groups)) - final_algo_groups = [result_files[i] for i in range(len(result_files)) - if (algo_group_files[0][i] - in split_algo_groups[0]) - and (algo_group_files[1][i] - in split_algo_groups[1])] + final_algo_groups = [ + result_files[i] + for i in range(len(result_files)) + if (algo_group_files[0][i] in split_algo_groups[0]) + and (algo_group_files[1][i] in split_algo_groups[1]) + ] final_results = final_results + final_algo_groups final_results = set(final_results) @@ -433,13 +446,13 @@ def main(): parser.add_argument( "--groups", help="plot only comma separated groups of parameters", - default="base" + default="base", ) parser.add_argument( "--algo-groups", "--algo-groups", - help="add comma separated . to plot. \ - Example usage: \"--algo-groups=raft_cagra.large,hnswlib.large\"", + help='add comma separated . to plot. \ + Example usage: "--algo-groups=raft_cagra.large,hnswlib.large"', ) parser.add_argument( "-k", 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 9a14c378a7..f7717fc5e0 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 @@ -14,11 +14,12 @@ # limitations under the License. import argparse -from importlib import import_module import itertools import json import os import subprocess +from importlib import import_module + import yaml @@ -222,12 +223,12 @@ def main(): parser.add_argument( "--groups", help="run only comma separated groups of parameters", - default="base" + default="base", ) parser.add_argument( "--algo-groups", - help="add comma separated . to run. \ - Example usage: \"--algo-groups=raft_cagra.large,hnswlib.large\"", + help='add comma separated . to run. \ + Example usage: "--algo-groups=raft_cagra.large,hnswlib.large"', ) parser.add_argument( "-f", @@ -264,9 +265,7 @@ def main(): if args.dataset_configuration: dataset_conf_f = args.dataset_configuration else: - dataset_conf_f = os.path.join( - scripts_path, "conf", "datasets.yaml" - ) + dataset_conf_f = os.path.join(scripts_path, "conf", "datasets.yaml") with open(dataset_conf_f, "r") as f: dataset_conf_all = yaml.safe_load(f) @@ -288,15 +287,17 @@ def main(): conf_file["search_basic_param"]["batch_size"] = batch_size algos_conf_fs = os.listdir(os.path.join(scripts_path, "conf", "algos")) - algos_conf_fs = [os.path.join(scripts_path, "conf", "algos", f) - for f in algos_conf_fs] + algos_conf_fs = [ + os.path.join(scripts_path, "conf", "algos", f) for f in algos_conf_fs + ] conf_filedir = os.path.join(scripts_path, "conf", "algos") if args.configuration: if os.path.isdir(args.configuration): conf_filedir = args.configuration - algos_conf_fs = algos_conf_fs + \ - [os.path.join(args.configuration, f) - for f in os.listdir(args.configuration)] + algos_conf_fs = algos_conf_fs + [ + os.path.join(args.configuration, f) + for f in os.listdir(args.configuration) + ] elif os.path.isfile(args.configuration): conf_filedir = os.path.normpath(args.configuration).split(os.sep) algos_conf_fs = algos_conf_fs + [args.configuration] @@ -308,8 +309,9 @@ def main(): filter_algo_groups = True if args.algo_groups else False allowed_algo_groups = None if filter_algo_groups: - allowed_algo_groups = [algo_group.split(".") for algo_group - in args.algo_groups.split(",")] + allowed_algo_groups = [ + algo_group.split(".") for algo_group in args.algo_groups.split(",") + ] allowed_algo_groups = list(zip(*allowed_algo_groups)) algos_conf = dict() for algo_f in algos_conf_fs: @@ -326,14 +328,14 @@ def main(): def add_algo_group(group_list): if algo["name"] not in algos_conf: - algos_conf[algo["name"]] = {"groups" : {}} + algos_conf[algo["name"]] = {"groups": {}} for group in algo["groups"].keys(): if group in group_list: - algos_conf[algo["name"]]["groups"][group] = \ - algo["groups"][group] + algos_conf[algo["name"]]["groups"][group] = algo[ + "groups" + ][group] if "validators" in algo: - algos_conf[algo["name"]]["validators"] = \ - algo["validators"] + algos_conf[algo["name"]]["validators"] = algo["validators"] if insert_algo: add_algo_group(named_groups) @@ -344,8 +346,9 @@ def add_algo_group(group_list): for algo in algos_conf.keys(): validate_algorithm(algos_yaml, algo, gpu_present) for group in algos_conf[algo]["groups"].keys(): - executable = find_executable(algos_yaml, algo, group, k, - batch_size) + executable = find_executable( + algos_yaml, algo, group, k, batch_size + ) if executable not in executables_to_run: executables_to_run[executable] = {"index": []} build_params = algos_conf[algo]["groups"][group]["build"] @@ -377,8 +380,10 @@ def add_algo_group(group_list): if "validators" in algos_conf[algo]: if "build" in algos_conf[algo]["validators"]: - importable = \ - algos_conf[algo]["validators"]["build"].split(".") + importable = algos_conf[algo]["validators"][ + "build" + ] + importable = importable.split(".") module = ".".join(importable[:-1]) func = importable[-1] validator = import_module(module) @@ -387,8 +392,9 @@ def add_algo_group(group_list): continue index["name"] = index_name - index["file"] = os.path.join(args.dataset_path, args.dataset, - "index", index_name) + index["file"] = os.path.join( + args.dataset_path, args.dataset, "index", index_name + ) index["search_params"] = [] all_search_params = itertools.product(*search_param_lists) for search_params in all_search_params: @@ -397,8 +403,9 @@ def add_algo_group(group_list): 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"]["search"] + importable = algos_conf[algo]["validators"][ + "search" + ] importable = importable.split(".") module = ".".join(importable[:-1]) func = importable[-1] From 9ff0fbbded8aa7d3c8e05c8dd2c33b6683b7582a Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 30 Oct 2023 13:31:10 -0700 Subject: [PATCH 22/49] more style fixes --- .../src/raft-ann-bench/data_export/__main__.py | 2 +- python/raft-ann-bench/src/raft-ann-bench/plot/__main__.py | 2 +- python/raft-ann-bench/src/raft-ann-bench/run/__main__.py | 6 ++---- 3 files changed, 4 insertions(+), 6 deletions(-) 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 97572e30fd..fe2670411e 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 @@ -45,7 +45,7 @@ def convert_json_to_csv_build(dataset, dataset_path): filename = filepath[-1].split("-")[0] + ".csv" write.to_csv( os.path.join(f"{os.sep}".join(filepath[:-1]), filename), - index=False + index=False, ) 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 84d11f3f14..ef81768f4d 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 @@ -352,7 +352,7 @@ def load_all_results( k, batch_size, method, - index_key + index_key, ): results_path = os.path.join(dataset_path, "result", method) result_files = os.listdir(results_path) 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 f7717fc5e0..85a1e633ec 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 @@ -332,7 +332,7 @@ def add_algo_group(group_list): for group in algo["groups"].keys(): if group in group_list: algos_conf[algo["name"]]["groups"][group] = algo[ - "groups" + "groups" ][group] if "validators" in algo: algos_conf[algo["name"]]["validators"] = algo["validators"] @@ -380,9 +380,7 @@ def add_algo_group(group_list): if "validators" in algos_conf[algo]: if "build" in algos_conf[algo]["validators"]: - importable = algos_conf[algo]["validators"][ - "build" - ] + importable = algos_conf[algo]["validators"]["build"] importable = importable.split(".") module = ".".join(importable[:-1]) func = importable[-1] From cf086c24a5f7f887da1115a2d92e93550c05b5bb Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 30 Oct 2023 17:08:06 -0700 Subject: [PATCH 23/49] fix when --configuration is a file --- python/raft-ann-bench/src/raft-ann-bench/run/__main__.py | 1 + 1 file changed, 1 insertion(+) 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 85a1e633ec..06dc948829 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 @@ -300,6 +300,7 @@ def main(): ] elif os.path.isfile(args.configuration): conf_filedir = os.path.normpath(args.configuration).split(os.sep) + conf_filedir = os.path.join(*conf_filedir[:-1]) algos_conf_fs = algos_conf_fs + [args.configuration] filter_algos = True if args.algorithms else False From 0b6cc09dccad6e292f62c3f6b3923e8c036608f9 Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 30 Oct 2023 17:16:28 -0700 Subject: [PATCH 24/49] don't read json confs --- python/raft-ann-bench/src/raft-ann-bench/run/__main__.py | 2 ++ 1 file changed, 2 insertions(+) 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 06dc948829..80f5a86592 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 @@ -289,6 +289,7 @@ def main(): algos_conf_fs = os.listdir(os.path.join(scripts_path, "conf", "algos")) algos_conf_fs = [ os.path.join(scripts_path, "conf", "algos", f) for f in algos_conf_fs + if ".json" not in f ] conf_filedir = os.path.join(scripts_path, "conf", "algos") if args.configuration: @@ -297,6 +298,7 @@ def main(): algos_conf_fs = algos_conf_fs + [ os.path.join(args.configuration, f) for f in os.listdir(args.configuration) + if ".json" not in f ] elif os.path.isfile(args.configuration): conf_filedir = os.path.normpath(args.configuration).split(os.sep) From 227486c2656c5dfc543e9b047e71ef214dc981e4 Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 30 Oct 2023 17:16:47 -0700 Subject: [PATCH 25/49] add nvtx dependency --- conda/environments/all_cuda-118_arch-aarch64.yaml | 1 + conda/environments/all_cuda-118_arch-x86_64.yaml | 1 + conda/environments/all_cuda-120_arch-aarch64.yaml | 1 + conda/environments/all_cuda-120_arch-x86_64.yaml | 1 + conda/environments/bench_ann_cuda-118_arch-aarch64.yaml | 1 + conda/environments/bench_ann_cuda-118_arch-x86_64.yaml | 1 + conda/environments/bench_ann_cuda-120_arch-aarch64.yaml | 1 + conda/environments/bench_ann_cuda-120_arch-x86_64.yaml | 1 + dependencies.yaml | 5 +++++ 9 files changed, 13 insertions(+) diff --git a/conda/environments/all_cuda-118_arch-aarch64.yaml b/conda/environments/all_cuda-118_arch-aarch64.yaml index 27894d27e7..762d436028 100644 --- a/conda/environments/all_cuda-118_arch-aarch64.yaml +++ b/conda/environments/all_cuda-118_arch-aarch64.yaml @@ -12,6 +12,7 @@ dependencies: - clang-tools=16.0.6 - clang==16.0.6 - cmake>=3.26.4 +- cuda-nvtx=11.8 - cuda-profiler-api=11.8.86 - cuda-python>=11.7.1,<12.0a0 - cuda-version=11.8 diff --git a/conda/environments/all_cuda-118_arch-x86_64.yaml b/conda/environments/all_cuda-118_arch-x86_64.yaml index 96e5227186..48cd34c6ca 100644 --- a/conda/environments/all_cuda-118_arch-x86_64.yaml +++ b/conda/environments/all_cuda-118_arch-x86_64.yaml @@ -12,6 +12,7 @@ dependencies: - clang-tools=16.0.6 - clang==16.0.6 - cmake>=3.26.4 +- cuda-nvtx=11.8 - cuda-profiler-api=11.8.86 - cuda-python>=11.7.1,<12.0a0 - cuda-version=11.8 diff --git a/conda/environments/all_cuda-120_arch-aarch64.yaml b/conda/environments/all_cuda-120_arch-aarch64.yaml index e4ba4281c2..9108fb6215 100644 --- a/conda/environments/all_cuda-120_arch-aarch64.yaml +++ b/conda/environments/all_cuda-120_arch-aarch64.yaml @@ -14,6 +14,7 @@ dependencies: - cmake>=3.26.4 - cuda-cudart-dev - cuda-nvcc +- cuda-nvtx-dev - cuda-profiler-api - cuda-python>=12.0,<13.0a0 - cuda-version=12.0 diff --git a/conda/environments/all_cuda-120_arch-x86_64.yaml b/conda/environments/all_cuda-120_arch-x86_64.yaml index d0430e10f6..8f1fbf6744 100644 --- a/conda/environments/all_cuda-120_arch-x86_64.yaml +++ b/conda/environments/all_cuda-120_arch-x86_64.yaml @@ -14,6 +14,7 @@ dependencies: - cmake>=3.26.4 - cuda-cudart-dev - cuda-nvcc +- cuda-nvtx-dev - cuda-profiler-api - cuda-python>=12.0,<13.0a0 - cuda-version=12.0 diff --git a/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml index 579a8a0ceb..b5fc4e3bd5 100644 --- a/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml @@ -12,6 +12,7 @@ dependencies: - clang-tools=16.0.6 - clang==16.0.6 - cmake>=3.26.4 +- cuda-nvtx=11.8 - cuda-profiler-api=11.8.86 - cuda-version=11.8 - cudatoolkit diff --git a/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml index 2c92ad0a99..b868f26e15 100644 --- a/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml @@ -12,6 +12,7 @@ dependencies: - clang-tools=16.0.6 - clang==16.0.6 - cmake>=3.26.4 +- cuda-nvtx=11.8 - cuda-profiler-api=11.8.86 - cuda-version=11.8 - cudatoolkit diff --git a/conda/environments/bench_ann_cuda-120_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-120_arch-aarch64.yaml index 99ec02859d..4a3818fe5d 100644 --- a/conda/environments/bench_ann_cuda-120_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-120_arch-aarch64.yaml @@ -14,6 +14,7 @@ dependencies: - cmake>=3.26.4 - cuda-cudart-dev - cuda-nvcc +- cuda-nvtx-dev - cuda-profiler-api - cuda-version=12.0 - cxx-compiler diff --git a/conda/environments/bench_ann_cuda-120_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-120_arch-x86_64.yaml index 1fbefc5c50..3d6f8c4ec1 100644 --- a/conda/environments/bench_ann_cuda-120_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-120_arch-x86_64.yaml @@ -14,6 +14,7 @@ dependencies: - cmake>=3.26.4 - cuda-cudart-dev - cuda-nvcc +- cuda-nvtx-dev - cuda-profiler-api - cuda-version=12.0 - cxx-compiler diff --git a/dependencies.yaml b/dependencies.yaml index 263a76f9c8..aba81d7ed9 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -246,6 +246,7 @@ dependencies: cuda: "12.0" packages: - cuda-version=12.0 + - cuda-nvtx-dev - cuda-cudart-dev - cuda-profiler-api - libcublas-dev @@ -257,6 +258,7 @@ dependencies: packages: - cuda-version=11.8 - cudatoolkit + - cuda-nvtx=11.8 - cuda-profiler-api=11.8.86 - libcublas-dev=11.11.3.6 - libcublas=11.11.3.6 @@ -271,6 +273,7 @@ dependencies: packages: - cuda-version=11.5 - cudatoolkit + - cuda-nvtx=11.5 - cuda-profiler-api>=11.4.240,<=11.8.86 # use any `11.x` version since pkg is missing several CUDA/arch packages - libcublas-dev>=11.7.3.1,<=11.7.4.6 - libcublas>=11.7.3.1,<=11.7.4.6 @@ -285,6 +288,7 @@ dependencies: packages: - cuda-version=11.4 - cudatoolkit + - &cudanvtx114 cuda-nvtx=11.4 - cuda-profiler-api>=11.4.240,<=11.8.86 # use any `11.x` version since pkg is missing several CUDA/arch packages - &libcublas_dev114 libcublas-dev>=11.5.2.43,<=11.6.5.2 - &libcublas114 libcublas>=11.5.2.43,<=11.6.5.2 @@ -299,6 +303,7 @@ dependencies: packages: - cuda-version=11.2 - cudatoolkit + - *cudanvtx114 - cuda-profiler-api>=11.4.240,<=11.8.86 # use any `11.x` version since pkg is missing several CUDA/arch packages # The NVIDIA channel doesn't publish pkgs older than 11.4 for these libs, # so 11.2 uses 11.4 packages (the oldest available). From d80af8bdca8b509f7b60482a3eae7495d71879a7 Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 30 Oct 2023 17:19:48 -0700 Subject: [PATCH 26/49] fix docs again --- docs/source/raft_ann_benchmarks.md | 62 +++++++++++++++++++----------- 1 file changed, 40 insertions(+), 22 deletions(-) diff --git a/docs/source/raft_ann_benchmarks.md b/docs/source/raft_ann_benchmarks.md index 32be9f15c7..72a164c34f 100644 --- a/docs/source/raft_ann_benchmarks.md +++ b/docs/source/raft_ann_benchmarks.md @@ -4,28 +4,46 @@ This project provides a benchmark program for various ANN search implementations ## Table of Contents -- [RAFT ANN Benchmarks](#raft-ann-benchmarks) - - [Table of Contents](#table-of-contents) - - [Installing the benchmarks](#installing-the-benchmarks) - - [Conda](#conda) - - [Docker](#docker) - - [How to run the benchmarks](#how-to-run-the-benchmarks) - - [Step 1: Prepare Dataset](#step-1-prepare-dataset) - - [Step 2: Build and Search Index](#step-2-build-and-search-index) - - [Step 3: Data Export](#step-3-data-export) - - [Step 4: Plot Results](#step-4-plot-results) - - [Running the benchmarks](#running-the-benchmarks) - - [End to end: small-scale benchmarks (\<1M to 10M)](#end-to-end-small-scale-benchmarks-1m-to-10m) - - [End to end: large-scale benchmarks (\>10M vectors)](#end-to-end-large-scale-benchmarks-10m-vectors) - - [Running with Docker containers](#running-with-docker-containers) - - [End-to-end run on GPU](#end-to-end-run-on-gpu) - - [End-to-end run on CPU](#end-to-end-run-on-cpu) - - [Manually run the scripts inside the container](#manually-run-the-scripts-inside-the-container) - - [Evaluating the results](#evaluating-the-results) - - [Creating and customizing dataset configurations](#creating-and-customizing-dataset-configurations) - - [Adding a new ANN algorithm](#adding-a-new-ann-algorithm) - - [Implementation and Configuration](#implementation-and-configuration) - - [Adding a CMake Target](#adding-a-cmake-target) +- [Installing the benchmarks](#installing-the-benchmarks) + - [Conda](#conda) + - [Docker](#docker) +- [How to run the benchmarks](#how-to-run-the-benchmarks) + - [Step 1: prepare dataset](#step-1-prepare-dataset) + - [Step 2: build and search index](#step-2-build-and-search-index) + - [Step 3: data export](#step-3-data-export) + - [Step 4: plot results](#step-4-plot-results) +- [Running the benchmarks](#running-the-benchmarks) + - [End to end: small-scale (<1M to 10M)](#end-to-end-small-scale-benchmarks-1m-to-10m) + - [End to end: large-scale (>10M)](#end-to-end-large-scale-benchmarks-10m-vectors) + - [Running with Docker containers](#running-with-docker-containers) + - [Evaluating the results](#evaluating-the-results) +- [Creating and customizing dataset configurations](#creating-and-customizing-dataset-configurations) +- [Adding a new ANN algorithm](#adding-a-new-ann-algorithm) +- [Parameter tuning guide](https://docs.rapids.ai/api/raft/nightly/ann_benchmarks_param_tuning/) +- [Wiki-all RAG/LLM Dataset](https://docs.rapids.ai/api/raft/nightly/wiki_all_dataset/) +## Installing the benchmarks# RAFT ANN Benchmarks + +This project provides a benchmark program for various ANN search implementations. It's especially suitable for comparing GPU implementations as well as comparing GPU against CPU. + +## Table of Contents + +- [Installing the benchmarks](#installing-the-benchmarks) + - [Conda](#conda) + - [Docker](#docker) +- [How to run the benchmarks](#how-to-run-the-benchmarks) + - [Step 1: prepare dataset](#step-1-prepare-dataset) + - [Step 2: build and search index](#step-2-build-and-search-index) + - [Step 3: data export](#step-3-data-export) + - [Step 4: plot results](#step-4-plot-results) +- [Running the benchmarks](#running-the-benchmarks) + - [End to end: small-scale (<1M to 10M)](#end-to-end-small-scale-benchmarks-1m-to-10m) + - [End to end: large-scale (>10M)](#end-to-end-large-scale-benchmarks-10m-vectors) + - [Running with Docker containers](#running-with-docker-containers) + - [Evaluating the results](#evaluating-the-results) +- [Creating and customizing dataset configurations](#creating-and-customizing-dataset-configurations) +- [Adding a new ANN algorithm](#adding-a-new-ann-algorithm) +- [Parameter tuning guide](https://docs.rapids.ai/api/raft/nightly/ann_benchmarks_param_tuning/) +- [Wiki-all RAG/LLM Dataset](https://docs.rapids.ai/api/raft/nightly/wiki_all_dataset/) ## Installing the benchmarks From 33349ce02b4fba87585c681a17a1b87350af9aef Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 30 Oct 2023 17:21:02 -0700 Subject: [PATCH 27/49] again fix docs --- docs/source/raft_ann_benchmarks.md | 23 ----------------------- 1 file changed, 23 deletions(-) diff --git a/docs/source/raft_ann_benchmarks.md b/docs/source/raft_ann_benchmarks.md index 72a164c34f..305cb836e3 100644 --- a/docs/source/raft_ann_benchmarks.md +++ b/docs/source/raft_ann_benchmarks.md @@ -4,29 +4,6 @@ This project provides a benchmark program for various ANN search implementations ## Table of Contents -- [Installing the benchmarks](#installing-the-benchmarks) - - [Conda](#conda) - - [Docker](#docker) -- [How to run the benchmarks](#how-to-run-the-benchmarks) - - [Step 1: prepare dataset](#step-1-prepare-dataset) - - [Step 2: build and search index](#step-2-build-and-search-index) - - [Step 3: data export](#step-3-data-export) - - [Step 4: plot results](#step-4-plot-results) -- [Running the benchmarks](#running-the-benchmarks) - - [End to end: small-scale (<1M to 10M)](#end-to-end-small-scale-benchmarks-1m-to-10m) - - [End to end: large-scale (>10M)](#end-to-end-large-scale-benchmarks-10m-vectors) - - [Running with Docker containers](#running-with-docker-containers) - - [Evaluating the results](#evaluating-the-results) -- [Creating and customizing dataset configurations](#creating-and-customizing-dataset-configurations) -- [Adding a new ANN algorithm](#adding-a-new-ann-algorithm) -- [Parameter tuning guide](https://docs.rapids.ai/api/raft/nightly/ann_benchmarks_param_tuning/) -- [Wiki-all RAG/LLM Dataset](https://docs.rapids.ai/api/raft/nightly/wiki_all_dataset/) -## Installing the benchmarks# RAFT ANN Benchmarks - -This project provides a benchmark program for various ANN search implementations. It's especially suitable for comparing GPU implementations as well as comparing GPU against CPU. - -## Table of Contents - - [Installing the benchmarks](#installing-the-benchmarks) - [Conda](#conda) - [Docker](#docker) From 879c744d6de9bde75d9321da7cd457411059302a Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 30 Oct 2023 17:38:05 -0700 Subject: [PATCH 28/49] fix style --- python/raft-ann-bench/src/raft-ann-bench/run/__main__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 80f5a86592..d11cd85387 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 @@ -288,7 +288,8 @@ def main(): algos_conf_fs = os.listdir(os.path.join(scripts_path, "conf", "algos")) algos_conf_fs = [ - os.path.join(scripts_path, "conf", "algos", f) for f in algos_conf_fs + os.path.join(scripts_path, "conf", "algos", f) + for f in algos_conf_fs if ".json" not in f ] conf_filedir = os.path.join(scripts_path, "conf", "algos") From d2ccf2a47b53033ad3ac41b6efdace060f878f9e Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 30 Oct 2023 18:09:42 -0700 Subject: [PATCH 29/49] fix wike datasets config --- .../src/raft-ann-bench/run/conf/datasets.yaml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 33f07b9f22..f422d302c5 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 @@ -85,15 +85,15 @@ distance: euclidean - name: wiki_all_1M, - base_file: wiki_all_1M/base.88M.fbin, + base_file: wiki_all_1M/base.1MM.fbin, query_file: wiki_all_1M/queries.fbin, - groundtruth_neighbors_file: wiki_all_1M/groundtruth.88M.neighbors.ibin, + groundtruth_neighbors_file: wiki_all_1M/groundtruth.1M.neighbors.ibin, distance: euclidean - name: wiki_all_10M, - base_file: wiki_all_10M/base.88M.fbin, + base_file: wiki_all_10M/base.10M.fbin, query_file: wiki_all_10M/queries.fbin, - groundtruth_neighbors_file: wiki_all_10M/groundtruth.88M.neighbors.ibin, + groundtruth_neighbors_file: wiki_all_10M/groundtruth.10M.neighbors.ibin, distance: euclidean - name: wiki_all_88M, From 124d091244f4a5b325f76b678a3b200f805489bc Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Mon, 30 Oct 2023 23:14:46 -0400 Subject: [PATCH 30/49] Changing FAISS M to M_ratio. Adding build validator for ivf-pq. Adding `dim` to datasets.yaml. Adding new yaml configs for FAISS. Fixing yaml config for RAFT --- .../ann/src/faiss/faiss_cpu_benchmark.cpp | 2 +- cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h | 10 ++- .../ann/src/faiss/faiss_gpu_benchmark.cu | 2 +- cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h | 5 +- cpp/bench/ann/src/raft/raft_benchmark.cu | 10 +-- docs/source/ann_benchmarks_param_tuning.md | 74 +++++++++---------- .../src/raft-ann-bench/run/__main__.py | 7 +- .../run/conf/algos/faiss_gpu_ivf_flat.yaml | 10 +++ .../run/conf/algos/faiss_gpu_ivf_pq.yaml | 12 +++ .../run/conf/algos/raft_ivf_flat.yaml | 4 +- .../run/conf/algos/raft_ivf_pq.yaml | 9 ++- .../src/raft-ann-bench/run/conf/datasets.yaml | 17 +++++ .../src/raft-ann-bench/validators/__init__.py | 5 ++ 13 files changed, 106 insertions(+), 61 deletions(-) create mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_flat.yaml create mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yaml diff --git a/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp b/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp index 97d1bbf307..8a0a3ffd37 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 = conf.at("M"); + param.M_ratio = conf.at("M_ratio"); 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 028a444530..626e52b086 100644 --- a/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h +++ b/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h @@ -229,7 +229,7 @@ template class FaissCpuIVFPQ : public FaissCpu { public: struct BuildParam : public FaissCpu::BuildParam { - int M; + int M_ratio; int bitsPerCode; bool usePrecomputed; }; @@ -237,8 +237,12 @@ 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, param.M, param.bitsPerCode, this->metric_type_); + this->index_ = std::make_unique(this->quantizer_.get(), + dim, + param.nlist, + dim / param.M_ratio, + 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 8b04ba1980..8de8973f16 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 = conf.at("M"); + param.M_ratio = conf.at("M_ratio"); 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 38eeddf813..8db8e29ef7 100644 --- a/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h +++ b/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h @@ -263,7 +263,7 @@ template class FaissGpuIVFPQ : public FaissGpu { public: struct BuildParam : public FaissGpu::BuildParam { - int M; + int M_ratio; bool useFloat16; bool usePrecomputed; }; @@ -274,11 +274,12 @@ class FaissGpuIVFPQ : public FaissGpu { config.useFloat16LookupTables = param.useFloat16; config.usePrecomputedTables = param.usePrecomputed; config.device = this->device_; + this->index_ = std::make_unique(&(this->gpu_resource_), dim, param.nlist, - param.M, + dim / param.M_ratio, 8, // FAISS only supports bitsPerCode=8 this->metric_type_, config); diff --git a/cpp/bench/ann/src/raft/raft_benchmark.cu b/cpp/bench/ann/src/raft/raft_benchmark.cu index 3b9bcc7e15..fa20c5c223 100644 --- a/cpp/bench/ann/src/raft/raft_benchmark.cu +++ b/cpp/bench/ann/src/raft/raft_benchmark.cu @@ -272,13 +272,5 @@ REGISTER_ALGO_INSTANCE(std::uint8_t); #ifdef ANN_BENCH_BUILD_MAIN #include "../common/benchmark.hpp" -int main(int argc, char** argv) -{ - rmm::mr::cuda_memory_resource cuda_mr; - // Construct a resource that uses a coalescing best-fit pool allocator - rmm::mr::pool_memory_resource pool_mr{&cuda_mr}; - rmm::mr::set_current_device_resource( - &pool_mr); // Updates the current device resource pointer to `pool_mr` - return raft::bench::ann::run_main(argc, argv); -} +int main(int argc, char** argv) { return raft::bench::ann::run_main(argc, argv); } #endif diff --git a/docs/source/ann_benchmarks_param_tuning.md b/docs/source/ann_benchmarks_param_tuning.md index 075d82a135..90a5e54e32 100644 --- a/docs/source/ann_benchmarks_param_tuning.md +++ b/docs/source/ann_benchmarks_param_tuning.md @@ -15,34 +15,34 @@ IVF-flat uses an inverted-file index, which partitions the vectors into a series IVF-flat is a simple algorithm which won't save any space, but it provides competitive search times even at higher levels of recall. -| Parameter | Type | Required | Data Type | Default | Description | -|-----------------------|------------------|----------|----------------------------|----------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `nlists` | `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. | -| `niter` | `build_param` | N | Positive Integer >0 | 20 | 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. | +| 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. | +| `niter` | `build_param` | N | Positive Integer >0 | 20 | 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. | | `dataset_memory_type` | `build_param` | N | ["device", "host", "mmap"] | "device" | What memory type should the dataset reside? | -| `query_memory_type` | `search_params` | N | ["device", "host", "mmap"] | "device | What memory type should the queries reside? | -| `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. | +| `query_memory_type` | `search_params` | N | ["device", "host", "mmap"] | "device | What memory type should the queries reside? | +| `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. | ### `raft_ivf_pq` IVF-pq is an inverted-file index, which partitions the vectors into a series of clusters, or lists, in a similar way to IVF-flat above. The difference is that IVF-PQ uses product quantization to also compress the vectors, giving the index a smaller memory footprint. Unfortunately, higher levels of compression can also shrink recall, which a refinement step can improve when the original vectors are still available. -| Parameter | Type | Required | Data Type | Default | Description | -|-------------------------|----------------|---|----------------------------------|---------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `nlists` | `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. | -| `niter` | `build_param` | N | Positive Integer >0 | 20 | Number of k-means iterations to use when training the clusters. | -| `ratio` | `build_param` | N | Positive Integer >0 | 2 | `1/ratio` is the number of training points which should be used to train the clusters. | -| `pq_dim` | `build_param` | N | Positive Integer. Multiple of 8. | 0 | Dimensionality of the vector after product quantization. When 0, a heuristic is used to select this value. `pq_dim` * `pq_bits` must be a multiple of 8. | -| `pq_bits` | `build_param` | N | Positive Integer. [4-8] | 8 | Bit length of the vector element after quantization. | -| `codebook_kind` | `build_param` | N | ["cluster", "subspace"] | "subspace" | Type of codebook. See the [API docs](https://docs.rapids.ai/api/raft/nightly/cpp_api/neighbors_ivf_pq/#_CPPv412codebook_gen) for more detail | -| `dataset_memory_type` | `build_param` | N | ["device", "host", "mmap"] | "device" | What memory type should the dataset reside? | -| `query_memory_type` | `search_params` | N | ["device", "host", "mmap"] | "device | What memory type should the queries reside? | -| `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. | +| 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. | +| `niter` | `build_param` | N | Positive Integer >0 | 20 | Number of k-means iterations to use when training the clusters. | +| `ratio` | `build_param` | N | Positive Integer >0 | 2 | `1/ratio` is the number of training points which should be used to train the clusters. | +| `pq_dim` | `build_param` | N | Positive Integer. Multiple of 8. | 0 | Dimensionality of the vector after product quantization. When 0, a heuristic is used to select this value. `pq_dim` * `pq_bits` must be a multiple of 8. | +| `pq_bits` | `build_param` | N | Positive Integer. [4-8] | 8 | Bit length of the vector element after quantization. | +| `codebook_kind` | `build_param` | N | ["cluster", "subspace"] | "subspace" | Type of codebook. See the [API docs](https://docs.rapids.ai/api/raft/nightly/cpp_api/neighbors_ivf_pq/#_CPPv412codebook_gen) for more detail | +| `dataset_memory_type` | `build_param` | N | ["device", "host", "mmap"] | "device" | What memory type should the dataset reside? | +| `query_memory_type` | `search_params` | N | ["device", "host", "mmap"] | "device | What memory type should the queries reside? | +| `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. | | `internalDistanceDtype` | `search_params` | N | [`float`, `half`] | `half` | The precision to use for the distance computations. Lower precision can increase performance at the cost of accuracy. | -| `smemLutDtype` | `search_params` | N | [`float`, `half`, `fp8`] | `half` | The precision to use for the lookup table in shared memory. Lower precision can increase performance at the cost of accuracy. | -| `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. | +| `smemLutDtype` | `search_params` | N | [`float`, `half`, `fp8`] | `half` | The precision to use for the lookup table in shared memory. Lower precision can increase performance at the cost of accuracy. | +| `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. | ### `raft_cagra` @@ -86,13 +86,13 @@ IVF-pq is an inverted-file index, which partitions the vectors into a series of | Parameter | Type | Required | Data Type | Default | Description | |------------------|----------------|----------|----------------------------------|---------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `nlists` | `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. | +| `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. | +| `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. | -| `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. | +| `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` @@ -108,26 +108,26 @@ Use FAISS flat index on the CPU, which performs an exact search using brute-forc Use FAISS IVF-Flat index on CPU | Parameter | Type | Required | Data Type | Default | Description | -|-----------|----------------|----------|---------------------|---------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `nlists` | `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. | -| `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. | -| `numThreads` | `search_params` | N | Positive Integer >0 | 1 | Number of threads to use for queries. | +|----------|----------------|----------|---------------------|---------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `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. | +| `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. | +| `numThreads` | `search_params` | N | Positive Integer >0 | 1 | Number of threads to use for queries. | ### `faiss_cpu_ivf_pq` Use FAISS IVF-PQ index on CPU -| Parameter | Type | Required | Data Type | Default | Description | -|------------------|----------------|----------|------------------------------------|---------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `nlists` | `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. | +| `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. | ## HNSW 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 80f5a86592..afc5d53556 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 @@ -288,7 +288,8 @@ def main(): algos_conf_fs = os.listdir(os.path.join(scripts_path, "conf", "algos")) algos_conf_fs = [ - os.path.join(scripts_path, "conf", "algos", f) for f in algos_conf_fs + os.path.join(scripts_path, "conf", "algos", f) + for f in algos_conf_fs if ".json" not in f ] conf_filedir = os.path.join(scripts_path, "conf", "algos") @@ -389,7 +390,9 @@ def add_algo_group(group_list): func = importable[-1] validator = import_module(module) build_validator = getattr(validator, func) - if not build_validator(index["build_param"]): + if not build_validator( + index["build_param"], conf_file["dataset"]["dims"] + ): continue index["name"] = index_name 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 new file mode 100644 index 0000000000..6542bbab4c --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_flat.yaml @@ -0,0 +1,10 @@ +name: faiss_gpu_ivf_flat +groups: + base: + build: + nlists: [1024, 2048, 4096, 8192, 16000, 32000] + ratio: [1, 10, 25] + useFloat16: [True, False] + search: + numProbes: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] + refine_ratio: [1, 2, 4, 10] \ No newline at end of file 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 new file mode 100644 index 0000000000..7e453d506e --- /dev/null +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yaml @@ -0,0 +1,12 @@ +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] + search: + numProbes: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] + refine_ratio: [1, 2, 4, 10] \ No newline at end of file diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yaml index 6c476af39e..c36a26514d 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yaml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_ivf_flat.yaml @@ -2,8 +2,8 @@ name: raft_ivf_flat groups: base: build: - nlist: [1024, 16384] - ratio: [1, 2] + nlist: [1024, 2048, 4096, 8192, 16384, 32000, 64000] + ratio: [1, 2, 4] niter: [20, 25] search: nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] \ No newline at end of file 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 107e0ac640..82a4f9a7f1 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,13 +1,14 @@ 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 groups: base: build: - nlist: [1024] - pq_dim: [128, 64] - pq_bits: [8, 6] - ratio: [1] + nlist: [1024, 2048, 4096, 8192, 16000, 32000] + pq_dim: [384, 256, 128, 64, 32] + pq_bits: [8, 6, 5, 4] + ratio: [1, 10, 25] niter: [25] search: nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] 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 33f07b9f22..28f29a5c92 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,6 +1,7 @@ - name: bigann-100M 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 distance: euclidean @@ -8,12 +9,14 @@ - name: deep-1B base_file: deep-1B/base.1B.fbin query_file: deep-1B/query.public.10K.fbin + dims: 96 groundtruth_neighbors_file: deep-1B/groundtruth.neighbors.ibin distance: inner_product - name: bigann-100M 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 distance: euclidean @@ -21,6 +24,7 @@ - name: deep-image-96-inner base_file: deep-image-96-inner/base.fbin query_file: deep-image-96-inner/query.fbin + dims: 96 groundtruth_neighbors_file: deep-image-96-inner/groundtruth.neighbors.ibin distance: euclidean @@ -31,72 +35,85 @@ distance: euclidean - name: gist-960-euclidean + dims: 960 base_file: gist-960-euclidean/base.fbin query_file: gist-960-euclidean/query.fbin distance: euclidean - name: glove-50-angular + dims: 50 base_file: glove-50-angular/base.fbin query_file: glove-50-angular/query.fbin distance: euclidean - name: glove-50-inner + dims: 50 base_file: glove-50-inner/base.fbin query_file: glove-50-inner/query.fbin distance: euclidean - name: glove-100-angular + dims: 100 base_file: glove-100-angular/base.fbin query_file: glove-100-angular/query.fbin distance: euclidean - name: glove-100-inner + dims: 100 base_file: glove-100-inner/base.fbin query_file: glove-100-inner/query.fbin distance: euclidean - name: lastfm-65-angular + dims: 65 base_file: lastfm-65-angular/base.fbin query_file: lastfm-65-angular/query.fbin distance: euclidean - name: mnist-784-euclidean + dims: 784 base_file: mnist-784-euclidean/base.fbin query_file: mnist-784-euclidean/query.fbin groundtruth_neighbors_file: mnist-784-euclidean/groundtruth.neighbors.ibin distance: euclidean - name: nytimes-256-angular + dims: 256 base_file: nytimes-256-angular/base.fbin query_file: nytimes-256-angular/query.fbin groundtruth_neighbors_file: nytimes-256-angular/groundtruth.neighbors.ibin distance: euclidean - name: nytimes-256-inner + dims: 256 base_file: nytimes-256-inner/base.fbin query_file: nytimes-256-inner/query.fbin groundtruth_neighbors_file: nytimes-256-inner/groundtruth.neighbors.ibin distance: euclidean - name: sift-128-euclidean + dims: 128 base_file: sift-128-euclidean/base.fbin query_file: sift-128-euclidean/query.fbin groundtruth_neighbors_file: sift-128-euclidean/groundtruth.neighbors.ibin distance: euclidean - name: wiki_all_1M, + dims: 784 base_file: wiki_all_1M/base.88M.fbin, query_file: wiki_all_1M/queries.fbin, groundtruth_neighbors_file: wiki_all_1M/groundtruth.88M.neighbors.ibin, distance: euclidean - name: wiki_all_10M, + dims: 784 base_file: wiki_all_10M/base.88M.fbin, query_file: wiki_all_10M/queries.fbin, groundtruth_neighbors_file: wiki_all_10M/groundtruth.88M.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, diff --git a/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py b/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py index 04448425ee..81fc181d41 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py @@ -16,6 +16,11 @@ DTYPE_SIZES = {"float": 4, "half": 2, "fp8": 1} +def raft_ivf_pq_build_validator(params, dims): + if "pq_dim" in params: + return params["pq_dim"] <= dims + + def raft_ivf_pq_search_validator(params, k, batch_size): if "internalDistanceDtype" in params and "smemLutDtype" in params: return ( From 5538b05b7f31fa2347e1f97da4fe6809a79847e2 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 00:22:19 -0400 Subject: [PATCH 31/49] More work on configs and validators --- .../src/raft-ann-bench/run/__main__.py | 9 +++++++-- .../raft-ann-bench/run/conf/algos/raft_ivf_pq.yaml | 9 +++++---- .../src/raft-ann-bench/run/conf/datasets.yaml | 1 + .../src/raft-ann-bench/validators/__init__.py | 14 ++++++++++---- 4 files changed, 23 insertions(+), 10 deletions(-) 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 afc5d53556..477c289666 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 @@ -243,7 +243,7 @@ def main(): "--search-mode", help="run search in 'latency' (measure individual batches) or " "'throughput' (pipeline batches and measure end-to-end) mode", - default="throughput", + default="latency", ) args = parser.parse_args() @@ -415,7 +415,12 @@ def add_algo_group(group_list): func = importable[-1] validator = import_module(module) search_validator = getattr(validator, func) - if search_validator(search_dict, k, batch_size): + if search_validator( + search_dict, + index["build_param"], + k, + batch_size, + ): index["search_params"].append(search_dict) executables_to_run[executable]["index"].append(index) 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 82a4f9a7f1..2e1912c6b0 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 @@ -5,12 +5,13 @@ validators: groups: base: build: - nlist: [1024, 2048, 4096, 8192, 16000, 32000] - pq_dim: [384, 256, 128, 64, 32] + nlist: [1024, 2048, 4096, 8192] + pq_dim: [64, 32] pq_bits: [8, 6, 5, 4] ratio: [1, 10, 25] niter: [25] search: - nprobe: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] + nprobe: [1, 5, 10, 50, 100, 200, 500] internalDistanceDtype: ["float", "half"] - smemLutDtype: ["float", "fp8", "half"] \ No newline at end of file + 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 b72bdb23c9..23476cc056 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 @@ -29,6 +29,7 @@ distance: euclidean - name: fashion-mnist-784-euclidean + dims: 784 base_file: fashion-mnist-784-euclidean/base.fbin query_file: fashion-mnist-784-euclidean/query.fbin groundtruth_neighbors_file: fashion-mnist-784-euclidean/groundtruth.neighbors.ibin diff --git a/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py b/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py index 81fc181d41..03bf707e6b 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py @@ -19,16 +19,22 @@ def raft_ivf_pq_build_validator(params, dims): if "pq_dim" in params: return params["pq_dim"] <= dims + return True -def raft_ivf_pq_search_validator(params, k, batch_size): +def raft_ivf_pq_search_validator(params, build_params, k, batch_size): + ret = True if "internalDistanceDtype" in params and "smemLutDtype" in params: - return ( + 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"] + return ret -def raft_cagra_search_validator(params, k, batch_size): + +def raft_cagra_search_validator(params, build_params, k, batch_size): if "itopk" in params: return params["itopk"] >= k From 32948f7e934e9918da1ac1d447b778c9965d3a1a Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 10:36:28 -0400 Subject: [PATCH 32/49] Many fixes around --- cpp/bench/ann/src/common/benchmark.hpp | 13 ++++++- cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h | 2 +- cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h | 38 ++++++++++++++----- docs/source/ann_benchmarks_param_tuning.md | 20 +++++----- .../src/raft-ann-bench/run/__main__.py | 17 +++++---- .../run/conf/algos/faiss_gpu_ivf_pq.yaml | 18 +++++++-- .../run/conf/algos/raft_ivf_pq.yaml | 4 +- 7 files changed, 77 insertions(+), 35 deletions(-) diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index 10a256bd63..bff3cfa5e1 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -216,7 +216,10 @@ void bench_search(::benchmark::State& state, } catch (const std::exception& e) { state.SkipWithError("Failed to create an algo: " + std::string(e.what())); } - algo->set_search_param(*search_param); + + /** + * It's important that we guarantee search dataset will always be set before search params. + */ 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); @@ -230,6 +233,14 @@ void bench_search(::benchmark::State& state, "Exception: " + std::string(ex.what())); 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; + } } } diff --git a/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h b/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h index 626e52b086..2c5f9439ac 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; } diff --git a/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h b/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h index 8db8e29ef7..b6506479a7 100644 --- a/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h +++ b/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h @@ -142,7 +142,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_; @@ -217,7 +217,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 +248,6 @@ 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; - } } void save(const std::string& file) const override @@ -268,6 +269,12 @@ class FaissGpuIVFPQ : public FaissGpu { bool usePrecomputed; }; + struct SearchParam : public FaissGpu::SearchParam { + int nprobe; + float refine_ratio = 1.0; + auto needs_dataset() const -> bool override { return true; } + }; + FaissGpuIVFPQ(Metric metric, int dim, const BuildParam& param) : FaissGpu(metric, dim, param) { faiss::gpu::GpuIndexIVFPQConfig config; @@ -285,6 +292,11 @@ class FaissGpuIVFPQ : public FaissGpu { config); } + void set_search_dataset(const T* dataset, size_t nrow) + { + this->index_refine_ = std::make_unique(this->index_.get(), dataset); + } + void set_search_param(const typename FaissGpu::AnnSearchParam& param) override { auto search_param = dynamic_cast::SearchParam&>(param); @@ -297,7 +309,6 @@ class FaissGpuIVFPQ : public FaissGpu { 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; } } @@ -321,6 +332,12 @@ class FaissGpuIVFSQ : public FaissGpu { std::string quantizer_type; }; + struct SearchParam : public FaissGpu::SearchParam { + int nprobe; + float refine_ratio = 1.0; + auto needs_dataset() const -> bool override { return true; } + }; + FaissGpuIVFSQ(Metric metric, int dim, const BuildParam& param) : FaissGpu(metric, dim, param) { faiss::ScalarQuantizer::QuantizerType qtype; @@ -339,6 +356,10 @@ class FaissGpuIVFSQ : public FaissGpu { &(this->gpu_resource_), dim, param.nlist, qtype, this->metric_type_, true, config); } + void set_search_dataset(const T* dataset, size_t nrow) + { + this->index_refine_ = std::make_unique(this->index_.get(), dataset); + } void set_search_param(const typename FaissGpu::AnnSearchParam& param) override { auto search_param = dynamic_cast::SearchParam&>(param); @@ -351,7 +372,6 @@ class FaissGpuIVFSQ : public FaissGpu { 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; } } 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/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..462b5097c6 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 @@ -128,14 +128,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"] @@ -401,6 +401,7 @@ def add_algo_group(group_list): ) index["search_params"] = [] all_search_params = itertools.product(*search_param_lists) + for search_params in all_search_params: search_dict = dict() for i in range(len(search_params)): @@ -415,15 +416,15 @@ def add_algo_group(group_list): func = importable[-1] validator = import_module(module) search_validator = getattr(validator, func) - if search_validator( - search_dict, - index["build_param"], - k, - batch_size, - ): + if search_validator(search_dict, 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_pq.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yaml index 7e453d506e..53a3d0ba1d 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,21 @@ name: faiss_gpu_ivf_pq groups: base: build: - nlist: [1024, 2048, 4096, 8192, 16000, 32000, 64000, 100000] + nlist: [1024, 2048, 4096, 8192] M_ratio: [2, 4] - ratio: [1, 10, 25] + ratio: [10, 25] usePrecomputed: [True, False] useFloat16: [True, 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, 2, 4] + test: + build: + nlist: [1024] + M_ratio: [2] + ratio: [10] + useFloat16: [False] + search: + nprobe: [1, 5, 10, 50, 100, 200] + refine_ratio: [1] + 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..dd76510ef7 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 @@ -8,10 +8,10 @@ groups: 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] + nprobe: [1, 5, 10, 50, 100, 200] internalDistanceDtype: ["float", "half"] smemLutDtype: ["float", "fp8", "half"] refine_ratio: [1, 2, 4] \ No newline at end of file From 03c2578cf70ab065ba5d02eca5e0e8f8e8b37056 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 10:43:59 -0400 Subject: [PATCH 33/49] Properly calling valiator --- python/raft-ann-bench/src/raft-ann-bench/run/__main__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 462b5097c6..28714c5037 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 @@ -416,7 +416,9 @@ def add_algo_group(group_list): func = importable[-1] validator = import_module(module) search_validator = getattr(validator, func) - if search_validator(search_dict, k, batch_size): + if search_validator( + search_dict, build_params, k, batch_size + ): index["search_params"].append(search_dict) else: index["search_params"].append(search_dict) From aa0fada0fc7ed12310b8029d8280bee6e3bf8f40 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 10:46:11 -0400 Subject: [PATCH 34/49] More fixes --- python/raft-ann-bench/src/raft-ann-bench/run/__main__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 28714c5037..d0400eeb58 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 @@ -417,7 +417,10 @@ def add_algo_group(group_list): validator = import_module(module) search_validator = getattr(validator, func) if search_validator( - search_dict, build_params, k, batch_size + search_dict, + index["build_param"], + k, + batch_size, ): index["search_params"].append(search_dict) else: From b11326c338dccd52f3dbd312753fcf07ccbaa87b Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 10:56:43 -0400 Subject: [PATCH 35/49] More fixes --- .../raft-ann-bench/src/raft-ann-bench/validators/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py b/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py index 03bf707e6b..4891e19908 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/validators/__init__.py @@ -27,11 +27,11 @@ def raft_ivf_pq_search_validator(params, build_params, k, batch_size): 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 = build_params["nlist"] >= params["nprobe"] return ret From df043bb6dd82a1281ed34c4acc8c3749e87a6ca0 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 15:08:08 -0400 Subject: [PATCH 36/49] Pushing for second set of eyes. --- cpp/bench/ann/src/common/ann_types.hpp | 2 +- cpp/bench/ann/src/common/benchmark.hpp | 25 ++++++----- .../ann/src/faiss/faiss_cpu_benchmark.cpp | 2 +- cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h | 10 ++--- .../ann/src/faiss/faiss_gpu_benchmark.cu | 2 +- cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h | 42 ++++++++++++------- .../run/conf/algos/faiss_gpu_ivf_pq.yaml | 9 ++-- 7 files changed, 50 insertions(+), 42 deletions(-) diff --git a/cpp/bench/ann/src/common/ann_types.hpp b/cpp/bench/ann/src/common/ann_types.hpp index 2c1105a272..f89743ac81 100644 --- a/cpp/bench/ann/src/common/ann_types.hpp +++ b/cpp/bench/ann/src/common/ann_types.hpp @@ -124,7 +124,7 @@ class ANN : public AnnBase { // 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. - virtual void set_search_dataset(const T* /*dataset*/, size_t /*nrow*/){}; + virtual void set_search_dataset(const T* /*dataset*/, size_t /*nrow*/) { printf("Setting \n"); }; }; } // namespace raft::bench::ann diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index bff3cfa5e1..ae36779c86 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -172,6 +172,7 @@ void bench_search(::benchmark::State& state, std::ptrdiff_t batch_offset = 0; std::size_t queries_processed = 0; + printf("Starting benchmark search\n"); double total_time = 0; const auto& sp_json = index.search_params[search_param_ix]; @@ -201,6 +202,7 @@ void bench_search(::benchmark::State& state, index_file = index.file; } + printf("Loading index from file\n"); std::unique_ptr::AnnSearchParam> search_param; ANN* algo; try { @@ -217,14 +219,23 @@ void bench_search(::benchmark::State& state, state.SkipWithError("Failed to create an algo: " + std::string(e.what())); } - /** - * It's important that we guarantee search dataset will always be set before search params. - */ + printf("Set search params\n"); + 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; + } + + printf("Setting search dataset\n"); 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); + + printf("AFTER!\n"); if (search_param->needs_dataset()) { try { + printf("About to set search datast\n"); algo->set_search_dataset(dataset->base_set(current_algo_props->dataset_memory_type), dataset->base_set_size()); } catch (const std::exception& ex) { @@ -233,14 +244,6 @@ void bench_search(::benchmark::State& state, "Exception: " + std::string(ex.what())); 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; - } } } 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 2c5f9439ac..755fe9f197 100644 --- a/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h +++ b/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h @@ -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 b6506479a7..c94450d480 100644 --- a/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h +++ b/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h @@ -84,6 +84,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 { @@ -109,6 +110,8 @@ class FaissGpu : public ANN { virtual void set_search_param(const FaissGpu::AnnSearchParam& param) {} + virtual void set_search_dataset(const T* dataset, size_t nrow) {} + // 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, @@ -142,7 +145,7 @@ class FaissGpu : public ANN { mutable faiss::gpu::StandardGpuResources gpu_resource_; std::unique_ptr index_; - std::unique_ptr index_refine_{nullptr}; + std::unique_ptr index_refine_; faiss::MetricType metric_type_; int nlist_; int device_; @@ -194,7 +197,14 @@ 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 (index_refine_->k_factor > 1) { + printf("Using refine!\n"); + index_refine_->search( + batch_size, queries, k, distances, reinterpret_cast(neighbors)); + } else { + index_->search(batch_size, queries, k, distances, reinterpret_cast(neighbors)); + } stream_wait(stream); } @@ -264,17 +274,11 @@ template class FaissGpuIVFPQ : public FaissGpu { public: struct BuildParam : public FaissGpu::BuildParam { - int M_ratio; + int M; bool useFloat16; bool usePrecomputed; }; - struct SearchParam : public FaissGpu::SearchParam { - int nprobe; - float refine_ratio = 1.0; - auto needs_dataset() const -> bool override { return true; } - }; - FaissGpuIVFPQ(Metric metric, int dim, const BuildParam& param) : FaissGpu(metric, dim, param) { faiss::gpu::GpuIndexIVFPQConfig config; @@ -286,19 +290,21 @@ 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); } - void set_search_dataset(const T* dataset, size_t nrow) + void set_search_dataset(const T* dataset, size_t nrow) override { - this->index_refine_ = std::make_unique(this->index_.get(), dataset); + printf("Setting search ataset for refine\n"); + dataset_ = dataset; } void set_search_param(const typename FaissGpu::AnnSearchParam& param) override { + printf("Setting ivfpq search params\n"); auto search_param = dynamic_cast::SearchParam&>(param); int nprobe = search_param.nprobe; assert(nprobe <= nlist_); @@ -309,6 +315,8 @@ class FaissGpuIVFPQ : public FaissGpu { 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->dataset_); this->index_refine_.get()->k_factor = search_param.refine_ratio; } } @@ -321,6 +329,8 @@ class FaissGpuIVFPQ : public FaissGpu { { this->template load_(file); } + + const T* dataset_; }; // TODO: Enable this in cmake @@ -356,10 +366,7 @@ class FaissGpuIVFSQ : public FaissGpu { &(this->gpu_resource_), dim, param.nlist, qtype, this->metric_type_, true, config); } - void set_search_dataset(const T* dataset, size_t nrow) - { - this->index_refine_ = std::make_unique(this->index_.get(), dataset); - } + void set_search_dataset(const T* dataset, size_t nrow) override { this->dataset_ = dataset; } void set_search_param(const typename FaissGpu::AnnSearchParam& param) override { auto search_param = dynamic_cast::SearchParam&>(param); @@ -372,6 +379,7 @@ class FaissGpuIVFSQ : public FaissGpu { 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(), dataset_); this->index_refine_.get()->k_factor = search_param.refine_ratio; } } @@ -386,6 +394,8 @@ class FaissGpuIVFSQ : public FaissGpu { this->template load_( file); } + + const T* dataset_; }; template 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 53a3d0ba1d..3a17eacc2d 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 @@ -3,15 +3,15 @@ groups: base: build: nlist: [1024, 2048, 4096, 8192] - M_ratio: [2, 4] + M: [8, 16] ratio: [10, 25] - usePrecomputed: [True, False] - useFloat16: [True, False] + usePrecomputed: [False] + useFloat16: [False] search: nprobe: [1, 5, 10, 50, 100, 200] refine_ratio: [1, 2, 4] test: - build: + build: nlist: [1024] M_ratio: [2] ratio: [10] @@ -19,4 +19,3 @@ groups: search: nprobe: [1, 5, 10, 50, 100, 200] refine_ratio: [1] - From b71ef73ca48e55563bc049cd6425605522cff43e Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 20:31:22 -0400 Subject: [PATCH 37/49] Turning off faiss refinement for the time being. --- cpp/bench/ann/src/common/ann_types.hpp | 4 +- cpp/bench/ann/src/common/benchmark.cpp | 1 + cpp/bench/ann/src/common/benchmark.hpp | 20 +++---- cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h | 60 ++++++++++--------- .../run/conf/algos/faiss_gpu_ivf_flat.yaml | 10 ++-- .../run/conf/algos/faiss_gpu_ivf_pq.yaml | 2 +- 6 files changed, 48 insertions(+), 49 deletions(-) diff --git a/cpp/bench/ann/src/common/ann_types.hpp b/cpp/bench/ann/src/common/ann_types.hpp index f89743ac81..e964a81efa 100644 --- a/cpp/bench/ann/src/common/ann_types.hpp +++ b/cpp/bench/ann/src/common/ann_types.hpp @@ -120,11 +120,11 @@ 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. - virtual void set_search_dataset(const T* /*dataset*/, size_t /*nrow*/) { printf("Setting \n"); }; + virtual void set_search_dataset(const T* /*dataset*/, size_t /*nrow*/){}; }; } // namespace raft::bench::ann diff --git a/cpp/bench/ann/src/common/benchmark.cpp b/cpp/bench/ann/src/common/benchmark.cpp index 6424a36471..62c91b0791 100644 --- a/cpp/bench/ann/src/common/benchmark.cpp +++ b/cpp/bench/ann/src/common/benchmark.cpp @@ -88,6 +88,7 @@ template std::unique_ptr::AnnSearchParam> create_search_param( const std::string& algo, const nlohmann::json& conf) { + printf("INside create_search_param\n"); static auto fname = get_fun_name(reinterpret_cast(&create_search_param)); auto handle = load_lib(algo); auto fun_addr = dlsym(handle, fname.c_str()); diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index ae36779c86..3a930d288e 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -172,7 +172,6 @@ void bench_search(::benchmark::State& state, std::ptrdiff_t batch_offset = 0; std::size_t queries_processed = 0; - printf("Starting benchmark search\n"); double total_time = 0; const auto& sp_json = index.search_params[search_param_ix]; @@ -202,7 +201,6 @@ void bench_search(::benchmark::State& state, index_file = index.file; } - printf("Loading index from file\n"); std::unique_ptr::AnnSearchParam> search_param; ANN* algo; try { @@ -217,25 +215,15 @@ 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())); - } - - printf("Set search params\n"); - 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; } - printf("Setting search dataset\n"); 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); - printf("AFTER!\n"); if (search_param->needs_dataset()) { try { - printf("About to set search datast\n"); algo->set_search_dataset(dataset->base_set(current_algo_props->dataset_memory_type), dataset->base_set_size()); } catch (const std::exception& ex) { @@ -245,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_gpu_wrapper.h b/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h index c94450d480..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 @@ -102,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_)); } @@ -110,7 +114,7 @@ class FaissGpu : public ANN { virtual void set_search_param(const FaissGpu::AnnSearchParam& param) {} - virtual void set_search_dataset(const T* dataset, size_t nrow) {} + 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 @@ -126,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; } @@ -145,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_; @@ -153,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 @@ -198,12 +205,23 @@ 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"); - if (index_refine_->k_factor > 1) { - printf("Using refine!\n"); - index_refine_->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)); + index_->search(batch_size, + queries, + k, + distances, + reinterpret_cast(neighbors), + this->search_params_.get()); } stream_wait(stream); } @@ -258,6 +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); + this->refine_ratio_ = search_param.refine_ratio; } void save(const std::string& file) const override @@ -296,19 +315,12 @@ class FaissGpuIVFPQ : public FaissGpu { config); } - void set_search_dataset(const T* dataset, size_t nrow) override - { - printf("Setting search ataset for refine\n"); - dataset_ = dataset; - } - void set_search_param(const typename FaissGpu::AnnSearchParam& param) override { - printf("Setting ivfpq search params\n"); 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; @@ -329,8 +341,6 @@ class FaissGpuIVFPQ : public FaissGpu { { this->template load_(file); } - - const T* dataset_; }; // TODO: Enable this in cmake @@ -342,12 +352,6 @@ class FaissGpuIVFSQ : public FaissGpu { std::string quantizer_type; }; - struct SearchParam : public FaissGpu::SearchParam { - int nprobe; - float refine_ratio = 1.0; - auto needs_dataset() const -> bool override { return true; } - }; - FaissGpuIVFSQ(Metric metric, int dim, const BuildParam& param) : FaissGpu(metric, dim, param) { faiss::ScalarQuantizer::QuantizerType qtype; @@ -366,7 +370,6 @@ class FaissGpuIVFSQ : public FaissGpu { &(this->gpu_resource_), dim, param.nlist, qtype, this->metric_type_, true, config); } - void set_search_dataset(const T* dataset, size_t nrow) override { this->dataset_ = dataset; } void set_search_param(const typename FaissGpu::AnnSearchParam& param) override { auto search_param = dynamic_cast::SearchParam&>(param); @@ -377,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(), dataset_); + this->index_refine_ = + std::make_unique(this->index_.get(), this->dataset_); this->index_refine_.get()->k_factor = search_param.refine_ratio; } } @@ -394,8 +398,6 @@ class FaissGpuIVFSQ : public FaissGpu { this->template load_( file); } - - const T* dataset_; }; template 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 3a17eacc2d..9c8b99c21f 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 @@ -9,7 +9,7 @@ groups: useFloat16: [False] search: nprobe: [1, 5, 10, 50, 100, 200] - refine_ratio: [1, 2, 4] + refine_ratio: [1] test: build: nlist: [1024] From 2ca91336aa7f95830e9236d97a8499045c9b4914 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 20:36:48 -0400 Subject: [PATCH 38/49] Adding hnswlib configuration --- .../src/raft-ann-bench/run/conf/algos/hnswlib.yaml | 8 ++++++++ 1 file changed, 8 insertions(+) create mode 100644 python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/hnswlib.yaml 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 From c7d7cc4bf47f2a90e40a7c401ad7ba3e671ee4c4 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 20:37:59 -0400 Subject: [PATCH 39/49] Removing test group --- .../run/conf/algos/faiss_gpu_ivf_pq.yaml | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) 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 9c8b99c21f..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 @@ -9,13 +9,4 @@ groups: useFloat16: [False] search: nprobe: [1, 5, 10, 50, 100, 200] - refine_ratio: [1] - test: - build: - nlist: [1024] - M_ratio: [2] - ratio: [10] - useFloat16: [False] - search: - nprobe: [1, 5, 10, 50, 100, 200] - refine_ratio: [1] + refine_ratio: [1] \ No newline at end of file From 1da90a14e5f65a26bd6c46568d53d3fce2715091 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 20:41:56 -0400 Subject: [PATCH 40/49] validator->constraints --- .../{validators => constraints}/__init__.py | 6 ++--- .../src/raft-ann-bench/run/__main__.py | 26 ++++++++++--------- 2 files changed, 17 insertions(+), 15 deletions(-) rename python/raft-ann-bench/src/raft-ann-bench/{validators => constraints}/__init__.py (85%) 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 85% 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 4891e19908..cfc0246e0c 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,13 +16,13 @@ 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 = ( @@ -35,6 +35,6 @@ def raft_ivf_pq_search_validator(params, build_params, k, batch_size): 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/run/__main__.py b/python/raft-ann-bench/src/raft-ann-bench/run/__main__.py index 35c44697c8..2e85b148d1 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 @@ -338,8 +338,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 +384,15 @@ 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 not build_constraints( index["build_param"], conf_file["dataset"]["dims"] ): continue @@ -405,17 +407,17 @@ 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, From 4ba835bdf4187faae87b685ede8f992a591d58fc Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 20:44:40 -0400 Subject: [PATCH 41/49] Throwing error when `dims` not specified in datasets.yaml --- python/raft-ann-bench/src/raft-ann-bench/run/__main__.py | 5 +++++ 1 file changed, 5 insertions(+) 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 2e85b148d1..de201b7089 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 @@ -392,6 +392,11 @@ def add_algo_group(group_list): func = importable[-1] validator = import_module(module) 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"] ): From 89590d241ca1dd79bda1a0841365d303d91d4300 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 20:47:31 -0400 Subject: [PATCH 42/49] Updating datasets in the docs --- docs/source/raft_ann_benchmarks.md | 1 + 1 file changed, 1 insertion(+) 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 ``` From 824bce7c1a6fc1d6445105f19795f9b150f34cfc Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 20:51:29 -0400 Subject: [PATCH 43/49] Updating bigann-1B --- .../src/raft-ann-bench/run/conf/datasets.yaml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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..78bab3e92e 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 From ac15453bb793c2d1bf7a40680ae9c37bf7aa562f Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 21:26:35 -0400 Subject: [PATCH 44/49] Fixing bad validator --- .../src/raft-ann-bench/constraints/__init__.py | 4 ++-- .../src/raft-ann-bench/run/conf/algos/raft_cagra.yaml | 4 ++-- .../src/raft-ann-bench/run/conf/algos/raft_ivf_pq.yaml | 8 ++++---- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/python/raft-ann-bench/src/raft-ann-bench/constraints/__init__.py b/python/raft-ann-bench/src/raft-ann-bench/constraints/__init__.py index cfc0246e0c..b61d185356 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/constraints/__init__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/constraints/__init__.py @@ -27,11 +27,11 @@ def raft_ivf_pq_search_constraints(params, build_params, k, batch_size): 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 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..ebeab918f1 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,6 +1,6 @@ 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: 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 dd76510ef7..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,7 +1,7 @@ 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: @@ -12,6 +12,6 @@ groups: niter: [25] search: nprobe: [1, 5, 10, 50, 100, 200] - internalDistanceDtype: ["float", "half"] + internalDistanceDtype: ["float"] smemLutDtype: ["float", "fp8", "half"] refine_ratio: [1, 2, 4] \ No newline at end of file From 657b0db264ab6dfddb6d4aece2d2a3b22c34a16c Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 21:38:15 -0400 Subject: [PATCH 45/49] One additional ix --- cpp/bench/ann/src/common/benchmark.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/cpp/bench/ann/src/common/benchmark.cpp b/cpp/bench/ann/src/common/benchmark.cpp index 62c91b0791..6424a36471 100644 --- a/cpp/bench/ann/src/common/benchmark.cpp +++ b/cpp/bench/ann/src/common/benchmark.cpp @@ -88,7 +88,6 @@ template std::unique_ptr::AnnSearchParam> create_search_param( const std::string& algo, const nlohmann::json& conf) { - printf("INside create_search_param\n"); static auto fname = get_fun_name(reinterpret_cast(&create_search_param)); auto handle = load_lib(algo); auto fun_addr = dlsym(handle, fname.c_str()); From 67380ed86a1fceb1dfbadbdfbb314f72be7456d6 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 22:09:10 -0400 Subject: [PATCH 46/49] Fixing docs and datasets.yaml --- docs/source/wiki_all_dataset.md | 8 ++++---- .../src/raft-ann-bench/run/conf/datasets.yaml | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) 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/run/conf/datasets.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml index 78bab3e92e..ed8336534d 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 @@ -100,21 +100,21 @@ distance: euclidean - name: wiki_all_1M, - dims: 784 + 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 + 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 + 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, From dfe33c19605c01638a13881d71755424e741d2c4 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 23:22:43 -0400 Subject: [PATCH 47/49] Removing commas --- .../src/raft-ann-bench/run/conf/datasets.yaml | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) 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 ed8336534d..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 @@ -99,23 +99,23 @@ groundtruth_neighbors_file: sift-128-euclidean/groundtruth.neighbors.ibin distance: euclidean -- name: wiki_all_1M, +- 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, + 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: 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, + 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, +- 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, + 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 From c97351889f942ab0e860220443a2767767539f6a Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 31 Oct 2023 20:57:29 -0700 Subject: [PATCH 48/49] fix plotting --- .../src/raft-ann-bench/data_export/__main__.py | 2 ++ python/raft-ann-bench/src/raft-ann-bench/plot/__main__.py | 6 ++---- python/raft-ann-bench/src/raft-ann-bench/run/__main__.py | 5 +---- 3 files changed, 5 insertions(+), 8 deletions(-) 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 de201b7089..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: From bbce6fae04efd84a6177f089160689c65df58b56 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Wed, 1 Nov 2023 00:09:25 -0400 Subject: [PATCH 49/49] Using nn-descent for cagra graph build algo --- .../src/raft-ann-bench/run/conf/algos/raft_cagra.yaml | 1 + 1 file changed, 1 insertion(+) 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 ebeab918f1..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 @@ -6,6 +6,7 @@ groups: build: graph_degree: [32, 64] intermediate_graph_degree: [64, 96] + graph_build_algo: ["NN_DESCENT"] search: itopk: [32, 64, 128]