feat($profiling): support @cpu_profile & @mem_profile decorators

python_boilerplate/common/profiling.py

@@ -1,8 +1,10 @@
 import functools
+import os
 import time
 from datetime import timedelta
 from typing import Callable
 
+import psutil
 from loguru import logger
 
 
@@ -42,3 +44,95 @@ def wrapped(*arg, **kwarg):
         return wrapped
 
     return elapsed_time_wrapper
+
+
+def get_memory_usage():
+    process = psutil.Process(os.getpid())
+    mem_info = process.memory_info()
+    return mem_info.rss
+
+
+def get_cpu_usage():
+    """
+    Getting cpu_percent non-blocking (percentage since last call)
+    """
+    cpu_usage = psutil.cpu_percent()
+    return cpu_usage
+
+
+def mem_profile(level="INFO"):
+    """
+    The decorator to monitor the memory usage of a function.
+
+    Usage:
+
+    * decorate the function with `@mem_profile()` to profile the function with INFO log
+    >>> @mem_profile()
+    >>> def some_function():
+    >>>    pass
+
+    * decorate the function with `@mem_profile("DEBUG")` to profile the function with DEBUG log
+    >>> @mem_profile("DEBUG")
+    >>> def some_function():
+    >>>    pass
+
+    https://stackoverflow.com/questions/12295974/python-decorators-just-syntactic-sugar
+
+    :param level: logging level, default is "INFO". Available values: ["TRACE", "DEBUG", "INFO", "WARNING", "ERROR"]
+    """
+
+    def mem_profile_wrapper(func: Callable):
+        @functools.wraps(func)
+        def wrapped(*arg, **kwarg):
+            mem_before = get_memory_usage()
+            return_value = func(*arg, **kwarg)
+            mem_after = get_memory_usage()
+            logger.log(
+                level,
+                f"{func.__qualname__}() -> Mem before: {mem_before}, mem after: {mem_after}. "
+                f"Consumed memory: {(mem_after - mem_before) / (1024 * 1024):.2f} MB",
+            )
+            return return_value
+
+        return wrapped
+
+    return mem_profile_wrapper
+
+
+def cpu_profile(level="INFO"):
+    """
+    The decorator to monitor the CPU usage of a function.
+
+    Usage:
+
+    * decorate the function with `@mem_profile()` to profile the function with INFO log
+    >>> @cpu_profile()
+    >>> def some_function():
+    >>>    pass
+
+    * decorate the function with `@cpu_profile("DEBUG")` to profile the function with DEBUG log
+    >>> @mem_profile("DEBUG")
+    >>> def some_function():
+    >>>    pass
+
+    https://stackoverflow.com/questions/12295974/python-decorators-just-syntactic-sugar
+
+    :param level: logging level, default is "INFO". Available values: ["TRACE", "DEBUG", "INFO", "WARNING", "ERROR"]
+    """
+
+    def cpu_profile_wrapper(func: Callable):
+        @functools.wraps(func)
+        def wrapped(*arg, **kwarg):
+            cpu_before = get_cpu_usage()
+            return_value = func(*arg, **kwarg)
+            cpu_after = get_cpu_usage()
+            logger.log(
+                level,
+                f"{func.__qualname__}() -> CPU before: {cpu_before}, CPU after: {cpu_after}, "
+                f"delta: {(cpu_after - cpu_before):.2f}",
+            )
+            return return_value
+
+        return wrapped
+
+    return cpu_profile_wrapper

python_boilerplate/demo/pandas_usage.py

@@ -10,7 +10,7 @@
 from python_boilerplate.__main__ import startup
 from python_boilerplate.common.asynchronization import async_function
 from python_boilerplate.common.common_function import get_data_dir, get_resources_dir
-from python_boilerplate.common.profiling import elapsed_time
+from python_boilerplate.common.profiling import cpu_profile, elapsed_time, mem_profile
 from python_boilerplate.common.trace import async_trace
 
 # 10 minutes to pandas https://pandas.pydata.org/pandas-docs/stable/user_guide/10min.html#minutes-to-pandas
@@ -38,6 +38,8 @@ def pandas_data_structure_date_range() -> DatetimeIndex:
 
 @async_trace
 @elapsed_time("INFO")
+@mem_profile("INFO")
+@cpu_profile("INFO")
 def look_for_sony_published_games() -> DataFrame:
     all_columns = set(video_games)
     selected_columns = {

tests/demo/test_multithread_and_thread_pool_usage.py

@@ -1,12 +1,16 @@
 import pandas as pd
 
+from python_boilerplate.common.profiling import cpu_profile, elapsed_time, mem_profile
 from python_boilerplate.demo.multithread_and_thread_pool_usage import (
     COLUMN_NAMES,
     ROW_ARRAY,
     async_generate_data_frame,
 )
 
 
+@elapsed_time("INFO")
+@mem_profile("INFO")
+@cpu_profile("INFO")
 def test_async_generate_data_frame():
     data_frame = async_generate_data_frame()
     assert data_frame is not None

tests/demo/test_pandas_usage.py

@@ -5,6 +5,7 @@
 from loguru import logger
 from pandas import DatetimeIndex, Series
 
+from python_boilerplate.common.profiling import cpu_profile, elapsed_time, mem_profile
 from python_boilerplate.demo.pandas_usage import (
     data_generation,
     look_for_sony_published_games,
@@ -17,13 +18,19 @@
 # https://pandas.pydata.org/pandas-docs/stable/user_guide/10min.html#minutes-to-pandas
 
 
+@elapsed_time("INFO")
+@mem_profile("INFO")
+@cpu_profile("INFO")
 def test_pandas_data_structure_series() -> None:
     series: Series = pandas_data_structure_series()
     assert series.dtype == np.dtype("float64")
     assert len(series) == 6
     logger.info(f"series:\n{series}")
 
 
+@elapsed_time("INFO")
+@mem_profile("INFO")
+@cpu_profile("INFO")
 def test_pandas_data_structure_date_range() -> None:
     dates: DatetimeIndex = pandas_data_structure_date_range()
     assert dates.dtype == np.dtype("datetime64[ns]")
@@ -33,6 +40,9 @@ def test_pandas_data_structure_date_range() -> None:
     logger.info(f"dates:\n{dates}")
 
 
+@elapsed_time("INFO")
+@mem_profile("INFO")
+@cpu_profile("INFO")
 def test_pandas_reading_csv() -> None:
     """
     Read a csv file.
@@ -55,13 +65,19 @@ def test_pandas_reading_csv() -> None:
     assert video_games.iloc[1]["Release.Console"] == "Sony PSP"
 
 
+@elapsed_time("INFO")
+@mem_profile("INFO")
+@cpu_profile("INFO")
 def test_look_for_sony_published_games():
     sony_published_games = look_for_sony_published_games()
     assert sony_published_games is not None
     assert len(sony_published_games) == 9
     assert Path(sony_published_video_games_path).exists(), "CSV file NOT exists!"
 
 
+@elapsed_time("INFO")
+@mem_profile("INFO")
+@cpu_profile("INFO")
 def test_data_generation():
     try:
         data_generation()

tests/demo/test_pgp_usage.py

@@ -4,6 +4,7 @@
 from loguru import logger
 from pgpy.constants import HashAlgorithm
 
+from python_boilerplate.common.profiling import cpu_profile, elapsed_time, mem_profile
 from python_boilerplate.demo.pgp_usage import (
     alice_public_key,
     alice_secret_key,
@@ -12,6 +13,9 @@
 )
 
 
+@elapsed_time("INFO")
+@mem_profile("INFO")
+@cpu_profile("INFO")
 def test_when_alice_sends_message_to_bob():
     plaintext = "Hello World!"
     signature = pgpy.PGPMessage.new(plaintext)
@@ -34,6 +38,9 @@ def test_when_alice_sends_message_to_bob():
     assert verified is not None
 
 
+@elapsed_time("INFO")
+@mem_profile("INFO")
+@cpu_profile("INFO")
 def test_when_bob_sends_message_to_alice():
     plaintext = "Hello World!"
     signature = pgpy.PGPMessage.new(plaintext)
@@ -56,6 +63,9 @@ def test_when_bob_sends_message_to_alice():
     assert verified is not None
 
 
+@elapsed_time("INFO")
+@mem_profile("INFO")
+@cpu_profile("INFO")
 def test_load_armor_ciphertext_and_then_verify():
     message_text = textwrap.dedent(
         """\