import os import logging import yaml YAML_PATH = "./cicd/configs" LOG_FILE = "temp_log" logger = logging.getLogger(__name__) class Dumper(yaml.Dumper): def increase_indent(self, flow=False, *args, **kwargs): return super().increase_indent(flow=flow, indentless=False) def get_submitted_yaml_path(uid): if not os.path.exists(f"{YAML_PATH}/submitted"): os.makedirs(f"{YAML_PATH}/submitted") if not os.path.exists(f"{YAML_PATH}/{uid}_config.yaml"): logger.error(f"config.yaml does not exist for {uid}") os.system(f"cp config.yaml {YAML_PATH}/{uid}_config.yaml") if not os.path.exists(f"{YAML_PATH}/submitted/{uid}_config.yaml"): os.system(f"cp {YAML_PATH}/{uid}_config.yaml {YAML_PATH}/submitted/{uid}_config.yaml") return f"{YAML_PATH}/submitted/{uid}_config.yaml" def get_yaml_path(uid): if not os.path.exists(YAML_PATH): os.makedirs(YAML_PATH) if not os.path.exists(f"{YAML_PATH}/{uid}_config.yaml"): os.system(f"cp config.yaml {YAML_PATH}/{uid}_config.yaml") return f"{YAML_PATH}/{uid}_config.yaml" # read scanners from yaml file # return a list of scanners def read_scanners(uid): scanners = [] with open(get_yaml_path(uid), "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) scanners = config.get("detectors", []) return scanners # convert a list of scanners to yaml file def write_scanners(scanners, uid): with open(get_yaml_path(uid), "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) if config: config["detectors"] = scanners # save scanners to detectors in yaml with open(get_yaml_path(uid), "w") as f: yaml.dump(config, f, Dumper=Dumper) # read model_type from yaml file def read_inference_type(uid): inference_type = "" with open(get_yaml_path(uid), "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) inference_type = config.get("inference_type", "") return inference_type # write model_type to yaml file def write_inference_type(use_inference, inference_token, uid): with open(get_yaml_path(uid), "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) if use_inference: config["inference_type"] = "hf_inference_api" config["inference_token"] = inference_token else: config["inference_type"] = "hf_pipeline" # FIXME: A quick and temp fix for missing token config["inference_token"] = "" # save inference_type to inference_type in yaml with open(get_yaml_path(uid), "w") as f: yaml.dump(config, f, Dumper=Dumper) # read column mapping from yaml file def read_column_mapping(uid): column_mapping = {} with open(get_yaml_path(uid), "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) if config: column_mapping = config.get("column_mapping", dict()) if column_mapping is None: column_mapping = {} return column_mapping # write column mapping to yaml file def write_column_mapping(mapping, uid): with open(get_yaml_path(uid), "r") as f: config = yaml.load(f, Loader=yaml.FullLoader) if config is None: return if mapping is None and "column_mapping" in config.keys(): del config["column_mapping"] else: config["column_mapping"] = mapping with open(get_yaml_path(uid), "w") as f: # yaml Dumper will by default sort the keys yaml.dump(config, f, Dumper=Dumper, sort_keys=False) # convert column mapping dataframe to json def convert_column_mapping_to_json(df, label=""): column_mapping = {} column_mapping[label] = [] for _, row in df.iterrows(): column_mapping[label].append(row.tolist()) return column_mapping def get_log_file_with_uid(uid): try: print(f"Loading {uid}.log") with open(f"./tmp/{uid}.log", "a") as file: return file.read() except Exception: return "Log file does not exist" def get_logs_file(): try: with open(LOG_FILE, "r") as file: return file.read() except Exception: return "Log file does not exist" def write_log_to_user_file(task_id, log): with open(f"./tmp/{task_id}.log", "a") as f: f.write(log)