From f4c70ba6c3d2fe0ee54d0924ee30d83504b539a5 Mon Sep 17 00:00:00 2001
From: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>
Date: Fri, 11 Aug 2023 10:46:46 +0800
Subject: [PATCH] [Feature] Support filtering specified levels message (#187)

* Support filtering message

* minor fix
---
 opencompass/utils/logging.py | 73 ++++++++++++++++++++++++++++++++++--
 1 file changed, 69 insertions(+), 4 deletions(-)

diff --git a/opencompass/utils/logging.py b/opencompass/utils/logging.py
index 0631fa7c..a576ad3c 100644
--- a/opencompass/utils/logging.py
+++ b/opencompass/utils/logging.py
@@ -1,13 +1,78 @@
+import logging
+import os
+
 from mmengine.logging import MMLogger
 
+_nameToLevel = {
+    'CRITICAL': logging.CRITICAL,
+    'FATAL': logging.FATAL,
+    'ERROR': logging.ERROR,
+    'WARN': logging.WARNING,
+    'WARNING': logging.WARNING,
+    'INFO': logging.INFO,
+    'DEBUG': logging.DEBUG,
+    'NOTSET': logging.NOTSET,
+}
 
-def get_logger(log_level='INFO') -> MMLogger:
+
+def get_logger(log_level='INFO', filter_duplicate_level=None) -> MMLogger:
     """Get the logger for OpenCompass.
 
     Args:
         log_level (str): The log level. Default: 'INFO'. Choices are 'DEBUG',
             'INFO', 'WARNING', 'ERROR', 'CRITICAL'.
     """
-    return MMLogger.get_instance('OpenCompass',
-                                 logger_name='OpenCompass',
-                                 log_level=log_level)
+    logger = MMLogger.get_instance('OpenCompass',
+                                   logger_name='OpenCompass',
+                                   log_level=log_level)
+
+    if filter_duplicate_level is None:
+        # export OPENCOMPASS_FILTER_DUPLICATE_LEVEL=error
+        # export OPENCOMPASS_FILTER_DUPLICATE_LEVEL=error,warning
+        filter_duplicate_level = os.getenv(
+            'OPENCOMPASS_FILTER_DUPLICATE_LEVEL', None)
+
+    if filter_duplicate_level:
+        logger.addFilter(
+            FilterDuplicateMessage('OpenCompass', filter_duplicate_level))
+
+    return logger
+
+
+class FilterDuplicateMessage(logging.Filter):
+    """Filter the repeated message.
+
+    Args:
+        name (str): name of the filter.
+    """
+
+    def __init__(self, name, filter_duplicate_level):
+        super().__init__(name)
+        self.seen: set = set()
+
+        if isinstance(filter_duplicate_level, str):
+            filter_duplicate_level = filter_duplicate_level.split(',')
+
+        self.filter_duplicate_level = []
+        for level in filter_duplicate_level:
+            _level = level.strip().upper()
+            if _level not in _nameToLevel:
+                raise ValueError(f'Invalid log level: {_level}')
+            self.filter_duplicate_level.append(_nameToLevel[_level])
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        """Filter the repeated error message.
+
+        Args:
+            record (LogRecord): The log record.
+
+        Returns:
+            bool: Whether to output the log record.
+        """
+        if record.levelno not in self.filter_duplicate_level:
+            return True
+
+        if record.msg not in self.seen:
+            self.seen.add(record.msg)
+            return True
+        return False