public
Last active

memit: magic memory usage benching for IPython

  • Download Gist
magic_memit.py
Python
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
# Author: Vlad Niculae <vlad@vene.ro>
# Makes use of memory_profiler from Fabian Pedregosa
# available at https://github.com/fabianp/memory_profiler
 
from IPython.core.magic import Magics, line_magic, magics_class
 
 
class MemMagics(Magics):
@line_magic
def memit(self, line='', setup='pass'):
"""Measure memory usage of a Python statement
 
Usage, in line mode:
%memit [-ir<R>t<T>] statement
 
Options:
-r<R>: repeat the loop iteration <R> times and take the best result.
Default: 3
 
-i: run the code in the current environment, without forking a new
process. This is required on some MacOS versions of Accelerate if your
line contains a call to `np.dot`.
 
-t<T>: timeout after <T> seconds. Unused if `-i` is active.
Default: None
 
Examples
--------
::
 
In [1]: import numpy as np
 
In [2]: %memit np.zeros(1e7)
maximum of 3: 76.402344 MB per loop
 
In [3]: %memit np.ones(1e6)
maximum of 3: 7.820312 MB per loop
 
In [4]: %memit -r 10 np.empty(1e8)
maximum of 10: 0.101562 MB per loop
 
In [5]: memit -t 3 while True: pass;
Subprocess timed out.
Subprocess timed out.
Subprocess timed out.
ERROR: all subprocesses exited unsuccessfully. Try again with the
`-i` option.
maximum of 3: -inf MB per loop
 
"""
opts, stmt = self.parse_options(line, 'r:t:i', posix=False,
strict=False)
repeat = int(getattr(opts, 'r', 3))
if repeat < 1:
repeat == 1
timeout = int(getattr(opts, 't', 0))
if timeout <= 0:
timeout = None
run_in_place = hasattr(opts, 'i')
 
# Don't depend on multiprocessing:
try:
import multiprocessing as pr
from multiprocessing.queues import SimpleQueue
q = SimpleQueue()
except ImportError:
class ListWithPut(list):
"Just a list where the `append` method is aliased to `put`."
def put(self, x):
self.append(x)
q = ListWithPut()
print ('WARNING: cannot import module `multiprocessing`. Forcing '
'the `-i` option.')
run_in_place = True
 
ns = self.shell.user_ns
 
def _get_usage(q, stmt, setup='pass', ns={}):
from memory_profiler import memory_usage as _mu
try:
exec setup in ns
_mu0 = _mu()[0]
exec stmt in ns
_mu1 = _mu()[0]
q.put(_mu1 - _mu0)
except Exception as e:
q.put(float('-inf'))
raise e
 
if run_in_place:
for _ in xrange(repeat):
_get_usage(q, stmt, ns=ns)
else:
# run in consecutive subprocesses
at_least_one_worked = False
for _ in xrange(repeat):
p = pr.Process(target=_get_usage, args=(q, stmt, 'pass', ns))
p.start()
p.join(timeout=timeout)
if p.exitcode == 0:
at_least_one_worked = True
else:
p.terminate()
if p.exitcode == None:
print 'Subprocess timed out.'
else:
print 'Subprocess exited with code %d.' % p.exitcode
q.put(float('-inf'))
 
if not at_least_one_worked:
print ('ERROR: all subprocesses exited unsuccessfully. Try '
'again with the `-i` option.')
 
usages = [q.get() for _ in xrange(repeat)]
usage = max(usages)
print u"maximum of %d: %f MB per loop" % (repeat, usage)
 
 
if __name__ == '__main__':
ip = get_ipython()
ip.register_magics(MemMagics)

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.