Elastic client update iteration 2 [WIP]

algos.yaml

@@ -672,14 +672,6 @@ float:
             - {"n_neighbors": 60, "diversify_prob": 0.0,
                "pruning_degree_multiplier":[2.0, 3.0], "leaf_size": 48}
           query-args: [[0.0, 0.04, 0.08, 0.12, 0.16, 0.20, 0.24, 0.28, 0.32, 0.36]]
-    elasticsearch:
-      docker-tag: ann-benchmarks-elasticsearch
-      module: ann_benchmarks.algorithms.elasticsearch
-      constructor: ElasticsearchScriptScoreQuery
-      base-args: [ "@metric", "@dimension" ]
-      run-groups:
-        empty:
-          args: []
     elastiknn-l2lsh:
       docker-tag: ann-benchmarks-elastiknn
       module: ann_benchmarks.algorithms.elastiknn
@@ -974,10 +966,44 @@ float:
       docker-tag: ann-benchmarks-elasticsearch
       module: ann_benchmarks.algorithms.elasticsearch
       constructor: ElasticsearchScriptScoreQuery
-      base-args: [ "@metric", "@dimension" ]
+      base-args: [ "@metric", "@dimension", "@connection" ]
       run-groups:
-        empty:
-          args: []
+        M-4:
+          arg-groups:
+            - {"m": 4,  "ef_construction": 500, "type": "hnsw"}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-8:
+          arg-groups:
+            - {"m": 8,  "ef_construction": 500, "type": "hnsw"}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-12:
+          arg-groups:
+            - {"m": 12,  "ef_construction": 500, "type": "hnsw"}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-16:
+          arg-groups:
+            - {"m": 16,  "ef_construction": 500, "type": "hnsw"}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-24:
+          arg-groups:
+            - {"m": 24,  "ef_construction": 500, "type": "hnsw"}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-36:
+          arg-groups:
+            - {"m": 36,  "ef_construction": 500, "type": "hnsw"}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-48:
+          arg-groups:
+            - {"m": 48,  "ef_construction": 500, "type": "hnsw"}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-64:
+          arg-groups:
+            - {"m": 64,  "ef_construction": 500, "type": "hnsw"}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-96:
+          arg-groups:
+            - {"m": 96,  "ef_construction": 500, "type": "hnsw"}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
     opensearchknn:
       docker-tag: ann-benchmarks-opensearchknn
       module: ann_benchmarks.algorithms.opensearchknn

ann_benchmarks/algorithms/elasticsearch.py

@@ -3,36 +3,38 @@
 Note that this requires X-Pack, which is not included in the OSS version of Elasticsearch.
 """
 import logging
+import os
 from time import sleep
 from urllib.error import URLError
-from urllib.request import Request, urlopen
 
-from elasticsearch import Elasticsearch
+from elasticsearch import Elasticsearch, BadRequestError
 from elasticsearch.helpers import bulk
-
+import urllib3
 from ann_benchmarks.algorithms.base import BaseANN
 
 # Configure the elasticsearch logger.
 # By default, it writes an INFO statement for every request.
 logging.getLogger("elasticsearch").setLevel(logging.WARN)
 
+# Disable InsecureRequestWarning
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
 # Uncomment these lines if you want to see timing for every HTTP request and its duration.
 # logging.basicConfig(level=logging.INFO)
 # logging.getLogger("elasticsearch").setLevel(logging.INFO)
 
-def es_wait():
+def es_wait(es):
     print("Waiting for elasticsearch health endpoint...")
-    req = Request("http://localhost:9200/_cluster/health?wait_for_status=yellow&timeout=1s")
     for i in range(30):
         try:
-            res = urlopen(req)
-            if res.getcode() == 200:
+            res = es.cluster.health(wait_for_status='green', timeout='1s')
+            if not res['timed_out']: # then status is OK
                 print("Elasticsearch is ready")
                 return
         except URLError:
             pass
         sleep(1)
-    raise RuntimeError("Failed to connect to local elasticsearch")
+    raise RuntimeError("Failed to connect to elasticsearch server")
 
 
 class ElasticsearchScriptScoreQuery(BaseANN):
@@ -42,61 +44,75 @@ class ElasticsearchScriptScoreQuery(BaseANN):
     - Dense vector queries: https://www.elastic.co/guide/en/elasticsearch/reference/master/query-dsl-script-score-query.html
     """
 
-    def __init__(self, metric: str, dimension: int):
-        self.name = f"elasticsearch-script-score-query_metric={metric}_dimension={dimension}"
-        self.metric = metric
+    def __init__(self, metric: str, dimension: int, conn_params, method_param):
+        self.name = f"elasticsearch-script-score-query_metric={metric}_dimension={dimension}_params{method_param}"
+        self.metric = {"euclidean": 'l2_norm', "angular": 'cosine'}[metric]
+        self.method_param = method_param
         self.dimension = dimension
