from __future__ import annotations import time from dataclasses import dataclass import datetime import yaml import os from .preprocess.task import DefaultTask as PreProcessTask from .preprocess.config import TaskConfig as PreProcessTaskConfig from .diagnosis.diagnosis import Diagnosis as DiagnosisTask from .diagnosis.config import TaskConfig as DiagnosisTaskConfig from ..dataloader.localloader import LocalFileLoader from ..dataloader.s3loader import * from ..utils.utils import * class EngineConfig: def __init__(self, engine_id=None, cache_path=None, o11y_auth0_client_id=None, o11y_auth0_client_secret=None, clinic_token=None, bucket_name=None): self.engine_id = engine_id if engine_id is not None else os.environ["ENGINE_ID"] self.cache_path = cache_path if cache_path is not None else os.environ["CACHE_PATH"] self.o11y_auth0_client_id = o11y_auth0_client_id if o11y_auth0_client_id is not None else os.environ[ "O11Y_AUTH0_CLIENT_ID"] self.o11y_auth0_client_secret = o11y_auth0_client_secret if o11y_auth0_client_secret is not None else os.environ[ "O11Y_AUTH0_CLIENT_SECRET"] self.clinic_token = clinic_token if clinic_token is not None else os.environ[ "CLINIC_TOKEN"] self.bucket_name = bucket_name if bucket_name is not None else os.environ["BUCKET_NAME"] self.run_once = os.environ["RUN_ONCE"] == 'True' if "RUN_ONCE" in os.environ else False self.skip_preprocess = os.environ["SKIP_PREPROCESS"] == 'True' if "SKIP_PREPROCESS" in os.environ else False self.fetch_timestemp = int( os.environ["FETCH_TIMESTEMP"]) if "FETCH_TIMESTEMP" in os.environ else None self.skip_slowlog = os.environ["SKIP_SLOWLOG"] == 'True' if "SKIP_SLOWLOG" in os.environ else False self.checkpoint_time = int(os.environ["CHECKPOINT"]) if "CHECKPOINT" in os.environ else None self.cache_save_interval_day = int(os.environ["CACHE_SAVE_INTERVAL_DAY"]) if "CACHE_SAVE_INTERVAL_DAY" in os.environ else 7 self.catalog_name = os.environ["CATALOG_NAME"] if "CATALOG_NAME" in os.environ else "drhouse/diag-catalog/catalog.yaml" @dataclass class PreProcessTaskInfo: task_id: str task_success: bool task_retry_times: int task_start_time: int task_end_time: int task_day_start: int task_day_end: int skip_tiflash: bool task_info: dict @dataclass class DiagnosisTaskInfo: task_id: str task_info: dict task_info_org: dict task_success: bool task_retry_times: int task_day_start: int task_day_end: int checkpoint_time:datetime.datetime class Manager: def __init__(self, remoteloader, config: EngineConfig) -> None: self.remoteloader = remoteloader self.config = config def load_checkpoint(self): file_name = "{}-checkpoint.yaml".format(self.config.engine_id) file_path = os.path.join(get_checkpoint_prefix(), file_name) # if chekpoint_time is not None, return chekpoint_time if self.config.checkpoint_time is not None: return datetime.datetime.fromtimestamp(self.config.checkpoint_time) # check if checkpoint file exists, if not, load from s3, if not, create a new one cache = LocalFileLoader(self.config.cache_path) if not cache.file_exists(file_path): if not self.remoteloader.file_exists(file_path): # create a new checkpoint file with tomorrow's midnight midnight = self.get_tomorrow_midnight(datetime.date.today()) checkpoint = {"checkpoint_time": int(midnight.timestamp())} cache.save_file(file_path, yaml.dump(checkpoint)) # upload checkpoint to s3 file_full_path_name = os.path.join(self.config.cache_path, file_path) key = os.path.join(get_checkpoint_prefix(), file_name) self.remoteloader.upload_file(file_full_path_name, key) else: f = self.remoteloader.get_file(file_path) checkpoint = yaml.safe_load(f) cache.save_file(file_path, yaml.dump(checkpoint)) # load checkpoint from local file f = cache.get_file(file_path) checkpoint = yaml.safe_load(f) return datetime.datetime.fromtimestamp(checkpoint["checkpoint_time"]) def save_checkpoint(self, checkpoint: datetime.datetime): # save checkpoint to local file file_name = "{}-checkpoint.yaml".format(self.config.engine_id) file_path = os.path.join(get_checkpoint_prefix(), file_name) cache = LocalFileLoader(self.config.cache_path) content = {"checkpoint_time": int(checkpoint.timestamp())} cache.save_file(file_path, yaml.dump(content)) # upload checkpoint to s3 file_full_path_name = os.path.join(self.config.cache_path, file_path) if os.path.exists(file_full_path_name): key = os.path.join(get_checkpoint_prefix(), file_name) self.remoteloader.upload_file(file_full_path_name, key) return def clean_up_cache(self): # skip,if run_once is True if self.config.run_once: return # clean up cache file modify time older than cache_save_interval_day cache = LocalFileLoader(self.config.cache_path) file_list = cache.get_file_list("./") last_time = int((datetime.datetime.now() - datetime.timedelta(days=self.config.cache_save_interval_day)).timestamp()) for file in file_list: # check file last modify time and delete it if it is older than cache_save_interval_day if os.path.getmtime(file) < last_time: print("delete file: ", file) os.remove(file) def do_recover_preprocess_task(self, task, task_info: dict, checkpoint_time: datetime.datetime,midnight:datetime.datetime): # if checkpoint_time equal midnight，then we check checkpoint_time-8 to checkpoint_time-1 # if checkpoint_time less than midnight，then we check checkpoint_time-8 to checkpoint_time # if create_time no longger than checkpoint_time-8,then we check create_time to checkpoint_time or -1 start_time = checkpoint_time - datetime.timedelta(days=8) create_time = task_info["create_time"] if create_time > start_time: _,_,create_day = self.get_timestemp(create_time) start_time = datetime.datetime.fromtimestamp(create_day) if checkpoint_time == midnight: end_time = checkpoint_time - datetime.timedelta(days=1) else: end_time = checkpoint_time # convert start_time and end_time to day_start and day_end,each task's day_start and day_end length is 1 day while start_time <= end_time: _,ck_start,ck_end = self.get_timestemp(start_time) task_id = self.get_cluster_id(task) meta_info = task_info['metas'] skip_tiflash = True if self.pick_topology( meta_info, "tiflash") is None else False pretask=PreProcessTaskInfo( task_id, False, 0, ck_start, ck_end, ck_start, ck_end, skip_tiflash, meta_info) success=self._run_preprocess_task(pretask) if not success: return False start_time = start_time + datetime.timedelta(days=1) return True def do_recover_diagnosis_task(self, task,task_info: dict, checkpoint_time: datetime.datetime,midnight:datetime.datetime): # if checkpoint_time equal midnight, return # if checkpoint_time less than midnight，do checkpoint_time diagnosis task if checkpoint_time == midnight: return True create_time = task_info["create_time"] day7_time = datetime.datetime.now() - datetime.timedelta(days=7) if create_time > day7_time: return True # convert start_time and end_time to day_start and day_end,each task's day_start and day_end length is 1 day _,ck_start,ck_end = self.get_timestemp(checkpoint_time) task_id = self.get_cluster_id(task) task_meta = task_info['metas'] diagtask = DiagnosisTaskInfo( task_id, task_meta, task_info,False, 0, ck_start, ck_end,checkpoint_time) success=self._run_diagnosis_task(diagtask) if not success: print("do recover diagnosis task failed {} checkpoint_time: {}".format(task,checkpoint_time.timestamp())) return False return True def do_recover_task(self,checkpoint_time:datetime.datetime): # get lasted checkpoint time midnight= self.get_tomorrow_midnight(datetime.date.today()) #if checkpoint set by user, we do not need to recover if self.config.checkpoint_time is not None: print("checkpoint time is set by user, we do not need to recover") return failed_task_list = [] while True: # get all task list which need to be recover task_list, task_infos,success = self.get_task_list() if not success: print("sleep 5 min for next process") time.sleep(300) continue # we should do once task first for each task for task in task_list: print("do recover task: ", task) # do recover preprocess task preprocess_success=self.do_recover_preprocess_task(task, task_infos[task],checkpoint_time, midnight) if not preprocess_success: print("do recover preprocess task failed") failed_task_list.append(task) continue # do recover diagnosis task diagnosis_success=self.do_recover_diagnosis_task(task,task_infos[task],checkpoint_time, midnight) if not diagnosis_success: print("do recover diagnosis task failed") failed_task_list.append(task) continue # if failed task list is not empty, we should do it again if len(failed_task_list) != 0: failed_task_list = [] # sleep 1min for next process print("sleep 1 min for next process") time.sleep(60) continue # if checkpoint time is now time, break if midnight == checkpoint_time: break # update checkpoint time checkpoint_time = checkpoint_time + datetime.timedelta(days=1) self.save_checkpoint(checkpoint_time) def start(self): preprocess_task = {} diagnosis_task = [] # we should save this checkpoint and load it when restart checkpoint_time = self.load_checkpoint() # if checkpoint to now not have repoter we should do recover task self.do_recover_task(checkpoint_time) while True: # sleep 15 min if not self.config.run_once: print("now sleep 15 min for next process") time.sleep(15 * 60) # get now time now_time, today_start_timestamp, today_end_timestamp = self.get_timestemp() if self.config.fetch_timestemp is None else self.get_timestemp(datetime.datetime.fromtimestamp(self.config.fetch_timestemp)) # get s3 task list task_list, task_infos,task_list_success = self.get_task_list() if not task_list_success: print("get task list failed sleep and wait next process") continue if not self.config.skip_preprocess: # build preprocess task preprocess_task = self.build_preprocess_task( task_list, task_infos, preprocess_task, now_time, today_start_timestamp, today_end_timestamp) # run preprocess task tasks, success = self.run_preprocess_tasks(preprocess_task) if not success: continue preprocess_task = tasks # do diagnosis task if now_time > checkpoint_time if now_time < int(time.mktime(checkpoint_time.timetuple())): continue # get checkpoint ts _, ck_start, ck_end = self.get_timestemp(checkpoint_time) # build diagnosis task tasks = self.build_diagnosis_task( task_list, task_infos, diagnosis_task, ck_start, ck_end,checkpoint_time) # run diagnosis task failed_tasks, success = self.run_diagnosis_tasks(tasks) if not success: diagnosis_task = failed_tasks continue if self.config.run_once: break # update checkpoint checkpoint_time = datetime.datetime.fromtimestamp(today_end_timestamp) self.save_checkpoint(checkpoint_time) self.clean_up_cache() def get_timestemp(self, checkpoint_time: datetime.datetime = None): # get now time if checkpoint_time is None: now_time = datetime.datetime.now() # get 0 timestemp today_start_time = now_time.replace( hour=0, minute=0, second=0, microsecond=0) # get 24 timestemp today_end_time = (now_time + datetime.timedelta(days=1) ).replace(hour=0, minute=0, second=0, microsecond=0) else: # checkpoint time alway one day's 0:0:0 so it's one days' end time now_time = checkpoint_time today_end_time = checkpoint_time today_start_time = checkpoint_time - datetime.timedelta(days=1) # to timestamp today_start_ts = int(time.mktime(today_start_time.timetuple())) today_end_ts = int(time.mktime(today_end_time.timetuple())) now_time_ts = int(time.mktime(now_time.timetuple())) return now_time_ts, today_start_ts, today_end_ts def get_task_list(self): try: # get catalog yaml file catalog = self.remoteloader.get_file(self.config.catalog_name) catalog = yaml.load(catalog, yaml.FullLoader) # get task list which need to be executed if self.config.engine_id not in catalog['tasks']: return [], {},True task_list = catalog['tasks'][self.config.engine_id] task_infos = {} for task in task_list: # load task info from yamlfile task_info = self.remoteloader.get_file(task) tmp = yaml.load(task_info, yaml.FullLoader) # check task list, if cluster creation less than 7 days, remove it from task list create_time = datetime.datetime.strptime(tmp['metas']['createTime'], "%Y-%m-%d %H:%M:%S") tmp["create_time"]=create_time task_infos[task] = tmp return task_list, task_infos,True except Exception as e: # we should return false if we can't get task list print("get task list failed: " + str(e)) return [], {},False def build_preprocess_task(self, task_list: list[str], task_infos: dict, pre_task_dict: dict[str, PreProcessTaskInfo | list[PreProcessTaskInfo]], now_time, today_start_timestamp, today_end_timestamp): task_dict = {} task_dict['extend'] = [] for task in task_list: task_id = self.get_cluster_id(task) meta_info = task_infos[task]['metas'] skip_tiflash = True if self.pick_topology( meta_info, "tiflash") is None else False old = pre_task_dict.get(task_id) if old is None: task_dict[task_id] = PreProcessTaskInfo(task_id, False, 0, today_start_timestamp, now_time, today_start_timestamp, today_end_timestamp, skip_tiflash, meta_info) else: if now_time > old.task_day_end: task_dict[task_id] = PreProcessTaskInfo( task_id, False, old.task_retry_times, today_start_timestamp, now_time, today_start_timestamp, today_end_timestamp, skip_tiflash, meta_info) task_dict["extend"].append(PreProcessTaskInfo( task_id, False, old.task_retry_times, old.task_day_start, old.task_day_end, old.task_day_start, old.task_day_end, skip_tiflash, meta_info)) else: task_dict[task_id] = PreProcessTaskInfo( task_id, False, old.task_retry_times, old.task_day_start, now_time, old.task_day_start, old.task_day_end, skip_tiflash, meta_info) for extend_task in pre_task_dict["extend"] if pre_task_dict.get("extend") else []: # retry extend task if not sucess,and ensure all in task_list if not extend_task.task_success and extend_task.task_id in task_list: task_dict['extend'].append(extend_task) return task_dict def run_preprocess_tasks(self, task_dict: dict[str, PreProcessTaskInfo | list[PreProcessTaskInfo]]): all_success = True for task_id, task_info in task_dict.items(): # skip extend task here,extend task will be run in next_loop if task_id == "extend": continue success = self._run_preprocess_task(task_info) if not success: all_success = False extend_task_list = task_dict.get("extend") for extend_task in extend_task_list: success = self._run_preprocess_task(extend_task) if not success: all_success = False return task_dict, all_success def _run_preprocess_task(self, task: PreProcessTaskInfo): try: PreProcessTask(PreProcessTaskConfig(task.task_id, self.config.o11y_auth0_client_id, self.config.o11y_auth0_client_secret, self.config.clinic_token, task.task_start_time, task.task_end_time, self.config.cache_path, self.build_metrics_sqlite_name(task.task_id, task.task_day_start, task.task_day_end), self.build_slowlog_sqlite_name(task.task_id, task.task_day_start, task.task_day_end), self.build_events_sqlite_name(task.task_id, task.task_day_start, task.task_day_end), task.skip_tiflash, self.config.skip_slowlog, task.task_info), self.remoteloader).execute() # set task success when task finished task.task_retry_times = 0 task.task_success = True except Exception as e: print("run_preprocess_task failed: " + str(e)) # set task retry times when task failed task.task_retry_times += 1 task.task_success = False return task.task_success def build_diagnosis_task(self, task_list: list[str], task_infos: dict, failed_task_dict: dict[str, DiagnosisTaskInfo], today_start_timestamp, today_end_timestamp,checkpoint_time:datetime.datetime): task_dict = {} if failed_task_dict is not None and len(failed_task_dict) > 0: # if failed task dict is not None,add failed task to task dict task_dict.update(failed_task_dict) else: for task in task_list: task_id = self.get_cluster_id(task) task_meta = task_infos[task]['metas'] create_time = task_infos[task]['create_time'] day7_time = datetime.datetime.now() - datetime.timedelta(days=7) if create_time > day7_time: print("cluster create time less than 7 days, skip diagnosis task: " + task_id) continue task_dict[task_id] = DiagnosisTaskInfo( task_id, task_meta,task_infos[task], False, 0, today_start_timestamp, today_end_timestamp,checkpoint_time) return task_dict def run_diagnosis_tasks(self, task_dict: dict[str, DiagnosisTaskInfo]): all_success = True failed_task = {} for task_id, task_info in task_dict.items(): # we should make sure we have 7days data before run diagnosis task midnight= self.get_tomorrow_midnight(datetime.date.today()) if not self.config.run_once: preprocess_success=self.do_recover_preprocess_task(task_id,task_info.task_info_org,task_info.checkpoint_time,midnight) if not preprocess_success: print("do_recover_preprocess_task before diagnosis failed: " + task_id) all_success = False failed_task[task_id] = task_info continue success = self._run_diagnosis_task(task_info) if not success: all_success = False failed_task[task_id] = task_info return failed_task, all_success def _run_diagnosis_task(self, task: DiagnosisTaskInfo): try: DiagnosisTask(DiagnosisTaskConfig(task.task_id, task.task_day_start, task.task_day_end, self.config.cache_path, self.config.bucket_name, task.task_info), self.remoteloader).execute() # set task success when task finished task.task_success = True except Exception as e: # set task retry times when task failed print('get error when diagnosis', e) task.task_retry_times += 1 task.task_success = False return task.task_success def get_cluster_id(self, file_path: str): # get cluster id from file path example: drhouse/diag-task/cluster_id.yaml cluster_id = file_path.split('/')[-1].split('.')[0] return cluster_id def build_metrics_sqlite_name(self, task_id, start_ts, end_ts): # build metrics sqlite name from task id return "{}_metrics_{}_{}.sqlite".format(task_id, start_ts, end_ts) def build_slowlog_sqlite_name(self, task_id, start_ts, end_ts): # build slowlog sqlite name from task id return "{}_slowlog_{}_{}.sqlite".format(task_id, start_ts, end_ts) def build_events_sqlite_name(self, task_id, start_ts, end_ts): # build slowlog sqlite name from task id return "{}_events_{}_{}.sqlite".format(task_id, start_ts, end_ts) def pick_topology(self, meta_info: dict, name: str): for item in meta_info['topology']: if item['name'] == name: return item return None def get_tomorrow_midnight(self,today:datetime.date): tomorrow = today + datetime.timedelta(days=1) midnight = datetime.datetime(tomorrow.year, tomorrow.month, tomorrow.day) return midnight