Add logger helper functions from detectron2 (licensed as Apache 2.0)

sdks/python/apache_beam/utils/logger.py

@@ -0,0 +1,137 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You 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.
+#
+
+"""Helper functions for easier logging.
+
+This module provides a few convenient logging methods, some of which
+were adopted from
+https://github.com/abseil/abseil-py/blob/master/absl/logging/__init__.py
+in
+https://github.com/facebookresearch/detectron2/blob/main/detectron2/utils/logger.py
+"""
+import logging
+import os
+import sys
+import time
+from collections import Counter
+from types import FrameType
+from typing import Optional
+from typing import Union
+
+
+def _find_caller() -> tuple[str, tuple]:
+  """
+    Returns:
+        str: module name of the caller
+        tuple: a hashable key to be used to identify different callers
+    """
+  frame: Optional[FrameType] = sys._getframe(2)
+  while frame:
+    code = frame.f_code
+    if os.path.join("utils", "logger.") not in code.co_filename:
+      mod_name = frame.f_globals["__name__"]
+      if mod_name == "__main__":
+        mod_name = "apache_beam"
+      return mod_name, (code.co_filename, frame.f_lineno, code.co_name)
+    frame = frame.f_back
+
+  # To appease mypy. Code returns earlier in practice.
+  return "unknown", ("unknown", 0, "unknown")
+
+
+_LOG_COUNTER = Counter()
+_LOG_TIMER = {}
+
+
+def log_first_n(
+    lvl: int,
+    msg: str,
+    *args,
+    n: int = 1,
+    name: Optional[str] = None,
+    key: Union[str, tuple[str]] = "caller") -> None:
+  """
+    Log only for the first n times.
+
+    Args:
+        lvl (int): the logging level
+        msg (str):
+        n (int):
+        name (str): name of the logger to use. Will use the caller's module
+          by default.
+        key (str or tuple[str]): the string(s) can be one of "caller" or
+            "message", which defines how to identify duplicated logs.
+            For example, if called with `n=1, key="caller"`, this function
+            will only log the first call from the same caller, regardless of
+            the message content.
+            If called with `n=1, key="message"`, this function will log the
+            same content only once, even if they are called from different
+            places. If called with `n=1, key=("caller", "message")`, this
+            function will not log only if the same caller has logged the same
+            message before.
+    """
+  key_tuple = (key, ) if isinstance(key, str) else key
+  assert len(key_tuple) > 0
+
+  caller_module, caller_key = _find_caller()
+  hash_key: tuple = ()
+  if "caller" in key_tuple:
+    hash_key = hash_key + caller_key
+  if "message" in key_tuple:
+    hash_key = hash_key + (msg, )
+
+  _LOG_COUNTER[hash_key] += 1
+  if _LOG_COUNTER[hash_key] <= n:
+    logging.getLogger(name or caller_module).log(lvl, msg, *args)
+
+
+def log_every_n(
+    lvl: int, msg: str, *args, n: int = 1, name: Optional[str] = None) -> None:
+  """
+    Log once per n times.
+
+    Args:
+        lvl (int): the logging level
+        msg (str):
+        n (int):
+        name (str): name of the logger to use. Will use the caller's module
+          by default.
+    """
+  caller_module, key = _find_caller()
+  _LOG_COUNTER[key] += 1
+  if n == 1 or _LOG_COUNTER[key] % n == 1:
+    logging.getLogger(name or caller_module).log(lvl, msg, *args)
+
+
+def log_every_n_seconds(
+    lvl: int, msg: str, *args, n: int = 1, name: Optional[str] = None) -> None:
+  """
+    Log no more than once per n seconds.
+
+    Args:
+        lvl (int): the logging level
+        msg (str):
+        n (int):
+        name (str): name of the logger to use. Will use the caller's module
+          by default.
+    """
+  caller_module, key = _find_caller()
+  last_logged = _LOG_TIMER.get(key, None)
+  current_time = time.time()
+  if last_logged is None or current_time - last_logged >= n:
+    logging.getLogger(name or caller_module).log(lvl, msg, *args)
+    _LOG_TIMER[key] = current_time

sdks/python/apache_beam/utils/logger_test.py

