improve distributed logging

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/llmcompressor/logger.py

@@ -8,8 +8,6 @@
 
 - `LLM_COMPRESSOR_LOG_DISABLED`: Disable logging.
     Default: `False`.
-- `LLM_COMPRESSOR_CLEAR_LOGGERS`: Clear existing loggers from loguru.
-    Default: `True`.
 - `LLM_COMPRESSOR_LOG_LEVEL`: Log level for console logging.
     Default: `None`. Options: `DEBUG`, `INFO`, `WARNING`, `ERROR`, `CRITICAL`.
 - `LLM_COMPRESSOR_LOG_FILE`: Path to the log file for file logging.
@@ -25,7 +23,6 @@
     configure_logger(
         config=LoggerConfig(
             disabled=False,
-            clear_loggers=True,
             console_log_level="DEBUG",
             log_file=None,
             log_file_level=None,
@@ -65,9 +62,7 @@ class LoggerConfig:
 LOGGER_CONFIG = LoggerConfig()
 
 
-def configure_logger(
-    logger_config: LoggerConfig = LOGGER_CONFIG, clear_loggers: bool = False
-):
+def configure_logger(logger_config: LoggerConfig = LOGGER_CONFIG):
     """
     Configure the logger for LLM Compressor.
 
@@ -97,8 +92,12 @@ def configure_logger(
     logger.enable("llmcompressor")
 
     # reset logger configuration
-    if clear_loggers:
-        logger.remove()
+    logger.remove()
+
+    # initialize metric logger on loguru
+    logger_levels = logger._core.levels.keys()
+    if not logger_config.metrics_disabled and "METRIC" not in logger_levels:
+        logger.level("METRIC", no=38, color="<yellow>", icon="📈")
 
     # set format (optionally adding rank)
     format = "{time:YYYY-MM-DDTHH:mm:ss.SSSS} | {function} | {level} - {message}"
@@ -107,7 +106,6 @@ def configure_logger(
         format = "[Rank {extra[rank]}] " + format
 
     if logger_config.console_log_level:
-        # log as a human readable string with the time, function, level, and message
         logger.add(
             sys.stdout,
             level=logger_config.console_log_level.upper(),
@@ -123,16 +121,9 @@ def configure_logger(
             log_file,
             level=log_file_level.upper(),
             serialize=True,
-            format=format,
             filter=support_log_once,
         )
 
-    if logger_config.metrics_disabled or "METRIC" in logger._core.levels.keys():
-        return
-
-    # initialize metric logger on loguru
-    logger.level("METRIC", no=38, color="<yellow>", icon="📈")
-
     # set global value for later calls
     global LOGGER_CONFIG
     LOGGER_CONFIG = logger_config
@@ -164,7 +155,7 @@ def support_log_once(record: Dict[str, Any]) -> bool:
 
 def configure_distributed_logger(logger_config: LoggerConfig = LOGGER_CONFIG):
     logger_config.rank = dist.get_rank()
-    configure_logger(logger_config, clear_loggers=True)
+    configure_logger(logger_config)
 
 
 # invoke logger setup on import with default values enabling console logging with INFO

src/llmcompressor/utils/metric_logging.py

@@ -54,14 +54,14 @@ def __exit__(self, _exc_type, _exc_val, _exc_tb):
             patch.log("METRIC", f"error {self._loss:.2f}")
 
         for device_id in _get_visible_devices():
-            used_memory = torch.cuda.max_memory_allocated(device_id)
-            max_memory = torch.cuda.get_device_properties(device_id).total_memory
+            used_memory = torch.cuda.max_memory_allocated(device_id) / 1e9
+            max_memory = torch.cuda.get_device_properties(device_id).total_memory / 1e9
             perc_used = 100 * used_memory / max_memory
             patch.log(
                 "METRIC",
                 (
                     f"GPU {device_id} | usage: {perc_used:.2f}%"
-                    f" | total memory: {max_memory:.1f} GB"
+                    f" | total memory: {max_memory:.1f} Gb"
                 ),
             )
 

tests/unit/test_logger.py

@@ -26,7 +26,7 @@ def test_default_logger_settings(capsys):
 def test_configure_logger_console_settings(capsys):
     # Test configuring the logger to change console log level
     config = LoggerConfig(console_log_level="DEBUG")
-    configure_logger(config=config)
+    configure_logger(config)
     logger.info("Info message")
     logger.debug("Debug message")
 
@@ -39,7 +39,7 @@ def test_configure_logger_file_settings(tmp_path):
     # Test configuring the logger to log to a file
     log_file = tmp_path / "test.log"
     config = LoggerConfig(log_file=str(log_file), log_file_level="DEBUG")
-    configure_logger(config=config)
+    configure_logger(config)
     logger.info("Info message")
     logger.debug("Debug message")
 
@@ -55,7 +55,7 @@ def test_configure_logger_console_and_file(capsys, tmp_path):
     config = LoggerConfig(
         console_log_level="ERROR", log_file=str(log_file), log_file_level="INFO"
     )
-    configure_logger(config=config)
+    configure_logger(config)
     logger.info("Info message")
     logger.error("Error message")
 
@@ -75,7 +75,7 @@ def test_environment_variable_override(monkeypatch, capsys, tmp_path):
     monkeypatch.setenv("LLM_COMPRESSOR_LOG_FILE", str(tmp_path / "env_test.log"))
     monkeypatch.setenv("LLM_COMPRESSOR_LOG_FILE_LEVEL", "DEBUG")
 
-    configure_logger(config=LoggerConfig())
+    configure_logger(LoggerConfig())
     logger.info("Info message")
     logger.error("Error message")
     logger.debug("Debug message")
@@ -96,7 +96,7 @@ def test_environment_variable_disable_logging(monkeypatch, capsys):
     # Test environment variable to disable logging
     monkeypatch.setenv("LLM_COMPRESSOR_LOG_DISABLED", "true")
 
-    configure_logger(config=LoggerConfig())
+    configure_logger(LoggerConfig())
     logger.info("Info message")
     logger.error("Error message")
 
@@ -109,7 +109,7 @@ def test_environment_variable_enable_logging(monkeypatch, capsys):
     # Test environment variable to enable logging
     monkeypatch.setenv("LLM_COMPRESSOR_LOG_DISABLED", "false")
 
-    configure_logger(config=LoggerConfig())
+    configure_logger(LoggerConfig())
     logger.info("Info message")
     logger.error("Error message")