Merge branch 'mhaseebtariq-master'

Author: fabianp

Makefile

@@ -14,3 +14,4 @@ test:
 	$(PYTHON) test/test_import.py
 	$(PYTHON) test/test_memory_usage.py
 	$(PYTHON) test/test_precision_import.py
+	$(PYTHON) test/test_exception.py

README.rst

@@ -447,6 +447,8 @@ cleanup.
 
 `Benjamin Bengfort <https://github.com/bbengfort>`_ added support for tracking the usage of individual child processes and plotting them.
 
+`Muhammad Haseeb Tariq <https://github.com/mhaseebtariq>`_ fixed issue #152, which made the whole interpreter hang on functions that launched an exception.
+
 =========
  License
 =========

memory_profiler.py

@@ -16,6 +16,8 @@
 import inspect
 import subprocess
 import logging
+import traceback
+from signal import SIGKILL
 
 
 # TODO: provide alternative when multiprocessing is not available
@@ -324,21 +326,34 @@ def memory_usage(proc=-1, interval=.1, timeout=None, timestamps=False,
             raise ValueError
 
         while True:
+            exit_block = False
             child_conn, parent_conn = Pipe()  # this will store MemTimer's results
             p = MemTimer(os.getpid(), interval, child_conn, backend,
                          timestamps=timestamps,
                          max_usage=max_usage,
                          include_children=include_children)
             p.start()
             parent_conn.recv()  # wait until we start getting memory
-            returned = f(*args, **kw)
-            parent_conn.send(0)  # finish timing
-            ret = parent_conn.recv()
-            n_measurements = parent_conn.recv()
-            if retval:
-                ret = ret, returned
+
+            # When there is an exception in the "proc" - the (spawned) monitoring processes don't get killed.
+            # Therefore, the whole process hangs indefinitely. Here, we are ensuring that the process gets killed!
+            try:
+                returned = f(*args, **kw)
+                parent_conn.send(0)  # finish timing
+                ret = parent_conn.recv()
+                n_measurements = parent_conn.recv()
+                if retval:
+                    ret = ret, returned
+            except Exception:
+                if has_psutil:
+                    parent = psutil.Process(os.getpid())
+                    for child in parent.children(recursive=True):
+                        os.kill(child.pid, SIGKILL)
+                p.join(0)
+                raise
+
             p.join(5 * interval)
-            if n_measurements > 4 or interval < 1e-6:
+            if exit_block or n_measurements > 4 or interval < 1e-6:
                 break
             interval /= 10.
     elif isinstance(proc, subprocess.Popen):
@@ -1108,7 +1123,7 @@ def exec_with_profiler(filename, profiler, backend):
         execfile(filename, ns, ns)
 else:
     def exec_with_profiler(filename, profiler, backend):
-        choose_backend(backend)
+        _backend = choose_backend(backend)
         if _backend == 'tracemalloc' and has_tracemalloc:
             tracemalloc.start()
         builtins.__dict__['profile'] = profiler

mprof

@@ -351,8 +351,9 @@ def read_mprofile_file(filename):
 def plot_file(filename, index=0, timestamps=True, children=True, options=None):
     try:
         import pylab as pl
-    except ImportError:
+    except ImportError as e:
         print("matplotlib is needed for plotting.")
+        print(e)
         sys.exit(1)
     import numpy as np  # pylab requires numpy anyway
     mprofile = read_mprofile_file(filename)
@@ -455,11 +456,6 @@ def plot_action():
             raise OptionValueError("'%s' option must contain two numbers separated with a comma" % value)
         setattr(parser.values, option.dest, newvalue)
 
-    try:
-        import pylab as pl
-    except ImportError:
-        print("matplotlib is needed for plotting.")
-        sys.exit(1)
     desc = """Plots using matplotlib the data file `file.dat` generated
 using `mprof run`. If no .dat file is given, it will take the most recent
 such file in the current directory."""
@@ -477,8 +473,21 @@ such file in the current directory."""
                       type="str", action="callback",
                       callback=get_comma_separated_args,
                       help="Plot a time-subset of the data. E.g. to plot between 0 and 20.5 seconds: --window 0,20.5")
+    parser.add_option("--backend", 
+                      help="Specify the Matplotlib backend to use")
     (options, args) = parser.parse_args()
 
+    try:
+        if options.backend is not None:
+            import matplotlib
+            matplotlib.use(options.backend)
+
+        import pylab as pl
+    except ImportError as e:
+        print("matplotlib is needed for plotting.")
+        print(e)
+        sys.exit(1)
+
     profiles = glob.glob("mprofile_??????????????.dat")
     profiles.sort()
 

test/test_exception.py

@@ -0,0 +1,12 @@
+
+# make sure that memory profiler does not hang on exception
+from memory_profiler import memory_usage
+
+def foo():
+    raise NotImplementedError('Error')
+
+try:
+    out = memory_usage((foo, tuple(), {}), timeout=1)
+except NotImplementedError:
+    pass
+print('Success')