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

redisbench_admin/run/ann/pkg/algosP.yaml

@@ -0,0 +1,67 @@
+float:
+  any:
+    bruteforce:
+      docker-tag: ann-benchmarks-sklearn
+      module: ann_benchmarks.algorithms.bruteforce
+      constructor: BruteForce
+      base-args: ["@metric"]
+      run-groups:
+        empty:
+          args: []
+    bruteforce-blas:
+      docker-tag: ann-benchmarks-sklearn
+      module: ann_benchmarks.algorithms.bruteforce
+      constructor: BruteForceBLAS
+      base-args: ["@metric"]
+      run-groups:
+        empty:
+          args: []
+  angular:
+    pp-bruteforce-lo:
+      module: ann_benchmarks.algorithms.subprocess
+      docker-tag: ann-benchmarks-subprocess
+      constructor: FloatSubprocess
+      base-args: [["protocol/bf-runner"]]
+      run-groups:
+        jf-linear:
+          args: {"point-type": "float", "distance": "angular"}
+    pp-bruteforce-hi:
+      module: ann_benchmarks.algorithms.subprocess
+      docker-tag: ann-benchmarks-subprocess
+      constructor: FloatSubprocessPrepared
+      base-args: [["protocol/bf-runner"]]
+      run-groups:
+        jf-linear:
+          args: {"point-type": "float", "distance": "angular"}
+    pp-bruteforce-blas-lo:
+      module: ann_benchmarks.algorithms.subprocess
+      docker-tag: ann-benchmarks-subprocess
+      constructor: FloatSubprocess
+      base-args: [["protocol/bf-runner"]]
+      run-groups:
+        jf-linear:
+          args: {"point-type": "float", "distance": "angular", "fast": 1}
+    pp-bruteforce-blas-hi:
+      module: ann_benchmarks.algorithms.subprocess
+      docker-tag: ann-benchmarks-subprocess
+      constructor: FloatSubprocessPrepared
+      base-args: [["protocol/bf-runner"]]
+      run-groups:
+        jf-linear:
+          args: {"point-type": "float", "distance": "angular", "fast": 1}
+    pp-bruteforce-batch:
+      module: ann_benchmarks.algorithms.subprocess
+      docker-tag: ann-benchmarks-subprocess
+      constructor: FloatSubprocessBatch
+      base-args: [["protocol/bf-runner"]]
+      run-groups:
+        jf-linear:
+          args: {"point-type": "float", "distance": "angular"}
+    pp-bruteforce-blas-batch:
+      module: ann_benchmarks.algorithms.subprocess
+      docker-tag: ann-benchmarks-subprocess
+      constructor: FloatSubprocessBatch
+      base-args: [["protocol/bf-runner"]]
+      run-groups:
+        jf-linear:
+          args: {"point-type": "float", "distance": "angular", "fast": 1}

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

@@ -0,0 +1,2 @@
+from __future__ import absolute_import
+# from ann_benchmarks.main import *

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

@@ -0,0 +1,26 @@
+from __future__ import absolute_import
+import annoy
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class Annoy(BaseANN):
+    def __init__(self, metric, n_trees):
+        self._n_trees = n_trees
+        self._search_k = None
+        self._metric = metric
+
+    def fit(self, X):
+        self._annoy = annoy.AnnoyIndex(X.shape[1], metric=self._metric)
+        for i, x in enumerate(X):
+            self._annoy.add_item(i, x.tolist())
+        self._annoy.build(self._n_trees)
+
+    def set_query_arguments(self, search_k):
+        self._search_k = search_k
+
+    def query(self, v, n):
+        return self._annoy.get_nns_by_vector(v.tolist(), n, self._search_k)
+
+    def __str__(self):
+        return 'Annoy(n_trees=%d, search_k=%d)' % (self._n_trees,
+                                                   self._search_k)

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

@@ -0,0 +1,22 @@
+from __future__ import absolute_import
+import sklearn.neighbors
+import sklearn.preprocessing
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class BallTree(BaseANN):
+    def __init__(self, metric, leaf_size=20):
+        self._leaf_size = leaf_size
+        self._metric = metric
+        self.name = 'BallTree(leaf_size=%d)' % self._leaf_size
+
+    def fit(self, X):
+        if self._metric == 'angular':
+            X = sklearn.preprocessing.normalize(X, axis=1, norm='l2')
+        self._tree = sklearn.neighbors.BallTree(X, leaf_size=self._leaf_size)
+
+    def query(self, v, n):
+        if self._metric == 'angular':
+            v = sklearn.preprocessing.normalize([v], axis=1, norm='l2')[0]
+        dist, ind = self._tree.query([v], k=n)
+        return ind[0]

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

