Added cmd-level threads and memory profiling

Author: pintohutch

nipype/interfaces/base.py

@@ -751,17 +751,8 @@ def __init__(self, **inputs):
                             self.__class__.__name__)
         self.inputs = self.input_spec(**inputs)
         self.estimated_memory = 1
-        self._real_memory = 0
         self.num_threads = 1
 
-    @property
-    def real_memory(self):
-        return self._real_memory
-
-    @real_memory.setter
-    def real_memory(self, value):
-        self._real_memory = value
-
     @classmethod
     def help(cls, returnhelp=False):
         """ Prints class help
@@ -1269,7 +1260,8 @@ def _process(drain=0):
                     stream.read(drain)
         while proc.returncode is None:
             if mem_prof:
-                mem_mb = max([mem_mb, _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.Process(proc.pid).num_threads())
                 time.sleep(interval)
             proc.poll()
             _process()
@@ -1288,8 +1280,8 @@ def _process(drain=0):
     if output == 'allatonce':
         if mem_prof:
             while proc.returncode is None:
-                mem_mb = max([mem_mb, _get_memory(proc.pid, include_children=True)])
-                num_threads = max([num_threads, psutil.Proc(proc.pid).num_threads()])
+                mem_mb = max(mem_mb, _get_memory(proc.pid, include_children=True))
+                num_threads = max(num_threads, psutil.Process(proc.pid).num_threads())
                 time.sleep(interval)
                 proc.poll()
         stdout, stderr = proc.communicate()
@@ -1299,8 +1291,8 @@ def _process(drain=0):
     if output == 'file':
         if mem_prof:
             while proc.returncode is None:
-                mem_mb = max([mem_mb, _get_memory(proc.pid, include_children=True)])
-                num_threads = max([num_threads, psutil.Proc(proc.pid).num_threads()])
+                mem_mb = max(mem_mb, _get_memory(proc.pid, include_children=True))
+                num_threads = max(num_threads, psutil.Process(proc.pid).num_threads())
                 time.sleep(interval)
                 proc.poll()
         ret_code = proc.wait()
@@ -1312,8 +1304,8 @@ def _process(drain=0):
     if output == 'none':
         if mem_prof:
             while proc.returncode is None:
-                mem_mb = max([mem_mb, _get_memory(proc.pid, include_children=True)])
-                num_threads = max([num_threads, psutil.Proc(proc.pid).num_threads()])
+                mem_mb = max(mem_mb, _get_memory(proc.pid, include_children=True))
+                num_threads = max(num_threads, psutil.Process(proc.pid).num_threads())
                 time.sleep(interval)
                 proc.poll()
         proc.communicate()
@@ -1322,7 +1314,7 @@ def _process(drain=0):
         result['merged'] = ''
 
     setattr(runtime, 'cmd_memory', mem_mb/1024.0)
-    setattr(runtime, 'num_threads', num_threads)
+    setattr(runtime, 'cmd_threads', num_threads)
     runtime.stderr = '\n'.join(result['stderr'])
     runtime.stdout = '\n'.join(result['stdout'])
     runtime.merged = result['merged']

nipype/pipeline/plugins/base.py

@@ -419,8 +419,6 @@ def _task_finished_cb(self, jobid, result=None):
             if result == None:
                 if self._taskresult.has_key(jobid):
                     result = self._taskresult[jobid].get()
-                else:
-                    result = {'real_memory' : 'nokey'}
             self._status_callback(self.procs[jobid], 'end', result)
         # Update job and worker queues
         self.proc_pending[jobid] = False

nipype/pipeline/plugins/callback_log.py

@@ -4,11 +4,12 @@
 def log_nodes_cb(node, status, result=None):
     logger = logging.getLogger('callback')
     try:
-        real_mem1 = result['real_memory']
-        real_mem2 = result['real_memory2']
+        node_mem = result['node_memory']
+        cmd_mem = result['cmd_memory']
         run_seconds = result['run_seconds']
+        cmd_threads = result['cmd_threads']
     except Exception as exc:
-        real_mem1 = real_mem2 = run_seconds = 'N/A'
+        node_mem = cmd_mem = run_seconds = cmd_threads = 'N/A'
     if status == 'start':
         message  = '{"name":' + '"' + node.name + '"' + ',"id":' + '"' +\
         node._id + '"' + ',"start":' + '"' +str(datetime.datetime.now()) +\
@@ -19,16 +20,20 @@ def log_nodes_cb(node, status, result=None):
 
     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) + ',"real_memory1":' + str(real_mem1) + ',"real_memory2":' + str(real_mem2) + ',"run_seconds":' + str(run_seconds) + '}'
+        node._id + '"' + ',"finish":' + '"' + str(datetime.datetime.now()) + \
+        '"' + ',"estimate memory":' + str(node._interface.estimated_memory) + \
+        ',"num_threads":' + str(node._interface.num_threads) + \
+        ',"cmd-level threads":' + str(cmd_threads) + \
+        ',"node-level memory":' + str(node_mem) + \
+        ',"cmd-level memory":' + str(cmd_mem) + \
+        ',"run_seconds":' + str(run_seconds) + '}'
 
         logger.debug(message)
 
     else:
         message  = '{"name":' + '"' + node.name + '"' + ',"id":' + '"' + \
         node._id + '"' + ',"finish":' + '"' + str(datetime.datetime.now()) +\
-        '"' + ',"memory":' + str(node._interface.estimated_memory) + ',"num_threads":' \
+        '"' + ',"estimate memory":' + str(node._interface.estimated_memory) + ',"num_threads":' \
         + str(node._interface.num_threads) + ',"error":"True"}'
 
         logger.debug(message)

nipype/pipeline/plugins/multiproc.py

@@ -13,25 +13,37 @@
 from .base import (DistributedPluginBase, report_crash)
 
 
+# Run node
 def run_node(node, updatehash, plugin_args=None):
-    result = dict(result=None, traceback=None)
+    """docstring
+    """
+
+    # Import packages
     try:
-        run_memory = plugin_args['memory_profile']
+        runtime_profile = plugin_args['runtime_profile']
         import memory_profiler
+        import datetime
     except KeyError:
-        run_memory = False
+        runtime_profile = False
     except ImportError:
-        run_memory = False
-    if run_memory:
-        import datetime
+        runtime_profile = False
+
+    # Init variables
+    result = dict(result=None, traceback=None)
+
+    # If we're profiling the run
+    if runtime_profile:
+        # Init function tuple
         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['node_memory'] = mem_mb[0]/1024.0
         result['cmd_memory'] = retval.runtime.get('cmd_memory')
+        result['cmd_threads'] = retval.runtime.get('cmd_threads')
         result['run_seconds'] = runtime
+    # Otherwise, execute node.run as normal
     else:
         try:
             result['result'] = node.run(updatehash=updatehash)
@@ -141,15 +153,15 @@ class ResourceMultiProcPlugin(MultiProcPlugin):
     the number of threads and memory of the system is used.
 
     System consuming nodes should be tagged:
-    memory_consuming_node.interface.memory = 8 #Gb
+    memory_consuming_node.interface.estimated_memory = 8 #Gb
     thread_consuming_node.interface.num_threads = 16
 
     The default number of threads and memory for a node is 1. 
 
     Currently supported options are:
 
-    - num_thread: maximum number of threads to be executed in parallel
-    - memory: maximum memory that can be used at once.
+    - num_threads: maximum number of threads to be executed in parallel
+    - estimated_memory: maximum memory that can be used at once.
 
     """
 
@@ -198,7 +210,7 @@ def _send_procs_to_workers(self, updatehash=False, graph=None):
         for jobid in jobids:
             busy_memory+= self.procs[jobid]._interface.estimated_memory
             busy_processors+= self.procs[jobid]._interface.num_threads
-                
+
         free_memory = self.memory - busy_memory
         free_processors = self.processors - busy_processors
 
@@ -222,7 +234,7 @@ def _send_procs_to_workers(self, updatehash=False, graph=None):
             if self.procs[jobid]._interface.estimated_memory <= free_memory and self.procs[jobid]._interface.num_threads <= free_processors:
                 logger.info('Executing: %s ID: %d' %(self.procs[jobid]._id, jobid))
                 executing_now.append(self.procs[jobid])
-                
+
                 if isinstance(self.procs[jobid], MapNode):
                     try:
                         num_subnodes = self.procs[jobid].num_subnodes()