Updated validatoes

rapidsai · Oct 25, 2023 · 48f192a · 48f192a
1 parent b49679c
commit 48f192a
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
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
@@ -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
@@ -1,5 +1,6 @@
 name: raft_ivf_pq
-validator: raft-ann-bench.validators.raft_ivf_pq_validator
+validators:
+  search: raft-ann-bench.validators.raft_cagra_build_validator
 base:
   build:
     nlist: [1024]

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
@@ -13,13 +13,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 DTYPE_SIZES = {"float": 4, "half": 2, "fp8": 1}
 
 
-def ivf_pq_validator(params):
+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