Adding search-threads option

python/raft-ann-bench/src/raft-ann-bench/run/__main__.py

@@ -82,6 +82,7 @@ def run_build_and_search(
     search,
     k,
     batch_size,
+    search_threads,
     mode="throughput",
 ):
     for executable, ann_executable_path, algo in executables_to_run.keys():
@@ -115,6 +116,9 @@ def run_build_and_search(
             ]
             if force:
                 cmd = cmd + ["--overwrite"]
+
+            if search_threads:
+                cmd = cmd + ["--threads=%s" % search_threads]
             cmd = cmd + [temp_conf_filepath]
             subprocess.run(cmd, check=True)
 
@@ -128,7 +132,7 @@ def run_build_and_search(
                 "--benchmark_counters_tabular=true",
                 "--override_kv=k:%s" % k,
                 "--override_kv=n_queries:%s" % batch_size,
-                "--benchmark_min_warmup_time=0.01",
+                "--benchmark_min_warmup_time=1",
                 "--benchmark_out_format=json",
                 "--mode=%s" % mode,
                 "--benchmark_out="
@@ -243,6 +247,16 @@ def main():
         default="latency",
     )
 
+    parser.add_argument(
+        "-t",
+        "--search-threads",
+        help="specify the number threads to use for throughput benchmark."
+        " Power of 2 values between 'min' and 'max' will be used. If only "
+        "'min' is specified, then a single test is run with 'min' threads. "
+        "By default min=1, max=<num hyper threads>.",
+        default=None,
+    )
+
     args = parser.parse_args()
 
     # If both build and search are not provided,