monitor_peak_memory

mufasa.utils.memory.monitor_peak_memory(output_attr=None, key=None, unit='GB')

Decorator to monitor and record the peak memory usage of a function.

This decorator measures the peak memory usage, including intermediate peaks, during the execution of a function. It supports multi-threaded tasks and can store results in an instance attribute or print them.

Parameters:

• output_attr (str or None, optional) – The name of the instance attribute to store the peak memory usage. If None, the memory usage is printed. If specified and a key is provided, the attribute is treated as a dictionary. Default is None.

• key (str or None, optional) – The dictionary key to use when storing memory usage in the attribute specified by output_attr. If None, the value is stored directly in the attribute. Default is None.

• unit ({'KB', 'MB', 'GB', 'TB'}, default='GB') – The unit for reporting memory usage. If an unrecognized unit is given, a warning is logged, and the default (‘GB’) is used.

Returns:

A decorator that monitors peak memory usage for the decorated function.

Return type:

callable

Raises:

ValueError – If output_attr is specified as a non-dictionary attribute and a key is also provided.

Examples

Monitor memory usage and print results:

>>> @monitor_peak_memory()
... def my_function(self):
...     # Function logic here
...     pass

Store memory usage in an instance attribute:

>>> @monitor_peak_memory(output_attr='memory_usage', key='task1')
... def my_function(self):
...     # Function logic here
...     pass
>>> self.memory_usage['task1']  # Access the stored peak memory

Use a custom unit for memory usage:

>>> @monitor_peak_memory(unit='MB')
... def my_function(self):
...     pass