diff --git a/parl/utils/logger.py b/parl/utils/logger.py index c894ea777..f2459d94b 100644 --- a/parl/utils/logger.py +++ b/parl/utils/logger.py @@ -137,89 +137,92 @@ def set_level(level): _logger.setLevel(level) -def set_dir(dirname): - global LOG_DIR, _FILE_HANDLER, _logger - if _FILE_HANDLER: - # unload and close the old file handler, so that we may safely delete the logger directory - _logger.removeHandler(_FILE_HANDLER) - _FILE_HANDLER.close() - del _FILE_HANDLER - - shutil.rmtree(dirname, ignore_errors=True) - _makedirs(dirname) - LOG_DIR = dirname - _set_file(os.path.join(dirname, 'log.log')) - -def auto_set_dir(action=None): - """Set the global logging directory automatically. The default path is "./train_log/{scriptname}". "scriptname" is the name of the main python file currently running" - - Note: This function references `https://github.com/tensorpack/tensorpack/blob/master/tensorpack/utils/logger.py#L93` - - Args: - dir_name(str): log directory - action(str): an action of ["k","d","q"] to be performed - when the directory exists. Will ask user by default. - "d": delete the directory. Note that the deletion may fail when - the directory is used by tensorboard. - "k": keep the directory. This is useful when you resume from a - previous training and want the directory to look as if the - training was not interrupted. - Note that this option does not load old models or any other - old states for you. It simply does nothing. +def get_dir(): + return LOG_DIR - Returns: - dirname(str): log directory used in the global logging directory. - """ +def get_default_dirname(): + """Determine the default directory name based on the script's name.""" mod = sys.modules['__main__'] if hasattr(mod, '__file__'): basename = os.path.basename(mod.__file__) else: basename = '' - dirname = os.path.join('train_log', basename[:basename.rfind('.')]) - dirname = os.path.normpath(dirname) + return os.path.normpath(os.path.join('train_log', basename[:basename.rfind('.')])) + +def dir_nonempty(dirname): + """Check if directory exists and is nonempty (ignore hidden files).""" + return os.path.isdir(dirname) and len([x for x in os.listdir(dirname) if x[0] != '.']) + +def handle_existing_dir(dirname, action=None): + """Handle user actions when the directory exists and is non-empty.""" + while not action: + _logger.warning(f"Log directory {dirname} exists!") + action = input("Select Action: k (keep) / d (delete) / n (new) / q (quit):").lower().strip() + + if action == 'd': + shutil.rmtree(dirname, ignore_errors=True) + if dir_nonempty(dirname): + shutil.rmtree(dirname, ignore_errors=False) + elif action == 'n': + dirname = dirname + _get_time_str() + _logger.info(f"Use a new log directory {dirname}") + elif action == 'k': + pass + else: + raise OSError(f"Directory {dirname} exits!") + return dirname - global LOG_DIR, _FILE_HANDLER +def set_dir(dirname=None, action=None): + """ + Set the global logging directory. If dirname is not provided, it will be set automatically based on the script's name. + + Args: + dirname (str, optional): Desired log directory. If not provided, it will be set automatically. + action (str, optional): An action of ["k","d","n","q"] to be performed when the directory exists. + "d": delete the directory. + "k": keep the directory. + "n": use a new log directory with a timestamp. + "q": quit without doing anything. + Will ask the user by default if the directory exists and is non-empty. + + Returns: + str: The log directory used in the global logging directory. + """ + if not dirname: + dirname = get_default_dirname() + + global LOG_DIR, _FILE_HANDLER, _logger if _FILE_HANDLER: - # unload and close the old file handler, so that we may safely delete the logger directory _logger.removeHandler(_FILE_HANDLER) + _FILE_HANDLER.close() del _FILE_HANDLER - def dir_nonempty(dirname): - # If directory exists and nonempty (ignore hidden files), prompt for action - return os.path.isdir(dirname) and len( - [x for x in os.listdir(dirname) if x[0] != '.']) - if dir_nonempty(dirname): - if not action: - _logger.warning("""\ -Log directory {} exists! Use 'd' to delete it. """.format(dirname)) - _logger.warning("""\ -If you're resuming from a previous run, you can choose to keep it. -Press any other key to exit. """) - while not action: - action = input("Select Action: k (keep) / d (delete) / q (quit):" - ).lower().strip() - act = action - if act == 'd': - shutil.rmtree(dirname, ignore_errors=True) - if dir_nonempty(dirname): - shutil.rmtree(dirname, ignore_errors=False) - elif act == 'n': - dirname = dirname + _get_time_str() - info("Use a new log directory {}".format(dirname)) # noqa: F821 - elif act == 'k': - pass - else: - raise OSError("Directory {} exits!".format(dirname)) + dirname = handle_existing_dir(dirname, action) + LOG_DIR = dirname _makedirs(dirname) _set_file(os.path.join(dirname, 'log.log')) return dirname +def auto_set_dir(action=None): + """ + Automatically set the global logging directory based on the name of the main script being executed. + + Args: + action (str, optional): An action of ["k","d","n","q"] to be performed when the directory exists. + "d": delete the directory. + "k": keep the directory. + "n": use a new log directory with a timestamp. + "q": quit without doing anything. + Will ask the user by default if the directory exists and is non-empty. -def get_dir(): - return LOG_DIR + Returns: + str: The log directory used in the global logging directory. + """ + dirname = get_default_dirname() + return set_dir(dirname, action) def add_stdout_handler(): diff --git a/parl/utils/tensorboard.py b/parl/utils/tensorboard.py index d4c4d2d1d..0453b4765 100644 --- a/parl/utils/tensorboard.py +++ b/parl/utils/tensorboard.py @@ -19,7 +19,7 @@ __all__ = [] _writer = None -_WRITTER_METHOD = ['add_scalar', 'add_histogram', 'close', 'flush'] +_WRITTER_METHOD = ['add_scalar', 'add_histogram', 'add_text', 'close', 'flush'] def create_file_after_first_call(func_name):