-        self.index = f"es-ssq-{metric}-{dimension}"
-        self.es = Elasticsearch(["http://localhost:9200"])
+        self.timeout = 60 * 60
+        h = conn_params['host'] if conn_params['host'] is not None else 'localhost'
+        p = conn_params['port'] if conn_params['port'] is not None else '9200'
+        u = conn_params['user'] if conn_params['user'] is not None else 'elastic'
+        a = conn_params['auth'] if conn_params['auth'] is not None else ''
+        self.index = "ann_benchmark"
+        self.es = Elasticsearch("{}:{}".format(h,p), request_timeout=self.timeout, basic_auth=(u, a), verify_certs=False)
+        self.es.info()
+        self.shards = conn_params['shards']
+        self.replicas = int(os.getenv('ES_REPLICA_COUNT', "0"))
+        self.bulk_size = int(os.getenv('ES_BULK_SIZE', "500"))
         self.batch_res = []
-        if self.metric == "euclidean":
-            self.script = "1 / (1 + l2norm(params.query_vec, \"vec\"))"
-        elif self.metric == "angular":
-            self.script = "1.0 + cosineSimilarity(params.query_vec, \"vec\")"
-        else:
-            raise NotImplementedError(f"Not implemented for metric {self.metric}")
-        es_wait()
+        es_wait(self.es)
 
