Started adding in logic for num_threads and changed names of real memory stats keys

Author: pintohutch

nipype/interfaces/base.py

@@ -1248,7 +1248,8 @@ def run_command(runtime, output=None, timeout=0.01, redirect_x=False):
     outfile = os.path.join(runtime.cwd, 'stdout.nipype')
 
     # Init variables for memory profiling
-    ret = -1
+    mem_mb = -1
+    num_threads = -1
     interval = 0.1
 
     if output == 'stream':
@@ -1268,7 +1269,7 @@ def _process(drain=0):
                     stream.read(drain)
         while proc.returncode is None:
             if mem_prof:
-                ret = max([ret, _get_memory(proc.pid, include_children=True)])
+                mem_mb = max([mem_mb, _get_memory(proc.pid, include_children=True)])
                 time.sleep(interval)
             proc.poll()
             _process()
@@ -1287,7 +1288,8 @@ def _process(drain=0):
     if output == 'allatonce':
         if mem_prof:
             while proc.returncode is None:
-                ret = max([ret, _get_memory(proc.pid, include_children=True)])
+                mem_mb = max([mem_mb, _get_memory(proc.pid, include_children=True)])
+                num_threads = max([num_threads, psutil.Proc(proc.pid).num_threads()])
                 time.sleep(interval)
                 proc.poll()
         stdout, stderr = proc.communicate()
@@ -1297,7 +1299,8 @@ def _process(drain=0):
     if output == 'file':
         if mem_prof:
             while proc.returncode is None:
-                ret = max([ret, _get_memory(proc.pid, include_children=True)])
+                mem_mb = max([mem_mb, _get_memory(proc.pid, include_children=True)])
+                num_threads = max([num_threads, psutil.Proc(proc.pid).num_threads()])
                 time.sleep(interval)
                 proc.poll()
         ret_code = proc.wait()
@@ -1309,15 +1312,17 @@ def _process(drain=0):
     if output == 'none':
         if mem_prof:
             while proc.returncode is None:
-                ret = max([ret, _get_memory(proc.pid, include_children=True)])
+                mem_mb = max([mem_mb, _get_memory(proc.pid, include_children=True)])
+                num_threads = max([num_threads, psutil.Proc(proc.pid).num_threads()])
                 time.sleep(interval)
                 proc.poll()
         proc.communicate()
         result['stdout'] = []
         result['stderr'] = []
         result['merged'] = ''
 
-    setattr(runtime, 'real_memory2', ret/1024.0)
+    setattr(runtime, 'cmd_memory', mem_mb/1024.0)
+    setattr(runtime, 'num_threads', num_threads)
     runtime.stderr = '\n'.join(result['stderr'])
     runtime.stdout = '\n'.join(result['stdout'])
     runtime.merged = result['merged']

nipype/interfaces/utility.py

@@ -442,12 +442,15 @@ def _run_interface(self, runtime):
             if isdefined(value):
                 args[name] = value
 
-        # mem stuff
-        import memory_profiler
-        proc = (function_handle, (), args)
-        mem_mb, out = memory_profiler.memory_usage(proc=proc, retval=True, include_children=True, max_usage=True)
-        setattr(runtime, 'real_memory2', mem_mb[0]/1024.0)
-        #out = function_handle(**args)
+        # Record memory of function_handle
+        try:
+            import memory_profiler
+            proc = (function_handle, (), args)
+            mem_mb, out = memory_profiler.memory_usage(proc=proc, retval=True, include_children=True, max_usage=True)
+            setattr(runtime, 'cmd_memory', mem_mb[0]/1024.0)
+        # If no memory_profiler package, run without recording memory
+        except ImportError:
+            out = function_handle(**args)
 
         if len(self._output_names) == 1:
             self._out[self._output_names[0]] = out

nipype/pipeline/plugins/multiproc.py

@@ -17,18 +17,20 @@ def run_node(node, updatehash, plugin_args=None):
     result = dict(result=None, traceback=None)
     try:
         run_memory = plugin_args['memory_profile']
-    except Exception:
+        import memory_profiler
+    except KeyError:
+        run_memory = False
+    except ImportError:
         run_memory = False
     if run_memory:
-        import memory_profiler
         import datetime
         proc = (node.run, (), {'updatehash' : updatehash})
         start = datetime.datetime.now()
         mem_mb, retval = memory_profiler.memory_usage(proc=proc, retval=True, include_children=True, max_usage=True)
         runtime = (datetime.datetime.now() - start).total_seconds()
         result['result'] = retval
-        result['real_memory'] = mem_mb[0]/1024.0
-        result['real_memory2'] = retval.runtime.get('real_memory2')
+        result['node_memory'] = mem_mb[0]/1024.0
+        result['cmd_memory'] = retval.runtime.get('cmd_memory')
         result['run_seconds'] = runtime
     else:
         try: