redisbench-admin 0.11.63__py3-none-any.whl → 0.11.65__py3-none-any.whl

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/subprocess.py

@@ -0,0 +1,246 @@
+from __future__ import absolute_import
+from os.path import basename
+import shlex
+from types import MethodType
+import psutil
+import subprocess
+from ann_benchmarks.data import \
+    bit_unparse_entry, int_unparse_entry, float_unparse_entry
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class SubprocessStoppedError(Exception):
+    def __init__(self, code):
+        super(Exception, self).__init__(code)
+        self.code = code
+
+
+class Subprocess(BaseANN):
+    def _raw_line(self):
+        return shlex.split(
+            self._get_program_handle().stdout.readline().strip())
+
+    def _line(self):
+        line = self._raw_line()
+#       print("<- %s" % (" ".join(line)))
+        while len(line) < 1 or line[0] != "epbprtv0":
+            line = self._raw_line()
+        return line[1:]
+
+    @staticmethod
+    def _quote(token):
+        return "'" + str(token).replace("'", "'\\'") + "'"
+
+    def _write(self, string):
+        #       print("-> %s" % string)
+        self._get_program_handle().stdin.write(string + "\n")
+
+    # Called immediately before transitioning from query mode to training mode
+    def _configuration_hook(self):
+        pass
+
+    def _get_program_handle(self):
+        if self._program:
+            self._program.poll()
+            if self._program.returncode:
+                raise SubprocessStoppedError(self._program.returncode)
+        else:
+            self._program = subprocess.Popen(
+                self._args,
+                bufsize=1,  # line buffering
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                universal_newlines=True)
+
+            for key, value in iter(self._params.items()):
+                self._write("%s %s" %
+                            (Subprocess._quote(key), Subprocess._quote(value)))
+                assert self._line()[0] == "ok", """\
+assigning value '%s' to option '%s' failed""" % (value, key)
+            self._configuration_hook()
+
+            self._write("")
+            assert self._line()[0] == "ok", """\
+transitioning to training mode failed"""
+        return self._program
+
+    def __init__(self, args, encoder, params):
+        self.name = "Subprocess(program = %s, %s)" % \
+            (basename(args[0]), str(params))
+        self._program = None
+        self._args = args
+        self._encoder = encoder
+        self._params = params
+
+    def get_memory_usage(self):
+        if not self._program:
+            self._get_program_handle()
+        return psutil.Process(pid=self._program.pid).memory_info().rss / 1024
+
+    def fit(self, X):
+        for entry in X:
+            d = Subprocess._quote(self._encoder(entry))
+            self._write(d)
+            assert self._line()[0] == "ok", """\
+encoded training point '%s' was rejected""" % d
+        self._write("")
+        assert self._line()[0] == "ok", """\
+transitioning to query mode failed"""
+
+    def query(self, v, n):
+        d = Subprocess._quote(self._encoder(v))
+        self._write("%s %d" % (d, n))
+        return self._handle_query_response()
+
+    def _handle_query_response(self):
+        status = self._line()
+        if status[0] == "ok":
+            count = int(status[1])
+            return self._collect_query_response_lines(count)
+        else:
+            assert status[0] == "fail", """\
+query neither succeeded nor failed"""
+            return []
+
+    def _collect_query_response_lines(self, count):
+        results = []
+        i = 0
+        while i < count:
+            line = self._line()
+            results.append(int(line[0]))
+            i += 1
+        return results
+
+    def done(self):
+        if self._program:
+            self._program.poll()
+            if not self._program.returncode:
+                self._program.terminate()
+
+
+class PreparedSubprocess(Subprocess):
+    def __init__(self, args, encoder, params):
+        super(PreparedSubprocess, self).__init__(args, encoder, params)
+        self._result_count = None
+
+    def _configuration_hook(self):
+        self._write("frontend prepared-queries 1")
+        assert self._line()[0] == "ok", """\
+enabling prepared queries mode failed"""
+
+    def query(self, v, n):
+        self.prepare_query(v, n)
+        self.run_prepared_query()
+        return self.get_prepared_query_results()
+
+    def prepare_query(self, v, n):
+        d = Subprocess._quote(self._encoder(v))
+        self._write("%s %d" % (d, n))
+        assert self._line()[0] == "ok", """\
+preparing the query '%s' failed""" % d
+
+    def run_prepared_query(self):
+        self._write("query")
+        status = self._line()
+        if status[0] == "ok":
+            self._result_count = int(status[1])
+        else:
+            assert status[0] == "fail", """\
+query neither succeeded nor failed"""
+            self._result_count = 0
+
+    def get_prepared_query_results(self):
+        if self._result_count:
+            try:
+                return self._collect_query_response_lines(self._result_count)
+            finally:
+                self._result_count = 0
+        else:
+            return []
+
+
+class BatchSubprocess(Subprocess):
+    def __init__(self, args, encoder, params):
+        super(BatchSubprocess, self).__init__(args, encoder, params)
+        self._qp_count = None
+
+    def _configuration_hook(self):
+        self._write("frontend batch-queries 1")
+        assert self._line()[0] == "ok", """\
+enabling batch queries mode failed"""
+
+    def query(self, v, n):
+        self.prepare_batch_query([v], n)
+        self.run_batch_query()
+        return self.get_batch_results()[0]
+
+    def prepare_batch_query(self, X, n):
+        d = " ".join(map(lambda p: Subprocess._quote(self._encoder(p)), X))
+        self._qp_count = len(X)
+        self._write("%s %d" % (d, n))
+        assert self._line()[0] == "ok", """\
+preparing the batch query '%s' failed""" % d
+
+    def run_batch_query(self):
+        self._write("query")
+        status = self._line()
+        assert status[0] == "ok", """\
+batch query failed completely"""
+
+    def get_batch_results(self):
+        results = []
+        i = 0
+        while i < self._qp_count:
+            #           print("%d/%d" % (i, self._qp_count))
+            status = self._line()
+            if status[0] == "ok":
+                rc = int(status[1])
+                results.append(self._collect_query_response_lines(rc))
+            else:
+                results.append([])
+            i += 1
+        return results
+
+
+def BitSubprocess(args, params):
+    return Subprocess(args, bit_unparse_entry, params)
+
+
+def BitSubprocessPrepared(args, params):
+    return PreparedSubprocess(args, bit_unparse_entry, params)
+
+
+def FloatSubprocess(args, params):
+    return Subprocess(args, float_unparse_entry, params)
+
+
+def FloatSubprocessPrepared(args, params):
+    return PreparedSubprocess(args, float_unparse_entry, params)
+
+
+def FloatSubprocessBatch(args, params):
+    return BatchSubprocess(args, float_unparse_entry, params)
+
+
+def IntSubprocess(args, params):
+    return Subprocess(args, int_unparse_entry, params)
+
+
+def QueryParamWrapper(constructor, args, params):
+    r = constructor(args, params)
+
+    def _do(self, original=r._configuration_hook):
+        original()
+        self._write("frontend query-parameters 1")
+        assert self._line()[0] == "ok", """\
+enabling query parameter support failed"""
+    r._configuration_hook = MethodType(_do, r)
+
+    def _sqa(self, *args):
+        self._write("query-params %s set" %
+                    (" ".join(map(Subprocess._quote, args))))
+        assert self._line()[0] == "ok", """\
+reconfiguring query parameters failed"""
+        print(args)
+    r.set_query_arguments = MethodType(_sqa, r)
+    return r

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/vald.py

@@ -0,0 +1,149 @@
+from __future__ import absolute_import
+
+import atexit
+import subprocess
+import urllib.error
+import urllib.request
+
+import grpc
+import yaml
+from ann_benchmarks.algorithms.base import BaseANN
+
+from vald.v1.vald import insert_pb2_grpc, search_pb2_grpc
+from vald.v1.agent.core import agent_pb2_grpc
+from vald.v1.payload import payload_pb2
+
+
+default_server_config = {
+    'version': 'v0.0.0',
+    'logging': {
+        'logger': 'glg',
+        'level': 'info',
+        'format': 'raw'
+    },
+    'server_config': {
+        'servers': [
+            {
+                'name': 'agent-grpc',
+                'host': '127.0.0.1',
+                'port': 8082,
+                'mode': 'GRPC',
+                'probe_wait_time': '3s',
+                #'grpc': {
+                #    'bidirectional_stream_concurrency': 1
+                #},
+                "network": "unix",
+                "socket_path": "/var/run/vald.sock"
+            }
+        ],
+        'health_check_servers': [
+            {
+                'name': 'readiness',
+                'host': '127.0.0.1',
+                'port': 3001,
+                'mode': '',
+                'probe_wait_time': '3s',
+                'http': {
+                    'shutdown_duration': '5s',
+                    'handler_timeout': '',
+                    'idle_timeout': '',
+                    'read_header_timeout': '',
+                    'read_timeout': '',
+                    'write_timeout': ''
+                }
+            }
+        ],
+        'startup_strategy': ['agent-grpc', 'readiness'],
+        'shutdown_strategy': ['readiness', 'agent-grpc'],
+        'full_shutdown_duration': '600s',
+        'tls': {
+            'enabled': False,
+        }
+    },
+    'ngt': {
+        'enable_in_memory_mode': True,
+        'default_pool_size': 10000,
+        'default_epsilon': 0.01,
+        'default_radius': -1.0,
+        #'vqueue': {
+        #    'insert_buffer_size': 100,
+        #    'insert_buffer_pool_size': 1000,
+        #    'delete_buffer_size': 100,
+        #    'delete_buffer_pool_size': 1000
+        #}
+    }
+}
+
+grpc_opts = [
+    ('grpc.keepalive_time_ms', 1000 * 10),
+    ('grpc.keepalive_timeout_ms', 1000 * 10),
+    ('grpc.max_connection_idle_ms', 1000 * 50)
+]
+
+metrics = {'euclidean': 'l2', 'angular': 'cosine'}
+
+
+class Vald(BaseANN):
+    def __init__(self, metric, object_type, params):
+        self._param = default_server_config
+        self._ngt_config = {
+            'distance_type': metrics[metric],
+            'object_type': object_type,
+            'search_edge_size': int(params['searchedge']),
+            'creation_edge_size': int(params['edge']),
+            'bulk_insert_chunk_size': int(params['bulk'])
+        }
+        #self._address = 'localhost:8082'
+        self._address = 'unix:///var/run/vald.sock'
+
+    def fit(self, X):
+        dim = len(X[0])
+        self._ngt_config['dimension'] = dim
+        self._param['ngt'].update(self._ngt_config)
+        with open('config.yaml', 'w') as f:
+            yaml.dump(self._param, f)
+
+        cfg = payload_pb2.Insert.Config(skip_strict_exist_check=True)
+        vectors = [
+            payload_pb2.Insert.Request(
+                vector=payload_pb2.Object.Vector(id=str(i), vector=x.tolist()),
+                config=cfg) for i, x in enumerate(X[:100])]
+
+        p = subprocess.Popen(['/go/bin/ngt', '-f', 'config.yaml'])
+        atexit.register(lambda: p.kill())
+
+        while True:
+            try:
+                with urllib.request.urlopen('http://localhost:3001/readiness') as response:
+                    if response.getcode() == 200:
+                        break
+            except (urllib.error.HTTPError, urllib.error.URLError):
+                pass
+
+        channel = grpc.insecure_channel(self._address, grpc_opts)
+        istub = insert_pb2_grpc.InsertStub(channel)
+        for _ in istub.StreamInsert(iter(vectors)):
+            pass
+
+        astub = agent_pb2_grpc.AgentStub(channel)
+        astub.CreateIndex(
+            payload_pb2.Control.CreateIndexRequest(
+                pool_size=10000))
+
+    def set_query_arguments(self, epsilon):
+        self._epsilon = epsilon - 1.0
+        channel = grpc.insecure_channel(self._address, grpc_opts)
+        self._stub = search_pb2_grpc.SearchStub(channel)
+
+    def query(self, v, n):
+        cfg = payload_pb2.Search.Config(num=n, radius=-1.0, epsilon=self._epsilon, timeout=3000000)
+        response = self._stub.Search(payload_pb2.Search.Request(vector=v.tolist(), config=cfg))
+        return [int(result.id) for result in response.results]
+
+    def __str__(self):
+        return 'Vald(%d, %d, %d, %1.3f)' % (
+            self._ngt_config['creation_edge_size'],
+            self._ngt_config['search_edge_size'],
+            self._ngt_config['bulk_insert_chunk_size'],
+            self._epsilon + 1.0
+        )

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/vecsim-hnsw.py

