Source code for pycbc.pool

""" Tools for creating pools of worker processes
import multiprocessing.pool
import functools
from multiprocessing import TimeoutError, cpu_count
import types
import signal
import atexit
import logging

[docs]def is_main_process(): """ Check if this is the main control process and may handle one time tasks """ try: from mpi4py import MPI comm = MPI.COMM_WORLD rank = comm.Get_rank() return rank == 0 except (ImportError, ValueError, RuntimeError): return True
# Allow the pool to be interrupted, need to disable the children processes # from intercepting the keyboard interrupt def _noint(init, *args): signal.signal(signal.SIGINT, signal.SIG_IGN) if init is not None: return init(*args) _process_lock = None _numdone = None def _lockstep_fcn(values): """ Wrapper to ensure that all processes execute together """ numrequired, fcn, args = values with _process_lock: _numdone.value += 1 # yep this is an ugly busy loop, do something better please # when we care about the performance of this call and not just the # guarantee it provides (ok... maybe never) while 1: if _numdone.value == numrequired: return fcn(args) def _shutdown_pool(p): p.terminate() p.join()
[docs]class BroadcastPool(multiprocessing.pool.Pool): """ Multiprocessing pool with a broadcast method """ def __init__(self, processes=None, initializer=None, initargs=(), **kwds): global _process_lock global _numdone _process_lock = multiprocessing.Lock() _numdone = multiprocessing.Value('i', 0) noint = functools.partial(_noint, initializer) super(BroadcastPool, self).__init__(processes, noint, initargs, **kwds) atexit.register(_shutdown_pool, self) def __len__(self): return len(self._pool)
[docs] def broadcast(self, fcn, args): """ Do a function call on every worker. Parameters ---------- fcn: funtion Function to call. args: tuple The arguments for """ results =, [(len(self), fcn, args)] * len(self)) _numdone.value = 0 return results
[docs] def allmap(self, fcn, args): """ Do a function call on every worker with different arguments Parameters ---------- fcn: funtion Function to call. args: tuple The arguments for """ results =, [(len(self), fcn, arg) for arg in args]) _numdone.value = 0 return results
[docs] def map(self, func, items, chunksize=None): """ Catch keyboard interuppts to allow the pool to exit cleanly. Parameters ---------- func: function Function to call items: list of tuples Arguments to pass chunksize: int, Optional Number of calls for each process to handle at once """ results = self.map_async(func, items, chunksize) while True: try: return results.get(1800) except TimeoutError: pass except KeyboardInterrupt: self.terminate() self.join() raise KeyboardInterrupt
def _dummy_broadcast(self, f, args):, [args] * self.size)
[docs]class SinglePool(object):
[docs] def broadcast(self, fcn, args): return, [args])
[docs] def map(self, f, items): return [f(a) for a in items]
# This is single core, so imap and map # would not behave differently. This is defined # so that the general pool interfaces can use # imap irrespective of the pool type. imap = map
[docs]def use_mpi(require_mpi=False, log=True): """ Get whether MPI is enabled and if so the current size and rank """ use_mpi = False try: from mpi4py import MPI comm = MPI.COMM_WORLD size = comm.Get_size() rank = comm.Get_rank() if size > 1: use_mpi = True if log:'Running under mpi with size: %s, rank: %s', size, rank) except ImportError as e: if require_mpi: print(e) raise ValueError("Failed to load mpi, ensure mpi4py is installed") if not use_mpi: size = rank = 0 return use_mpi, size, rank
[docs]def choose_pool(processes, mpi=False): """ Get processing pool """ do_mpi, size, rank = use_mpi(require_mpi=mpi) if do_mpi: try: import schwimmbad pool = schwimmbad.choose_pool(mpi=do_mpi, processes=(size - 1)) pool.broadcast = types.MethodType(_dummy_broadcast, pool) atexit.register(pool.close) if processes:'NOTE: that for MPI process size determined by ' 'MPI launch size, not the processes argument') if do_mpi and not mpi:'NOTE: using MPI as this process was launched' 'under MPI') except ImportError: raise ValueError("Failed to start up an MPI pool, " "install mpi4py / schwimmbad") elif processes == 1: pool = SinglePool() else: if processes == -1: processes = cpu_count() pool = BroadcastPool(processes) pool.size = processes if size: pool.size = size return pool