141 lines
5.3 KiB
Python
141 lines
5.3 KiB
Python
import numpy as np
|
|
import gc
|
|
import inspect
|
|
from alloc_hook import NumpyAllocHook
|
|
|
|
class AllocationTracker:
|
|
def __init__(self, threshold=0):
|
|
'''track numpy allocations of size threshold bytes or more.'''
|
|
|
|
self.threshold = threshold
|
|
|
|
# The total number of bytes currently allocated with size above
|
|
# threshold
|
|
self.total_bytes = 0
|
|
|
|
# We buffer requests line by line and move them into the allocation
|
|
# trace when a new line occurs
|
|
self.current_line = None
|
|
self.pending_allocations = []
|
|
|
|
self.blocksizes = {}
|
|
|
|
# list of (lineinfo, bytes allocated, bytes freed, # allocations, #
|
|
# frees, maximum memory usage, long-lived bytes allocated)
|
|
self.allocation_trace = []
|
|
|
|
self.numpy_hook = NumpyAllocHook(self.hook)
|
|
|
|
def __enter__(self):
|
|
self.numpy_hook.__enter__()
|
|
|
|
def __exit__(self, type, value, traceback):
|
|
self.check_line_changed() # forces pending events to be handled
|
|
self.numpy_hook.__exit__()
|
|
|
|
def hook(self, inptr, outptr, size):
|
|
# minimize the chances that the garbage collector kicks in during a
|
|
# cython __dealloc__ call and causes a double delete of the current
|
|
# object. To avoid this fully the hook would have to avoid all python
|
|
# api calls, e.g. by being implemented in C like python 3.4's
|
|
# tracemalloc module
|
|
gc_on = gc.isenabled()
|
|
gc.disable()
|
|
if outptr == 0: # it's a free
|
|
self.free_cb(inptr)
|
|
elif inptr != 0: # realloc
|
|
self.realloc_cb(inptr, outptr, size)
|
|
else: # malloc
|
|
self.alloc_cb(outptr, size)
|
|
if gc_on:
|
|
gc.enable()
|
|
|
|
def alloc_cb(self, ptr, size):
|
|
if size >= self.threshold:
|
|
self.check_line_changed()
|
|
self.blocksizes[ptr] = size
|
|
self.pending_allocations.append(size)
|
|
|
|
def free_cb(self, ptr):
|
|
size = self.blocksizes.pop(ptr, 0)
|
|
if size:
|
|
self.check_line_changed()
|
|
self.pending_allocations.append(-size)
|
|
|
|
def realloc_cb(self, newptr, oldptr, size):
|
|
if (size >= self.threshold) or (oldptr in self.blocksizes):
|
|
self.check_line_changed()
|
|
oldsize = self.blocksizes.pop(oldptr, 0)
|
|
self.pending_allocations.append(size - oldsize)
|
|
self.blocksizes[newptr] = size
|
|
|
|
def get_code_line(self):
|
|
# first frame is this line, then check_line_changed(), then 2 callbacks,
|
|
# then actual code.
|
|
try:
|
|
return inspect.stack()[4][1:]
|
|
except Exception:
|
|
return inspect.stack()[0][1:]
|
|
|
|
def check_line_changed(self):
|
|
line = self.get_code_line()
|
|
if line != self.current_line and (self.current_line is not None):
|
|
# move pending events into the allocation_trace
|
|
max_size = self.total_bytes
|
|
bytes_allocated = 0
|
|
bytes_freed = 0
|
|
num_allocations = 0
|
|
num_frees = 0
|
|
before_size = self.total_bytes
|
|
for allocation in self.pending_allocations:
|
|
self.total_bytes += allocation
|
|
if allocation > 0:
|
|
bytes_allocated += allocation
|
|
num_allocations += 1
|
|
else:
|
|
bytes_freed += -allocation
|
|
num_frees += 1
|
|
max_size = max(max_size, self.total_bytes)
|
|
long_lived = max(self.total_bytes - before_size, 0)
|
|
self.allocation_trace.append((self.current_line, bytes_allocated,
|
|
bytes_freed, num_allocations,
|
|
num_frees, max_size, long_lived))
|
|
# clear pending allocations
|
|
self.pending_allocations = []
|
|
# move to the new line
|
|
self.current_line = line
|
|
|
|
def write_html(self, filename):
|
|
with open(filename, "w") as f:
|
|
f.write('<HTML><HEAD><script src="sorttable.js"></script></HEAD><BODY>\n')
|
|
f.write('<TABLE class="sortable" width=100%>\n')
|
|
f.write("<TR>\n")
|
|
cols = "event#,lineinfo,bytes allocated,bytes freed,#allocations,#frees,max memory usage,long lived bytes".split(',')
|
|
for header in cols:
|
|
f.write(" <TH>{0}</TH>".format(header))
|
|
f.write("\n</TR>\n")
|
|
for idx, event in enumerate(self.allocation_trace):
|
|
f.write("<TR>\n")
|
|
event = [idx] + list(event)
|
|
for col, val in zip(cols, event):
|
|
if col == 'lineinfo':
|
|
# special handling
|
|
try:
|
|
filename, line, module, code, index = val
|
|
val = "{0}({1}): {2}".format(filename, line, code[index])
|
|
except Exception:
|
|
# sometimes this info is not available (from eval()?)
|
|
val = str(val)
|
|
f.write(" <TD>{0}</TD>".format(val))
|
|
f.write("\n</TR>\n")
|
|
f.write("</TABLE></BODY></HTML>\n")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
tracker = AllocationTracker(1000)
|
|
with tracker:
|
|
for i in range(100):
|
|
np.zeros(i * 100)
|
|
np.zeros(i * 200)
|
|
tracker.write_html("allocations.html")
|