@@ -0,0 +1,43 @@
+from __future__ import absolute_import
+import os
+from VecSim import *
+import numpy as np
+from ann_benchmarks.constants import INDEX_DIR
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class VecSimHnsw(BaseANN):
+    def __init__(self, metric, method_param):
+        self.metric = {'angular': VecSimMetric_Cosine, 'euclidean': VecSimMetric_L2}[metric]
+        self.method_param = method_param
+        # print(self.method_param,save_index,query_param)
+        self.ef = None
+        self.name = 'VecSim-hnsw (%s)' % (self.method_param)
+
+    def fit(self, X):
+        hnswparams = HNSWParams()
+        hnswparams.M =self.method_param['M']
+        hnswparams.efConstruction = self.method_param['efConstruction']
+        hnswparams.initialCapacity = len(X)
+        hnswparams.dim = len(X[0])
+        hnswparams.type = VecSimType_FLOAT32
+        hnswparams.metric = self.metric
+        hnswparams.multi = False
+
+        self.index = HNSWIndex(hnswparams)
+
+        for i, vector in enumerate(X):
+            self.index.add_vector(vector, i)
+
+    def set_query_arguments(self, ef):
+        self.ef = ef
+        self.index.set_ef(ef)
+
+    def query(self, v, n):
+        return self.index.knn_query(np.expand_dims(v, axis=0), k=n)[0][0]
+
+    def freeIndex(self):
+        del self.index
+
+    def __str__(self):
+        return f"{self.name}, efRuntime: {self.ef}"

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/vespa.py

