Added real memory recording to plugn

Author: ccraddock

nipype/pipeline/plugins/callback_log.py

@@ -7,16 +7,16 @@ def log_nodes_cb(node, status):
     if status == 'start':
         message  = '{"name":' + '"' + node.name + '"' + ',"id":' + '"' +\
         node._id + '"' + ',"start":' + '"' +str(datetime.datetime.now()) +\
-        '"' + ',"memory":' + str(node._interface.estimated_memory) + ',"num_threads":' \
-        + str(node._interface.num_threads) +  '}'
+        '"' + ',"estimate memory":' + str(node._interface.estimated_memory) + ',"num_threads":' \
+        + str(node._interface.num_threads) + '}'
 
         logger.debug(message)
 
     elif status == 'end':
         message  = '{"name":' + '"' + node.name + '"' + ',"id":' + '"' + \
         node._id + '"' + ',"finish":' + '"' + str(datetime.datetime.now()) +\
         '"' + ',"memory":' + str(node._interface.estimated_memory) + ',"num_threads":' \
-        + str(node._interface.num_threads) +  '}'
+        + str(node._interface.num_threads) + ',"real memory":' str(node._interface.real_memory) + '}'
 
         logger.debug(message)
 
@@ -26,4 +26,4 @@ def log_nodes_cb(node, status):
         '"' + ',"memory":' + str(node._interface.estimated_memory) + ',"num_threads":' \
         + str(node._interface.num_threads) + ',"error":"True"}'
 
-        logger.debug(message)
+        logger.debug(message)

nipype/pipeline/plugins/multiproc.py

@@ -13,8 +13,18 @@
 from .base import (DistributedPluginBase, report_crash)
 
 
-def run_node(node, updatehash):
+def run_node(node, updatehash, plugin_args=None):
     result = dict(result=None, traceback=None)
+    try:
+        run_memory = plugin_args['memory_profile']
+    except Exception:
+        run_memory = False
+    if run_memory:
+        import memory_profiler
+        proc = (node.run(), (), {'updatehash' : updatehash})
+        mem_mb, retval = memory_profiler.memory_usage(proc, max_usage=True, retval=True)
+        result['result'] = retval
+        node._interface.real_memory = mem_mb[0]/1024.0
     try:
         result['result'] = node.run(updatehash=updatehash)
     except:
@@ -160,8 +170,9 @@ def _submit_job(self, node, updatehash=False):
                 node.inputs.terminal_output = 'allatonce'
         except:
             pass
-        self._taskresult[self._taskid] = self.pool.apply_async(run_node, (node,
-                                                                updatehash,), callback=release_lock)
+        self._taskresult[self._taskid] = self.pool.apply_async(run_node,
+                                                               (node, updatehash, self.plugin_args),
+                                                               callback=release_lock)
         return self._taskid
 
     def _send_procs_to_workers(self, updatehash=False, graph=None):
@@ -263,4 +274,4 @@ def _send_procs_to_workers(self, updatehash=False, graph=None):
             else:
                 break
 
-        logger.debug('No jobs waiting to execute')
+        logger.debug('No jobs waiting to execute')