Restore performance optimizations from PR #16 (85a6bc7)

Author: fcostaoliveira

engine/base_client/client.py

@@ -40,8 +40,9 @@ def save_search_results(
     ):
         now = datetime.now()
         timestamp = now.strftime("%Y-%m-%d-%H-%M-%S")
+        pid = os.getpid()  # Get the current process ID
         experiments_file = (
-            f"{self.name}-{dataset_name}-search-{search_id}-{timestamp}.json"
+            f"{self.name}-{dataset_name}-search-{search_id}-{pid}-{timestamp}.json"
         )
         result_path = RESULTS_DIR / experiments_file
         with open(result_path, "w") as out:
@@ -99,7 +100,8 @@ def run_experiment(
         reader = dataset.get_reader(execution_params.get("normalize", False))
 
         if skip_if_exists:
-            glob_pattern = f"{self.name}-{dataset.config.name}-search-*-*.json"
+            pid = os.getpid()  # Get the current process ID
+            glob_pattern = f"{self.name}-{dataset.config.name}-search-*-{pid}-*.json"
             existing_results = list(RESULTS_DIR.glob(glob_pattern))
             if len(existing_results) == len(self.searchers):
                 print(
@@ -137,8 +139,9 @@ def run_experiment(
             print("Experiment stage: Search")
             for search_id, searcher in enumerate(self.searchers):
                 if skip_if_exists:
+                    pid = os.getpid()  # Get the current process ID
                     glob_pattern = (
-                        f"{self.name}-{dataset.config.name}-search-{search_id}-*.json"
+                        f"{self.name}-{dataset.config.name}-search-{search_id}-{pid}-*.json"
                     )
                     existing_results = list(RESULTS_DIR.glob(glob_pattern))
                     print("Pattern", glob_pattern, "Results:", existing_results)

engine/base_client/search.py

@@ -1,6 +1,6 @@
 import functools
 import time
-from multiprocessing import get_context
+from multiprocessing import Process, Queue
 from typing import Iterable, List, Optional, Tuple
 from itertools import islice
 
@@ -106,40 +106,54 @@ def search_all(
             used_queries = queries_list
 
         if parallel == 1:
+            # Single-threaded execution
             start = time.perf_counter()
-            precisions, latencies = list(
-                zip(*[search_one(query) for query in tqdm.tqdm(used_queries)])
-            )
+            results = [search_one(query) for query in tqdm.tqdm(used_queries)]
+            total_time = time.perf_counter() - start
         else:
-            ctx = get_context(self.get_mp_start_method())
+            # Dynamically calculate chunk size
+            chunk_size = max(1, len(used_queries) // parallel)
+            query_chunks = list(chunked_iterable(used_queries, chunk_size))
 
-            def process_initializer():
-                """Initialize each process before starting the search."""
+            # Function to be executed by each worker process
+            def worker_function(chunk, result_queue):
                 self.__class__.init_client(
                     self.host,
                     distance,
                     self.connection_params,
                     self.search_params,
                 )
                 self.setup_search()
+                results = process_chunk(chunk, search_one)
+                result_queue.put(results)
 
-            # Dynamically chunk the generator
-            query_chunks = list(chunked_iterable(used_queries, max(1, len(used_queries) // parallel)))
-
-            with ctx.Pool(
-                processes=parallel,
-                initializer=process_initializer,
-            ) as pool:
-                if parallel > 10:
-                    time.sleep(15)  # Wait for all processes to start
-                start = time.perf_counter()
-                results = pool.starmap(
-                    process_chunk,
-                    [(chunk, search_one) for chunk in query_chunks],
-                )
-                precisions, latencies = zip(*[result for chunk in results for result in chunk])
+            # Create a queue to collect results
+            result_queue = Queue()
+
+            # Create and start worker processes
+            processes = []
+            for chunk in query_chunks:
+                process = Process(target=worker_function, args=(chunk, result_queue))
+                processes.append(process)
+                process.start()
+
+            # Start measuring time for the critical work
+            start = time.perf_counter()
 
-        total_time = time.perf_counter() - start
+            # Collect results from all worker processes
+            results = []
+            for _ in processes:
+                results.extend(result_queue.get())
+
+            # Wait for all worker processes to finish
+            for process in processes:
+                process.join()
+
+            # Stop measuring time for the critical work
+            total_time = time.perf_counter() - start
+
+        # Extract precisions and latencies (outside the timed section)
+        precisions, latencies = zip(*results)
 
         self.__class__.delete_client()
 
@@ -179,3 +193,8 @@ def chunked_iterable(iterable, size):
 def process_chunk(chunk, search_one):
     """Process a chunk of queries using the search_one function."""
     return [search_one(query) for query in chunk]
+
+
+def process_chunk_wrapper(chunk, search_one):
+    """Wrapper to process a chunk of queries."""
+    return process_chunk(chunk, search_one)