Optimize Towhee logging module (#2640)

Signed-off-by: Kaiyuan Hu <kaiyuan.hu@zilliz.com>

tests/unittests/utils/test_log.py

@@ -0,0 +1,90 @@
+# Copyright 2021 Zilliz. 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.
+import re
+import logging
+import unittest
+from pathlib import Path
+from io import StringIO
+
+from towhee.config import LogConfig
+from towhee.utils.log import engine_log, enable_log_config
+
+
+class TestLog(unittest.TestCase):
+    """
+    Test towhee logs.
+    """
+    def test_log(self):
+        f = StringIO()
+
+        with self.assertRaises(ValueError):
+            config = LogConfig(mode='test')
+
+        with self.assertRaises(ValueError):
+            config = LogConfig(rotate_mode='test')
+
+        config = LogConfig(mode='console', rotate_mode='time', stream=f)
+        root = logging.getLogger()
+
+        self.assertEqual(root.level, 30)
+        self.assertEqual(engine_log.level, 0)
+
+        # Default stream log
+        enable_log_config()
+        engine_log.warning('test_warning')
+        pattern = r'\d*-\d*-\d* \d*:\d*:\d*,\d* - \d* - test_log.py-test_log:\d* - WARNING: test_warning'
+        self.assertTrue(re.search(pattern, f.getvalue()) is not None)
+        self.assertEqual(root.level, 30)
+
+        # Debug level stream log
+        config.stream_level = 'DEBUG'
+        enable_log_config(config)
+        engine_log.debug('test_debug')
+        pattern = r'\d*-\d*-\d* \d*:\d*:\d*,\d* - \d* - test_log.py-test_log:\d* - DEBUG: test_debug'
+        self.assertTrue(re.search(pattern, f.getvalue()) is not None)
+        self.assertEqual(root.level, 10)
+
+        f.close()
+
+        # Default file log
+        config.mode = 'file'
+        enable_log_config(config)
+        engine_log.info('test_info')
+        pattern = r'\d*-\d*-\d* \d*:\d*:\d*,\d* - \d* - test_log.py-test_log:\d* - INFO: test_info'
+        with open('towhee.log', encoding='utf-8') as file:
+            self.assertTrue(re.search(pattern, file.read()) is not None)
+        self.assertEqual(root.level, 20)
+        self.assertIsInstance(root.handlers[-1], logging.handlers.TimedRotatingFileHandler)
+        Path('towhee.log').unlink()
+
+        # Error level file log
+        config.mode = 'file'
+        config.rotate_mode = 'size'
+        config.file_level = 'ERROR'
+        enable_log_config(config)
+        engine_log.error('test_error')
+        pattern = r'\d*-\d*-\d* \d*:\d*:\d*,\d* - \d* - test_log.py-test_log:\d* - ERROR: test_error'
+        with open('towhee.log', encoding='utf-8') as file:
+            self.assertTrue(re.search(pattern, file.read()) is not None)
+        self.assertEqual(root.level, 40)
+        self.assertIsInstance(root.handlers[-1], logging.handlers.RotatingFileHandler)
+        Path('towhee.log').unlink()
+
+        with self.assertRaises(ValueError):
+            config.mode = 'test'
+            enable_log_config()
+
+        with self.assertRaises(ValueError):
+            config.rotate_mode = 'test'
+            enable_log_config()

towhee/config/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2021 Zilliz. 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.
+
+from .log_config import LogConfig

towhee/config/log_config.py

@@ -0,0 +1,69 @@
+# Copyright 2021 Zilliz. 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.
+from typing import Optional, Any
+
+from pydantic import BaseModel, validator
+from towhee.utils.singleton import singleton
+
+
+@singleton
+class LogConfig(BaseModel):
+    """
+    Towhee logger config.
+
+    Args:
+        mode (`Optional[str]`):
+            The mode of Logger, decide which Handler to use, 'file' or 'console', defaults to 'console'.
+        rotate_mode (`optional[str]`):
+            The mode of rotating files, 'time' or 'size', defaults to 'time'.
+        filename (`Optional[str]`):
+            Path for log files, defaults to 'towhee.log'.
+        rotate_when (`Optional[str]`):
+            The type of TimedRotatingFileHandler interval, supports 'S', 'M', 'H', 'D', 'W0'-'W6', 'midnight', deafults to 'D'.
+        rotate_interval (`Optional[int]`):
+            The interval value of timed rotating, defaults to 1.
+        backup_count (`Optional[int]`):
+            The number of log files to keep, defaults to 10.
+        stream_level (`Optional[str]`):
+            The log level in console mode, defaults to 'WARNING'.
+        file_level (`Optional[str]`):
+            The log level in file mode, defautls to 'INFO'.
+        utc (`Optional[bool]`):
+            Whether to use utc time, defaults to False.
+        stream (`Optional[Any]`):
+            The stream to write logs into in console mode, defeaults to None and sys.stderr is used.
+    """
+    mode: Optional[str] = 'console'
+    rotate_mode: Optional[str] = 'time'
+    filename: Optional[str] = 'towhee.log'
+    file_max_size: Optional[int] = 100 * 1000 * 1000
+    rotate_when: Optional[str] = 'D'
+    rotate_interval: Optional[int] = 1
+    backup_count: Optional[int] = 10
+    stream_level: Optional[str] = 'WARNING'
+    file_level: Optional[str] = 'INFO'
+    utc: Optional[bool] = False
+    stream: Optional[Any] = None
+
+    @validator('mode')
+    def mode_range(cls, mode):  # pylint: disable=no-self-argument
+        if mode not in ['console', 'file']:
+            raise ValueError(f'value should be either \'console\' or \'file\', not \'{mode}\'.')
+        return mode
+
+    @validator('rotate_mode')
+    def rotate_mode_range(cls, rotate_mode):  # pylint: disable=no-self-argument
+        if rotate_mode not in ['time', 'size']:
+            raise ValueError(f'value should be either \'time\' or \'size\', not \'{rotate_mode}\'.')
+        return rotate_mode

towhee/utils/log.py

@@ -11,13 +11,90 @@
 # 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.
-
-
 import logging
+from towhee.config.log_config import LogConfig
 
 FORMAT = '%(asctime)s - %(thread)d - %(filename)s-%(module)s:%(lineno)s - %(levelname)s: %(message)s'
-logging.basicConfig(format=FORMAT)
 
+# TODO: for service `access_log`
+# TODO: multi progress log
+# TODO: user-defined log
+logging.basicConfig(format=FORMAT)
+root_log = logging.getLogger()
 engine_log = logging.getLogger('towhee.engine')
 trainer_log = logging.getLogger('towhee.trainer')
 models_log = logging.getLogger('towhee.models')
+
+formatter = logging.Formatter(FORMAT)
+
+
+def _get_time_file_hdlr(config):
+    time_rotating_handler = logging.handlers.TimedRotatingFileHandler(
+        filename=config.filename,
+        when=config.rotate_when,
+        interval=config.rotate_interval,
+        backupCount=config.backup_count,
+        utc=config.utc,
+        encoding='utf-8',
+        delay=True
+    )
+
+    time_rotating_handler.setFormatter(formatter)
+    time_rotating_handler.setLevel(config.file_level)
+    time_rotating_handler.set_name('time_rotating_file')
+
+    return time_rotating_handler
+
+
+def _get_size_file_hdlr(config):
+    size_rotating_handler = logging.handlers.RotatingFileHandler(
+        filename=config.filename,
+        maxBytes=config.file_max_size,
+        backupCount=config.backup_count,
+        encoding='utf-8',
+        delay=True
+    )
+    size_rotating_handler.setFormatter(formatter)
+    size_rotating_handler.setLevel(config.file_level)
+    size_rotating_handler.set_name('size_rotating_file')
+
+    return size_rotating_handler
+
+
+def _get_console_hdlr(config):
+    console_handler = logging.StreamHandler(config.stream)
+    console_handler.setFormatter(formatter)
+    console_handler.setLevel(config.stream_level)
+    console_handler.set_name('console')
+
+    return console_handler
+
+
+def _config_logger(logger, handler):
+    for hdlr in logger.handlers:
+        logger.removeHandler(hdlr)
+
+    logger.addHandler(handler)
+
+
+def enable_log_config(config: 'towhee.config.LogConfig' = None):
+    if not config:
+        config = LogConfig()
+
+    time_rotating_handler = _get_time_file_hdlr(config)
+    size_rotating_handler = _get_size_file_hdlr(config)
+    console_handler = _get_console_hdlr(config)
+
+    if config.mode == 'console':
+        _config_logger(root_log, console_handler)
+        root_log.setLevel(config.stream_level)
+    elif config.mode == 'file':
+        if config.rotate_mode == 'size':
+            _config_logger(root_log, size_rotating_handler)
+        elif config.rotate_mode == 'time':
+            _config_logger(root_log, time_rotating_handler)
+        else:
+            raise ValueError(f'value should be either \'time\' or \'size\', not \'{config.rotate_mode}\'.')
+        root_log.setLevel(config.file_level)
+    else:
+        raise ValueError(f'value should be either \'console\' or \'file\', not \'{config.mode}\'.')