@@ -0,0 +1,108 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You 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 logging
+import unittest
+from unittest.mock import patch
+
+import pytest
+
+from apache_beam.utils.logger import _LOG_COUNTER
+from apache_beam.utils.logger import _LOG_TIMER
+from apache_beam.utils.logger import log_every_n
+from apache_beam.utils.logger import log_every_n_seconds
+from apache_beam.utils.logger import log_first_n
+
+
+@pytest.mark.no_xdist
+class TestLogFirstN(unittest.TestCase):
+  def setUp(self):
+    _LOG_COUNTER.clear()
+    _LOG_TIMER.clear()
+
+  @patch('apache_beam.utils.logger.logging.getLogger')
+  def test_log_first_n_once(self, mock_get_logger):
+    mock_logger = mock_get_logger.return_value
+    for _ in range(5):
+      log_first_n(logging.INFO, "Test message %s", "arg", n=1)
+    mock_logger.log.assert_called_once_with(
+        logging.INFO, "Test message %s", "arg")
+
+  @patch('apache_beam.utils.logger.logging.getLogger')
+  def test_log_first_n_multiple(self, mock_get_logger):
+    mock_logger = mock_get_logger.return_value
+    for _ in range(5):
+      log_first_n(logging.INFO, "Test message %s", "arg", n=3)
+    self.assertEqual(mock_logger.log.call_count, 3)
+    mock_logger.log.assert_called_with(logging.INFO, "Test message %s", "arg")
+
+  @patch('apache_beam.utils.logger.logging.getLogger')
+  def test_log_first_n_with_different_callers(self, mock_get_logger):
+    mock_logger = mock_get_logger.return_value
+    for _ in range(5):
+      log_first_n(logging.INFO, "Test message", n=2)
+
+    # call from another "caller" (another line)
+    for _ in range(5):
+      log_first_n(logging.INFO, "Test message", n=2)
+
+    self.assertEqual(mock_logger.log.call_count, 4)
+
+  @patch('apache_beam.utils.logger.logging.getLogger')
+  def test_log_first_n_with_message_key(self, mock_get_logger):
+    mock_logger = mock_get_logger.return_value
+    log_first_n(logging.INFO, "Test message", n=1, key="message")
+    log_first_n(logging.INFO, "Test message", n=1, key="message")
+    self.assertEqual(mock_logger.log.call_count, 1)
+
+  @patch('apache_beam.utils.logger.logging.getLogger')
+  def test_log_first_n_with_caller_and_message_key(self, mock_get_logger):
+    mock_logger = mock_get_logger.return_value
+    for message in ["Test message", "Another message"]:
+      for _ in range(5):
+        log_first_n(logging.INFO, message, n=1, key=("caller", "message"))
+    self.assertEqual(mock_logger.log.call_count, 2)
+
+  @patch('apache_beam.utils.logger.logging.getLogger')
+  def test_log_every_n_multiple(self, mock_get_logger):
+    mock_logger = mock_get_logger.return_value
+    for _ in range(9):
+      log_every_n(logging.INFO, "Test message", n=2)
+
+    self.assertEqual(mock_logger.log.call_count, 5)
+
+  @patch('apache_beam.utils.logger.logging.getLogger')
+  @patch('apache_beam.utils.logger.time.time')
+  def test_log_every_n_seconds_always(self, mock_time, mock_get_logger):
+    mock_logger = mock_get_logger.return_value
+    for i in range(3):
+      mock_time.return_value = i
+      log_every_n_seconds(logging.INFO, "Test message", n=0)
+    self.assertEqual(mock_logger.log.call_count, 3)
+
+  @patch('apache_beam.utils.logger.logging.getLogger')
+  @patch('apache_beam.utils.logger.time.time')
+  def test_log_every_n_seconds_multiple(self, mock_time, mock_get_logger):
+    mock_logger = mock_get_logger.return_value
+    for i in range(4):
+      mock_time.return_value = i
+      log_every_n_seconds(logging.INFO, "Test message", n=2)
+    self.assertEqual(mock_logger.log.call_count, 2)
+
+
+if __name__ == '__main__':
+  unittest.main()