RF: Futures-based MultiProc

Author: effigies

nipype/pipeline/plugins/multiproc.py

@@ -11,7 +11,8 @@
 
 # Import packages
 import os
-from multiprocessing import Process, Pool, cpu_count, pool
+from multiprocessing import cpu_count
+from concurrent.futures import ProcessPoolExecutor
 from traceback import format_exception
 import sys
 from logging import INFO
@@ -74,25 +75,6 @@ def run_node(node, updatehash, taskid):
     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, not sending more jobs at once
@@ -139,8 +121,6 @@ def __init__(self, plugin_args=None):
         self._cwd = os.getcwd()
 
         # Read in options or set defaults.
-        non_daemon = self.plugin_args.get('non_daemon', True)
-        maxtasks = self.plugin_args.get('maxtasksperchild', 10)
         self.processors = self.plugin_args.get('n_procs', cpu_count())
         self.memory_gb = self.plugin_args.get(
             'memory_gb',  # Allocate 90% of system memory
@@ -149,30 +129,19 @@ def __init__(self, plugin_args=None):
                                                        True)
 
         # Instantiate different thread pools for non-daemon processes
-        logger.debug('[MultiProc] Starting in "%sdaemon" mode (n_procs=%d, '
-                     'mem_gb=%0.2f, cwd=%s)', 'non' * int(non_daemon),
+        logger.debug('[MultiProc] Starting (n_procs=%d, '
+                     'mem_gb=%0.2f, cwd=%s)',
                      self.processors, self.memory_gb, self._cwd)
 
-        NipypePool = NonDaemonPool if non_daemon else Pool
-        try:
-            self.pool = NipypePool(
-                processes=self.processors,
-                maxtasksperchild=maxtasks,
-                initializer=os.chdir,
-                initargs=(self._cwd,)
-            )
-        except TypeError:
-            # Python < 3.2 does not have maxtasksperchild
-            # When maxtasksperchild is not set, initializer is not to be
-            # called
-            self.pool = NipypePool(processes=self.processors)
+        self.pool = ProcessPoolExecutor(max_workers=self.processors)
 
         self._stats = None
 
     def _async_callback(self, args):
         # Make sure runtime is not left at a dubious working directory
         os.chdir(self._cwd)
-        self._taskresult[args['taskid']] = args
+        result = args.result()
+        self._taskresult[result['taskid']] = result
 
     def _get_result(self, taskid):
         return self._taskresult.get(taskid)
@@ -187,9 +156,9 @@ def _submit_job(self, node, updatehash=False):
         if getattr(node.interface, 'terminal_output', '') == 'stream':
             node.interface.terminal_output = 'allatonce'
 
-        self._task_obj[self._taskid] = self.pool.apply_async(
-            run_node, (node, updatehash, self._taskid),
-            callback=self._async_callback)
+        result_future = self.pool.submit(run_node, node, updatehash, self._taskid)
+        result_future.add_done_callback(self._async_callback)
+        self._task_obj[self._taskid] = result_future
 
         logger.debug('[MultiProc] Submitted task %s (taskid=%d).',
                      node.fullname, self._taskid)
@@ -218,7 +187,7 @@ def _prerun_check(self, graph):
                 raise RuntimeError('Insufficient resources available for job')
 
     def _postrun_check(self):
-        self.pool.close()
+        self.pool.shutdown()
 
     def _check_resources(self, running_tasks):
         """