Add --queries option to specify number of queries to run

Author: fcostaoliveira

engine/base_client/client.py

@@ -90,6 +90,7 @@ def run_experiment(
         parallels: [int] = [],
         upload_start_idx: int = 0,
         upload_end_idx: int = -1,
+        num_queries: int = -1,
     ):
         execution_params = self.configurator.execution_params(
             distance=dataset.config.distance, vector_size=dataset.config.vector_size
@@ -161,7 +162,7 @@ def run_experiment(
                     )
 
                     search_stats = searcher.search_all(
-                        dataset.config.distance, reader.read_queries()
+                        dataset.config.distance, reader.read_queries(), num_queries
                     )
                     # ensure we specify the client count in the results
                     search_params["parallel"] = client_count

engine/base_client/search.py

@@ -2,11 +2,10 @@
 import time
 from multiprocessing import get_context
 from typing import Iterable, List, Optional, Tuple
-import itertools
 
 import numpy as np
 import tqdm
-import os 
+import os
 
 from dataset_reader.base_reader import Query
 
@@ -62,6 +61,7 @@ def search_all(
         self,
         distance,
         queries: Iterable[Query],
+        num_queries: int = -1,
     ):
         parallel = self.search_params.get("parallel", 1)
         top = self.search_params.get("top", None)
@@ -72,13 +72,38 @@ def search_all(
         self.setup_search()
 
         search_one = functools.partial(self.__class__._search_one, top=top)
-        used_queries = queries
 
+        # Convert queries to a list for potential reuse
+        queries_list = list(queries)
 
+        # Handle MAX_QUERIES environment variable
         if MAX_QUERIES > 0:
-            used_queries = itertools.islice(queries, MAX_QUERIES)
+            queries_list = queries_list[:MAX_QUERIES]
             print(f"Limiting queries to [0:{MAX_QUERIES-1}]")
 
+        # Handle num_queries parameter
+        if num_queries > 0:
+            # If we need more queries than available, cycle through the list
+            if num_queries > len(queries_list) and len(queries_list) > 0:
+                print(f"Requested {num_queries} queries but only {len(queries_list)} are available.")
+                print(f"Extending queries by cycling through the available ones.")
+                # Calculate how many complete cycles and remaining items we need
+                complete_cycles = num_queries // len(queries_list)
+                remaining = num_queries % len(queries_list)
+
+                # Create the extended list
+                extended_queries = []
+                for _ in range(complete_cycles):
+                    extended_queries.extend(queries_list)
+                extended_queries.extend(queries_list[:remaining])
+
+                used_queries = extended_queries
+            else:
+                used_queries = queries_list[:num_queries]
+                print(f"Using {num_queries} queries")
+        else:
+            used_queries = queries_list
+
         if parallel == 1:
             start = time.perf_counter()
             precisions, latencies = list(

run.py

@@ -26,6 +26,7 @@ def run(
     timeout: float = 86400.0,
     upload_start_idx: int = 0,
     upload_end_idx: int = -1,
+    queries: int = typer.Option(-1, help="Number of queries to run. If the available queries are fewer, they will be reused."),
 ):
     """
     Example:
@@ -68,6 +69,7 @@ def run(
                         parallels,
                         upload_start_idx,
                         upload_end_idx,
+                        queries,
                     )
                 client.delete_client()