Simple Memory Profiling for Python

Sometimes I find it useful to see how much memory does the python process take after we created some huge object. Usually objects are created and returned by functions, therefore I decided to make a decorator:

import time
import os
import psutil


def elapsed_since(start):
    return time.strftime("%H:%M:%S", time.gmtime(time.time() - start))


def get_process_memory():
    process = psutil.Process(os.getpid())
    return process.get_memory_info().rss


def profile(func):
    def wrapper(*args, **kwargs):
        mem_before = get_process_memory()
        start = time.time()
        result = func(*args, **kwargs)
        elapsed_time = elapsed_since(start)
        mem_after = get_process_memory()
        print("{}: memory before: {:,}, after: {:,}, consumed: {:,}; exec time: {}".format(
            func.__name__,
            mem_before, mem_after, mem_after - mem_before,
            elapsed_time))
        return result
    return wrapper

Let us see a sample of usage. Imagine we have a function that creates and returns some list:

from utils import profile

@profile
def list_create(n):
    print("inside a dummy list create")
    x = [1] * n
    return x

list_create(20000)

The result will look like this:

inside a dummy list_create
list_create: memory before: 45,899,776, after: 46,174,208, consumed: 274,432; exec time: 00:00:00

If you have some function func_x which calls func_y, but you don’t want to decorate the func_y (or you can’t do this because it is NOT your function but rather some library function), you can decorate just its call, like this:

@profile
def list_create_2(n):
    print("inside a dummy list_create_2")
    for i in range(1,10):
        x = [1,2,3] * n
        # x.sort()  - instead of doing this, do what is on the next line
        profile(list.sort)(x)
    return x

l = list_create_2(2000000)

Output:

inside a dummy list_create_2
sort: memory before: 94,121,984, after: 98,111,488, consumed: 3,989,504; exec time: 00:00:00
sort: memory before: 98,111,488, after: 110,002,176, consumed: 11,890,688; exec time: 00:00:00
sort: memory before: 110,014,464, after: 110,014,464, consumed: 0; exec time: 00:00:00
sort: memory before: 110,014,464, after: 110,014,464, consumed: 0; exec time: 00:00:00
sort: memory before: 110,014,464, after: 110,014,464, consumed: 0; exec time: 00:00:00
sort: memory before: 110,014,464, after: 110,014,464, consumed: 0; exec time: 00:00:00
sort: memory before: 110,014,464, after: 110,014,464, consumed: 0; exec time: 00:00:00
sort: memory before: 110,014,464, after: 110,014,464, consumed: 0; exec time: 00:00:00
sort: memory before: 110,014,464, after: 110,014,464, consumed: 0; exec time: 00:00:00
list_create_2: memory before: 46,006,272, after: 110,014,464, consumed: 64,008,192; exec time: 00:00:07

Leave a Comment

Your email address will not be published. Required fields are marked *

You may use these HTML tags and attributes: <a href="" title=""> <abbr title=""> <acronym title=""> <b> <blockquote cite=""> <cite> <code> <del datetime=""> <em> <i> <q cite=""> <s> <strike> <strong>