simplecloud/VidChain-exercise

1import datetime2import logging3import logging.handlers4import os5import sys6import time7 8import requests9 10from vtimellm.constants import LOGDIR11 12server_error_msg = "**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**"13moderation_msg = "YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN."14 15handler = None16 17 18def build_logger(logger_name, logger_filename):19    global handler20 21    formatter = logging.Formatter(22        fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",23        datefmt="%Y-%m-%d %H:%M:%S",24    )25 26    # Set the format of root handlers27    if not logging.getLogger().handlers:28        logging.basicConfig(level=logging.INFO)29    logging.getLogger().handlers[0].setFormatter(formatter)30 31    # Redirect stdout and stderr to loggers32    stdout_logger = logging.getLogger("stdout")33    stdout_logger.setLevel(logging.INFO)34    sl = StreamToLogger(stdout_logger, logging.INFO)35    sys.stdout = sl36 37    stderr_logger = logging.getLogger("stderr")38    stderr_logger.setLevel(logging.ERROR)39    sl = StreamToLogger(stderr_logger, logging.ERROR)40    sys.stderr = sl41 42    # Get logger43    logger = logging.getLogger(logger_name)44    logger.setLevel(logging.INFO)45 46    # Add a file handler for all loggers47    if handler is None:48        os.makedirs(LOGDIR, exist_ok=True)49        filename = os.path.join(LOGDIR, logger_filename)50        handler = logging.handlers.TimedRotatingFileHandler(51            filename, when='D', utc=True)52        handler.setFormatter(formatter)53 54        for name, item in logging.root.manager.loggerDict.items():55            if isinstance(item, logging.Logger):56                item.addHandler(handler)57 58    return logger59 60 61class StreamToLogger(object):62    """63    Fake file-like stream object that redirects writes to a logger instance.64    """65    def __init__(self, logger, log_level=logging.INFO):66        self.terminal = sys.stdout67        self.logger = logger68        self.log_level = log_level69        self.linebuf = ''70 71    def __getattr__(self, attr):72        return getattr(self.terminal, attr)73 74    def write(self, buf):75        temp_linebuf = self.linebuf + buf76        self.linebuf = ''77        for line in temp_linebuf.splitlines(True):78            # From the io.TextIOWrapper docs:79            #   On output, if newline is None, any '\n' characters written80            #   are translated to the system default line separator.81            # By default sys.stdout.write() expects '\n' newlines and then82            # translates them so this is still cross platform.83            if line[-1] == '\n':84                self.logger.log(self.log_level, line.rstrip())85            else:86                self.linebuf += line87 88    def flush(self):89        if self.linebuf != '':90            self.logger.log(self.log_level, self.linebuf.rstrip())91        self.linebuf = ''92 93 94def disable_torch_init():95    """96    Disable the redundant torch default initialization to accelerate model creation.97    """98    import torch99    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)100    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)101 102 103def violates_moderation(text):104    """105    Check whether the text violates OpenAI moderation API.106    """107    url = "https://api.openai.com/v1/moderations"108    headers = {"Content-Type": "application/json",109               "Authorization": "Bearer " + os.environ["OPENAI_API_KEY"]}110    text = text.replace("\n", "")111    data = "{" + '"input": ' + f'"{text}"' + "}"112    data = data.encode("utf-8")113    try:114        ret = requests.post(url, headers=headers, data=data, timeout=5)115        flagged = ret.json()["results"][0]["flagged"]116    except requests.exceptions.RequestException as e:117        flagged = False118    except KeyError as e:119        flagged = False120 121    return flagged122 123 124def pretty_print_semaphore(semaphore):125    if semaphore is None:126        return "None"127    return f"Semaphore(value={semaphore._value}, locked={semaphore.locked()})"128 129def get_gpu_status():130    """131    Check the gpu stats and return the valid number of available gpus132    """133    from gpustat.core import GPUStatCollection134 135    gpus_stats = GPUStatCollection.new_query()136 137    info = gpus_stats.jsonify()["gpus"]138    gpu_list = []139 140    mem_ratio_threshold = 0.1  #141    util_ratio_threshold = 10  #142    for idx, each in enumerate(info):143        mem_ratio = each["memory.used"] / each["memory.total"]144        util_ratio = each["utilization.gpu"]145        if mem_ratio < mem_ratio_threshold and util_ratio < util_ratio_threshold:146            gpu_list.append(idx)147    print("Scan GPUs to get {} free GPU ({})".format(len(gpu_list), gpu_list))148    return gpu_list149 150def check_gpu_status(gpu_option):151    if gpu_option == 'cuda':152        env_var = 'CUDA_VISIBLE_DEVICES'153    elif gpu_option == 'gpu_vis':154        env_var = 'gpu_vis'155    print(f'gpu option is {env_var}')156    gpu_list = [int(x) for x in os.environ[env_var].split(',')]157    if len(gpu_list) == 0:158        assert False, 'Please specify the gpu_vis in the environment variable with export.'159    available_gpu = get_gpu_status()160    while gpu_list != available_gpu:161        print("No available GPU, waiting for 1 minutes until {} get freed. Current time : {}".format(gpu_list, time.ctime()))162        time.sleep(60)163        available_gpu = get_gpu_status()164    print("GPU is available now. Current time : {}".format(time.ctime()))