import os import logging import pdb import time import re import random import traceback import json from multiprocessing import Process import numpy as np from yaml import full_load, dump from client import MilvusClient import parser from runner import Runner from milvus_metrics.api import report from milvus_metrics.models import Env, Hardware, Server, Metric import utils logger = logging.getLogger("milvus_benchmark.k8s_runner") namespace = "milvus" DELETE_INTERVAL_TIME = 5 # INSERT_INTERVAL = 100000 INSERT_INTERVAL = 50000 timestamp = int(time.time()) default_path = "/var/lib/milvus" class K8sRunner(Runner): """run docker mode""" def __init__(self): super(K8sRunner, self).__init__() self.name = utils.get_unique_name() self.host = None self.ip = None self.hostname = None self.env_value = None def init_env(self, server_config, server_host, image_type, image_tag): self.hostname = server_host # update values helm_path = os.path.join(os.getcwd(), "../milvus-helm") values_file_path = helm_path+"/values.yaml" if not os.path.exists(values_file_path): raise Exception("File %s not existed" % values_file_path) utils.update_values(values_file_path, server_host, server_config) try: logger.debug("Start install server") self.host, self.ip = utils.helm_install_server(helm_path, image_tag, image_type, self.name, namespace) except Exception as e: logger.error("Helm install server failed: %s" % str(e)) logger.error(traceback.format_exc()) self.clean_up() return False # for debugging # self.host = "192.168.1.101" if not self.host: logger.error("Helm install server failed") self.clean_up() return False return True def clean_up(self): logger.debug(self.name) utils.helm_del_server(self.name, namespace) def report_wrapper(self, milvus_instance, env_value, hostname, collection_info, index_info, search_params): metric = Metric() metric.set_run_id(timestamp) metric.env = Env(env_value) metric.env.OMP_NUM_THREADS = 0 metric.hardware = Hardware(name=hostname) server_version = milvus_instance.get_server_version() server_mode = milvus_instance.get_server_mode() commit = milvus_instance.get_server_commit() metric.server = Server(version=server_version, mode=server_mode, build_commit=commit) metric.collection = collection_info metric.index = index_info metric.search = search_params return metric def run(self, run_type, collection): logger.debug(run_type) logger.debug(collection) collection_name = collection["collection_name"] milvus_instance = MilvusClient(collection_name=collection_name, ip=self.ip) self.env_value = milvus_instance.get_server_config() # ugly implemention self.env_value.pop("logs") if run_type == "insert_performance": (data_type, collection_size, index_file_size, dimension, metric_type) = parser.collection_parser(collection_name) ni_per = collection["ni_per"] build_index = collection["build_index"] if milvus_instance.exists_collection(): milvus_instance.delete() time.sleep(10) index_info = {} search_params = {} milvus_instance.create_collection(collection_name, dimension, index_file_size, metric_type) if build_index is True: index_type = collection["index_type"] index_param = collection["index_param"] index_info = { "index_type": index_type, "index_param": index_param } milvus_instance.create_index(index_type, index_param) logger.debug(milvus_instance.describe_index()) res = self.do_insert(milvus_instance, collection_name, data_type, dimension, collection_size, ni_per) logger.info(res) milvus_instance.flush() collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, search_params) metric.metrics = { "type": run_type, "value": { "total_time": res["total_time"], "qps": res["qps"], "ni_time": res["ni_time"] } } report(metric) if build_index is True: logger.debug("Start build index for last file") milvus_instance.create_index(index_type, index_param) logger.debug(milvus_instance.describe_index()) if run_type == "insert_flush_performance": (data_type, collection_size, index_file_size, dimension, metric_type) = parser.collection_parser(collection_name) ni_per = collection["ni_per"] if milvus_instance.exists_collection(): milvus_instance.delete() time.sleep(10) index_info = {} search_params = {} milvus_instance.create_collection(collection_name, dimension, index_file_size, metric_type) res = self.do_insert(milvus_instance, collection_name, data_type, dimension, collection_size, ni_per) logger.info(res) logger.debug(milvus_instance.count()) start_time = time.time() milvus_instance.flush() end_time = time.time() logger.debug(milvus_instance.count()) collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, search_params) metric.metrics = { "type": run_type, "value": { "flush_time": round(end_time - start_time, 1) } } report(metric) elif run_type == "build_performance": (data_type, collection_size, index_file_size, dimension, metric_type) = parser.collection_parser(collection_name) index_type = collection["index_type"] index_param = collection["index_param"] collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } index_info = { "index_type": index_type, "index_param": index_param } if not milvus_instance.exists_collection(): logger.error("Table name: %s not existed" % collection_name) return search_params = {} start_time = time.time() # drop index logger.debug("Drop index") milvus_instance.drop_index() start_mem_usage = milvus_instance.get_mem_info()["memory_used"] milvus_instance.create_index(index_type, index_param) logger.debug(milvus_instance.describe_index()) logger.debug(milvus_instance.count()) end_time = time.time() end_mem_usage = milvus_instance.get_mem_info()["memory_used"] metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, search_params) metric.metrics = { "type": "build_performance", "value": { "build_time": round(end_time - start_time, 1), "start_mem_usage": start_mem_usage, "end_mem_usage": end_mem_usage, "diff_mem": end_mem_usage - start_mem_usage } } report(metric) elif run_type == "delete_performance": (data_type, collection_size, index_file_size, dimension, metric_type) = parser.collection_parser(collection_name) ni_per = collection["ni_per"] search_params = {} collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } if not milvus_instance.exists_collection(): logger.error("Table name: %s not existed" % collection_name) return length = milvus_instance.count() logger.info(length) index_info = milvus_instance.describe_index() logger.info(index_info) ids = [i for i in range(length)] loops = int(length / ni_per) milvus_instance.preload_collection() start_mem_usage = milvus_instance.get_mem_info()["memory_used"] start_time = time.time() for i in range(loops): delete_ids = ids[i*ni_per : i*ni_per+ni_per] logger.debug("Delete %d - %d" % (delete_ids[0], delete_ids[-1])) milvus_instance.delete_vectors(delete_ids) milvus_instance.flush() logger.debug("Table row counts: %d" % milvus_instance.count()) logger.debug("Table row counts: %d" % milvus_instance.count()) milvus_instance.flush() end_time = time.time() end_mem_usage = milvus_instance.get_mem_info()["memory_used"] logger.debug("Table row counts: %d" % milvus_instance.count()) metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, search_params) metric.metrics = { "type": "delete_performance", "value": { "delete_time": round(end_time - start_time, 1), "start_mem_usage": start_mem_usage, "end_mem_usage": end_mem_usage, "diff_mem": end_mem_usage - start_mem_usage } } report(metric) elif run_type == "search_performance": (data_type, collection_size, index_file_size, dimension, metric_type) = parser.collection_parser(collection_name) run_count = collection["run_count"] top_ks = collection["top_ks"] nqs = collection["nqs"] search_params = collection["search_params"] collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } # fro debugging # time.sleep(3600) if not milvus_instance.exists_collection(): logger.error("Table name: %s not existed" % collection_name) return logger.info(milvus_instance.count()) index_info = milvus_instance.describe_index() logger.info(index_info) milvus_instance.preload_collection() logger.info("Start warm up query") res = self.do_query(milvus_instance, collection_name, [1], [1], 2, search_param=search_params[0]) logger.info("End warm up query") for search_param in search_params: logger.info("Search param: %s" % json.dumps(search_param)) res = self.do_query(milvus_instance, collection_name, top_ks, nqs, run_count, search_param) headers = ["Nq/Top-k"] headers.extend([str(top_k) for top_k in top_ks]) logger.info("Search param: %s" % json.dumps(search_param)) utils.print_table(headers, nqs, res) for index_nq, nq in enumerate(nqs): for index_top_k, top_k in enumerate(top_ks): search_param_group = { "nq": nq, "topk": top_k, "search_param": search_param } search_time = res[index_nq][index_top_k] metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, search_param_group) metric.metrics = { "type": "search_performance", "value": { "search_time": search_time } } report(metric) elif run_type == "search_ids_stability": (data_type, collection_size, index_file_size, dimension, metric_type) = parser.collection_parser(collection_name) search_params = collection["search_params"] during_time = collection["during_time"] ids_length = collection["ids_length"] ids = collection["ids"] collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } if not milvus_instance.exists_collection(): logger.error("Table name: %s not existed" % collection_name) return logger.info(milvus_instance.count()) index_info = milvus_instance.describe_index() logger.info(index_info) g_top_k = int(collection["top_ks"].split("-")[1]) l_top_k = int(collection["top_ks"].split("-")[0]) g_id = int(ids.split("-")[1]) l_id = int(ids.split("-")[0]) g_id_length = int(ids_length.split("-")[1]) l_id_length = int(ids_length.split("-")[0]) milvus_instance.preload_collection() start_mem_usage = milvus_instance.get_mem_info()["memory_used"] logger.debug(start_mem_usage) start_time = time.time() while time.time() < start_time + during_time * 60: search_param = {} top_k = random.randint(l_top_k, g_top_k) ids_num = random.randint(l_id_length, g_id_length) ids_param = [random.randint(l_id_length, g_id_length) for _ in range(ids_num)] for k, v in search_params.items(): search_param[k] = random.randint(int(v.split("-")[0]), int(v.split("-")[1])) logger.debug("Query top-k: %d, ids_num: %d, param: %s" % (top_k, ids_num, json.dumps(search_param))) result = milvus_instance.query_ids(top_k, ids_param, search_param=search_param) end_mem_usage = milvus_instance.get_mem_info()["memory_used"] metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, {}) metric.metrics = { "type": "search_ids_stability", "value": { "during_time": during_time, "start_mem_usage": start_mem_usage, "end_mem_usage": end_mem_usage, "diff_mem": end_mem_usage - start_mem_usage } } report(metric) # for sift/deep datasets # TODO: enable elif run_type == "accuracy": (data_type, collection_size, index_file_size, dimension, metric_type) = parser.collection_parser(collection_name) search_params = collection["search_params"] # mapping to search param list search_params = self.generate_combinations(search_params) top_ks = collection["top_ks"] nqs = collection["nqs"] collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } if not milvus_instance.exists_collection(): logger.error("Table name: %s not existed" % collection_name) return logger.info(milvus_instance.count()) index_info = milvus_instance.describe_index() logger.info(index_info) milvus_instance.preload_collection() true_ids_all = self.get_groundtruth_ids(collection_size) for search_param in search_params: for top_k in top_ks: for nq in nqs: total = 0 search_param_group = { "nq": nq, "topk": top_k, "search_param": search_param } logger.info("Query params: %s" % json.dumps(search_param_group)) result_ids, result_distances = self.do_query_ids(milvus_instance, collection_name, top_k, nq, search_param=search_param) acc_value = self.get_recall_value(true_ids_all[:nq, :top_k].tolist(), result_ids) logger.info("Query accuracy: %s" % acc_value) metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, search_param_group) metric.metrics = { "type": "accuracy", "value": { "acc": acc_value } } report(metric) elif run_type == "ann_accuracy": hdf5_source_file = collection["source_file"] collection_name = collection["collection_name"] index_file_sizes = collection["index_file_sizes"] index_types = collection["index_types"] index_params = collection["index_params"] top_ks = collection["top_ks"] nqs = collection["nqs"] search_params = collection["search_params"] # mapping to search param list search_params = self.generate_combinations(search_params) # mapping to index param list index_params = self.generate_combinations(index_params) data_type, dimension, metric_type = parser.parse_ann_collection_name(collection_name) collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } dataset = utils.get_dataset(hdf5_source_file) if milvus_instance.exists_collection(collection_name): logger.info("Re-create collection: %s" % collection_name) milvus_instance.delete() time.sleep(DELETE_INTERVAL_TIME) true_ids = np.array(dataset["neighbors"]) for index_file_size in index_file_sizes: milvus_instance.create_collection(collection_name, dimension, index_file_size, metric_type) logger.info(milvus_instance.describe()) insert_vectors = self.normalize(metric_type, np.array(dataset["train"])) # Insert batch once # milvus_instance.insert(insert_vectors) loops = len(insert_vectors) // INSERT_INTERVAL + 1 for i in range(loops): start = i*INSERT_INTERVAL end = min((i+1)*INSERT_INTERVAL, len(insert_vectors)) tmp_vectors = insert_vectors[start:end] if start < end: if not isinstance(tmp_vectors, list): milvus_instance.insert(tmp_vectors.tolist(), ids=[i for i in range(start, end)]) else: milvus_instance.insert(tmp_vectors, ids=[i for i in range(start, end)]) milvus_instance.flush() logger.info("Table: %s, row count: %s" % (collection_name, milvus_instance.count())) if milvus_instance.count() != len(insert_vectors): logger.error("Table row count is not equal to insert vectors") return for index_type in index_types: for index_param in index_params: logger.debug("Building index with param: %s" % json.dumps(index_param)) milvus_instance.create_index(index_type, index_param=index_param) logger.info(milvus_instance.describe_index()) logger.info("Start preload collection: %s" % collection_name) milvus_instance.preload_collection() index_info = { "index_type": index_type, "index_param": index_param } logger.debug(index_info) for search_param in search_params: for nq in nqs: query_vectors = self.normalize(metric_type, np.array(dataset["test"][:nq])) for top_k in top_ks: search_param_group = { "nq": len(query_vectors), "topk": top_k, "search_param": search_param } logger.debug(search_param_group) if not isinstance(query_vectors, list): result = milvus_instance.query(query_vectors.tolist(), top_k, search_param=search_param) else: result = milvus_instance.query(query_vectors, top_k, search_param=search_param) result_ids = result.id_array acc_value = self.get_recall_value(true_ids[:nq, :top_k].tolist(), result_ids) logger.info("Query ann_accuracy: %s" % acc_value) metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, search_param_group) metric.metrics = { "type": "ann_accuracy", "value": { "acc": acc_value } } report(metric) elif run_type == "search_stability": (data_type, collection_size, index_file_size, dimension, metric_type) = parser.collection_parser(collection_name) search_params = collection["search_params"] during_time = collection["during_time"] collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } if not milvus_instance.exists_collection(): logger.error("Table name: %s not existed" % collection_name) return logger.info(milvus_instance.count()) index_info = milvus_instance.describe_index() logger.info(index_info) g_top_k = int(collection["top_ks"].split("-")[1]) g_nq = int(collection["nqs"].split("-")[1]) l_top_k = int(collection["top_ks"].split("-")[0]) l_nq = int(collection["nqs"].split("-")[0]) milvus_instance.preload_collection() start_mem_usage = milvus_instance.get_mem_info()["memory_used"] logger.debug(start_mem_usage) start_row_count = milvus_instance.count() logger.debug(milvus_instance.describe_index()) logger.info(start_row_count) start_time = time.time() while time.time() < start_time + during_time * 60: search_param = {} top_k = random.randint(l_top_k, g_top_k) nq = random.randint(l_nq, g_nq) for k, v in search_params.items(): search_param[k] = random.randint(int(v.split("-")[0]), int(v.split("-")[1])) query_vectors = [[random.random() for _ in range(dimension)] for _ in range(nq)] logger.debug("Query nq: %d, top-k: %d, param: %s" % (nq, top_k, json.dumps(search_param))) result = milvus_instance.query(query_vectors, top_k, search_param=search_param) end_mem_usage = milvus_instance.get_mem_info()["memory_used"] metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, {}) metric.metrics = { "type": "search_stability", "value": { "during_time": during_time, "start_mem_usage": start_mem_usage, "end_mem_usage": end_mem_usage, "diff_mem": end_mem_usage - start_mem_usage } } report(metric) elif run_type == "stability": (data_type, collection_size, index_file_size, dimension, metric_type) = parser.collection_parser(collection_name) search_params = collection["search_params"] insert_xb = collection["insert_xb"] insert_interval = collection["insert_interval"] delete_xb = collection["delete_xb"] during_time = collection["during_time"] collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } if not milvus_instance.exists_collection(): logger.error("Table name: %s not existed" % collection_name) return logger.info(milvus_instance.count()) index_info = milvus_instance.describe_index() logger.info(index_info) g_top_k = int(collection["top_ks"].split("-")[1]) g_nq = int(collection["nqs"].split("-")[1]) l_top_k = int(collection["top_ks"].split("-")[0]) l_nq = int(collection["nqs"].split("-")[0]) milvus_instance.preload_collection() start_mem_usage = milvus_instance.get_mem_info()["memory_used"] start_row_count = milvus_instance.count() logger.debug(milvus_instance.describe_index()) logger.info(start_row_count) start_time = time.time() i = 0 ids = [] insert_vectors = [[random.random() for _ in range(dimension)] for _ in range(insert_xb)] query_vectors = [[random.random() for _ in range(dimension)] for _ in range(10000)] while time.time() < start_time + during_time * 60: i = i + 1 for j in range(insert_interval): top_k = random.randint(l_top_k, g_top_k) nq = random.randint(l_nq, g_nq) search_param = {} for k, v in search_params.items(): search_param[k] = random.randint(int(v.split("-")[0]), int(v.split("-")[1])) logger.debug("Query nq: %d, top-k: %d, param: %s" % (nq, top_k, json.dumps(search_param))) result = milvus_instance.query(query_vectors[0:nq], top_k, search_param=search_param) count = milvus_instance.count() insert_ids = [(count+x) for x in range(len(insert_vectors))] ids.extend(insert_ids) status, res = milvus_instance.insert(insert_vectors, ids=insert_ids) logger.debug("%d, row_count: %d" % (i, milvus_instance.count())) milvus_instance.delete_vectors(ids[-delete_xb:]) milvus_instance.flush() milvus_instance.compact() end_mem_usage = milvus_instance.get_mem_info()["memory_used"] end_row_count = milvus_instance.count() metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, {}) metric.metrics = { "type": "stability", "value": { "during_time": during_time, "start_mem_usage": start_mem_usage, "end_mem_usage": end_mem_usage, "diff_mem": end_mem_usage - start_mem_usage, "row_count_increments": end_row_count - start_row_count } } report(metric) else: logger.warning("Run type not defined") return logger.debug("Test finished")