@@ -0,0 +1,47 @@
+from ann_benchmarks.algorithms.base import BaseANN
+from vespa_ann_benchmark import DistanceMetric, HnswIndexParams, HnswIndex
+import time
+
+# Class using the Vespa implementation of an HNSW index for nearest neighbor
+# search over data points in a high dimensional vector space.
+#
+# To use nearest neighbor search in a Vespa application,
+# see https://docs.vespa.ai/en/approximate-nn-hnsw.html for more details.
+class VespaHnswBase(BaseANN):
+    def __init__(self, enable_normalize, metric, dimension, param):
+        if metric not in ('angular', 'euclidean'):
+            raise NotImplementedError(
+                "VespaHnsw doesn't support metric %s" % metric)
+        self.metric = {'angular': DistanceMetric.Angular, 'euclidean': DistanceMetric.Euclidean}[metric]
+        normalize = False
+        if self.metric == DistanceMetric.Angular and enable_normalize:
+            normalize = True
+            self.metric = DistanceMetric.InnerProduct
+        self.param = param
+        self.neighbors_to_explore_at_insert = param.get("efConstruction", 200)
+        self.max_links_per_node = param.get("M", 8)
+        self.dimension = dimension
+        self.neighbors_to_explore = 200
+        self.name = 'VespaHnsw()'
+        self.index = HnswIndex(dimension, HnswIndexParams(self.max_links_per_node, self.neighbors_to_explore_at_insert, self.metric, False), normalize)
+
+    def fit(self, X):
+        for i, x in enumerate(X):
+            self.index.set_vector(i, x)
+
+    def set_query_arguments(self, ef):
+        print("VespaHnsw: ef = " + str(ef))
+        self.neighbors_to_explore = ef
+
+    def query(self, v, n):
+        return [index for index, _ in self.query_with_distances(v, n)]
+
+    def query_with_distances(self, v, n):
+        return self.index.find_top_k(n, v, self.neighbors_to_explore)
+
+class VespaHnsw(VespaHnswBase):
+    def __init__(self, metric, dimension, param):
+        super().__init__(True, metric, dimension, param)
+
+    def __str__(self):
+        return 'VespaHnsw ({}, ef: {})'.format(self.param, self.neighbors_to_explore)

redisbench_admin/run/ann/pkg/ann_benchmarks/constants.py

@@ -0,0 +1 @@
+INDEX_DIR = 'indices'

redisbench_admin/run/ann/pkg/ann_benchmarks/data.py

@@ -0,0 +1,48 @@
+from __future__ import absolute_import
+import numpy
+
+
+def float_parse_entry(line):
+    return [float(x) for x in line.strip().split()]
+
+
+def float_unparse_entry(entry):
+    return " ".join(map(str, entry))
+
+
+def int_parse_entry(line):
+    return frozenset([int(x) for x in line.strip().split()])
+
+
+def int_unparse_entry(entry):
+    return " ".join(map(str, map(int, entry)))
+
+
+def bit_parse_entry(line):
+    return [bool(int(x)) for x in list(line.strip()
+                                       .replace(" ", "")
+                                       .replace("\t", ""))]
+
+
+def bit_unparse_entry(entry):
+    return " ".join(map(lambda el: "1" if el else "0", entry))
+
+
+type_info = {
+    "float": {
+        "type": numpy.float,
+        "parse_entry": float_parse_entry,
+        "unparse_entry": float_unparse_entry,
+        "finish_entries": numpy.vstack
+    },
+    "bit": {
+        "type": numpy.bool_,
+        "parse_entry": bit_parse_entry,
+        "unparse_entry": bit_unparse_entry
+    },
+    "int": {
+        "type": numpy.object,
+        "parse_entry": int_parse_entry,
+        "unparse_entry": int_unparse_entry,
+    },
+}