More work on configs and validators

rapidsai · Oct 31, 2023 · 5538b05 · 5538b05
1 parent 683cebe
commit 5538b05
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 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
@@ -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
@@ -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"]
+      smemLutDtype: ["float", "fp8", "half"]
+      refine_ratio: [1, 2, 4]
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
@@ -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
@@ -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