Python Memory Alloc Trace

trace_mem.py

def trace_mem(nframe=6,top=8):
    '''
        naive memory trace
    '''
    import tracemalloc
    is_tracing = tracemalloc.is_tracing()
    if not is_tracing:
        # start tracing
        tracemalloc.start(nframe)
        return {}
    else:
        # read traced memory alloc
        current_mem, peak_mem = tracemalloc.get_traced_memory()
        overhead = tracemalloc.get_tracemalloc_memory()
        stats = tracemalloc.take_snapshot().statistics('traceback')[:top]
        # memory summary
        summary = {}
        summary['memory'] = int(current_mem // 1024)
        summary['peak'] = int(peak_mem // 1024)
        summary['overhead'] = int(overhead // 1024)
        summary['description'] = "traced memory: %d KiB  peak: %d KiB  overhead: %d KiB" % (
            int(current_mem // 1024), int(peak_mem // 1024), int(overhead // 1024)
        )

        # traceback
        out_lines = []
        for trace in stats:
            stacktrace = {}
            stacktrace['memory'] = int(trace.size // 1024)
            stacktrace['blocks'] = int(trace.count)
            stacktrace['stack'] = trace.traceback.format()
            out_lines.append(stacktrace)
        
        data = {}
        data['summary'] = summary
        data['traceback'] = out_lines

        # stop tracing
        tracemalloc.stop()

        return data

Use like this

import json
# start memory tracing
trace_mem(nframe=6, top=8)
## do your heavy and memory-bound job
# dump memory and stop tracing
print("\n--------mem--------\n", json.dumps(trace_mem(nframe=6, top=8), indent=4), "\n--------mem--------")