Multiclient tool

algos.yaml

@@ -240,14 +240,23 @@ float:
           # This run group produces 3 algorithm instances -- Annoy("angular",
           # 100), Annoy("angular", 200), and Annoy("angular", 400) -- each of
           # which will be used to run 12 different queries.
-    milvus:
+    milvus-hnsw:
       docker-tag: ann-benchmarks-milvus
       module: ann_benchmarks.algorithms.milvus
       constructor: Milvus
-      base-args: ["@metric"]
+      base-args: ["@metric", "@connection"]
+      run-groups:
+        milvus:
+          args: [['HNSW'], [{"M": 4,  "efConstruction": 500}, {"M": 32,  "efConstruction": 500}]]
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]  # ef 
+    milvus-ivf:
+      docker-tag: ann-benchmarks-milvus
+      module: ann_benchmarks.algorithms.milvus
+      constructor: Milvus
+      base-args: ["@metric", "@connection"]
       run-groups:
         milvus:
-          args: [['IVF_FLAT', 'IVF_SQ8'], [100, 300, 1000, 3000, 10000, 30000]]  # nlist
+          args: [['IVF_FLAT', 'IVF_SQ8'], [{"nlist": 100}, {"nlist": 300}, {"nlist": 1000}, {"nlist": 3000}, {"nlist": 10000}, {"nlist": 30000}]]  # nlist
           query-args: [[1, 3, 10, 30, 100, 300]]  # nprobe (should be <= nlist)
     nearpy:
       disabled: true

ann_benchmarks/algorithms/definitions.py

@@ -97,7 +97,7 @@ def get_unique_algorithms(definition_file):
 
 
 def get_definitions(definition_file, dimension, point_type="float",
-                    distance_metric="euclidean", count=10, conn_params=dict()):
+                    distance_metric="euclidean", count=10, conn_params={'host': None, 'port': None, 'auth': None, 'user': None, 'cluster': False}):
     definitions = _get_definitions(definition_file)
 
     algorithm_definitions = {}

ann_benchmarks/algorithms/milvus.py

@@ -1,45 +1,71 @@
 from __future__ import absolute_import
-import milvus
+from sqlite3 import paramstyle
+from pymilvus import (
+    connections,
+    utility,
+    FieldSchema,
+    CollectionSchema,
+    DataType,
+    IndexType,
+    Collection,
+)
 import numpy
 import sklearn.preprocessing
 from ann_benchmarks.algorithms.base import BaseANN
 
 
 class Milvus(BaseANN):
-    def __init__(self, metric, index_type, nlist):
-        self._nlist = nlist
+    def __init__(self, metric, conn_params, index_type, method_params):
+        self._host = conn_params['host']
+        self._port = conn_params['port'] # 19530
+        # connections.connect(host=conn_params['host'], port=conn_params['port'])
+        # fields = [
+        #     FieldSchema(name="pk", dtype=DataType.INT64, is_primary=True, auto_id=False),
+        #     FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=100)
+        # ]
+        # schema = CollectionSchema(fields)
+        # self._milvus = Collection('milvus', schema)
+        self._index_type = index_type
+        self._method_params = method_params
         self._nprobe = None
         self._metric = metric
-        self._milvus = milvus.Milvus()
-        self._milvus.connect(host='localhost', port='19530')
-        self._table_name = 'test01'
-        self._index_type = index_type
 
     def fit(self, X):
         if self._metric == 'angular':
-            X = sklearn.preprocessing.normalize(X, axis=1, norm='l2')
+            X = sklearn.preprocessing.normalize(X, axis=1)
+
+        # TODO: if we can set the dim later, mabe return this to the init func
+        connections.connect(host=self._host, port=self._port)
+        fields = [
+            FieldSchema(name="pk", dtype=DataType.INT64, is_primary=True, auto_id=False),
+            FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=len(X[0]))
+        ]
+        schema = CollectionSchema(fields)
+        self._milvus = Collection('milvus', schema)
 
-        self._milvus.create_table({'table_name': self._table_name, 'dimension': X.shape[1]})
-        vector_ids = [id for id in range(len(X))]
-        self._milvus.insert(table_name=self._table_name, records=X.tolist(), ids=vector_ids)
-        index_type = getattr(milvus.IndexType, self._index_type)  # a bit hacky but works
-        self._milvus.create_index(self._table_name, {'index_type': index_type, 'nlist': self._nlist})
+        self._milvus.insert([[id for id in range(len(X))], X.tolist()])
+        self._milvus.create_index('vector', {'index_type': self._index_type, 'metric_type':'L2', 'params':self._method_params})
+        self._milvus.load()
 
-    def set_query_arguments(self, nprobe):
-        if nprobe > self._nlist:
-            print('warning! nprobe > nlist')
-            nprobe = self._nlist
-        self._nprobe = nprobe
+    def set_query_arguments(self, param):
+        self._query_params = dict()
+        if 'IVF_' in self._index_type:
+            if param > self._method_params['nlist']:
+                print('warning! nprobe > nlist')
+                param = self._method_params['nlist']
+            self._query_params['nprobe'] = param
+        if 'HNSW' in self._index_type:
+            self._query_params['ef'] = param
 
     def query(self, v, n):
         if self._metric == 'angular':
             v /= numpy.linalg.norm(v)
         v = v.tolist()
