Fixed logging of real memory

Author: ccraddock

nipype/interfaces/base.py

@@ -1206,9 +1206,18 @@ def run_command(runtime, output=None, timeout=0.01, redirect_x=False):
 
     The returned runtime contains a merged stdout+stderr log with timestamps
     """
-    PIPE = subprocess.PIPE
 
+    # Import packages
+    try:
+        from memory_profiler import _get_memory
+        mem_prof = True
+    except:
+        mem_prof = False
+
+    # Init variables
+    PIPE = subprocess.PIPE
     cmdline = runtime.cmdline
+
     if redirect_x:
         exist_xvfb, _ = _exists_in_path('xvfb-run', runtime.environ)
         if not exist_xvfb:
@@ -1237,6 +1246,11 @@ def run_command(runtime, output=None, timeout=0.01, redirect_x=False):
     result = {}
     errfile = os.path.join(runtime.cwd, 'stderr.nipype')
     outfile = os.path.join(runtime.cwd, 'stdout.nipype')
+
+    # Init variables for memory profiling
+    ret = -1
+    interval = 0.1
+
     if output == 'stream':
         streams = [Stream('stdout', proc.stdout), Stream('stderr', proc.stderr)]
 
@@ -1252,8 +1266,10 @@ def _process(drain=0):
             else:
                 for stream in res[0]:
                     stream.read(drain)
-
         while proc.returncode is None:
+            if mem_prof:
+                ret = max([ret, _get_memory(proc.pid, include_children=True)])
+                time.sleep(interval)
             proc.poll()
             _process()
         _process(drain=1)
@@ -1267,23 +1283,41 @@ def _process(drain=0):
             result[stream._name] = [r[2] for r in rows]
         temp.sort()
         result['merged'] = [r[1] for r in temp]
+
     if output == 'allatonce':
+        if mem_prof:
+            while proc.returncode is None:
+                ret = max([ret, _get_memory(proc.pid, include_children=True)])
+                time.sleep(interval)
+                proc.poll()
         stdout, stderr = proc.communicate()
         result['stdout'] = stdout.split('\n')
         result['stderr'] = stderr.split('\n')
         result['merged'] = ''
     if output == 'file':
+        if mem_prof:
+            while proc.returncode is None:
+                ret = max([ret, _get_memory(proc.pid, include_children=True)])
+                time.sleep(interval)
+                proc.poll()
         ret_code = proc.wait()
         stderr.flush()
         stdout.flush()
         result['stdout'] = [line.strip() for line in open(outfile).readlines()]
         result['stderr'] = [line.strip() for line in open(errfile).readlines()]
         result['merged'] = ''
     if output == 'none':
+        if mem_prof:
+            while proc.returncode is None:
+                ret = max([ret, _get_memory(proc.pid, include_children=True)])
+                time.sleep(interval)
+                proc.poll()
         proc.communicate()
         result['stdout'] = []
         result['stderr'] = []
         result['merged'] = ''
+
+    setattr(runtime, 'real_memory2', ret/1024.0)
     runtime.stderr = '\n'.join(result['stderr'])
     runtime.stdout = '\n'.join(result['stdout'])
     runtime.merged = result['merged']

nipype/interfaces/utility.py

@@ -442,7 +442,12 @@ def _run_interface(self, runtime):
             if isdefined(value):
                 args[name] = value
 
-        out = function_handle(**args)
+        # 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)
 
         if len(self._output_names) == 1:
             self._out[self._output_names[0]] = out

nipype/pipeline/plugins/base.py

@@ -241,7 +241,8 @@ def run(self, graph, config, updatehash=False):
                             notrun.append(self._clean_queue(jobid, graph,
                                                             result=result))
                         else:
-                            self._task_finished_cb(jobid)
+                            print "DJC: Calling task finished for %s cb from DistributedPluginBase.run"%(str(taskid))
+                            self._task_finished_cb(jobid, result)
                             self._remove_node_dirs()
                         self._clear_task(taskid)
                     else:
@@ -379,6 +380,7 @@ def _send_procs_to_workers(self, updatehash=False, graph=None):
                                  )
                                ):
                                 continue_with_submission = False
+                                print "DJC: Calling task finised cb from DistributedPluginBase._send_procs_to_workers hash==true"
                                 self._task_finished_cb(jobid)
                                 self._remove_node_dirs()
                         except Exception:
@@ -395,6 +397,7 @@ def _send_procs_to_workers(self, updatehash=False, graph=None):
                                 self.procs[jobid].run()
                             except Exception:
                                 self._clean_queue(jobid, graph)
+                            print "DJC: Calling task finised cb from DistributedPluginBase._send_procs_to_workers continue_with_submission==true"
                             self._task_finished_cb(jobid)
                             self._remove_node_dirs()
                         else:
@@ -408,18 +411,23 @@ def _send_procs_to_workers(self, updatehash=False, graph=None):
             else:
                 break
 
-    def _task_finished_cb(self, jobid):
+    def _task_finished_cb(self, jobid, result=None):
         """ Extract outputs and assign to inputs of dependent tasks
 
         This is called when a job is completed.
         """
         logger.info('[Job finished] jobname: %s jobid: %d' %
                     (self.procs[jobid]._id, jobid))
         if self._status_callback:
-            print '!!!!!!!!!!!!!!!!!!!'
-            print self._taskresult
-            print self._taskresult.keys()
-            self._status_callback(self.procs[jobid], 'end', self._taskresult[self.taskresultid])
+            if result == None:
+                if self._taskresult.has_key(jobid):
+                    result = self._taskresult[jobid].get()
+                    print 'MMMM'
+                    print result['real_memory'], result['real_memory2']
+                else:
+                    print "DJC: %s not found, taskresult keys are: %s"%(str(jobid),":".join([str(k) for k in self._taskresult.keys()]))
+                    result = {'real_memory' : 'nokey'}
+            self._status_callback(self.procs[jobid], 'end', result)
         # Update job and worker queues
         self.proc_pending[jobid] = False
         # update the job dependency structure

nipype/pipeline/plugins/callback_log.py

@@ -2,8 +2,12 @@
 import logging
 
 def log_nodes_cb(node, status, result=None):
-    print 'status', status
     logger = logging.getLogger('callback')
+    try:
+        real_mem1 = result['real_memory']
+        real_mem2 = result['result'].runtime.get('real_memory2')
+    except Exception as exc:
+        real_mem1 = real_mem2 = 'N/A'
     if status == 'start':
         message  = '{"name":' + '"' + node.name + '"' + ',"id":' + '"' +\
         node._id + '"' + ',"start":' + '"' +str(datetime.datetime.now()) +\
@@ -16,7 +20,7 @@ def log_nodes_cb(node, status, result=None):
         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 memory":' + str(result['real_memory']) + '}'
+        + str(node._interface.num_threads) + ',"real memory1":' + str(real_mem1) + ',"real memory2":' + str(real_mem2) + '}'
 
         logger.debug(message)
 

nipype/pipeline/plugins/multiproc.py

@@ -22,12 +22,10 @@ def run_node(node, updatehash, plugin_args=None):
     if run_memory:
         import memory_profiler
         proc = (node.run, (), {'updatehash' : updatehash})
-        mem_mb, retval = memory_profiler.memory_usage(proc, max_usage=True, retval=True)
+        mem_mb, retval = memory_profiler.memory_usage(proc=proc, retval=True, include_children=True, max_usage=True)
         result['result'] = retval
-        result['real_memory'] = 100
-        print 'Just populated task result!!!!!!!!!!!!!!!!!!!'
-        print result
-        #node._interface.real_memory = mem_mb[0]/1024.0
+        result['real_memory'] = mem_mb[0]/1024.0
+        result['real_memory2'] = retval.runtime.get('real_memory2')
     else:
         try:
             result['result'] = node.run(updatehash=updatehash)
@@ -177,8 +175,6 @@ def _submit_job(self, node, updatehash=False):
         self._taskresult[self._taskid] = self.pool.apply_async(run_node,
                                                                (node, updatehash, self.plugin_args),
                                                                callback=release_lock)
-        print 'Printing on output!!!!!!!!!!'
-        print self._taskresult, self._taskid
         return self._taskid
 
     def _send_procs_to_workers(self, updatehash=False, graph=None):