Skip to content

Commit 227a632

Browse files
tvalentyntvalentyn
authored
Add logger helper functions from detectron2 (licensed as Apache 2.0) (#36432)
* Add logger helper functions from detectron2 (licensed as Apache 2.0) * Type hints * Allow *args. * Add some tests. * Add license * lint
1 parent 62df216 commit 227a632

2 files changed

Lines changed: 245 additions & 0 deletions

File tree

‎sdks/python/apache_beam/utils/logger.py‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one or more
3+
# contributor license agreements. See the NOTICE file distributed with
4+
# this work for additional information regarding copyright ownership.
5+
# The ASF licenses this file to You under the Apache License, Version 2.0
6+
# (the "License"); you may not use this file except in compliance with
7+
# the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
#
17+
18+
"""Helper functions for easier logging.
19+
20+
This module provides a few convenient logging methods, some of which
21+
were adopted from
22+
https://github.com/abseil/abseil-py/blob/master/absl/logging/__init__.py
23+
in
24+
https://github.com/facebookresearch/detectron2/blob/main/detectron2/utils/logger.py
25+
"""
26+
import logging
27+
import os
28+
import sys
29+
import time
30+
from collections import Counter
31+
from types import FrameType
32+
from typing import Optional
33+
from typing import Union
34+
35+
36+
def _find_caller() -> tuple[str, tuple]:
37+
"""
38+
Returns:
39+
str: module name of the caller
40+
tuple: a hashable key to be used to identify different callers
41+
"""
42+
frame: Optional[FrameType] = sys._getframe(2)
43+
while frame:
44+
code = frame.f_code
45+
if os.path.join("utils", "logger.") not in code.co_filename:
46+
mod_name = frame.f_globals["__name__"]
47+
if mod_name == "__main__":
48+
mod_name = "apache_beam"
49+
return mod_name, (code.co_filename, frame.f_lineno, code.co_name)
50+
frame = frame.f_back
51+
52+
# To appease mypy. Code returns earlier in practice.
53+
return "unknown", ("unknown", 0, "unknown")
54+
55+
56+
_LOG_COUNTER = Counter()
57+
_LOG_TIMER = {}
58+
59+
60+
def log_first_n(
61+
lvl: int,
62+
msg: str,
63+
*args,
64+
n: int = 1,
65+
name: Optional[str] = None,
66+
key: Union[str, tuple[str]] = "caller") -> None:
67+
"""
68+
Log only for the first n times.
69+
70+
Args:
71+
lvl (int): the logging level
72+
msg (str):
73+
n (int):
74+
name (str): name of the logger to use. Will use the caller's module
75+
by default.
76+
key (str or tuple[str]): the string(s) can be one of "caller" or
77+
"message", which defines how to identify duplicated logs.
78+
For example, if called with `n=1, key="caller"`, this function
79+
will only log the first call from the same caller, regardless of
80+
the message content.
81+
If called with `n=1, key="message"`, this function will log the
82+
same content only once, even if they are called from different
83+
places. If called with `n=1, key=("caller", "message")`, this
84+
function will not log only if the same caller has logged the same
85+
message before.
86+
"""
87+
key_tuple = (key, ) if isinstance(key, str) else key
88+
assert len(key_tuple) > 0
89+
90+
caller_module, caller_key = _find_caller()
91+
hash_key: tuple = ()
92+
if "caller" in key_tuple:
93+
hash_key = hash_key + caller_key
94+
if "message" in key_tuple:
95+
hash_key = hash_key + (msg, )
96+
97+
_LOG_COUNTER[hash_key] += 1
98+
if _LOG_COUNTER[hash_key] <= n:
99+
logging.getLogger(name or caller_module).log(lvl, msg, *args)
100+
101+
102+
def log_every_n(
103+
lvl: int, msg: str, *args, n: int = 1, name: Optional[str] = None) -> None:
104+
"""
105+
Log once per n times.
106+
107+
Args:
108+
lvl (int): the logging level
109+
msg (str):
110+
n (int):
111+
name (str): name of the logger to use. Will use the caller's module
112+
by default.
113+
"""
114+
caller_module, key = _find_caller()
115+
_LOG_COUNTER[key] += 1
116+
if n == 1 or _LOG_COUNTER[key] % n == 1:
117+
logging.getLogger(name or caller_module).log(lvl, msg, *args)
118+
119+
120+
def log_every_n_seconds(
121+
lvl: int, msg: str, *args, n: int = 1, name: Optional[str] = None) -> None:
122+
"""
123+
Log no more than once per n seconds.
124+
125+
Args:
126+
lvl (int): the logging level
127+
msg (str):
128+
n (int):
129+
name (str): name of the logger to use. Will use the caller's module
130+
by default.
131+
"""
132+
caller_module, key = _find_caller()
133+
last_logged = _LOG_TIMER.get(key, None)
134+
current_time = time.time()
135+
if last_logged is None or current_time - last_logged >= n:
136+
logging.getLogger(name or caller_module).log(lvl, msg, *args)
137+
_LOG_TIMER[key] = current_time

‎sdks/python/apache_beam/utils/logger_test.py‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one or more
3+
# contributor license agreements. See the NOTICE file distributed with
4+
# this work for additional information regarding copyright ownership.
5+
# The ASF licenses this file to You under the Apache License, Version 2.0
6+
# (the "License"); you may not use this file except in compliance with
7+
# the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
#
17+
18+
import logging
19+
import unittest
20+
from unittest.mock import patch
21+
22+
import pytest
23+
24+
from apache_beam.utils.logger import _LOG_COUNTER
25+
from apache_beam.utils.logger import _LOG_TIMER
26+
from apache_beam.utils.logger import log_every_n
27+
from apache_beam.utils.logger import log_every_n_seconds
28+
from apache_beam.utils.logger import log_first_n
29+
30+
31+
@pytest.mark.no_xdist
32+
class TestLogFirstN(unittest.TestCase):
33+
def setUp(self):
34+
_LOG_COUNTER.clear()
35+
_LOG_TIMER.clear()
36+
37+
@patch('apache_beam.utils.logger.logging.getLogger')
38+
def test_log_first_n_once(self, mock_get_logger):
39+
mock_logger = mock_get_logger.return_value
40+
for _ in range(5):
41+
log_first_n(logging.INFO, "Test message %s", "arg", n=1)
42+
mock_logger.log.assert_called_once_with(
43+
logging.INFO, "Test message %s", "arg")
44+
45+
@patch('apache_beam.utils.logger.logging.getLogger')
46+
def test_log_first_n_multiple(self, mock_get_logger):
47+
mock_logger = mock_get_logger.return_value
48+
for _ in range(5):
49+
log_first_n(logging.INFO, "Test message %s", "arg", n=3)
50+
self.assertEqual(mock_logger.log.call_count, 3)
51+
mock_logger.log.assert_called_with(logging.INFO, "Test message %s", "arg")
52+
53+
@patch('apache_beam.utils.logger.logging.getLogger')
54+
def test_log_first_n_with_different_callers(self, mock_get_logger):
55+
mock_logger = mock_get_logger.return_value
56+
for _ in range(5):
57+
log_first_n(logging.INFO, "Test message", n=2)
58+
59+
# call from another "caller" (another line)
60+
for _ in range(5):
61+
log_first_n(logging.INFO, "Test message", n=2)
62+
63+
self.assertEqual(mock_logger.log.call_count, 4)
64+
65+
@patch('apache_beam.utils.logger.logging.getLogger')
66+
def test_log_first_n_with_message_key(self, mock_get_logger):
67+
mock_logger = mock_get_logger.return_value
68+
log_first_n(logging.INFO, "Test message", n=1, key="message")
69+
log_first_n(logging.INFO, "Test message", n=1, key="message")
70+
self.assertEqual(mock_logger.log.call_count, 1)
71+
72+
@patch('apache_beam.utils.logger.logging.getLogger')
73+
def test_log_first_n_with_caller_and_message_key(self, mock_get_logger):
74+
mock_logger = mock_get_logger.return_value
75+
for message in ["Test message", "Another message"]:
76+
for _ in range(5):
77+
log_first_n(logging.INFO, message, n=1, key=("caller", "message"))
78+
self.assertEqual(mock_logger.log.call_count, 2)
79+
80+
@patch('apache_beam.utils.logger.logging.getLogger')
81+
def test_log_every_n_multiple(self, mock_get_logger):
82+
mock_logger = mock_get_logger.return_value
83+
for _ in range(9):
84+
log_every_n(logging.INFO, "Test message", n=2)
85+
86+
self.assertEqual(mock_logger.log.call_count, 5)
87+
88+
@patch('apache_beam.utils.logger.logging.getLogger')
89+
@patch('apache_beam.utils.logger.time.time')
90+
def test_log_every_n_seconds_always(self, mock_time, mock_get_logger):
91+
mock_logger = mock_get_logger.return_value
92+
for i in range(3):
93+
mock_time.return_value = i
94+
log_every_n_seconds(logging.INFO, "Test message", n=0)
95+
self.assertEqual(mock_logger.log.call_count, 3)
96+
97+
@patch('apache_beam.utils.logger.logging.getLogger')
98+
@patch('apache_beam.utils.logger.time.time')
99+
def test_log_every_n_seconds_multiple(self, mock_time, mock_get_logger):
100+
mock_logger = mock_get_logger.return_value
101+
for i in range(4):
102+
mock_time.return_value = i
103+
log_every_n_seconds(logging.INFO, "Test message", n=2)
104+
self.assertEqual(mock_logger.log.call_count, 2)
105+
106+
107+
if __name__ == '__main__':
108+
unittest.main()

0 commit comments

Comments
 (0)