Loguru logging standardization for LLM Compressor

src/llmcompressor/__init__.py

@@ -8,7 +8,7 @@
 
 # flake8: noqa
 
-from .logger import configure_logger, logger
+from .logger import LoggerConfig, configure_logger, logger
 from .version import (
     __version__,
     build_type,
@@ -31,6 +31,7 @@
     "version_build",
     "configure_logger",
     "logger",
+    "LoggerConfig",
 ]
 
 from llmcompressor.core.session_functions import (

src/llmcompressor/logger.py

@@ -5,11 +5,9 @@
 It supports console and file logging with options to configure via environment
 variables or direct function calls.
 
-By default, logging is disabled for this library to ensure it does not
-overwrite application logs. To enable logging, either set one of the
-environment variables or call configure_logger from within the application code.
-
 Environment Variables:
+    - 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
@@ -18,84 +16,104 @@
         (default: INFO if log file set else none).
 
 Usage:
-    from llmcompressor.metrics import configure_logger
+    from llmcompressor import logger, configure_logger, LoggerConfig
 
     # Configure metrics with default settings
-    configure_logger()
-
-    # Configure metrics with custom settings
     configure_logger(
-        console_log_level="DEBUG",
-        log_file="/path/to/logfile.log",
-        log_file_level="ERROR"
+        config=LoggerConfig(
+            disabled=False,
+            clear_loggers=True,
+            console_log_level="DEBUG",
+            log_file=None,
+            log_file_level=None,
+        )
     )
+
+    logger.debug("This is a debug message")
+    logger.info("This is an info message")
 """
 
 import os
 import sys
+from dataclasses import dataclass
 from typing import Optional
 
 from loguru import logger
 
-__all__ = ["configure_logger", "logger"]
+__all__ = ["LoggerConfig", "configure_logger", "logger"]
+
+
+@dataclass
+class LoggerConfig:
+    disabled: bool = False
+    clear_loggers: bool = True
+    console_log_level: Optional[str] = "INFO"
+    log_file: Optional[str] = None
+    log_file_level: Optional[str] = None
 
 
-def configure_logger(
-    console_log_level: Optional[str] = "INFO",
-    log_file: Optional[str] = None,
-    log_file_level: Optional[str] = None,
-):
+def configure_logger(config: Optional[LoggerConfig] = None):
     """
     Configure the metrics for LLM Compressor.
     This function sets up the console and file logging
     as per the specified or default parameters.
 
-    :param console_log_level: Log level for console output, defaults to "INFO"
-    :type console_log_level: Optional[str]
-    :param log_file: Path to the log file, defaults to "llm-compressor.log"
-        if log_file_level is set
-    :type log_file: Optional[str]
-    :param log_file_level: Log level for file output, defaults to "INFO"
-        if log_file is set
-    :type log_file_level: Optional[str]
+    Note: Environment variables take precedence over the function parameters.
+
+    :param config: The configuration for the logger to use.
+    :type config: LoggerConfig
     """
-    _logger_setup(True, console_log_level, log_file, log_file_level)
 
+    _ENV_CONFIG = LoggerConfig(
+        disabled=os.getenv("LLM_COMPRESSOR_LOG_DISABLED") == "true",
+        clear_loggers=os.getenv("LLM_COMPRESSOR_CLEAR_LOGGERS") == "true",
+        console_log_level=os.getenv("LLM_COMPRESSOR_LOG_LEVEL"),
+        log_file=os.getenv("LLM_COMPRESSOR_LOG_FILE"),
+        log_file_level=os.getenv("LLM_COMPRESSOR_LOG_FILE_LEVEL"),
+    )
 
-def _logger_setup(
-    api_request: bool,
-    console_log_level: Optional[str],
-    log_file: Optional[str],
-    log_file_level: Optional[str],
-):
-    enable_logging = api_request or console_log_level or log_file or log_file_level
+    if not config:
+        config = LoggerConfig()
+    # override from environment variables, if set
+    logger_config = LoggerConfig(
+        disabled=_ENV_CONFIG.disabled or config.disabled,
+        console_log_level=_ENV_CONFIG.console_log_level or config.console_log_level,
+        log_file=_ENV_CONFIG.log_file or config.log_file,
+        log_file_level=_ENV_CONFIG.log_file_level or config.log_file_level,
+    )
 
-    if not enable_logging:
+    if logger_config.disabled:
         logger.disable("llmcompressor")
         return
 
     logger.enable("llmcompressor")
-    logger.remove()
 
-    if console_log_level:
+    if logger_config.clear_loggers:
+        logger.remove()
+
+    if logger_config.console_log_level:
         # log as a human readable string with the time, function, level, and message
         logger.add(
             sys.stdout,
-            level=console_log_level.upper(),
+            level=logger_config.console_log_level.upper(),
             format="{time} | {function} | {level} - {message}",
         )
 
-    if log_file or log_file_level:
-        log_file = log_file or "llm-compressor.log"
-        log_file_level = log_file_level or "INFO"
+    if logger_config.log_file or logger_config.log_file_level:
+        log_file = logger_config.log_file or "llmcompressor.log"
+        log_file_level = logger_config.log_file_level or "INFO"
         # log as json to the file for easier parsing
         logger.add(log_file, level=log_file_level.upper(), serialize=True)
 
 
-# invoke the metrics setup on import if environment variables are set
-_logger_setup(
-    api_request=False,
-    console_log_level=os.getenv("LLM_COMPRESSOR_LOG_LEVEL"),
-    log_file=os.getenv("LLM_COMPRESSOR_LOG_FILE"),
-    log_file_level=os.getenv("LLM_COMPRESSOR_LOG_FILE_LEVEL"),
+# invoke logger setup on import with default values enabling console logging with INFO
+# and disabling file logging
+configure_logger(
+    config=LoggerConfig(
+        disabled=False,
+        clear_loggers=True,
+        console_log_level="INFO",
+        log_file=None,
+        log_file_level=None,
+    )
 )

src/llmcompressor/modifiers/obcq/base.py

@@ -1,8 +1,8 @@
-import logging
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
+from loguru import logger
 from torch.nn import Module
 from tqdm import tqdm
 
@@ -19,8 +19,6 @@
 
 __all__ = ["SparseGPTModifier"]
 
-_LOGGER = logging.getLogger(__name__)
-
 
 class SparseGPTModifier(Modifier):
     """
@@ -152,15 +150,15 @@ def initialize_compression(
         self._infer_mask_block_size()
 
         if self.sparsity_profile is not None and self.sparsity_profile.lower() == "owl":
-            _LOGGER.info(
+            logger.info(
                 "Inferring layer-wise sparsities from "
                 f"{len(dataloader)} calibration samples..."
             )
             self.sparsity = self._infer_layer_sparsity(dataloader)
         self._validate_layerwise_sparsity()
 
         for idx, (name, layer) in enumerate(self.compressible_layers_.items()):
-            _LOGGER.info(f"Preparing {name} for compression")
+            logger.info(f"Preparing {name} for compression")
             if isinstance(self.sparsity, Dict):
                 layer_sparsity = self.sparsity[name]
             elif isinstance(self.sparsity, List):
@@ -202,8 +200,9 @@ def apply_compression(
         :param dataloader: calibration data for WANDA
         """
         class_name = self.__class__.__name__.replace("PyTorch", "")
-        _LOGGER.info(
-            f"Running {class_name} calibration with " f"{len(dataloader)} samples..."
+        logger.info(
+            f"Running {class_name} calibration with "
+            f"{len(dataloader) if dataloader else 0} samples..."
         )
         if not self.sequential_update:
             # in non-sequential mode we run one forward batch for all modules
@@ -212,7 +211,7 @@ def apply_compression(
         num_layers = len(self.compressible_layers_)
         for idx, layer_compressor in enumerate(self.layer_compressors_):
             layer_sparsity = layer_compressor.args["sparsity"]
-            _LOGGER.info(
+            logger.info(
                 f"\n===== Compressing layer {idx+1}/{num_layers} "
                 f"to sparsity {layer_sparsity} ====="
             )
@@ -223,7 +222,7 @@ def apply_compression(
                 # want to compress, this will be really slow but allows compression in
                 # earlier layers to affect later layers
                 layer_compressor.pre_compress()
-                _LOGGER.info(f"Calibrating {layer_compressor.name}...")
+                logger.info(f"Calibrating {layer_compressor.name}...")
                 run_calibration_forward(self.model, dataloader, mask_padding=True)
             layer_compressor.compress()
             layer_compressor.post_compress()
@@ -239,8 +238,8 @@ def _validate_layerwise_sparsity(self):
 
         if len(target_layers) != len(self.sparsity):
             raise ValueError(
-                "Number of layer targets must match the number of "
-                f"sparsities. Got {len(target_layers)} layers and "
+                "Number of layer targets must match the number of sparsities. "
+                "Received {len(target_layers)} layers and "
                 f"{len(self.sparsity)} sparsities"
             )
 
@@ -320,9 +319,9 @@ def _infer_layer_sparsity(self, calibration_dataloader):
             )
             for k in outlier_ratios
         }
-        _LOGGER.info(f"OWL sparsities for sp={self.sparsity} are:")
+        logger.info(f"OWL sparsities for sp={self.sparsity} are:")
         for k in sparsities:
-            _LOGGER.info(f"Sparsity for {k}: {sparsities[k]}")
+            logger.info(f"Sparsity for {k}: {sparsities[k]}")
         return sparsities
 
 

src/llmcompressor/modifiers/obcq/utils/helpers.py

@@ -1,13 +1,12 @@
-import logging
 import operator
 from collections import defaultdict
 from math import ceil
 from typing import List, Optional
 
 import torch
+from loguru import logger
 from torch.nn.modules.sparse import Embedding
 
-_LOGGER = logging.getLogger(__name__)
 _DEFAULT_TARGET_IDS = ["attention_mask", "position_ids", "position_bias"]
 
 
@@ -148,7 +147,7 @@ def cache_attention_inputs(
 def ppl_eval_general(
     eval_logits, model, dataloader, dev, nsamples=None, max_samples_per_iteration=128
 ):
-    _LOGGER.info("Evaluating perplexity...")
+    logger.info("Evaluating perplexity...")
 
     if nsamples is None:
         nsamples = len(dataloader)
@@ -180,12 +179,16 @@ def ppl_eval_general(
         )
 
         number_tokens += labels.numel()
-        _LOGGER.info(torch.exp(neg_log_likelihood / number_tokens))
+        perplexity = torch.exp(neg_log_likelihood / number_tokens)
+        logger.debug(
+            f"Processed iteration {iteration} of {number_iterations} with perplexity: "
+            f"{perplexity}"
+        )
 
-    ppl = torch.exp(neg_log_likelihood / number_tokens)
-    _LOGGER.info(f"Perplexity: {ppl.item():3f}")
+    perplexity = torch.exp(neg_log_likelihood / number_tokens)
+    logger.info(f"Evaluated perplexity: {perplexity}")
 
-    return ppl.item()
+    return perplexity.item()
 
 
 def _get_pre_layer_modules(model_root, layers_name):

src/llmcompressor/modifiers/obcq/utils/sgpt_wrapper.py

@@ -8,16 +8,14 @@
     transformers = None
     transformers_err = err
 
-import logging
 import math
 
 import torch
 import torch.nn as nn
+from loguru import logger
 
 __all__ = ["SparseGptWrapper"]
 
-_LOGGER = logging.getLogger(__name__)
-
 
 class SparseGptWrapper(ModuleCompressionWrapper):
     """
@@ -194,8 +192,8 @@ def fasterprune(
             else:
                 W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])
 
-        _LOGGER.info("time %.2f" % (time.time() - tick))
-        _LOGGER.info("error %.2f" % torch.sum(Losses).item())
+        logger.info("time %.2f" % (time.time() - tick))
+        logger.info("error %.2f" % torch.sum(Losses).item())
 
         if isinstance(self.layer, transformers.Conv1D):
             W = W.t()