Python函数内存消耗测量需求：嵌套函数场景下的难题解决

核心需求 #

• 实现Python函数内存消耗日志记录，测量每个函数执行期间进程的峰值内存并存储，最终生成包含执行日期的日志文件
• 代码需部署，禁止使用memory-profiler的mprof这类命令行工具
• 无需修改目标函数代码，通过包装器实现内存测量

现有尝试的问题 #

psutil方案缺陷 #

使用psutil.Process.memory_info仅能获取实时内存数据，其中peak_wset不会自动重置，前序函数的高内存消耗会导致后续测量结果失真。例如测试中f函数执行前后rss、vms值无变化，差值为0；第二次执行f后peak_wset未更新，无法反映函数真实的内存峰值。

tracemalloc方案缺陷 #

tracemalloc功能接近需求，但无法准确测量嵌套函数的内存消耗：首次调用tracemalloc.start()后，后续调用不会重置测量上下文；若中途调用tracemalloc.stop()，则测量结果完全失效，不符合嵌套场景的需求。

测试代码 #

import tracemalloc
from typing import Callable
import numpy as np
import psutil

d: dict[str, dict[str, tuple]] = {}

def memory_mesure(func: Callable):
    def memory_wrapper(*args, **kwargs):
        tracemalloc.start()
        s = psutil.Process().memory_info()
        key = f"{func.__name__}_{len(d)}"
        d[key] = {}
        d[key]["start"] = (s.rss, s.vms, s.peak_wset)
        # d[key]["start"] = tracemalloc.get_traced_memory()
        func(*args, **kwargs)
        e = psutil.Process().memory_info()
        d[key]["end"] = (e.rss, e.vms, e.peak_wset)
        # d[key]["end"] = tracemalloc.get_traced_memory()
        # tracemalloc.stop()

    return memory_wrapper

@memory_mesure
def f(i):
    r = {}
    r[0] = np.zeros((i,))
    del r[0]
    r[1] = np.zeros((i,))
    del r[1]

@memory_mesure
def g():
    f(1000000000)
    f(100000)

g()

psutil测试输出结果 #

g_0
    start - (31166464, 257990656, 31166464)
    end - (31170560, 257990656, 31174656)
    => (4096, 0, 8192)
f_1
    start - (31170560, 257990656, 31170560)
    end - (31170560, 257990656, 31174656)
    => (0, 0, 4096)
f_2
    start - (31170560, 257990656, 31174656)
    end - (31170560, 257990656, 31174656)
    => (0, 0, 0)

预期结果 #

• g_0、f_1和f_2的内存消耗值均严格为正
• f_1的内存消耗值严格大于f_2（因为其创建的数组规模远大于f_2）
• g_0的内存消耗值大于f_1和f_2（因为g_0的执行包含两次f函数调用）

内容的提问来源于stack exchange，提问作者Paulo-99