trace-func-memory.py

import math
import os
import threading
import time

# pip install psutil
import psutil


def process_memory(pid):
    process = psutil.Process(pid)
    mem_info = process.memory_info()
    return mem_info.rss


def add_to_avg(avg, count, val):
    return (avg * count + val) / (count + 1)


def trace_memory(func, interval=0.5):
    def __f1(*args, **kwargs):
        t0 = time.time()
        x = True
        pid = os.getpid()
        start_mem = process_memory(pid)

        def __f():
            count = 0
            max_mem = 0
            avg_mem = 0

            while x:
                mem = process_memory(pid) - start_mem
                max_mem = max(max_mem, mem)
                avg_mem = add_to_avg(avg_mem, count, mem)
                time.sleep(interval)
                count += 1
            print("\nFunction {} Memory Usage: max:{} avg:{} over {}\n".format(func.__name__, sizeof_fmt(max_mem),
                                                                               sizeof_fmt(avg_mem),
                                                                               duration_fmt(
                                                                                   math.floor(time.time() - t0))))

        t = threading.Thread(target=__f)
        t.start()
        try:
            func(*args, **kwargs)
        except Exception as e:
            x = False
            t.join()
            raise e

    return __f1


def sizeof_fmt(num, suffix="B"):
    for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]:
        if abs(num) < 1024.0:
            return f"{num:3.1f}{unit}{suffix}"
        num /= 1024.0
    return f"{num:.1f}Yi{suffix}"
    
def duration_fmt(seconds):
    return '{0} hours, {1} minutes, {2} seconds'.format(seconds // 3600, seconds % 3600 // 60, seconds % 60)

#  example
if __name__ == "__main__":
    @trace_memory
    def test_trace_memory(ir):
        for i in range(0, ir):
            time.sleep(0.001)
        print("done")


    test_trace_memory(200)