Limit cache size

Author: trungleduc

src/sagemaker/config/config_utils.py

@@ -15,6 +15,7 @@
 These utils may be used inside or outside the config module.
 """
 from __future__ import absolute_import
+from collections import deque
 
 import logging
 import sys
@@ -200,26 +201,31 @@ def _log_sagemaker_config_merge(
         logger.debug("Skipped value because no value defined\n  config key = %s", config_key_path)
 
 
-def non_repeating_log_factory(logger: logging.Logger, method: str) -> Callable:
+def non_repeating_log_factory(logger: logging.Logger, method: str, cache_size=100) -> Callable:
     """Create log function that filters the repeated messages.
 
+    By default. It only keeps track of last 100 messages, if a repeated
+    message arrives after the ``cache_size`` messages, it will be displayed.
+
     Args:
         logger (logging.Logger): the logger to be used to dispatch the message.
         method (str): the log method, can be info, warning or debug.
+        cache_size (int): the number of last log messages to keep in cache.
+            Default to 100
 
     Returns:
         (Callable): the new log method
     """
     if method not in ["info", "warning", "debug"]:
         raise ValueError("Not supported logging method.")
 
-    _caches = set()
+    _caches = deque(maxlen=cache_size)
     log_method = getattr(logger, method)
 
     def new_log_method(msg, *args, **kwargs):
         key = f"{msg}:{args}"
         if key not in _caches:
             log_method(msg, *args, **kwargs)
-            _caches.add(key)
+            _caches.append(key)
 
     return new_log_method

tests/unit/sagemaker/config/test_config.py

@@ -458,3 +458,21 @@ def test_non_repeating_log_factory(method_name):
     log_function("foo")
 
     mock.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "method_name",
+    ["info", "warning", "debug"],
+)
+def test_non_repeating_log_factory_cache_size(method_name):
+    tmp_logger = logging.getLogger("test-logger")
+    mock = MagicMock()
+    setattr(tmp_logger, method_name, mock)
+
+    log_function = non_repeating_log_factory(tmp_logger, method_name, cache_size=2)
+    log_function("foo")
+    log_function("bar")
+    log_function("foo2")
+    log_function("foo")
+
+    assert mock.call_count == 4