import os import time import yaml from o11ydiagnosis.engine.log import log from o11ydiagnosis.subtasks import slowlogconverter from datasdk.utils import o11y from datasdk.config import config as o11yConfig data_prefix_template = "diagnosis/data/{}/" checkpoint_prefix_template = "diagnosis/checkpoints-deltalake/{}/" running_tasks = {} class Task: PROVIDER_AWS = 'aws' PROVIDER_GCP = 'gcp' def __init__(self, o11y_api, task_info, cache_prefix, remote_loader): self.cluster_id = task_info["meta"]["clusterID"] self.project_id = task_info["meta"]["projectID"] self.tenant_id = task_info["meta"]["tenantID"] self.o11y_cluster_id = task_info["meta"]["o11yClusterID"] self.serverless = task_info["meta"]["clusterType"].lower() == "serverless" self.interval = task_info["interval"] self.provider = task_info["meta"]["provider"] self.region = task_info["meta"]["region"] self.env = task_info["meta"]["env"] self.cache_prefix = cache_prefix self.remote_prefix = data_prefix_template.format(self.cluster_id) self.remote_loader = remote_loader self.o11y_api = o11y_api self.subtasks = [] for subtask_info in task_info["subtasks"]: if subtask_info["type"] == "slowlogconverter": self.subtasks.append(slowlogconverter.Config(self, task_info, subtask_info)) def pull(self, filelist): result = [] for file in filelist: if self.remote_loader.file_exists(file): content = self.remote_loader.get_file(file) fullpath = os.path.join(self.prefix, file) if not os.path.exists(os.path.dirname(fullpath)): os.makedirs(os.path.dirname(fullpath)) with open(fullpath, 'w') as f: f.write(content) result.append(True) else: result.append(False) return result def push(self, filelist): result = [] for file in filelist: cache_file_path = os.path.join(self.cache_prefix, file) remote_file_name = os.path.join(self.remote_prefix, file) self.remote_loader.upload_file(cache_file_path, remote_file_name) result.append(True) return result def exist_cache(self, filename): fullpath = self.get_cache_full_path(filename) return os.path.exists(fullpath) def get_cache_full_path(self, filename): return os.path.join(self.cache_prefix, filename) def create_cache(self, filename): fullpath = os.path.join(self.cache_prefix, filename) if not os.path.exists(os.path.dirname(fullpath)): os.makedirs(os.path.dirname(fullpath)) def save_cache(self, filename, content): fullpath = os.path.join(self.cache_prefix, filename) if not os.path.exists(os.path.dirname(fullpath)): os.makedirs(os.path.dirname(fullpath)) with open(fullpath, 'w') as f: f.write(content) class Manager: def __init__(self, engine_config, task_file: str) -> None: self.o11y_config = o11yConfig.Config(engine_config.env, engine_config.o11y_auth0_client_id, engine_config.o11y_auth0_client_secret) self.o11y_api = o11y.O11yAPI(self.o11y_config) self.running_task_key = task_file self.engine_config = engine_config self.remote_loader = self.engine_config.remote_loader if engine_config.use_local_loader == 'True': task_file = task_file.replace(self.remote_loader.prefix+"/", "") task_info_content = self.remote_loader.get_file(task_file) task_info = yaml.load(task_info_content, yaml.FullLoader) self.task_info = task_info cache_prefix = os.path.join(engine_config.cache_root, task_info["meta"]["clusterID"]) if not os.path.exists(cache_prefix): os.makedirs(cache_prefix) self.task = Task(self.o11y_api, task_info, cache_prefix, self.remote_loader) def start(self): now_start = self.load_checkpoint() if now_start is None: now_start = self.engine_config.start_timestemp # revise now_start if self.task_info["meta"]["createTime"] and now_start < self.task_info["meta"]["createTime"]: now_start = self.task_info["meta"]["createTime"] now_start = now_start // self.task.interval * self.task.interval while running_tasks[self.running_task_key]() and (self.engine_config.end_timestemp is None or now_start < self.engine_config.end_timestemp): if now_start > time.time(): time.sleep(10) continue else: #log.info(f"[{self.running_task_key}] execute at {now_start}") succeed, true_started = self.execute_subtasks(self.task.cluster_id, now_start) if succeed: try: # gcp storage limit 1/s update the same object or it will throw 429 error. if self.task.provider == "gcp": if true_started: time.sleep(1) self.save_checkpoint(now_start) else: self.save_checkpoint(now_start) except Exception: log.exception(f"[{self.task.cluster_id}] failed to save checkpoint") now_start += 30 log.info(f"[{self.running_task_key}] stop at {now_start}") def execute_subtasks(self, cluster_id, now_start): all_success = True true_started = False for subtask in self.task.subtasks: try: if subtask.is_timing(now_start) and subtask.ready(now_start): true_started = True subtask.start(now_start) subtask.close(now_start) except Exception: all_success = False log.exception(f"[{cluster_id}] failed to execute subtask") return all_success, true_started def load_checkpoint(self): file_name = "checkpoint.yaml" file_path = os.path.join(checkpoint_prefix_template.format(self.task.cluster_id), file_name) if not self.remote_loader.file_exists(file_path): return None else: f = self.remote_loader.get_file(file_path) checkpoint = yaml.safe_load(f) return int(checkpoint["checkpoint_time"]) def save_checkpoint(self, timestamp): file_name = "checkpoint.yaml" file_path = os.path.join(self.task.cache_prefix, file_name) checkpoint = {"checkpoint_time": timestamp} # save to cache path with open(file_path, "w") as f: yaml.dump(checkpoint, f) file_name = "checkpoint.yaml" remote_file_path = os.path.join(checkpoint_prefix_template.format(self.task.cluster_id), file_name) self.remote_loader.upload_file(file_path, remote_file_path)