import os import threading import logging import pdb import time import random import grpc from multiprocessing import Process from itertools import product import numpy as np import sklearn.preprocessing from milvus import DataType from client import MilvusClient import utils import parser logger = logging.getLogger("milvus_benchmark.runner") VECTORS_PER_FILE = 1000000 SIFT_VECTORS_PER_FILE = 100000 BINARY_VECTORS_PER_FILE = 2000000 MAX_NQ = 10001 FILE_PREFIX = "binary_" # FOLDER_NAME = 'ann_1000m/source_data' SRC_BINARY_DATA_DIR = '/test/milvus/raw_data/random/' SIFT_SRC_DATA_DIR = '/test/milvus/raw_data/sift1b/' DEEP_SRC_DATA_DIR = '/test/milvus/raw_data/deep1b/' BINARY_SRC_DATA_DIR = '/test/milvus/raw_data/binary/' SIFT_SRC_GROUNDTRUTH_DATA_DIR = SIFT_SRC_DATA_DIR + 'gnd' WARM_TOP_K = 1 WARM_NQ = 1 DEFAULT_DIM = 512 GROUNDTRUTH_MAP = { "1000000": "idx_1M.ivecs", "2000000": "idx_2M.ivecs", "5000000": "idx_5M.ivecs", "10000000": "idx_10M.ivecs", "20000000": "idx_20M.ivecs", "50000000": "idx_50M.ivecs", "100000000": "idx_100M.ivecs", "200000000": "idx_200M.ivecs", "500000000": "idx_500M.ivecs", "1000000000": "idx_1000M.ivecs", } def gen_file_name(idx, dimension, data_type): s = "%05d" % idx fname = FILE_PREFIX + str(dimension) + "d_" + s + ".npy" if data_type == "random": fname = SRC_BINARY_DATA_DIR+fname elif data_type == "sift": fname = SIFT_SRC_DATA_DIR+fname elif data_type == "deep": fname = DEEP_SRC_DATA_DIR+fname elif data_type == "binary": fname = BINARY_SRC_DATA_DIR+fname return fname def get_vectors_from_binary(nq, dimension, data_type): # use the first file, nq should be less than VECTORS_PER_FILE if nq > MAX_NQ: raise Exception("Over size nq") if data_type == "random": file_name = SRC_BINARY_DATA_DIR+'query_%d.npy' % dimension elif data_type == "sift": file_name = SIFT_SRC_DATA_DIR+'query.npy' elif data_type == "deep": file_name = DEEP_SRC_DATA_DIR+'query.npy' elif data_type == "binary": file_name = BINARY_SRC_DATA_DIR+'query.npy' data = np.load(file_name) vectors = data[0:nq].tolist() return vectors class Runner(object): def __init__(self): pass def gen_executors(self, operations): l = [] for name, operation in operations.items(): weight = operation["weight"] if "weight" in operation else 1 l.extend([name] * weight) random.shuffle(l) return l def get_vector_type(self, data_type): vector_type = '' if data_type in ["random", "sift", "deep", "glove"]: vector_type = DataType.FLOAT_VECTOR elif data_type in ["binary"]: vector_type = DataType.BINARY_VECTOR else: raise Exception("Data type: %s not defined" % data_type) return vector_type def get_vector_type_from_metric(self, metric_type): vector_type = '' if metric_type in ["hamming", "jaccard"]: vector_type = DataType.BINARY_VECTOR else: vector_type = DataType.FLOAT_VECTOR return vector_type def normalize(self, metric_type, X): if metric_type == "ip": logger.info("Set normalize for metric_type: %s" % metric_type) X = sklearn.preprocessing.normalize(X, axis=1, norm='l2') X = X.astype(np.float32) elif metric_type == "l2": X = X.astype(np.float32) elif metric_type in ["jaccard", "hamming", "sub", "super"]: tmp = [] for item in X: new_vector = bytes(np.packbits(item, axis=-1).tolist()) tmp.append(new_vector) X = tmp return X def generate_combinations(self, 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 do_insert(self, milvus, collection_name, data_type, dimension, size, ni): ''' @params: mivlus: server connect instance dimension: collection dimensionn # index_file_size: size trigger file merge size: row count of vectors to be insert ni: row count of vectors to be insert each time # store_id: if store the ids returned by call add_vectors or not @return: total_time: total time for all insert operation qps: vectors added per second ni_time: avarage insert operation time ''' bi_res = {} total_time = 0.0 qps = 0.0 ni_time = 0.0 if data_type == "random": if dimension == 512: vectors_per_file = VECTORS_PER_FILE elif dimension == 4096: vectors_per_file = 100000 elif dimension == 16384: vectors_per_file = 10000 elif data_type == "sift": vectors_per_file = SIFT_VECTORS_PER_FILE elif data_type in ["binary"]: vectors_per_file = BINARY_VECTORS_PER_FILE else: raise Exception("data_type: %s not supported" % data_type) if size % vectors_per_file or size % ni: raise Exception("Not invalid collection size or ni") i = 0 while i < (size // vectors_per_file): vectors = [] if vectors_per_file >= ni: file_name = gen_file_name(i, dimension, data_type) # logger.info("Load npy file: %s start" % file_name) data = np.load(file_name) # logger.info("Load npy file: %s end" % file_name) for j in range(vectors_per_file // ni): vectors = data[j*ni:(j+1)*ni].tolist() if vectors: # start insert vectors start_id = i * vectors_per_file + j * ni end_id = start_id + len(vectors) logger.debug("Start id: %s, end id: %s" % (start_id, end_id)) ids = [k for k in range(start_id, end_id)] entities = milvus.generate_entities(vectors, ids) ni_start_time = time.time() try: res_ids = milvus.insert(entities, ids=ids) except grpc.RpcError as e: if e.code() == grpc.StatusCode.UNAVAILABLE: logger.debug("Retry insert") def retry(): res_ids = milvus.insert(entities, ids=ids) t0 = threading.Thread(target=retry) t0.start() t0.join() logger.debug("Retry successfully") raise e assert ids == res_ids # milvus.flush() logger.debug(milvus.count()) ni_end_time = time.time() total_time = total_time + ni_end_time - ni_start_time i += 1 else: vectors.clear() loops = ni // vectors_per_file for j in range(loops): file_name = gen_file_name(loops*i+j, dimension, data_type) data = np.load(file_name) vectors.extend(data.tolist()) if vectors: start_id = i * vectors_per_file end_id = start_id + len(vectors) logger.info("Start id: %s, end id: %s" % (start_id, end_id)) ids = [k for k in range(start_id, end_id)] entities = milvus.generate_entities(vectors, ids) ni_start_time = time.time() try: res_ids = milvus.insert(entities, ids=ids) except grpc.RpcError as e: if e.code() == grpc.StatusCode.UNAVAILABLE: logger.debug("Retry insert") def retry(): res_ids = milvus.insert(entities, ids=ids) t0 = threading.Thread(target=retry) t0.start() t0.join() logger.debug("Retry successfully") raise e assert ids == res_ids # milvus.flush() logger.debug(milvus.count()) ni_end_time = time.time() total_time = total_time + ni_end_time - ni_start_time i += loops qps = round(size / total_time, 2) ni_time = round(total_time / (size / ni), 2) bi_res["total_time"] = round(total_time, 2) bi_res["qps"] = qps bi_res["ni_time"] = ni_time return bi_res def do_query(self, milvus, collection_name, vec_field_name, top_ks, nqs, run_count=1, search_param=None, filter_query=None): bi_res = [] (data_type, collection_size, index_file_size, dimension, metric_type) = parser.collection_parser(collection_name) base_query_vectors = get_vectors_from_binary(MAX_NQ, dimension, data_type) for nq in nqs: tmp_res = [] query_vectors = base_query_vectors[0:nq] for top_k in top_ks: avg_query_time = 0.0 min_query_time = 0.0 logger.info("Start query, query params: top-k: {}, nq: {}, actually length of vectors: {}".format(top_k, nq, len(query_vectors))) for i in range(run_count): logger.debug("Start run query, run %d of %s" % (i+1, run_count)) start_time = time.time() vector_query = {"vector": {vec_field_name: { "topk": top_k, "query": query_vectors, "metric_type": utils.metric_type_trans(metric_type), "params": search_param} }} query_res = milvus.query(vector_query, filter_query=filter_query) interval_time = time.time() - start_time if (i == 0) or (min_query_time > interval_time): min_query_time = interval_time logger.info("Min query time: %.2f" % min_query_time) tmp_res.append(round(min_query_time, 2)) bi_res.append(tmp_res) return bi_res def do_query_qps(self, milvus, query_vectors, top_k, search_param): start_time = time.time() result = milvus.query(query_vectors, top_k, search_param) end_time = time.time() return end_time - start_time def do_query_ids(self, milvus, collection_name, vec_field_name, top_k, nq, search_param=None, filter_query=None): (data_type, collection_size, index_file_size, dimension, metric_type) = parser.collection_parser(collection_name) base_query_vectors = get_vectors_from_binary(MAX_NQ, dimension, data_type) query_vectors = base_query_vectors[0:nq] logger.info("Start query, query params: top-k: {}, nq: {}, actually length of vectors: {}".format(top_k, nq, len(query_vectors))) vector_query = {"vector": {vec_field_name: { "topk": top_k, "query": query_vectors, "metric_type": utils.metric_type_trans(metric_type), "params": search_param} }} query_res = milvus.query(vector_query, filter_query=filter_query) result_ids = milvus.get_ids(query_res) return result_ids def do_query_acc(self, milvus, collection_name, top_k, nq, id_store_name, search_param=None): (data_type, collection_size, index_file_size, dimension, metric_type) = parser.collection_parser(collection_name) base_query_vectors = get_vectors_from_binary(MAX_NQ, dimension, data_type) vectors = base_query_vectors[0:nq] logger.info("Start query, query params: top-k: {}, nq: {}, actually length of vectors: {}".format(top_k, nq, len(vectors))) query_res = milvus.query(vectors, top_k, search_param=None) # if file existed, cover it if os.path.isfile(id_store_name): os.remove(id_store_name) with open(id_store_name, 'a+') as fd: for nq_item in query_res: for item in nq_item: fd.write(str(item.id)+'\t') fd.write('\n') # compute and print accuracy def compute_accuracy(self, flat_file_name, index_file_name): flat_id_list = []; index_id_list = [] logger.info("Loading flat id file: %s" % flat_file_name) with open(flat_file_name, 'r') as flat_id_fd: for line in flat_id_fd: tmp_list = line.strip("\n").strip().split("\t") flat_id_list.append(tmp_list) logger.info("Loading index id file: %s" % index_file_name) with open(index_file_name) as index_id_fd: for line in index_id_fd: tmp_list = line.strip("\n").strip().split("\t") index_id_list.append(tmp_list) if len(flat_id_list) != len(index_id_list): raise Exception("Flat index result length: not match, Acc compute exiting ..." % (len(flat_id_list), len(index_id_list))) # get the accuracy return self.get_recall_value(flat_id_list, index_id_list) def get_recall_value(self, true_ids, result_ids): """ Use the intersection length """ sum_radio = 0.0 for index, item in enumerate(result_ids): # tmp = set(item).intersection(set(flat_id_list[index])) tmp = set(true_ids[index]).intersection(set(item)) sum_radio = sum_radio + len(tmp) / len(item) # logger.debug(sum_radio) return round(sum_radio / len(result_ids), 3) """ Implementation based on: https://github.com/facebookresearch/faiss/blob/master/benchs/datasets.py """ def get_groundtruth_ids(self, collection_size): fname = GROUNDTRUTH_MAP[str(collection_size)] fname = SIFT_SRC_GROUNDTRUTH_DATA_DIR + "/" + fname a = np.fromfile(fname, dtype='int32') d = a[0] true_ids = a.reshape(-1, d + 1)[:, 1:].copy() return true_ids def get_fields(self, milvus, collection_name): fields = [] info = milvus.get_info(collection_name) for item in info["fields"]: fields.append(item["name"]) return fields # def get_filter_query(self, filter_query): # for filter in filter_query: