Merge pull request #550 from satra/fix/danmultiproc

Fix/danmultiproc

Author: satra

nipype/pipeline/plugins/multiproc.py

@@ -1,13 +1,16 @@
 # emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
 # vi: set ft=python sts=4 ts=4 sw=4 et:
 """Parallel workflow execution via multiprocessing
+
+Support for child processes running as non-daemons based on
+http://stackoverflow.com/a/8963618/1183453
 """
 
-from multiprocessing import Pool, cpu_count
+from multiprocessing import Process, Pool, cpu_count, pool
 from traceback import format_exception
 import sys
 
-from .base import (DistributedPluginBase, logger, report_crash)
+from .base import (DistributedPluginBase, report_crash)
 
 def run_node(node, updatehash):
     result = dict(result=None, traceback=None)
@@ -19,25 +22,49 @@ def run_node(node, updatehash):
         result['result'] = node.result
     return result
 
+class NonDaemonProcess(Process):
+    """A non-daemon process to support internal multiprocessing.
+    """
+    def _get_daemon(self):
+        return False
+
+    def _set_daemon(self, value):
+        pass
+
+    daemon = property(_get_daemon, _set_daemon)
+
+class NonDaemonPool(pool.Pool):
+    """A process pool with non-daemon processes.
+    """
+    Process = NonDaemonProcess
+
 class MultiProcPlugin(DistributedPluginBase):
     """Execute workflow with multiprocessing
 
     The plugin_args input to run can be used to control the multiprocessing
     execution. Currently supported options are:
 
     - n_procs : number of processes to use
+    - non_daemon : boolean flag to execute as non-daemon processes
 
     """
 
     def __init__(self, plugin_args=None):
         super(MultiProcPlugin, self).__init__(plugin_args=plugin_args)
         self._taskresult = {}
         self._taskid = 0
+        non_daemon = True
         n_procs = cpu_count()
         if plugin_args:
             if 'n_procs' in plugin_args:
                 n_procs = plugin_args['n_procs']
-        self.pool = Pool(processes=n_procs)
+            if 'non_daemon' in plugin_args:
+                non_daemon = plugin_args['non_daemon']
+        if non_daemon:
+            # run the execution using the non-daemon pool subclass
+            self.pool = NonDaemonPool(processes=n_procs)
+        else:
+            self.pool = Pool(processes=n_procs)
 
     def _get_result(self, taskid):
         if taskid not in self._taskresult:

nipype/pipeline/plugins/tests/test_multiproc_nondaemon.py

@@ -0,0 +1,144 @@
+import os
+from tempfile import mkdtemp
+from shutil import rmtree
+
+from nipype.testing import assert_equal, assert_true
+import nipype.pipeline.engine as pe
+from nipype.interfaces.utility import Function
+
+
+def mytestFunction(insum=0):
+    '''
+    Run a multiprocessing job and spawn child processes.
+    '''
+
+    # need to import here since this is executed as an external process
+    import multiprocessing
+    import tempfile
+    import time
+    import os
+
+    numberOfThreads = 2
+
+    # list of processes
+    t = [None] * numberOfThreads
+
+    # list of alive flags
+    a = [None] * numberOfThreads
+
+    # list of tempFiles
+    f = [None] * numberOfThreads
+
+    def dummyFunction(filename):
+        '''
+        This function writes the value 45 to the given filename.
+        '''
+        j = 0
+        for i in range(0, 10):
+          j += i
+
+        # j is now 45 (0+1+2+3+4+5+6+7+8+9)
+
+        with open(filename, 'w') as f:
+          f.write(str(j))
+
+    for n in xrange(numberOfThreads):
+
+      # mark thread as alive
+      a[n] = True
+
+      # create a temp file to use as the data exchange container
+      tmpFile = tempfile.mkstemp('.txt', 'test_engine_')[1]
+      f[n] = tmpFile # keep track of the temp file
+      t[n] = multiprocessing.Process(target=dummyFunction,
+                                     args=(tmpFile,))
+      # fire up the job
+      t[n].start()
+
+
+    # block until all processes are done
+    allDone = False
+    while not allDone:
+
+      time.sleep(1)
+
+      for n in xrange(numberOfThreads):
+
+        a[n] = t[n].is_alive()
+
+      if not any(a):
+        # if no thread is alive
+        allDone = True
+
+    # here, all processes are done
+
+    # read in all temp files and sum them up
+    total = insum
+    for file in f:
+      with open(file) as fd:
+        total += int(fd.read())
+      os.remove(file)
+
+    return total
+
+
+def run_multiproc_nondaemon_with_flag(nondaemon_flag):
+    '''
+    Start a pipe with two nodes using the multiproc plugin and passing the nondaemon_flag.
+    '''
+    
+    cur_dir = os.getcwd()
+    temp_dir = mkdtemp(prefix='test_engine_')
+    os.chdir(temp_dir)
+    
+    pipe = pe.Workflow(name='pipe')
+    
+    f1 = pe.Node(interface=Function(function=mytestFunction,
+                                    input_names=['insum'],
+                                    output_names=['sum_out']),
+                 name='f1')
+    f2 = pe.Node(interface=Function(function=mytestFunction,
+                                    input_names=['insum'],
+                                    output_names=['sum_out']),
+                 name='f2')
+
+    pipe.connect([(f1, f2, [('sum_out', 'insum')])])
+    pipe.base_dir = os.getcwd()
+    f1.inputs.insum = 0
+
+    pipe.config = {'execution': {'stop_on_first_crash': True}}
+    # execute the pipe using the MultiProc plugin with 2 processes and the non_daemon flag
+    # to enable child processes which start other multiprocessing jobs
+    execgraph = pipe.run(plugin="MultiProc",
+                         plugin_args={'n_procs': 2,
+                                      'non_daemon': nondaemon_flag})
+    
+    names = ['.'.join((node._hierarchy,node.name)) for node in execgraph.nodes()]
+    node = execgraph.nodes()[names.index('pipe.f2')]
+    result = node.get_output('sum_out')
+    os.chdir(cur_dir)
+    rmtree(temp_dir)
+    return result
+
+    
+def test_run_multiproc_nondaemon_false():
+    '''
+    This is the entry point for the test. Two times a pipe of several multiprocessing jobs gets
+    executed. First, without the nondaemon flag. Second, with the nondaemon flag.
+    
+    Since the processes of the pipe start child processes, the execution only succeeds when the
+    non_daemon flag is on.
+    '''
+    shouldHaveFailed = False
+    try:
+        # with nondaemon_flag = False, the execution should fail
+        run_multiproc_nondaemon_with_flag(False)
+    except:
+        shouldHaveFailed = True
+    yield assert_true, shouldHaveFailed
+
+def test_run_multiproc_nondaemon_true():
+    # with nondaemon_flag = True, the execution should succeed
+    result = run_multiproc_nondaemon_with_flag(True)
+    yield assert_equal, result, 180 # n_procs (2) * numberOfThreads (2) * 45 == 180
+