1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
|
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")
|