experiments with qdrant

Author: generall

benchmark/dataset.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 from dataclasses import dataclass
 from typing import Optional
 import tarfile
@@ -49,7 +50,8 @@ def download(self):
             else:
                 print(f"Moving: {tmp_path} -> {target_path}")
                 (DATASETS_DIR / self.config.path).parent.mkdir(exist_ok=True)
-                os.rename(tmp_path, target_path)
+                shutil.copy2(tmp_path, target_path)
+                os.remove(tmp_path)
 
     def get_reader(self) -> BaseReader:
         reader_class = READER_TYPE[self.config.type]

dataset_reader/ann_h5_reader.py

@@ -12,11 +12,12 @@ def __init__(self, path):
 
     def read_queries(self) -> Iterator[Query]:
         data = h5py.File(self.path)
-        for vector, expected_result in zip(data['test'], data['neighbors']):
+        for vector, expected_result, expected_scores in zip(data['test'], data['neighbors'], data['distances']):
             yield Query(
                 vector=vector.tolist(),
                 meta_conditions=None,
                 expected_result=expected_result.tolist(),
+                expected_scores=expected_scores.tolist()
             )
 
     def read_data(self) -> Iterator[Record]:

dataset_reader/base_reader.py

@@ -14,6 +14,7 @@ class Query:
     vector: List[float]
     meta_conditions: Optional[dict]
     expected_result: Optional[List[int]]
+    expected_scores: Optional[List[float]] = None
 
 
 class BaseReader:

engine/base_client/client.py

@@ -47,15 +47,18 @@ def save_upload_results(self, dataset_name: str, results: dict):
             )
 
     def run_experiment(self, dataset: Dataset):
+        print("Experiment stage: Configure")
         self.configurator.configure(
             distance=dataset.config.distance,
             vector_size=dataset.config.vector_size,
         )
 
         reader = dataset.get_reader()
+        print("Experiment stage: Upload")
         upload_stats = self.uploader.upload(reader.read_data())
         self.save_upload_results(dataset.config.name, upload_stats)
 
+        print("Experiment stage: Search")
         for search_id, searcher in enumerate(self.searchers):
             search_params = {**searcher.search_params}
             search_stats = searcher.search_all(reader.read_queries())

engine/base_client/upload.py

@@ -2,6 +2,8 @@
 from multiprocessing import get_context
 from typing import List, Optional, Iterable
 
+import tqdm
+
 from dataset_reader.base_reader import Record
 from engine.base_client.utils import iter_batches
 
@@ -30,7 +32,7 @@ def upload(
 
         if parallel == 1:
             self.init_client(self.host, self.connection_params, self.upload_params)
-            for ids, vectors, metadata in iter_batches(records, batch_size):
+            for ids, vectors, metadata in iter_batches(tqdm.tqdm(records), batch_size):
                 latencies.append(self._upload_batch(ids, vectors, metadata))
 
         else:
@@ -41,15 +43,22 @@ def upload(
                 initargs=(self.host, self.connection_params, self.upload_params),
             ) as pool:
                 latencies = pool.imap(
-                    self.__class__._upload_batch, iter_batches(records, batch_size)
+                    self.__class__._upload_batch, iter_batches(tqdm.tqdm(records), batch_size)
                 )
 
+        upload_time = time.perf_counter() - start
+
+        print("Upload time: {}".format(upload_time))
+
         post_upload_stats = self.post_upload()
 
+        total_time = time.perf_counter() - start
+
         return {
-            "latencies": latencies,
             "post_upload": post_upload_stats,
-            "total_time": time.perf_counter() - start
+            "upload_time": upload_time,
+            "total_time": total_time,
+            "latencies": latencies,
         }
 
     @classmethod

engine/clients/qdrant/upload.py

@@ -1,7 +1,8 @@
+import time
 from typing import Optional, List
 
 from qdrant_client import QdrantClient
-from qdrant_client.http.models import Batch
+from qdrant_client.http.models import Batch, CollectionStatus
 
 from engine.base_client.upload import BaseUploader
 from engine.clients.qdrant.config import QDRANT_COLLECTION_NAME
@@ -25,5 +26,21 @@ def upload_batch(
     ):
         cls.client.upsert(
             collection_name=QDRANT_COLLECTION_NAME,
-            points=Batch(ids=ids, vectors=vectors, payloads=metadata),
+            points=Batch(ids=ids, vectors=vectors, payloads=[payload or {} for payload in metadata]),
         )
+
+    @classmethod
+    def post_upload(cls):
+        cls.wait_collection_green()
+        return {}
+
+    @classmethod
+    def wait_collection_green(cls):
+        wait_time = 1.0
+        total = 0
+        collection_info = cls.client.get_collection(QDRANT_COLLECTION_NAME)
+        while collection_info.status != CollectionStatus.GREEN:
+            time.sleep(wait_time)
+            total += wait_time
+            collection_info = cls.client.get_collection(QDRANT_COLLECTION_NAME)
+        return total

experiments/configurations/qdrant-single-node.json

@@ -42,21 +42,5 @@
       {"parallel":  4},
       {"parallel":  100}
     ]
-  },
-  {
-    "name": "qdrant-distributed-3-node",
-    "engine": "qdrant-distributed",
-    "connection_params": {},
-    "collection_params": {
-      "hnsw_config": {
-        "m": 64
-      }
-    },
-    "search_params": [
-      {"parallel":  1},
-      {"parallel":  2},
-      {"parallel":  4},
-      {"parallel":  100}
-    ]
   }
 ]