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 d827c3d1e1..2b7b2728fe 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 @@ -16,6 +16,12 @@ DTYPE_SIZES = {"float": 4, "half": 2, "fp8": 1} +def raft_cagra_build_constraints(params, dims): + if "graph_degree" in params and "intermediate_graph_degree" in params: + return params["graph_degree"] <= params["intermediate_graph_degree"] + return True + + def raft_ivf_pq_build_constraints(params, dims): if "pq_dim" in params: return params["pq_dim"] <= dims @@ -36,8 +42,10 @@ def raft_ivf_pq_search_constraints(params, build_params, k, batch_size): def raft_cagra_search_constraints(params, build_params, k, batch_size): + ret = True if "itopk" in params: - return params["itopk"] >= k + ret = ret and params["itopk"] >= k + return ret def hnswlib_search_constraints(params, build_params, k, batch_size): 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 d8015da5c6..374458989a 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,5 +1,6 @@ name: raft_cagra constraints: + build: raft-ann-bench.constraints.raft_cagra_build_constraints search: raft-ann-bench.constraints.raft_cagra_search_constraints groups: base: 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 be63e1d535..188d24d20f 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 @@ -112,7 +112,7 @@ groundtruth_neighbors_file: wiki_all_1M/groundtruth.1M.neighbors.ibin distance: euclidean -- name: wiki_all_10M, +- name: wiki_all_10M dims: 768 base_file: wiki_all_10M/base.10M.fbin query_file: wiki_all_10M/queries.fbin