@@ -0,0 +1,36 @@
+from __future__ import absolute_import
+from multiprocessing.pool import ThreadPool
+import psutil
+
+
+class BaseANN(object):
+    def done(self):
+        pass
+
+    def get_memory_usage(self):
+        """Return the current memory usage of this algorithm instance
+        (in kilobytes), or None if this information is not available."""
+        # return in kB for backwards compatibility
+        return psutil.Process().memory_info().rss / 1024
+
+    def fit(self, X):
+        pass
+
+    def query(self, q, n):
+        return []  # array of candidate indices
+
+    def batch_query(self, X, n):
+        """Provide all queries at once and let algorithm figure out
+           how to handle it. Default implementation uses a ThreadPool
+           to parallelize query processing."""
+        pool = ThreadPool()
+        self.res = pool.map(lambda q: self.query(q, n), X)
+
+    def get_batch_results(self):
+        return self.res
+
+    def get_additional(self):
+        return {}
+
+    def __str__(self):
+        return self.name

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

@@ -0,0 +1,110 @@
+from __future__ import absolute_import
+import numpy
+import sklearn.neighbors
+from ann_benchmarks.distance import metrics as pd
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class BruteForce(BaseANN):
+    def __init__(self, metric):
+        if metric not in ('angular', 'euclidean', 'hamming'):
+            raise NotImplementedError(
+                "BruteForce doesn't support metric %s" % metric)
+        self._metric = metric
+        self.name = 'BruteForce()'
+
+    def fit(self, X):
+        metric = {'angular': 'cosine', 'euclidean': 'l2',
+                  'hamming': 'hamming'}[self._metric]
+        self._nbrs = sklearn.neighbors.NearestNeighbors(
+            algorithm='brute', metric=metric)
+        self._nbrs.fit(X)
+
+    def query(self, v, n):
+        return list(self._nbrs.kneighbors(
+            [v], return_distance=False, n_neighbors=n)[0])
+
+    def query_with_distances(self, v, n):
+        (distances, positions) = self._nbrs.kneighbors(
+            [v], return_distance=True, n_neighbors=n)
+        return zip(list(positions[0]), list(distances[0]))
+
+
+class BruteForceBLAS(BaseANN):
+    """kNN search that uses a linear scan = brute force."""
+
+    def __init__(self, metric, precision=numpy.float32):
+        if metric not in ('angular', 'euclidean', 'hamming', 'jaccard'):
+            raise NotImplementedError(
+                "BruteForceBLAS doesn't support metric %s" % metric)
+        elif metric == 'hamming' and precision != numpy.bool:
+            raise NotImplementedError(
+                "BruteForceBLAS doesn't support precision"
+                " %s with Hamming distances" % precision)
+        self._metric = metric
+        self._precision = precision
+        self.name = 'BruteForceBLAS()'
+
+    def fit(self, X):
+        """Initialize the search index."""
+        if self._metric == 'angular':
+            # precompute (squared) length of each vector
+            lens = (X ** 2).sum(-1)
+            # normalize index vectors to unit length
+            X /= numpy.sqrt(lens)[..., numpy.newaxis]
+            self.index = numpy.ascontiguousarray(X, dtype=self._precision)
+        elif self._metric == 'hamming':
+            # Regarding bitvectors as vectors in l_2 is faster for blas
+            X = X.astype(numpy.float32)
+            # precompute (squared) length of each vector
+            lens = (X ** 2).sum(-1)
+            self.index = numpy.ascontiguousarray(X, dtype=numpy.float32)
+            self.lengths = numpy.ascontiguousarray(lens, dtype=numpy.float32)
+        elif self._metric == 'euclidean':
+            # precompute (squared) length of each vector
+            lens = (X ** 2).sum(-1)
+            self.index = numpy.ascontiguousarray(X, dtype=self._precision)
+            self.lengths = numpy.ascontiguousarray(lens, dtype=self._precision)
+        elif self._metric == 'jaccard':
+            self.index = X
+        else:
+            # shouldn't get past the constructor!
+            assert False, "invalid metric"
+
+    def query(self, v, n):
+        return [index for index, _ in self.query_with_distances(v, n)]
+
+    def query_with_distances(self, v, n):
+        """Find indices of `n` most similar vectors from the index to query
+        vector `v`."""
+
+        if self._metric != 'jaccard':
+            # use same precision for query as for index
+            v = numpy.ascontiguousarray(v, dtype=self.index.dtype)
+
+        # HACK we ignore query length as that's a constant
+        # not affecting the final ordering
+        if self._metric == 'angular':
+            # argmax_a cossim(a, b) = argmax_a dot(a, b) / |a||b| = argmin_a -dot(a, b)  # noqa
+            dists = -numpy.dot(self.index, v)
+        elif self._metric == 'euclidean':
+            # argmin_a (a - b)^2 = argmin_a a^2 - 2ab + b^2 = argmin_a a^2 - 2ab  # noqa
+            dists = self.lengths - 2 * numpy.dot(self.index, v)
+        elif self._metric == 'hamming':
+            # Just compute hamming distance using euclidean distance
+            dists = self.lengths - 2 * numpy.dot(self.index, v)
+        elif self._metric == 'jaccard':
+            dists = [pd[self._metric]['distance'](v, e) for e in self.index]
+        else:
+            # shouldn't get past the constructor!
+            assert False, "invalid metric"
+        # partition-sort by distance, get `n` closest
+        nearest_indices = numpy.argpartition(dists, n)[:n]
+        indices = [idx for idx in nearest_indices if pd[self._metric]
+                   ["distance_valid"](dists[idx])]
+
+        def fix(index):
+            ep = self.index[index]
+            ev = v
+            return (index, pd[self._metric]['distance'](ep, ev))
+        return map(fix, indices)

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

