示例#1
0
class LoadBalanceChecker(Checker):
    """check loadbalance operations in a dependent thread"""
    def __init__(self, collection_name=None):
        super().__init__(collection_name=collection_name)
        self.utility_wrap = ApiUtilityWrapper()
        self.c_wrap.load(enable_traceback=enable_traceback)

    def keep_running(self):
        while True:
            c_name = self.c_wrap.name
            res, _ = self.c_wrap.get_replicas()
            # prepare load balance params
            # find a group which has multi nodes
            group_nodes = []
            for g in res.groups:
                if len(g.group_nodes) >= 2:
                    group_nodes = list(g.group_nodes)
                    break
            src_node_id = group_nodes[0]
            dst_node_ids = group_nodes[1:]
            res, _ = self.utility_wrap.get_query_segment_info(c_name)
            segment_distribution = cf.get_segment_distribution(res)
            sealed_segment_ids = segment_distribution[src_node_id]["sealed"]
            # load balance
            t0 = time.time()
            _, result = self.utility_wrap.load_balance(c_name, src_node_id,
                                                       dst_node_ids,
                                                       sealed_segment_ids)
            t1 = time.time()
            # get segments distribution after load balance
            time.sleep(3)
            res, _ = self.utility_wrap.get_query_segment_info(c_name)
            segment_distribution = cf.get_segment_distribution(res)
            sealed_segment_ids_after_load_banalce = segment_distribution[
                src_node_id]["sealed"]
            check_1 = len(
                set(sealed_segment_ids)
                & set(sealed_segment_ids_after_load_banalce)) == 0
            des_sealed_segment_ids = []
            for des_node_id in dst_node_ids:
                des_sealed_segment_ids += segment_distribution[des_node_id][
                    "sealed"]
            # assert sealed_segment_ids is subset of des_sealed_segment_ids
            check_2 = set(sealed_segment_ids).issubset(
                set(des_sealed_segment_ids))

            if result and (check_1 and check_2):
                self.rsp_times.append(t1 - t0)
                self.average_time = (
                    (t1 - t0) + self.average_time * self._succ) / (self._succ +
                                                                   1)
                self._succ += 1
                log.debug(
                    f"load balance success, time: {t1 - t0:.4f}, average_time: {self.average_time:.4f}"
                )
            else:
                self._fail += 1
            sleep(10)
示例#2
0
文件: checker.py 项目: avmi/milvus
class LoadBalanceChecker(Checker):
    """check loadbalance operations in a dependent thread"""

    def __init__(self, collection_name=None):
        if collection_name is None:
            collection_name = cf.gen_unique_str("LoadBalanceChecker_")
        super().__init__(collection_name=collection_name)
        self.utility_wrap = ApiUtilityWrapper()
        self.c_wrap.load()
        self.sealed_segment_ids = None
        self.dst_node_ids = None
        self.src_node_id = None

    @trace()
    def load_balance(self):
        res, result = self.utility_wrap.load_balance(
            self.c_wrap.name, self.src_node_id, self.dst_node_ids, self.sealed_segment_ids)
        return res, result

    def prepare(self):
        """prepare load balance params"""
        res, _ = self.c_wrap.get_replicas()
        # find a group which has multi nodes
        group_nodes = []
        for g in res.groups:
            if len(g.group_nodes) >= 2:
                group_nodes = list(g.group_nodes)
                break
        self.src_node_id = group_nodes[0]
        self.dst_node_ids = group_nodes[1:]
        res, _ = self.utility_wrap.get_query_segment_info(self.c_wrap.name)
        segment_distribution = cf.get_segment_distribution(res)
        self.sealed_segment_ids = segment_distribution[self.src_node_id]["sealed"]

    @exception_handler()
    def run_task(self):
        self.prepare()
        res, result = self.load_balance()
        return res, result

    def keep_running(self):
        while self._keep_running:
            self.run_task()
            sleep(constants.WAIT_PER_OP / 10)