main

関数のメモリ使用量をリアルタイムでプロットするコードを書いた。
アプリケーションのリソースを調査する過程で製作。
実際のところログで吐き出せばいいのでリアルタイムでプロットする意味はあまりないのだが、勉強になった。

リアルタイムプロットの参考
https://qiita.com/dendensho/items/79e9d2e3d4e8eb5061bc

from multiprocessing import Process, Value
import os
import psutil
import matplotlib.pyplot as plt

def memory_check(pid=os.getpid()):
    '''
    与えられたPIDの物理メモリと仮想メモリ使用量を返す。
    Windowsの場合は仮想メモリ=物理メモリ(共有部分を除く)+ページファイル
    '''
    unit = 1024**2
    mem_info = psutil.Process(pid=pid).memory_info()
    pmem = mem_info.rss / unit
    vmem = mem_info.vms / unit
    return pmem, vmem

def plotting(status, PID):
    '''グラフプロット関数
    status :
        0-5の整数値
        plotting関数のステータス
    PID :
        PID
    '''

    # 初期化
    times = [0 for i in range(100)]
    pmems = [0 for i in range(100)]
    vmems = [0 for i in range(100)]

    plt.ion()
    plt.figure()
    li_p, = plt.plot(times, pmems, label='Physical')
    li_v, = plt.plot(times, vmems, label='Virtual')

    plt.legend()

    plt.xlabel('time')
    plt.ylabel('used memory (MB)')
    plt.title('memory real time plot')

    current_time = 0
    status.value = 1 # 初期化終了フラグ
    # メインループ
    while True:
        if status.value == 1:
            # アプリケーション起動待ち
            plt.pause(0.1)
        elif status.value == 2:
            # リアルタイムプロット
            current_time += 0.1
            mem = memory_check(PID.value)

            times.append(current_time)
            times.pop(0)
            pmems.append(mem[0])
            pmems.pop(0)
            vmems.append(mem[1])
            vmems.pop(0)

            li_p.set_xdata(times)
            li_p.set_ydata(pmems)
            li_v.set_xdata(times)
            li_v.set_ydata(vmems)

            plt.xlim(min(times), max(times))
            plt.ylim(0, max(vmems)*1.1)
            plt.draw()

            plt.pause(0.1)
        elif status.value == 3:
            # グラフの更新を停止
            plt.ioff()
            status.value = 4
        elif status.value == 4:
            # 終了待ち
            plt.pause(0.1)
        elif status.value == 5:
            # グラフを閉じる
            plt.clf()
            plt.close()
            return 0
        else:
            raise ValueError("statusが異常です")

def realtime_mem_plot(func):
    def wrapper(*args):
        status = Value('i', 0)
        PID = Value('i', 0)
        p_plot = Process(target=plotting, args=(status, PID))
        p_plot.start()

        p_main = Process(target=func, args=(*args,))
        p_main.start()

        print('start plotting')
        parent = psutil.Process(os.getpid())
        children = parent.children()
        PID.value = children[-1].pid
        print(f'PID: {children[-1].pid}')

        status.value = 2

        p_main.join()
        status.value = 3

        input("Please press Enter key to close")

        status.value = 5
        p_plot.join()
    return wrapper

from time import sleep

# @realtime_mem_plot
def f(name):
    print(name)
    for j in range(10):
        k = []
        for i in range(100):
            k.append([0]*10000*j)
            sleep(0.01)
        del k
    return 10

if __name__=='__main__':
    realtime_mem_plot(f)('test')
    # f('test')

wrapper化不可能?

wrapper化すると便利かと思い試行錯誤するも、以下のエラーがでて止まっている。
(環境: Windows 10, Python 3.10.10)

Exception has occurred: PicklingError
Can't pickle <function f at 0x000001E06FFB7F40>: it's not the same object as __main__.f
  File "C:\Users\kubot\Documents\temp\realtime_mem.py", line 95, in wrapper
    p_main.start()
  File "C:\Users\kubot\Documents\temp\realtime_mem.py", line 129, in <module>
    f('test')
_pickle.PicklingError: Can't pickle <function f at 0x000001E06FFB7F40>: it's not the same object as __main__.f

wrapper化していない現行のコードだとなぜか実行可能

start_methodを'fork'にすれば行けるという記事もみたが ( https://medium.com/devopss-hole/python-multiprocessing-pickle-issue-e2d35ccf96a9 )、Windowsではこの方法が使えない・・・

改良点

plt.pause(0.1)が必ずしも0.1秒ちょうどとは限らないので実行時間と横軸がずれる
→timeを計測して横軸に設定する？
関数の返り値を受け取れるようにする
→wrapperの中にさらにwrapperを用いる構造にする
測定関数で子プロセスを呼び出したときの動作が未確認
→実際のところ、子プロセスのメモリ使用量が追えないだけだと予想

まとめ

wrapper化で便利ツール化できないかと試行するも、思わぬとろころでつまづいた…
解決策がみつかれば追記しようと思う

また、どなたかご存知の方がいらっしゃればコメントください。