@@ -0,0 +1,17 @@
+from __future__ import absolute_import
+from scipy.spatial import cKDTree
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class CKDTree(BaseANN):
+    def __init__(self, metric, leaf_size=20):
+        self._leaf_size = leaf_size
+        self._metric = metric
+        self.name = 'CKDTree(leaf_size=%d)' % self._leaf_size
+
+    def fit(self, X):
+        self._tree = cKDTree(X, leafsize=self._leaf_size)
+
+    def query(self, v, n):
+        dist, ind = self._tree.query([v], k=n)
+        return ind[0]

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

@@ -0,0 +1,29 @@
+from __future__ import absolute_import
+from datasketch import MinHashLSHForest, MinHash
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class DataSketch(BaseANN):
+    def __init__(self, metric, n_perm, n_rep):
+        if metric not in ('jaccard'):
+            raise NotImplementedError(
+                "Datasketch doesn't support metric %s" % metric)
+        self._n_perm = n_perm
+        self._n_rep = n_rep
+        self._metric = metric
+        self.name = 'Datasketch(n_perm=%d, n_rep=%d)' % (n_perm, n_rep)
+
+    def fit(self, X):
+        self._index = MinHashLSHForest(num_perm=self._n_perm, l=self._n_rep)
+        for i, x in enumerate(X):
+            m = MinHash(num_perm=self._n_perm)
+            for e in x:
+                m.update(str(e).encode('utf8'))
+            self._index.add(str(i), m)
+        self._index.index()
+
+    def query(self, v, n):
+        m = MinHash(num_perm=self._n_perm)
+        for e in v:
+            m.update(str(e).encode('utf8'))
+        return map(int, self._index.query(m, n))

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