-        status, results = self._milvus.search(table_name=self._table_name, query_records=[v], top_k=n, nprobe=self._nprobe)
+        results = self._milvus.search([v], 'vector', {'metric_type':'L2', 'params':self._query_params}, limit=n)
         if not results:
             return []  # Seems to happen occasionally, not sure why
         result_ids = [result.id for result in results[0]]
         return result_ids
 
     def __str__(self):
-        return 'Milvus(index_type=%s, nlist=%d, nprobe=%d)' % (self._index_type, self._nlist, self._nprobe)
+        return 'Milvus(index_type=%s, method_params=%s, query_params=%s)' % (self._index_type, str(self._method_params), str(self._nprobe))

ann_benchmarks/main.py

@@ -139,12 +139,12 @@ def main():
         '--host',
         metavar='NAME',
         help='host name or IP',
-        default="localhost")
+        default=None)
     parser.add_argument(
         '--port',
         type=positive_int,
         help='the port "host" is listening on',
-        default=6379)
+        default=None)
     parser.add_argument(
         '--auth', '-a',
         metavar='PASSWORD',

ann_benchmarks/results.py

@@ -16,13 +16,13 @@ def get_result_filename(dataset=None, count=None, definition=None,
         d.append(str(count))
     if definition:
         d.append(definition.algorithm + ('-batch' if batch_mode else ''))
-        if 'redisearch' in definition.algorithm:
-            prefix = re.sub(r'\W+', '_', json.dumps(query_arguments, sort_keys=True)).strip('_')
-            d.append(prefix + f'_client_{id}.hdf5')
-        else:
-            data = definition.arguments + query_arguments
-            d.append(re.sub(r'\W+', '_', json.dumps(data, sort_keys=True))
-                     .strip('_') + ".hdf5")
+        data = definition.arguments + query_arguments
+        for i in range(len(data)):
+            if isinstance(data[i], dict):
+                data[i] = {k:data[i][k] for k in data[i] if data[i][k] is not None}
+        data.append('client')
+        data.append(id)
+        d.append(re.sub(r'\W+', '_', json.dumps(data, sort_keys=True)).strip('_') + ".hdf5")
     return os.path.join(*d)
 
 

multirun.py

@@ -24,12 +24,12 @@
     parser.add_argument(
         '--host',
         help='host name or IP',
-        default='localhost')
+        default=None)
     parser.add_argument(
         '--port',
         type=positive_int,
         help='the port "host" is listening on',
-        default=6379)
+        default=None)
     parser.add_argument(
         '--auth', '-a',
         metavar='PASS',
@@ -60,20 +60,28 @@
         '--algorithm',
         metavar='ALGO',
         help='run redisearch with this algorithm',
-        default="hnsw")
+        default="redisearch-hnsw")
     parser.add_argument(
         '--cluster',
         action='store_true',
         help='working with a cluster')
 
     args = parser.parse_args()
+    isredis = True if 'redisearch' in args.algorithm else False
 
-    redis = RedisCluster if args.cluster else Redis
-    redis = redis(host=args.host, port=args.port, password=args.auth, username=args.user)
+    if isredis:
+        redis = RedisCluster if args.cluster else Redis
+        redis = redis(host=args.host, port=args.port, password=args.auth, username=args.user)
+    elif 'milvus' in args.algorithm:
+        if args.build_clients > 1:
+            print('milvus does not allow multi client build. running with one builder')
+            args.build_clients = 1
 
-    base = 'python run.py --local --algorithm redisearch-' + args.algorithm.lower() + ' -k ' + str(args.count) + \
-           ' --dataset ' + args.dataset + ' --host ' + str(args.host) + ' --port ' + str(args.port)
+    base = 'python run.py --local --algorithm ' + args.algorithm + ' -k ' + str(args.count) + \
+           ' --dataset ' + args.dataset
 
+    if args.host:   base += ' --host ' + str(args.host)
+    if args.port:   base += ' --port ' + str(args.port)
     if args.user:   base += ' --user ' + str(args.user)
     if args.auth:   base += ' --auth ' + str(args.auth)
     if args.force:  base += ' --force'
@@ -92,15 +100,17 @@
         print(f'total build time: {total_time}\n\n')
 
         fn = get_result_filename(args.dataset, args.count)
+        fn = path.join(fn, args.algorithm)
         if not path.isdir(fn):
             makedirs(fn)
         fn = path.join(fn, 'build_stats.hdf5')
         f = h5py.File(fn, 'w')
         f.attrs["build_time"] = total_time
-        if args.cluster:
-            f.attrs["index_size"] = -1 # TODO: get total size from all the shards
-        else:
-            f.attrs["index_size"] = redis.ft('ann_benchmark').info()['vector_index_sz_mb']*0x100000
+        if isredis:
+            if args.cluster:
+                f.attrs["index_size"] = -1 # TODO: get total size from all the shards
+            else:
+                f.attrs["index_size"] = redis.ft('ann_benchmark').info()['vector_index_sz_mb']*0x100000
         f.close()
 
     if args.test_clients > 0: