optimize nebula startup time

Author: camaze

.gitignore

@@ -19,7 +19,7 @@ zdatafront*
 *antgroup*
 *ipynb
 *log
-
+examples/mysql/db
 
 # frontend
 frontend/node_modules

muagent/db_handler/graph_db_handler/nebula_handler.py

@@ -63,13 +63,36 @@ def __init__(self,gb_config : GBConfig = None):
  self.nb_pw = '' or 'nebula'
  self.space_name = "client"
  else:
- logger.info('NebulaGraph容器启动中，等待20秒')
- time.sleep(20)
- self.connection_pool.init([(gb_config.extra_kwargs.get("host"), gb_config.extra_kwargs.get("port"))], config)
+ logger.info('NebulaGraph容器启动中，请等待')
+ 
+ if self.nebula_started(gb_config):
+ self.connection_pool.init([(gb_config.extra_kwargs.get("host"), gb_config.extra_kwargs.get("port"))], config)
+ 
  self.username = gb_config.extra_kwargs.get("username")
  self.nb_pw = gb_config.extra_kwargs.get("password")
  self.space_name = gb_config.extra_kwargs.get("space")
 
+ def nebula_started(self, gb_config: GBConfig):
+ '''
+ Continuously ping the nebula server until it has started
+ If it does not start within 40 seconds, raise an exception
+ '''
+ # 80 attempts at 0.5 seconds each = 40 seconds
+ max_attempts = 80 
+
+ host = gb_config.extra_kwargs.get('host')
+ port = gb_config.extra_kwargs.get('port')
+
+ for attempt in range(max_attempts):
+ if self.connection_pool.ping([host, port]):
+ logger.info(f"nebula server is ok, total waiting time: {attempt * 0.5}s")
+ return True
+ 
+ logger.info(f"ping nebula server.. waiting time: {attempt * 0.5}s")
+ time.sleep(0.5)
+
+ raise Exception('The nebula server did not start within 40 seconds.')
+
  def execute_cypher(self, cypher: str, space_name: str = '',ignore_log: bool = False, format_res: str = 'as_primitive', use_space_name: bool = True):
  '''
  @param space_name: space_name, if provided, will execute use space_name first
@@ -172,6 +195,11 @@ def show_space(self):
  def drop_space(self, space_name):
  cypher = f'DROP SPACE {space_name}'
  return self.execute_cypher(cypher)
+ 
+ def show_hosts(self):
+ cypher = 'SHOW HOSTS'
+ resp = self.execute_cypher(cypher, use_space_name=False)
+ return resp
 
  def create_tag(self, tag_name: str, prop_dict: dict = {}):
  '''

muagent/service/ekg_construct/ekg_construct_base.py

@@ -180,20 +180,19 @@ def init_gb(self, do_init: bool=None):
  initialize_space = self.initialize_space # True or False
  if initialize_space and self.gb_config.gb_type=="NebulaHandler":
  self.gb.add_hosts('storaged0', 9779)
- print('增加NebulaGraph Storage主机中，等待20秒')
- time.sleep(20)
+ self.waiting_host_initialize()
 
  if self.clear_history_data:
  # 初始化space
  self.gb.drop_space('client')
 
  self.gb.create_space('client')
+ self.waiting_space_initialize()
 
  # 创建node tags和edge types
  self.create_gb_tags_and_edgetypes()
+ self.waiting_tags_edgetypes_initialize()
 
- print('Node Tags和Edge Types初始化中，等待20秒......')
- time.sleep(20)
  else:
  self.gb = None
 
@@ -223,7 +222,54 @@ def init_sls(self, do_init: bool=None):
  self.embed_config, vb_config=self.vb_config)
  else:
  self.sls = None
+
+ def wait_for_condition(self, condition, timeout=40, wait_time=0.5, log_prefix=''):
+ max_attempts = int(timeout / wait_time)
+
+ total_time = 0
+
+ for _ in range(max_attempts):
+ if condition():
+ logger.info(f'[{log_prefix}] initialized successfully, total waiting time: {total_time}s')
+ return
+ 
+ logger.info(f'[{log_prefix}] have waited for {total_time}s...')
+ total_time += wait_time
+ time.sleep(wait_time)
+
+ raise Exception(f'[{log_prefix}] not initialized within {timeout} seconds')
 
+ def waiting_host_initialize(self):
+ self.wait_for_condition(
+ condition=lambda: any(
+ host.get('Host') == 'storaged0' and host.get('Status') == 'ONLINE'
+ for host in self.gb.show_hosts()
+ ),
+ log_prefix='host_initializing'
+ )
+
+ def waiting_tags_edgetypes_initialize(self):
+ '''
+ Make sure that tags and edgetypes are created properly
+ Refer to create_gb_tags_and_edgetypes for more details
+ '''
+ self.wait_for_condition(
+ condition=lambda: (
+ len(self.gb.show_tags()) == len(TYPE2SCHEMA.items()) - 1 and
+ len(self.gb.show_edge_type()) == (len(TYPE2SCHEMA.items()) - 1) ** 2 * 3
+ ),
+ log_prefix='tags_edgetypes_initializing'
+ )
+
+ def waiting_space_initialize(self):
+ self.wait_for_condition(
+ condition=lambda:any(
+ space.get('Name') == 'client'
+ for space in self.gb.show_space()
+ ),
+ log_prefix='space_initializing'
+ )
+
  def _get_local_graph(
  self, nodes: List[GNode], edges: List[GEdge], rootid
  ) -> Tuple[List[str], Graph]: