mem_usage.py

import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from clize import run
from subprocess import call
import time
import pandas as pd


def loop(pid, out=None):
    cmd = 'ps p {PID} -o %mem=,rss='.format(PID=pid)
    filename = 'mem_usage_{PID}'.format(PID=pid)
    if out is None:
        out = 'mem_usage_{PID}.png'.format(PID=pid)
    fd = open(filename, 'w')
    fd.close()
    while True:
        fd = open(filename, 'a')
        call(cmd, shell=True, stdout=fd)
        fd.close()
        draw(filename, out, title='PID : {PID}'.format(PID=pid))
        time.sleep(1)


def draw(filename, out, title=''):
    df = pd.read_table(
            filename,
            header=0,
            names=['%mem', 'rss'],
            sep=' ')
    df = df.reset_index()
    if len(df) == 0:
        return
    fig = plt.figure(figsize=(20, 5))
    gs = gridspec.GridSpec(1, 2)
    ax1 = plt.subplot(gs[0])
    ax2 = plt.subplot(gs[1])
    ax1.plot(df['%mem'], label='%mem', c='blue')
    ax1.legend(loc='best')
    ax2.plot(df['rss']/(1024.), label='Resident Set Size(MB)', c='green')
    ax2.legend(loc='best')
    ax1.set_title(title)
    ax2.set_title(title)
    plt.savefig(out)
    plt.show()
    plt.close(fig)


if __name__ == '__main__':
    run(loop)