-    def fit(self, X):
-        body = dict(settings=dict(number_of_shards=1, number_of_replicas=0))
-        mapping = dict(
+    def fit(self, X, offset=0, limit=None):
+        limit = limit if limit else len(X)
+        mappings = dict(
             properties=dict(
                 id=dict(type="keyword", store=True),
-                vec=dict(type="dense_vector", dims=self.dimension)
+                vec=dict(
+                    type="dense_vector",
+                    dims=self.dimension,
+                    similarity=self.metric,
+                    index=True,
+                    index_options=self.method_param
+                )
             )
         )
-        self.es.indices.create(self.index, body=body)
-        self.es.indices.put_mapping(mapping, self.index)
-
-        def gen():
-            for i, vec in enumerate(X):
-                yield { "_op_type": "index", "_index": self.index, "vec": vec.tolist(), 'id': str(i + 1) }
-
-        (_, errors) = bulk(self.es, gen(), chunk_size=500, max_retries=9)
-        assert len(errors) == 0, errors
-
-        self.es.indices.refresh(self.index)
-        self.es.indices.forcemerge(self.index, max_num_segments=1)
+        try:
+            self.es.indices.create(index=self.index, mappings=mappings, settings=dict(number_of_shards=self.shards, number_of_replicas=self.replicas))
+        except BadRequestError as e:
+            if 'resource_already_exists_exception' not in e.message: raise e
+        bulk_size = self.bulk_size
+        print(f'Using a bulk size of {bulk_size} vectors')
+        for bulk_array in [X[i: i+bulk_size] for i in range(0, len(X), bulk_size)]:
+            bulk_array_len = len(bulk_array)
+            print(f'inserting vectors {offset} to {offset+bulk_array_len}')
+            def gen():
+                for i, vec in enumerate(bulk_array):
+                    yield { "_op_type": "index", "_index": self.index, "vec": vec.tolist(), 'id': str(offset+i) }
+            (_, errors) = bulk(self.es, gen(), chunk_size=bulk_size, max_retries=9, refresh="wait_for")
+            assert len(errors) == 0, errors
+            offset += bulk_array_len
+
+        print('refreshing elastic index...')
+        self.es.indices.refresh(index=self.index)
+        print('finished refreshing elastic index...')
+
+    def set_query_arguments(self, ef):
+        self.ef = ef
 
     def query(self, q, n):
-        body = dict(
-            query=dict(
-                script_score=dict(
-                    query=dict(match_all=dict()),
-                    script=dict(
-                        source=self.script,
-                        params=dict(query_vec=q.tolist())
-                    )
-                )
-            )
-        )
-        res = self.es.search(index=self.index, body=body, size=n, _source=False, docvalue_fields=['id'],
-                             stored_fields="_none_", filter_path=["hits.hits.fields.id"])
-        return [int(h['fields']['id'][0]) - 1 for h in res['hits']['hits']]
+        knn = dict(field='vec', query_vector=q.tolist(), k=n, num_candidates=self.ef)
+        res = self.es.knn_search(index=self.index, knn=knn, source=False, docvalue_fields=['id'],
+                                 stored_fields="_none_", filter_path=["hits.hits.fields.id"])
+        return [int(h['fields']['id'][0]) for h in res['hits']['hits']]
 
     def batch_query(self, X, n):
         self.batch_res = [self.query(q, n) for q in X]
 
     def get_batch_results(self):
         return self.batch_res
 
+    def freeIndex(self):
+        print("Deleting elastic index named {}".format(self.index))
+        self.es.indices.delete(index=self.index)
+        print("Finished deleting elastic index named {}".format(self.index))

multirun.py

@@ -11,13 +11,6 @@
 from ann_benchmarks.results import get_result_filename
 from ann_benchmarks.algorithms.definitions import get_run_groups
 
-from redis import Redis
-from redis.cluster import RedisCluster
-
-from pymilvus import utility, connections
-
-import pinecone
-
 
 def aggregate_outputs(files, clients):
     different_attrs = set([f.split('client')[0] for f in files])
@@ -155,15 +148,33 @@ def aggregate_outputs(files, clients):
     print("Changing the workdir to {}".format(workdir))
     os.chdir(workdir)
 
-    isredis = True if 'redisearch' in args.algorithm else False
-    ismilvus = True if 'milvus' in args.algorithm else False
-    ispinecone = True if 'pinecone' in args.algorithm else False
+    # All supported algorithms that need spacial stuff
+    isredis = ismilvus = ispinecone = iselastic = False
+
+    if 'redisearch' in args.algorithm:
+        from redis import Redis
+        from redis.cluster import RedisCluster
+        isredis = True
+
+    elif 'milvus' in args.algorithm:
+        from pymilvus import utility, connections
+        ismilvus = True
+
+    elif 'pinecone' in args.algorithm:
+        import pinecone
+        ispinecone = True
+
+    elif 'elasticsearch' in args.algorithm:
+        from elasticsearch import Elasticsearch
+        from elastic_transport.client_utils import DEFAULT
+        iselastic = True
 
     if args.host is None:
         args.host = 'localhost'
     if args.port is None:
         if isredis: args.port = '6379'
         elif ismilvus: args.port = '19530'
+        elif iselastic: args.port = '9200'
 
     if isredis:
         redis = RedisCluster if args.cluster else Redis
@@ -172,6 +183,14 @@ def aggregate_outputs(files, clients):
         connections.connect(host=args.host, port=args.port)
     elif ispinecone:
         pinecone.init(api_key=args.auth)
+    elif iselastic:
+        args.user = args.user if args.user is not None else 'elastic'
+        args.auth = args.auth if args.auth is not None else os.environ.get('ELASTIC_PASSWORD', '')
+        try:
+            es = Elasticsearch([f'http://{args.host}:{args.port}'],  request_timeout=3600, basic_auth=(args.user, args.auth))
+            es.info()
+        except Exception:
+            es = Elasticsearch([f'https://{args.host}:{args.port}'], request_timeout=3600, basic_auth=(args.user, args.auth), ca_certs=os.environ.get('ELASTIC_CA', DEFAULT))
 
     if args.run_group is not None:
         run_groups = [args.run_group]
@@ -213,20 +232,23 @@ def on_created_or_modified(event):
         observer.start()
 
     for run_group in run_groups:
-        if isredis:
-            redis.flushall()
-        elif ismilvus:
-            if utility.has_collection('milvus'):
-                utility.drop_collection('milvus')
-        elif ispinecone:
-            for idx in pinecone.list_indexes():
-                pinecone.delete_index(idx)
-
         results_dict = {}
         curr_base_build = base_build + ' --run-group ' + run_group
         curr_base_test = base_test + ' --run-group ' + run_group
 
         if int(args.build_clients) > 0:
+            if isredis:
+                redis.flushall()
+            elif ismilvus:
+                if utility.has_collection('milvus'):
+                    utility.drop_collection('milvus')
+            elif ispinecone:
+                for idx in pinecone.list_indexes():
+                    pinecone.delete_index(idx)
+            elif iselastic:
+                for idx in es.indices.stats()['indices']:
+                    es.indices.delete(index=idx)
+
             clients = [Process(target=os.system, args=(curr_base_build + ' --client-id ' + str(i),)) for i in
                        range(1, int(args.build_clients) + 1)]
 
@@ -245,6 +267,8 @@ def on_created_or_modified(event):
                 if not args.cluster:  # TODO: get total size from all the shards
                     index_size = float(redis.ft('ann_benchmark').info()['vector_index_sz_mb']) * 1024
                 f.attrs["index_size"] = index_size
+            elif iselastic:
+                f.attrs["index_size"] = es.indices.stats(index='ann_benchmark')['indices']['ann_benchmark']['total']['store']['size_in_bytes']
             f.close()
             results_dict["build"] = {"total_clients": args.build_clients, "build_time": total_time,
                                      "vector_index_sz_mb": index_size}

requirements.txt

@@ -7,7 +7,8 @@ pyyaml==5.4
 psutil==5.6.6
 scipy==1.3.3
 scikit-learn==0.22.2
-jinja2==2.10
+jinja2==3.1.2
 pymilvus==2.0.2
 pinecone-client==2.0.11
 redis==4.3.2
+elasticsearch==8.3.1