Manual merge of s3_datasink and resource_multiproc branch for cpac run

Author: pintohutch

nipype/interfaces/base.py

@@ -764,6 +764,8 @@ def __init__(self, **inputs):
             raise Exception('No input_spec in class: %s' %
                             self.__class__.__name__)
         self.inputs = self.input_spec(**inputs)
+        self.estimated_memory = 1
+        self.num_threads = 1
 
     @classmethod
     def help(cls, returnhelp=False):
@@ -1202,14 +1204,43 @@ def _read(self, drain):
         self._lastidx = len(self._rows)
 
 
+def _get_num_threads(proc):
+    '''
+    '''
+
+    # Import packages
+    import psutil
+
+    # Init variables
+    num_threads = proc.num_threads()
+    try:
+        for child in proc.children():
+            num_threads = max(num_threads, child.num_threads(),
+                              len(child.children()), _get_num_threads(child))
+    except psutil.NoSuchProcess:
+        dummy = 1
+
+    return num_threads
+
+
 def run_command(runtime, output=None, timeout=0.01, redirect_x=False):
     """Run a command, read stdout and stderr, prefix with timestamp.
 
     The returned runtime contains a merged stdout+stderr log with timestamps
     """
-    PIPE = subprocess.PIPE
 
+    # Import packages
+    try:
+        from memory_profiler import _get_memory
+        import psutil
+        mem_proc = 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:
@@ -1238,6 +1269,12 @@ 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
+    mem_mb = -1
+    num_threads = -1
+    interval = 1
+
     if output == 'stream':
         streams = [Stream('stdout', proc.stdout), Stream('stderr', proc.stderr)]
 
@@ -1253,8 +1290,10 @@ def _process(drain=0):
             else:
                 for stream in res[0]:
                     stream.read(drain)
-
         while proc.returncode is None:
+            if mem_prof:
+                mem_mb = max(mem_mb, _get_memory(proc.pid, include_children=True))
+                num_threads = max(num_threads, _get_num_threads(psutil.Process(proc.pid)))
             proc.poll()
             _process()
         _process(drain=1)
@@ -1268,27 +1307,53 @@ 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:
+                mem_mb = max(mem_mb, _get_memory(proc.pid, include_children=True))
+                num_threads = max(num_threads, _get_num_threads(psutil.Process(proc.pid)))
+                proc.poll()
         stdout, stderr = proc.communicate()
         if stdout and isinstance(stdout, bytes):
-            stdout = stdout.decode()
+            try:
+                stdout = stdout.decode()
+            except UnicodeDecodeError:
+                stdout = stdout.decode("ISO-8859-1")
         if stderr and isinstance(stderr, bytes):
-            stderr = stderr.decode()
+            try:
+                stderr = stderr.decode()
+            except UnicodeDecodeError:
+                stdout = stdout.decode("ISO-8859-1")
+
         result['stdout'] = str(stdout).split('\n')
         result['stderr'] = str(stderr).split('\n')
         result['merged'] = ''
     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, _get_num_threads(psutil.Process(proc.pid)))
+                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:
+                mem_mb = max(mem_mb, _get_memory(proc.pid, include_children=True))
+                num_threads = max(num_threads, _get_num_threads(psutil.Process(proc.pid)))
+                proc.poll()
         proc.communicate()
         result['stdout'] = []
         result['stderr'] = []
         result['merged'] = ''
+
+    setattr(runtime, 'cmd_memory', mem_mb/1024.0)
+    setattr(runtime, 'cmd_threads', num_threads)
     runtime.stderr = '\n'.join(result['stderr'])
     runtime.stdout = '\n'.join(result['stdout'])
     runtime.merged = result['merged']

nipype/interfaces/utility.py

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

nipype/pipeline/engine/tests/test_engine.py

@@ -723,7 +723,7 @@ def func1(in1):
     # test running the workflow on default conditions
     error_raised = False
     try:
-        w1.run(plugin='MultiProc')
+        w1.run(plugin='ResourceMultiProc')
     except Exception as e:
         from nipype.pipeline.engine.base import logger
         logger.info('Exception: %s' % str(e))
@@ -737,7 +737,7 @@ def func1(in1):
     # test running the workflow on serial conditions
     error_raised = False
     try:
-        w1.run(plugin='MultiProc')
+        w1.run(plugin='ResourceMultiProc')
     except Exception as e:
         from nipype.pipeline.engine.base import logger
         logger.info('Exception: %s' % str(e))

nipype/pipeline/engine/tests/test_utils.py

@@ -214,7 +214,7 @@ def test_function3(arg):
 
     out_dir = mkdtemp()
 
-    for plugin in ('Linear',):  # , 'MultiProc'):
+    for plugin in ('Linear',):  # , 'ResourceMultiProc'):
         n1 = pe.Node(niu.Function(input_names=['arg1'],
                                   output_names=['out_file1', 'out_file2', 'dir'],
                                   function=test_function),

nipype/pipeline/plugins/__init__.py

@@ -9,11 +9,13 @@
 from .sge import SGEPlugin
 from .condor import CondorPlugin
 from .dagman import CondorDAGManPlugin
-from .multiproc import MultiProcPlugin
+from .multiproc import ResourceMultiProcPlugin
 from .ipython import IPythonPlugin
 from .somaflow import SomaFlowPlugin
 from .pbsgraph import PBSGraphPlugin
 from .sgegraph import SGEGraphPlugin
 from .lsf import LSFPlugin
 from .slurm import SLURMPlugin
 from .slurmgraph import SLURMGraphPlugin
+
+from .callback_log import log_nodes_cb

nipype/pipeline/plugins/base.py

@@ -20,7 +20,6 @@
 import numpy as np
 import scipy.sparse as ssp
 
-
 from ...utils.filemanip import savepkl, loadpkl
 from ...utils.misc import str2bool
 from ..engine.utils import (nx, dfs_preorder, topological_sort)
@@ -246,7 +245,7 @@ def run(self, graph, config, updatehash=False):
                             notrun.append(self._clean_queue(jobid, graph,
                                                             result=result))
                         else:
-                            self._task_finished_cb(jobid)
+                            self._task_finished_cb(jobid, result)
                             self._remove_node_dirs()
                         self._clear_task(taskid)
                     else:
@@ -265,10 +264,15 @@ def run(self, graph, config, updatehash=False):
                                             graph=graph)
             else:
                 logger.debug('Not submitting')
-            sleep(float(self._config['execution']['poll_sleep_duration']))
+            self._wait()
         self._remove_node_dirs()
         report_nodes_not_run(notrun)
 
+
+
+    def _wait(self):
+        sleep(float(self._config['execution']['poll_sleep_duration']))
+
     def _get_result(self, taskid):
         raise NotImplementedError
 
@@ -410,15 +414,18 @@ 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:
-            self._status_callback(self.procs[jobid], 'end')
+            if result == None:
+                if self._taskresult.has_key(jobid):
+                    result = self._taskresult[jobid].get()
+            self._status_callback(self.procs[jobid], 'end', result)
         # Update job and worker queues
         self.proc_pending[jobid] = False
         # update the job dependency structure