mirror of https://github.com/milvus-io/milvus.git
107 lines
4.2 KiB
Python
107 lines
4.2 KiB
Python
import time
|
|
import copy
|
|
import logging
|
|
from milvus_benchmark import parser
|
|
from milvus_benchmark.runners import utils
|
|
from milvus_benchmark.runners.base import BaseRunner
|
|
|
|
logger = logging.getLogger("milvus_benchmark.runners.build")
|
|
|
|
|
|
class BuildRunner(BaseRunner):
|
|
"""run build"""
|
|
name = "build_performance"
|
|
|
|
def __init__(self, env, metric):
|
|
super(BuildRunner, self).__init__(env, metric)
|
|
|
|
def extract_cases(self, collection):
|
|
collection_name = collection["collection_name"] if "collection_name" in collection else None
|
|
(data_type, collection_size, dimension, metric_type) = parser.collection_parser(collection_name)
|
|
ni_per = collection["ni_per"]
|
|
vector_type = utils.get_vector_type(data_type)
|
|
other_fields = collection["other_fields"] if "other_fields" in collection else None
|
|
collection_info = {
|
|
"dimension": dimension,
|
|
"metric_type": metric_type,
|
|
"dataset_name": collection_name,
|
|
"collection_size": collection_size,
|
|
"other_fields": other_fields,
|
|
"ni_per": ni_per
|
|
}
|
|
index_field_name = utils.get_default_field_name(vector_type)
|
|
index_type = collection["index_type"]
|
|
index_param = collection["index_param"]
|
|
index_info = {
|
|
"index_type": index_type,
|
|
"index_param": index_param
|
|
}
|
|
flush = True
|
|
if "flush" in collection and collection["flush"] == "no":
|
|
flush = False
|
|
self.init_metric(self.name, collection_info, index_info, search_info=None)
|
|
case_metric = copy.deepcopy(self.metric)
|
|
case_metric.set_case_metric_type()
|
|
case_metrics = list()
|
|
case_params = list()
|
|
case_metrics.append(case_metric)
|
|
case_param = {
|
|
"collection_name": collection_name,
|
|
"data_type": data_type,
|
|
"dimension": dimension,
|
|
"collection_size": collection_size,
|
|
"ni_per": ni_per,
|
|
"metric_type": metric_type,
|
|
"vector_type": vector_type,
|
|
"other_fields": other_fields,
|
|
"flush_after_insert": flush,
|
|
"index_field_name": index_field_name,
|
|
"index_type": index_type,
|
|
"index_param": index_param,
|
|
}
|
|
case_params.append(case_param)
|
|
return case_params, case_metrics
|
|
|
|
def prepare(self, **case_param):
|
|
collection_name = case_param["collection_name"]
|
|
self.milvus.set_collection(collection_name)
|
|
if not self.milvus.exists_collection():
|
|
logger.info("collection not exist")
|
|
logger.debug({"collection count": self.milvus.count()})
|
|
|
|
def run_case(self, case_metric, **case_param):
|
|
index_field_name = case_param["index_field_name"]
|
|
start_time = time.time()
|
|
self.milvus.create_index(index_field_name, case_param["index_type"], case_param["metric_type"],
|
|
index_param=case_param["index_param"])
|
|
build_time = round(time.time() - start_time, 2)
|
|
tmp_result = {"build_time": build_time}
|
|
return tmp_result
|
|
|
|
|
|
class InsertBuildRunner(BuildRunner):
|
|
"""run insert and build"""
|
|
name = "insert_build_performance"
|
|
|
|
def __init__(self, env, metric):
|
|
super(InsertBuildRunner, self).__init__(env, metric)
|
|
|
|
def prepare(self, **case_param):
|
|
collection_name = case_param["collection_name"]
|
|
dimension = case_param["dimension"]
|
|
vector_type = case_param["vector_type"]
|
|
other_fields = case_param["other_fields"]
|
|
self.milvus.set_collection(collection_name)
|
|
if self.milvus.exists_collection():
|
|
logger.debug("Start drop collection")
|
|
self.milvus.drop()
|
|
time.sleep(utils.DELETE_INTERVAL_TIME)
|
|
self.milvus.create_collection(dimension, data_type=vector_type, other_fields=other_fields)
|
|
self.insert(self.milvus, collection_name, case_param["data_type"], dimension,
|
|
case_param["collection_size"], case_param["ni_per"])
|
|
start_time = time.time()
|
|
self.milvus.flush()
|
|
flush_time = round(time.time() - start_time, 2)
|
|
logger.debug({"collection count": self.milvus.count()})
|
|
logger.debug({"flush_time": flush_time})
|