upgrade logger.utils

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():

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):