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)
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)