Python: profile total memory allocated with tracemalloc

tracemalloc is Python’s standard library module for tracking memory allocations. It has many bells and whistles for detailed analysis, allowing you to slice allocations by file and line or compare snapshots. But for simple purposes, displaying the total memory allocated is sufficient, which is what this recipe does:

import tracemalloc
from contextlib import contextmanager


@contextmanager
def print_memory_usage():
    tracemalloc.start()

    yield

    current, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    print(f"Peak memory allocation: {peak / 1024:,.2f}KiB")
    print(f"Final memory allocation: {current / 1024:,.2f}KiB")


with print_memory_usage():
    x = list(range(1_000))

Update (2024-08-30): Edited to also peak usage, from tracemalloc.get_traced_memory().

@contextmanager() makes print_memory_usage usable as a context manager or a function decorator. Both the peak memory and final memory usage are displayed, from tracemalloc.get_traced_memory(). I added output in kibibytes, which you may wish to customize.

For example:

$ python example.py
Peak memory allocation: 31.20KiB
Final memory allocation: 31.09KiB

So, it takes nearly 32,000 bytes to allocate the 1,000 byte list, with a little extra needed on the way. Python uses about ~32bytes per int, including the references in the list.

In [1]: import sys

In [2]: x = list(range(1_000))

In [3]: sys.getsizeof(x)
Out[3]: 8056

with print_memory_usage():
    import numpy

$ python example.py
Peak memory allocation: 6,721.72KiB
Final memory allocation: 6,658.15KiB

😸😸😸 Check out my new book on using GitHub effectively, Boost Your GitHub DX! 😸😸😸

One summary email a week, no spam, I pinky promise.

Related posts:

• Django: How to profile and improve startup time

Tags: python