mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-05 05:18:52 +08:00
127 lines
5.2 KiB
Python
127 lines
5.2 KiB
Python
import copy
|
|
import logging
|
|
import time
|
|
from operator import methodcaller
|
|
from yaml import full_load, dump
|
|
import threading
|
|
from milvus_benchmark import utils
|
|
from milvus_benchmark.runners import utils as runner_utils
|
|
from milvus_benchmark.chaos import utils as chaos_utils
|
|
from milvus_benchmark.runners.base import BaseRunner
|
|
from chaos.chaos_opt import ChaosOpt
|
|
from milvus_benchmark import config
|
|
from milvus_benchmark.chaos.chaos_mesh import PodChaos, NetworkChaos
|
|
|
|
logger = logging.getLogger("milvus_benchmark.runners.chaos")
|
|
|
|
kind_chaos_mapping = {
|
|
"PodChaos": PodChaos,
|
|
"NetworkChaos": NetworkChaos
|
|
}
|
|
|
|
assert_func_mapping = {
|
|
"fail": chaos_utils.assert_fail,
|
|
"pass": chaos_utils.assert_pass
|
|
}
|
|
|
|
|
|
class SimpleChaosRunner(BaseRunner):
|
|
"""run chaos"""
|
|
name = "simple_chaos"
|
|
|
|
def __init__(self, env, metric):
|
|
super(SimpleChaosRunner, self).__init__(env, metric)
|
|
|
|
async def async_call(self, func, **kwargs):
|
|
future = methodcaller(func, **kwargs)(self.milvus)
|
|
|
|
def run_step(self, interface_name, interface_params):
|
|
if interface_name == "create_collection":
|
|
collection_name = utils.get_unique_name("chaos")
|
|
self.data_type = interface_params["data_type"]
|
|
self.dimension = interface_params["dimension"]
|
|
self.milvus.set_collection(collection_name)
|
|
vector_type = runner_utils.get_vector_type(self.data_type)
|
|
self.milvus.create_collection(self.dimension, data_type=vector_type)
|
|
elif interface_name == "insert":
|
|
batch_size = interface_params["batch_size"]
|
|
collection_size = interface_params["collection_size"]
|
|
self.insert(self.milvus, self.milvus.collection_name, self.data_type, self.dimension, collection_size,
|
|
batch_size)
|
|
elif interface_name == "create_index":
|
|
metric_type = interface_params["metric_type"]
|
|
index_type = interface_params["index_type"]
|
|
index_param = interface_params["index_param"]
|
|
vector_type = runner_utils.get_vector_type(self.data_type)
|
|
field_name = runner_utils.get_default_field_name(vector_type)
|
|
self.milvus.create_index(field_name, index_type, metric_type, index_param=index_param)
|
|
elif interface_name == "flush":
|
|
self.milvus.flush()
|
|
|
|
def extract_cases(self, collection):
|
|
before_steps = collection["before"]
|
|
after = collection["after"] if "after" in collection else None
|
|
processing = collection["processing"]
|
|
case_metrics = []
|
|
case_params = [{
|
|
"before_steps": before_steps,
|
|
"after": after,
|
|
"processing": processing
|
|
}]
|
|
self.init_metric(self.name, {}, {}, None)
|
|
case_metric = copy.deepcopy(self.metric)
|
|
case_metric.set_case_metric_type()
|
|
case_metrics.append(case_metric)
|
|
return case_params, case_metrics
|
|
|
|
def prepare(self, **case_param):
|
|
steps = case_param["before_steps"]
|
|
for step in steps:
|
|
interface_name = step["interface_name"]
|
|
params = step["params"]
|
|
self.run_step(interface_name, params)
|
|
|
|
def run_case(self, case_metric, **case_param):
|
|
processing = case_param["processing"]
|
|
after = case_param["after"]
|
|
user_chaos = processing["chaos"]
|
|
kind = user_chaos["kind"]
|
|
spec = user_chaos["spec"]
|
|
metadata_name = config.NAMESPACE + "-" + kind.lower()
|
|
metadata = {"name": metadata_name}
|
|
process_assertion = processing["assertion"]
|
|
after_assertion = after["assertion"]
|
|
# load yaml from default template to generate stand chaos dict
|
|
chaos_mesh = kind_chaos_mapping[kind](config.DEFAULT_API_VERSION, kind, metadata, spec)
|
|
experiment_config = chaos_mesh.gen_experiment_config()
|
|
process_func = processing["interface_name"]
|
|
process_params = processing["params"] if "params" in processing else {}
|
|
after_func = after["interface_name"]
|
|
after_params = after["params"] if "params" in after else {}
|
|
logger.debug(chaos_mesh.kind)
|
|
chaos_opt = ChaosOpt(chaos_mesh.kind)
|
|
chaos_objects = chaos_opt.list_chaos_object()
|
|
if len(chaos_objects["items"]) != 0:
|
|
logger.debug(chaos_objects["items"])
|
|
chaos_opt.delete_chaos_object(chaos_mesh.metadata["name"])
|
|
# with open('./pod-newq.yaml', "w") as f:
|
|
# dump(experiment_config, f)
|
|
# f.close()
|
|
# concurrent inject chaos and run func
|
|
# logger.debug(experiment_config)
|
|
t_milvus = threading.Thread(target=assert_func_mapping[process_assertion], args=(process_func, self.milvus,), kwargs=process_params)
|
|
try:
|
|
t_milvus.start()
|
|
chaos_opt.create_chaos_object(experiment_config)
|
|
# processing assert exception
|
|
except Exception as e:
|
|
logger.info("exception {}".format(str(e)))
|
|
else:
|
|
chaos_opt.delete_chaos_object(chaos_mesh.metadata["name"])
|
|
# TODO retry connect milvus
|
|
time.sleep(15)
|
|
assert_func_mapping[after_assertion](after_func, self.milvus, **after_params)
|
|
finally:
|
|
chaos_opt.delete_all_chaos_object()
|
|
logger.info(chaos_opt.list_chaos_object())
|