Source code for lenstronomy.Sampling.Pool.pool
# coding: utf-8
"""
this file is taken from schwimmbad (https://github.com/adrn/schwimmbad) and an explicit fork by Aymeric Galan
to replace the multiprocessing with the multiprocess dependence as for multi-threading, multiprocessing is
not supporting dill (only pickle) which is required.
Tests show that the MPI mode works with Python 3.7.2 but not with Python 3.7.0 on a specific system due to mpi4py
dependencies and configurations.
Contributions by:
- Peter K. G. Williams
- JĂșlio Hoffimann Mendes
- Dan Foreman-Mackey
- Aymeric Galan
- Simon Birrer
Implementations of four different types of processing pools:
- MPIPool: An MPI pool.
- MultiPool: A multiprocessing for local parallelization.
- SerialPool: A serial pool, which uses the built-in `map` function
"""
__version__ = "0.3.0"
__author__ = "Adrian Price-Whelan <adrianmpw@gmail.com>"
# Standard library
import sys
import logging
log = logging.getLogger(__name__)
_VERBOSE = 5
# from schwimmbad.multiprocessing import MultiPool
# from schwimmbad.jl import JoblibPool
__all__ = ['choose_pool']
[docs]def choose_pool(mpi=False, processes=1, **kwargs):
"""
Extends the capabilities of the schwimmbad.choose_pool method.
It handles the `use_dill` parameters in kwargs, that would otherwise raise an error when processes > 1.
Any thread in the returned multiprocessing pool (e.g. processes > 1) also default
The requirement of schwimmbad relies on the master branch (as specified in requirements.txt).
The 'use_dill' functionality can raise if not following the requirement specified.
Choose between the different pools given options from, e.g., argparse.
Parameters
----------
mpi : bool, optional
Use the MPI processing pool, :class:`~schwimmbad.mpi.MPIPool`. By
default, ``False``, will use the :class:`~schwimmbad.serial.SerialPool`.
processes : int, optional
Use the multiprocessing pool,
:class:`~schwimmbad.multiprocessing.MultiPool`, with this number of
processes. By default, ``processes=1``, will use them:class:`~schwimmbad.serial.SerialPool`.
Any additional kwargs are passed in to the pool class initializer selected by the arguments.
"""
# Imports moved here to avoid crashing at import time if dependencies
# are missing
from lenstronomy.Sampling.Pool.multiprocessing import MultiPool
from schwimmbad.serial import SerialPool
from schwimmbad.mpi import MPIPool
if mpi:
if not MPIPool.enabled():
raise SystemError("Tried to run with MPI but MPIPool not enabled.")
try:
pool = MPIPool(**kwargs)
except:
raise ImportError('MPIPool of schwimmbad can not be generated. lenstronomy uses a specific branch of '
'schwimmbad specified in the requirements.txt. Make sure you are using the correct '
'version of schwimmbad. In particular the "use_dill" argument is not supported in the '
'pypi version 0.3.0.')
if not pool.is_master():
pool.wait()
sys.exit(0)
log.info("Running with MPI on {0} cores".format(pool.size))
return pool
elif processes != 1 and MultiPool.enabled():
if 'use_dill' in kwargs:
# schwimmbad MultiPool does not support dill so we remove this option from the kwargs
_ = kwargs.pop('use_dill')
log.info("Running with MultiPool on {0} cores".format(processes))
return MultiPool(processes=processes, **kwargs)
else:
log.info("Running with SerialPool")
return SerialPool(**kwargs)