From aca180c6d4faa17a6fa7accc5fc8492ae69a22c8 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Mon, 8 Sep 2025 20:59:58 -0400 Subject: [PATCH 1/5] support loguru Signed-off-by: Kyle Sayers --- src/compressed_tensors/__init__.py | 4 +- src/compressed_tensors/logger.py | 132 +++++++++++++++++++++++++++++ 2 files changed, 135 insertions(+), 1 deletion(-) create mode 100644 src/compressed_tensors/logger.py diff --git a/src/compressed_tensors/__init__.py b/src/compressed_tensors/__init__.py index 64a52ddac..c892e81a9 100644 --- a/src/compressed_tensors/__init__.py +++ b/src/compressed_tensors/__init__.py @@ -12,9 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +# flake8: noqa +# isort: off +from .logger import LoggerConfig, configure_logger, logger from .base import * -# flake8: noqa from .compressors import * from .config import * from .quantization import QuantizationConfig, QuantizationStatus diff --git a/src/compressed_tensors/logger.py b/src/compressed_tensors/logger.py new file mode 100644 index 000000000..02d99402b --- /dev/null +++ b/src/compressed_tensors/logger.py @@ -0,0 +1,132 @@ +# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Logger configuration for Compressed Tensors. +""" + +import os +import sys +from dataclasses import dataclass +from typing import Any, Dict, Optional + +from loguru import logger + + +__all__ = ["LoggerConfig", "configure_logger", "logger"] + + +# used by `support_log_once`` +_logged_once = set() + + +@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(config: Optional[LoggerConfig] = None) -> None: + """ + Configure the logger for Compressed Tensors. + This function sets up the console and file logging + as per the specified or default parameters. + + Note: Environment variables take precedence over the function parameters. + + :param config: The configuration for the logger to use. + :type config: LoggerConfig + """ + logger_config = config or LoggerConfig() + + # env vars get priority + if (disabled := os.getenv("COMPRESSED_TENSORS_LOG_DISABLED")) is not None: + logger_config.disabled = disabled.lower() == "true" + if (clear_loggers := os.getenv("COMPRESSED_TENSORS_CLEAR_LOGGERS")) is not None: + logger_config.clear_loggers = clear_loggers.lower() == "true" + if (console_log_level := os.getenv("COMPRESSED_TENSORS_LOG_LEVEL")) is not None: + logger_config.console_log_level = console_log_level.upper() + if (log_file := os.getenv("COMPRESSED_TENSORS_LOG_FILE")) is not None: + logger_config.log_file = log_file + if (log_file_level := os.getenv("COMPRESSED_TENSORS_LOG_FILE_LEVEL")) is not None: + logger_config.log_file_level = log_file_level.upper() + + if logger_config.disabled: + logger.disable("compressed_tensors") + return + + logger.enable("compressed_tensors") + + 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=logger_config.console_log_level.upper(), + format="{time} | {function} | {level} - {message}", + filter=support_log_once, + ) + + if logger_config.log_file or logger_config.log_file_level: + log_file = logger_config.log_file or "compressed_tensors.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, + filter=support_log_once, + ) + + +def support_log_once(record: Dict[str, Any]) -> bool: + """ + Support logging only once using `.bind(log_once=True)` + + ``` + logger.bind(log_once=False).info("This will log multiple times") + logger.bind(log_once=False).info("This will log multiple times") + logger.bind(log_once=True).info("This will only log once") + logger.bind(log_once=True).info("This will only log once") # skipped + ``` + """ + log_once = record["extra"].get("log_once", False) + level = getattr(record["level"], "name", "none") + message = str(level) + record["message"] + + if log_once and message in _logged_once: + return False + + if log_once: + _logged_once.add(message) + + return True + + +# 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, + ) +) From d634e529afbc46362b158d31439d0b829fcfa27a Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Mon, 8 Sep 2025 21:05:12 -0400 Subject: [PATCH 2/5] add loguru Signed-off-by: Kyle Sayers --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 58a1105f4..d60b8d8ce 100644 --- a/setup.py +++ b/setup.py @@ -88,7 +88,7 @@ def _setup_packages() -> List: ) def _setup_install_requires() -> List: - return ["torch>=1.7.0", "transformers", "pydantic>=2.0", "frozendict"] + return ["torch>=1.7.0", "transformers", "pydantic>=2.0", "frozendict", "loguru"] def _setup_extras() -> Dict: return { From 3629d072c68e99b9c42d3d39a7c19be18a445cb8 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Tue, 9 Sep 2025 10:10:19 -0400 Subject: [PATCH 3/5] fix quality Signed-off-by: Kyle Sayers --- .../test_compressors/quantized_compressors/test_nvfp4_quant.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_compressors/quantized_compressors/test_nvfp4_quant.py b/tests/test_compressors/quantized_compressors/test_nvfp4_quant.py index 496374a82..a53c1e60a 100644 --- a/tests/test_compressors/quantized_compressors/test_nvfp4_quant.py +++ b/tests/test_compressors/quantized_compressors/test_nvfp4_quant.py @@ -54,4 +54,4 @@ def test_pack_unpack_odd_dims(): ) with pytest.raises((ValueError, torch._dynamo.exc.Unsupported)): - _packed = pack_fp4_to_uint8(x) + _packed = pack_fp4_to_uint8(x) # noqa: F841 From 7b98675c0ada128de28f482e965f977bc2c571f3 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Tue, 9 Sep 2025 14:30:17 -0400 Subject: [PATCH 4/5] address Signed-off-by: Kyle Sayers --- src/compressed_tensors/logger.py | 22 +++++++--------------- 1 file changed, 7 insertions(+), 15 deletions(-) diff --git a/src/compressed_tensors/logger.py b/src/compressed_tensors/logger.py index 02d99402b..05c52cfdd 100644 --- a/src/compressed_tensors/logger.py +++ b/src/compressed_tensors/logger.py @@ -40,7 +40,7 @@ class LoggerConfig: log_file_level: Optional[str] = None -def configure_logger(config: Optional[LoggerConfig] = None) -> None: +def configure_logger(config: Optional[LoggerConfig] = None): """ Configure the logger for Compressed Tensors. This function sets up the console and file logging @@ -54,10 +54,10 @@ def configure_logger(config: Optional[LoggerConfig] = None) -> None: logger_config = config or LoggerConfig() # env vars get priority - if (disabled := os.getenv("COMPRESSED_TENSORS_LOG_DISABLED")) is not None: - logger_config.disabled = disabled.lower() == "true" - if (clear_loggers := os.getenv("COMPRESSED_TENSORS_CLEAR_LOGGERS")) is not None: - logger_config.clear_loggers = clear_loggers.lower() == "true" + if bool(os.getenv("COMPRESSED_TENSORS_LOG_DISABLED")): + logger_config.disabled = True + if bool(os.getenv("COMPRESSED_TENSORS_CLEAR_LOGGERS")): + logger_config.clear_loggers = True if (console_log_level := os.getenv("COMPRESSED_TENSORS_LOG_LEVEL")) is not None: logger_config.console_log_level = console_log_level.upper() if (log_file := os.getenv("COMPRESSED_TENSORS_LOG_FILE")) is not None: @@ -108,7 +108,7 @@ def support_log_once(record: Dict[str, Any]) -> bool: """ log_once = record["extra"].get("log_once", False) level = getattr(record["level"], "name", "none") - message = str(level) + record["message"] + message = hash(str(level) + record["message"]) if log_once and message in _logged_once: return False @@ -121,12 +121,4 @@ def support_log_once(record: Dict[str, Any]) -> bool: # 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, - ) -) +configure_logger(config=LoggerConfig()) From 7391fc4ef646fd88f3d264f2f5dcb939e34f509a Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Wed, 10 Sep 2025 09:34:17 -0400 Subject: [PATCH 5/5] [Utils] Improve type hints for `deprecated`, only log once (#455) * better type hints, warn once Signed-off-by: Kyle Sayers * remove unneeded import Signed-off-by: Kyle Sayers * use warnings Signed-off-by: Kyle Sayers --------- Signed-off-by: Kyle Sayers --- src/compressed_tensors/utils/helpers.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/compressed_tensors/utils/helpers.py b/src/compressed_tensors/utils/helpers.py index 712c4f837..ff2c6fc27 100644 --- a/src/compressed_tensors/utils/helpers.py +++ b/src/compressed_tensors/utils/helpers.py @@ -15,7 +15,7 @@ import contextlib import warnings from functools import wraps -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional, TypeVar import numpy import torch @@ -23,6 +23,9 @@ from transformers import AutoConfig +T = TypeVar("T", bound="Callable") # used by `deprecated` + + if TYPE_CHECKING: from compressed_tensors.compressors import ModelCompressor @@ -170,7 +173,9 @@ def getattr_chain(obj: Any, chain_str: str, *args, **kwargs) -> Any: return res -def deprecated(future_name: Optional[str] = None, message: Optional[str] = None): +def deprecated( + future_name: Optional[str] = None, message: Optional[str] = None +) -> Callable[[T], T]: """ Decorator to mark functions as deprecated @@ -178,7 +183,7 @@ def deprecated(future_name: Optional[str] = None, message: Optional[str] = None) :param message: Deprecation message, replaces default deprecation message """ - def decorator(func: Callable[[Any], Any]): + def decorator(func: T) -> T: nonlocal message if message is None: