import os import logging import pdb import time import re import random import traceback import json import csv import threading from multiprocessing import Process import numpy as np from milvus import DataType from yaml import full_load, dump import concurrent.futures import locust_user 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 helm_utils import utils logger = logging.getLogger("milvus_benchmark.k8s_runner") namespace = "milvus" default_port = 19530 DELETE_INTERVAL_TIME = 5 # INSERT_INTERVAL = 100000 INSERT_INTERVAL = 50000 BIG_FLUSH_INTERVAL = 3600 DEFAULT_FLUSH_INTERVAL = 1 timestamp = int(time.time()) default_path = "/var/lib/milvus" class K8sRunner(Runner): """run docker mode""" def __init__(self): super(K8sRunner, self).__init__() self.service_name = utils.get_unique_name() self.host = None self.port = default_port self.hostname = None self.env_value = None self.hardware = None self.deploy_mode = None def init_env(self, milvus_config, server_config, server_host, deploy_mode, image_type, image_tag): logger.debug("Tests run on server host:") logger.debug(server_host) self.hostname = server_host self.deploy_mode = deploy_mode if self.hostname: try: cpus = helm_utils.get_host_cpus(self.hostname) except Exception as e: logger.error(str(e)) cpus = 64 logger.debug(type(cpus)) if server_config: if "cpus" in server_config.keys(): cpus = min(server_config["cpus"], int(cpus)) else: server_config.update({"cpus": cpus}) else: server_config = {"cpus": cpus} self.hardware = Hardware(name=self.hostname, cpus=cpus) # update values helm_path = os.path.join(os.getcwd(), "../milvus-helm/charts/milvus") values_file_path = helm_path + "/values.yaml" if not os.path.exists(values_file_path): raise Exception("File %s not existed" % values_file_path) if milvus_config: helm_utils.update_values(values_file_path, deploy_mode, server_host, milvus_config, server_config) try: logger.debug("Start install server") self.host = helm_utils.helm_install_server(helm_path, deploy_mode, image_tag, image_type, self.service_name, namespace) except Exception as e: logger.error("Helm install server failed: %s" % (str(e))) logger.error(traceback.format_exc()) logger.error(self.hostname) self.clean_up() return False logger.debug(server_config) # for debugging if not self.host: logger.error("Helm install server failed") self.clean_up() return False return True def clean_up(self): logger.debug("Start clean up: %s" % self.service_name) helm_utils.helm_del_server(self.service_name, namespace) def report_wrapper(self, milvus_instance, env_value, hostname, collection_info, index_info, search_params, run_params=None, server_config=None): metric = Metric() metric.set_run_id(timestamp) metric.env = Env(env_value) metric.env.OMP_NUM_THREADS = 0 metric.hardware = self.hardware # TODO: removed # server_version = milvus_instance.get_server_version() # server_mode = milvus_instance.get_server_mode() # commit = milvus_instance.get_server_commit() server_version = "0.12.0" server_mode = self.deploy_mode metric.server = Server(version=server_version, mode=server_mode, build_commit=None) metric.collection = collection_info metric.index = index_info metric.search = search_params metric.run_params = run_params return metric def run(self, run_type, collection): logger.debug(run_type) logger.debug(collection) collection_name = collection["collection_name"] if "collection_name" in collection else None milvus_instance = MilvusClient(collection_name=collection_name, host=self.host) # TODO: removed # self.env_value = milvus_instance.get_server_config() # ugly implemention # self.env_value = utils.convert_nested(self.env_value) # self.env_value.pop("logs") # self.env_value.pop("network") self.env_value = collection if run_type == "insert_performance": (data_type, collection_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.drop() time.sleep(10) index_info = {} search_params = {} vector_type = self.get_vector_type(data_type) other_fields = collection["other_fields"] if "other_fields" in collection else None milvus_instance.create_collection(dimension, data_type=vector_type, other_fields=other_fields) 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 } index_field_name = utils.get_default_field_name(vector_type) milvus_instance.create_index(index_field_name, index_type, metric_type, index_param=index_param) logger.debug(milvus_instance.describe_index()) res = self.do_insert(milvus_instance, collection_name, data_type, dimension, collection_size, ni_per) flush_time = 0.0 if "flush" in collection and collection["flush"] == "no": logger.debug("No manual flush") else: start_time = time.time() milvus_instance.flush() flush_time = time.time() - start_time logger.debug(milvus_instance.count()) collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name, "other_fields": other_fields, "ni_per": ni_per } metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, search_params) total_time = res["total_time"] build_time = 0 if build_index is True: logger.debug("Start build index for last file") start_time = time.time() milvus_instance.create_index(index_field_name, index_type, metric_type, index_param=index_param) build_time = time.time() - start_time total_time = total_time + build_time metric.metrics = { "type": run_type, "value": { "total_time": total_time, "qps": res["qps"], "ni_time": res["ni_time"], "flush_time": flush_time, "build_time": build_time } } report(metric) elif run_type == "build_performance": (data_type, collection_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 = {} vector_type = self.get_vector_type(data_type) index_field_name = utils.get_default_field_name(vector_type) start_time = time.time() # drop index logger.debug("Drop index") milvus_instance.drop_index(index_field_name) # start_mem_usage = milvus_instance.get_mem_info()["memory_used"] # TODO: need to check milvus_instance.create_index(index_field_name, index_type, metric_type, index_param=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), } } report(metric) elif run_type == "delete_performance": (data_type, collection_size, dimension, metric_type) = parser.collection_parser( collection_name) ni_per = collection["ni_per"] auto_flush = collection["auto_flush"] if "auto_flush" in collection else True search_params = {} collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } if not milvus_instance.exists_collection(): logger.error(milvus_instance.show_collections()) 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.load_collection() # TODO: remove # start_mem_usage = milvus_instance.get_mem_info()["memory_used"] start_time = time.time() # if auto_flush is False: # milvus_instance.set_config("storage", "auto_flush_interval", BIG_FLUSH_INTERVAL) 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(delete_ids) logger.debug("Table row counts: %d" % milvus_instance.count()) logger.debug("Table row counts: %d" % milvus_instance.count()) start_flush_time = time.time() milvus_instance.flush() end_flush_time = time.time() end_time = time.time() # end_mem_usage = milvus_instance.get_mem_info()["memory_used"] logger.debug("Table row counts: %d" % milvus_instance.count()) # milvus_instance.set_config("storage", "auto_flush_interval", DEFAULT_FLUSH_INTERVAL) metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, search_params) delete_time = round(end_time - start_time, 1) metric.metrics = { "type": "delete_performance", "value": { "delete_time": delete_time, "qps": round(collection_size / delete_time, 1) } } if auto_flush is False: flush_time = round(end_flush_time - start_flush_time, 1) metric.metrics["value"].update({"flush_time": flush_time}) report(metric) elif run_type == "get_ids_performance": (data_type, collection_size, dimension, metric_type) = parser.collection_parser( collection_name) ids_length_per_segment = collection["ids_length_per_segment"] if not milvus_instance.exists_collection(): logger.error("Table name: %s not existed" % collection_name) return collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } search_params = {} logger.info(milvus_instance.count()) index_info = milvus_instance.describe_index() logger.info(index_info) for ids_num in ids_length_per_segment: segment_num, get_ids = milvus_instance.get_rand_ids_each_segment(ids_num) start_time = time.time() get_res = milvus_instance.get_entities(get_ids) total_time = time.time() - start_time avg_time = total_time / segment_num run_params = {"ids_num": ids_num} logger.info( "Segment num: %d, ids num per segment: %d, run_time: %f" % (segment_num, ids_num, total_time)) metric = self.report_wrapper(milvus_instance, self.env_value, self.hostname, collection_info, index_info, search_params, run_params=run_params) metric.metrics = { "type": run_type, "value": { "total_time": round(total_time, 1), "avg_time": round(avg_time, 1) } } report(metric) elif run_type == "search_performance": (data_type, collection_size, dimension, metric_type) = parser.collection_parser( collection_name) run_count = collection["run_count"] top_ks = collection["top_ks"] nqs = collection["nqs"] # filter_query = collection["filter"] if "filter" in collection else None filters = collection["filters"] if "filters" in collection else [] filter_query = [] search_params = collection["search_params"] fields = self.get_fields(milvus_instance, collection_name) collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name "fields": fields } if not milvus_instance.exists_collection(): logger.error("Table name: %s not existed" % collection_name) return vector_type = self.get_vector_type(data_type) vec_field_name = utils.get_default_field_name(vector_type) logger.info(milvus_instance.count()) index_info = milvus_instance.describe_index() logger.info(index_info) milvus_instance.load_collection() logger.info("Start warm up query") res = self.do_query(milvus_instance, collection_name, vec_field_name, [1], [1], 2, search_param=search_params[0], filter_query=filter_query) logger.info("End warm up query") for search_param in search_params: logger.info("Search param: %s" % json.dumps(search_param)) if not filters: filters.append(None) for filter in filters: filter_param = [] if isinstance(filter, dict) and "range" in filter: filter_query.append(eval(filter["range"])) filter_param.append(filter["range"]) if isinstance(filter, dict) and "term" in filter: filter_query.append(eval(filter["term"])) filter_param.append(filter["term"]) logger.info("filter param: %s" % json.dumps(filter_param)) res = self.do_query(milvus_instance, collection_name, vec_field_name, top_ks, nqs, run_count, search_param, filter_query=filter_query) 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, "filter": filter_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 == "locust_insert_stress": pass elif run_type in ["locust_search_performance", "locust_insert_performance", "locust_mix_performance"]: (data_type, collection_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.drop() time.sleep(10) index_info = {} search_params = {} vector_type = self.get_vector_type(data_type) index_field_name = utils.get_default_field_name(vector_type) milvus_instance.create_collection(dimension, data_type=vector_type, other_fields=None) vector_type = self.get_vector_type(data_type) vec_field_name = utils.get_default_field_name(vector_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_field_name, index_type, metric_type, index_param=index_param) logger.debug(milvus_instance.describe_index()) if run_type in ["locust_search_performance", "locust_mix_performance"]: res = self.do_insert(milvus_instance, collection_name, data_type, dimension, collection_size, ni_per) if "flush" in collection and collection["flush"] == "no": logger.debug("No manual flush") else: milvus_instance.flush() if build_index is True: logger.debug("Start build index for last file") milvus_instance.create_index(index_field_name, index_type, metric_type, _async=True, index_param=index_param) logger.debug(milvus_instance.describe_index()) logger.debug("Table row counts: %d" % milvus_instance.count()) milvus_instance.load_collection() logger.info("Start warm up query") for i in range(2): res = self.do_query(milvus_instance, collection_name, vec_field_name, [1], [1], 2, search_param={"nprobe": 16}) logger.info("End warm up query") real_metric_type = utils.metric_type_trans(metric_type) ### spawn locust requests task = collection["task"] connection_type = "single" connection_num = task["connection_num"] if connection_num > 1: connection_type = "multi" clients_num = task["clients_num"] hatch_rate = task["hatch_rate"] during_time = utils.timestr_to_int(task["during_time"]) task_types = task["types"] run_params = {"tasks": {}, "clients_num": clients_num, "spawn_rate": hatch_rate, "during_time": during_time} for task_type in task_types: run_params["tasks"].update({task_type["type"]: task_type["weight"] if "weight" in task_type else 1}) # . collect stats locust_stats = locust_user.locust_executor(self.host, self.port, collection_name, connection_type=connection_type, run_params=run_params) logger.info(locust_stats) 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": locust_stats} report(metric) elif run_type == "search_ids_stability": (data_type, collection_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.load_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, } } report(metric) # for sift/deep datasets # TODO: enable elif run_type == "accuracy": (data_type, collection_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.load_collection() true_ids_all = self.get_groundtruth_ids(collection_size) vector_type = self.get_vector_type(data_type) vec_field_name = utils.get_default_field_name(vector_type) for search_param in search_params: headers = ["Nq/Top-k"] res = [] for nq in nqs: for top_k in top_ks: tmp_res = [] search_param_group = { "nq": nq, "topk": top_k, "search_param": search_param, "metric_type": metric_type } logger.info("Query params: %s" % json.dumps(search_param_group)) result_ids = self.do_query_ids(milvus_instance, collection_name, vec_field_name, top_k, nq, search_param=search_param) # mem_used = milvus_instance.get_mem_info()["memory_used"] acc_value = self.get_recall_value(true_ids_all[:nq, :top_k].tolist(), result_ids) logger.info("Query accuracy: %s" % acc_value) tmp_res.append(acc_value) # logger.info("Memory usage: %s" % mem_used) 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) # logger.info("Memory usage: %s" % mem_used) res.append(tmp_res) 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) elif run_type == "ann_accuracy": hdf5_source_file = collection["source_file"] collection_name = collection["collection_name"] 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.drop() time.sleep(DELETE_INTERVAL_TIME) true_ids = np.array(dataset["neighbors"]) vector_type = self.get_vector_type_from_metric(metric_type) vec_field_name = utils.get_default_field_name(vector_type) real_metric_type = utils.metric_type_trans(metric_type) # re-create collection if milvus_instance.exists_collection(collection_name): milvus_instance.drop() time.sleep(DELETE_INTERVAL_TIME) milvus_instance.create_collection(dimension, data_type=vector_type) insert_vectors = self.normalize(metric_type, np.array(dataset["train"])) if len(insert_vectors) != dataset["train"].shape[0]: raise Exception("Row count of insert vectors: %d is not equal to dataset size: %d" % ( len(insert_vectors), dataset["train"].shape[0])) logger.debug("The row count of entities to be inserted: %d" % len(insert_vectors)) # 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)) if start < end: tmp_vectors = insert_vectors[start:end] ids = [i for i in range(start, end)] if not isinstance(tmp_vectors, list): entities = milvus_instance.generate_entities(tmp_vectors.tolist(), ids) res_ids = milvus_instance.insert(entities, ids=ids) else: entities = milvus_instance.generate_entities(tmp_vectors, ids) res_ids = milvus_instance.insert(entities, ids=ids) assert res_ids == ids milvus_instance.flush() res_count = milvus_instance.count() logger.info("Table: %s, row count: %d" % (collection_name, res_count)) if res_count != len(insert_vectors): raise Exception("Table row count is not equal to insert vectors") for index_type in index_types: for index_param in index_params: logger.debug("Building index with param: %s" % json.dumps(index_param)) if milvus_instance.get_config("cluster.enable") == "true": milvus_instance.create_index(vec_field_name, index_type, metric_type, _async=True, index_param=index_param) else: milvus_instance.create_index(vec_field_name, index_type, metric_type, index_param=index_param) logger.info(milvus_instance.describe_index()) logger.info("Start load collection: %s" % collection_name) milvus_instance.load_collection() logger.info("End load collection: %s" % collection_name) index_info = { "index_type": index_type, "index_param": index_param } logger.debug(index_info) warm_up = True for search_param in search_params: for nq in nqs: query_vectors = self.normalize(metric_type, np.array(dataset["test"][:nq])) if not isinstance(query_vectors, list): query_vectors = query_vectors.tolist() for top_k in top_ks: search_param_group = { "nq": len(query_vectors), "topk": top_k, "search_param": search_param, "metric_type": metric_type } logger.debug(search_param_group) vector_query = {"vector": {vec_field_name: { "topk": top_k, "query": query_vectors, "metric_type": real_metric_type, "params": search_param} }} for i in range(2): result = milvus_instance.query(vector_query) warm_up = False logger.info("End warm up") result = milvus_instance.query(vector_query) result_ids = milvus_instance.get_ids(result) 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, 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.load_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) vector_type = self.get_vector_type(data_type) vec_field_name = utils.get_default_field_name(vector_type) real_metric_type = utils.metric_type_trans(metric_type) 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))) vector_query = {"vector": {vec_field_name: { "topk": top_k, "query": query_vectors[:nq], "metric_type": real_metric_type, "params": search_param} }} milvus_instance.query(vector_query) # 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, } } report(metric) elif run_type == "loop_stability": # init data milvus_instance.clean_db() pull_interval = collection["pull_interval"] collection_num = collection["collection_num"] concurrent = collection["concurrent"] if "concurrent" in collection else False concurrent_num = collection_num dimension = collection["dimension"] if "dimension" in collection else 128 insert_xb = collection["insert_xb"] if "insert_xb" in collection else 100000 index_types = collection["index_types"] if "index_types" in collection else ['ivf_sq8'] index_param = {"nlist": 256} collection_names = [] milvus_instances_map = {} insert_vectors = [[random.random() for _ in range(dimension)] for _ in range(insert_xb)] ids = [i for i in range(insert_xb)] # initialize and prepare for i in range(collection_num): name = utils.get_unique_name(prefix="collection_%d_" % i) collection_names.append(name) metric_type = random.choice(["l2", "ip"]) # default float_vector milvus_instance = MilvusClient(collection_name=name, host=self.host) milvus_instance.create_collection(dimension, other_fields=None) index_type = random.choice(index_types) field_name = utils.get_default_field_name() milvus_instance.create_index(field_name, index_type, metric_type, index_param=index_param) logger.info(milvus_instance.describe_index()) insert_vectors = utils.normalize(metric_type, insert_vectors) entities = milvus_instance.generate_entities(insert_vectors, ids) res_ids = milvus_instance.insert(entities, ids=ids) milvus_instance.flush() milvus_instances_map.update({name: milvus_instance}) logger.info(milvus_instance.describe_index()) # loop time unit: min -> s pull_interval_seconds = pull_interval * 60 tasks = ["insert_rand", "query_rand", "flush"] i = 1 while True: logger.info("Loop time: %d" % i) start_time = time.time() while time.time() - start_time < pull_interval_seconds: if concurrent: threads = [] for name in collection_names: task_name = random.choice(tasks) task_run = getattr(milvus_instances_map[name], task_name) t = threading.Thread(target=task_run, args=()) threads.append(t) t.start() for t in threads: t.join() # with concurrent.futures.ThreadPoolExecutor(max_workers=concurrent_num) as executor: # future_results = {executor.submit(getattr(milvus_instances_map[mp[j][0]], mp[j][1])): j for j in range(concurrent_num)} # for future in concurrent.futures.as_completed(future_results): # future.result() else: tmp_collection_name = random.choice(collection_names) task_name = random.choice(tasks) logger.info(tmp_collection_name) logger.info(task_name) task_run = getattr(milvus_instances_map[tmp_collection_name], task_name) task_run() logger.debug("Restart server") helm_utils.restart_server(self.service_name, namespace) # new connection # for name in collection_names: # milvus_instance = MilvusClient(collection_name=name, host=self.host) # milvus_instances_map.update({name: milvus_instance}) time.sleep(30) i = i + 1 elif run_type == "stability": (data_type, collection_size, dimension, metric_type) = parser.collection_parser( collection_name) during_time = collection["during_time"] operations = collection["operations"] collection_info = { "dimension": dimension, "metric_type": metric_type, "dataset_name": collection_name } if not milvus_instance.exists_collection(): logger.error(milvus_instance.show_collections()) raise Exception("Table name: %s not existed" % collection_name) logger.info(milvus_instance.count()) index_info = milvus_instance.describe_index() logger.info(index_info) # start_mem_usage = milvus_instance.get_mem_info()["memory_used"] start_row_count = milvus_instance.count() logger.info(start_row_count) vector_type = self.get_vector_type(data_type) vec_field_name = utils.get_default_field_name(vector_type) real_metric_type = utils.metric_type_trans(metric_type) query_vectors = [[random.random() for _ in range(dimension)] for _ in range(10000)] if "insert" in operations: insert_xb = operations["insert"]["xb"] if "delete" in operations: delete_xb = operations["delete"]["xb"] if "query" in operations: g_top_k = int(operations["query"]["top_ks"].split("-")[1]) l_top_k = int(operations["query"]["top_ks"].split("-")[0]) g_nq = int(operations["query"]["nqs"].split("-")[1]) l_nq = int(operations["query"]["nqs"].split("-")[0]) search_params = operations["query"]["search_params"] i = 0 start_time = time.time() while time.time() < start_time + during_time * 60: i = i + 1 q = self.gen_executors(operations) for name in q: try: if name == "insert": insert_ids = random.sample(list(range(collection_size)), insert_xb) insert_vectors = [[random.random() for _ in range(dimension)] for _ in range(insert_xb)] entities = milvus_instance.generate_entities(insert_vectors, insert_ids) milvus_instance.insert(entities, ids=insert_ids) elif name == "delete": delete_ids = random.sample(list(range(collection_size)), delete_xb) milvus_instance.delete(delete_ids) elif name == "query": 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))) vector_query = {"vector": {vec_field_name: { "topk": top_k, "query": query_vectors[:nq], "metric_type": real_metric_type, "params": search_param} }} result = milvus_instance.query(vector_query) elif name in ["flush", "compact"]: func = getattr(milvus_instance, name) func() logger.debug(milvus_instance.count()) except Exception as e: logger.error(name) logger.error(str(e)) raise logger.debug("Loop time: %d" % i) # 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, "row_count_increments": end_row_count - start_row_count } } report(metric) elif run_type == "debug": time.sleep(7200) default_insert_vectors = [[random.random() for _ in range(128)] for _ in range(500000)] interval = 50000 for loop in range(1, 7): insert_xb = loop * interval insert_vectors = default_insert_vectors[:insert_xb] insert_ids = [i for i in range(insert_xb)] entities = milvus_instance.generate_entities(insert_vectors, insert_ids) for j in range(5): milvus_instance.insert(entities, ids=insert_ids) time.sleep(10) else: raise Exception("Run type not defined") logger.debug("All test finished")