@@ -0,0 +1,187 @@
+from __future__ import absolute_import
+from os import sep as pathsep
+import collections
+import importlib
+import os
+import sys
+import traceback
+import yaml
+from enum import Enum
+from itertools import product
+
+
+Definition = collections.namedtuple(
+    'Definition',
+    ['algorithm', 'run_group', 'constructor', 'module', 'docker_tag',
+     'arguments', 'query_argument_groups', 'disabled'])
+
+
+def instantiate_algorithm(definition):
+    print('Trying to instantiate %s.%s(%s)' %
+          (definition.module, definition.constructor, definition.arguments))
+    module = importlib.import_module(definition.module)
+    constructor = getattr(module, definition.constructor)
+    return constructor(*definition.arguments)
+
+
+class InstantiationStatus(Enum):
+    AVAILABLE = 0
+    NO_CONSTRUCTOR = 1
+    NO_MODULE = 2
+
+
+def algorithm_status(definition):
+    try:
+        module = importlib.import_module(definition.module)
+        if hasattr(module, definition.constructor):
+            return InstantiationStatus.AVAILABLE
+        else:
+            return InstantiationStatus.NO_CONSTRUCTOR
+    except ImportError:
+        return InstantiationStatus.NO_MODULE
+
+
+def _generate_combinations(args):
+    if isinstance(args, list):
+        args = [el if isinstance(el, list) else [el] for el in args]
+        return [list(x) for x in product(*args)]
+    elif isinstance(args, dict):
+        flat = []
+        for k, v in args.items():
+            if isinstance(v, list):
+                flat.append([(k, el) for el in v])
+            else:
+                flat.append([(k, v)])
+        return [dict(x) for x in product(*flat)]
+    else:
+        raise TypeError("No args handling exists for %s" % type(args).__name__)
+
+
+def _substitute_variables(arg, vs):
+    if isinstance(arg, dict):
+        return dict([(k, _substitute_variables(v, vs))
+                     for k, v in arg.items()])
+    elif isinstance(arg, list):
+        return [_substitute_variables(a, vs) for a in arg]
+    elif isinstance(arg, str) and arg in vs:
+        return vs[arg]
+    else:
+        return arg
+
+
+def _get_definitions(definition_file):
+    with open(definition_file, "r") as f:
+        return yaml.load(f, yaml.SafeLoader)
+
+
+def list_algorithms(definition_file):
+    definitions = _get_definitions(definition_file)
+
+    print('The following algorithms are supported...')
+    for point in definitions:
+        print('\t... for the point type "%s"...' % point)
+        for metric in definitions[point]:
+            print('\t\t... and the distance metric "%s":' % metric)
+            for algorithm in definitions[point][metric]:
+                print('\t\t\t%s' % algorithm)
+
+
+def get_unique_algorithms(definition_file):
+    definitions = _get_definitions(definition_file)
+    algos = set()
+    for point in definitions:
+        for metric in definitions[point]:
+            for algorithm in definitions[point][metric]:
+                algos.add(algorithm)
+    return list(sorted(algos))
+
+
+def get_run_groups(definition_file, algo = None):
+    definitions = _get_definitions(definition_file)
+    run_groups = set()
+    for point in definitions:
+        for metric in definitions[point]:
+            for algorithm in definitions[point][metric]:
+                if algo == None or algo == algorithm:
+                    for run_group in definitions[point][metric][algorithm]['run-groups'].keys():
+                        run_groups.add(run_group)
+    return list(sorted(run_groups))
+
+
+def get_definitions(definition_file, dimension, point_type="float",
+                    distance_metric="euclidean", count=10, conn_params={'host': None, 'port': None, 'auth': None, 'user': None, 'cluster': False, 'shards': 1}):
+    definitions = _get_definitions(definition_file)
+
+    algorithm_definitions = {}
+    if "any" in definitions[point_type]:
+        algorithm_definitions.update(definitions[point_type]["any"])
+    algorithm_definitions.update(definitions[point_type][distance_metric])
+
+    definitions = []
+    for (name, algo) in algorithm_definitions.items():
+        for k in ['docker-tag', 'module', 'constructor']:
+            if k not in algo:
+                raise Exception(
+                    'algorithm %s does not define a "%s" property' % (name, k))
+
+        base_args = []
+        if "base-args" in algo:
+            base_args = algo["base-args"]
+
+        for run_group_name, run_group in algo["run-groups"].items():
+            if "arg-groups" in run_group:
+                groups = []
+                for arg_group in run_group["arg-groups"]:
+                    if isinstance(arg_group, dict):
+                        # Dictionaries need to be expanded into lists in order
+                        # for the subsequent call to _generate_combinations to
+                        # do the right thing
+                        groups.append(_generate_combinations(arg_group))
+                    else:
+                        groups.append(arg_group)
+                args = _generate_combinations(groups)
+            elif "args" in run_group:
+                args = _generate_combinations(run_group["args"])
+            else:
+                assert False, "? what? %s" % run_group
+
+            if "query-arg-groups" in run_group:
+                groups = []
+                for arg_group in run_group["query-arg-groups"]:
+                    if isinstance(arg_group, dict):
+                        groups.append(_generate_combinations(arg_group))
+                    else:
+                        groups.append(arg_group)
+                query_args = _generate_combinations(groups)
+            elif "query-args" in run_group:
+                query_args = _generate_combinations(run_group["query-args"])
+            else:
+                query_args = []
+
+            for arg_group in args:
+                aargs = []
+                aargs.extend(base_args)
+                if isinstance(arg_group, list):
+                    aargs.extend(arg_group)
+                else:
+                    aargs.append(arg_group)
+
+                vs = {
+                    "@count": count,
+                    "@metric": distance_metric,
+                    "@dimension": dimension,
+                    "@connection": conn_params
+                }
+                aargs = [_substitute_variables(arg, vs) for arg in aargs]
+                definitions.append(Definition(
+                    algorithm=name,
+                    run_group = run_group_name,
+                    docker_tag=algo['docker-tag'],
+                    module=algo['module'],
+                    constructor=algo['constructor'],
+                    arguments=aargs,
+                    query_argument_groups=query_args,
+                    disabled=algo.get('disabled', False)
+                ))
+
+    return definitions