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)
def test_mpi():
with MPIPool() as pool:
all_tasks = [[random.random() for i in range(1000)]]
# test map alone
for tasks in all_tasks:
results = pool.map(_function, tasks)
for r1, r2 in zip(results, [_function(x) for x in tasks]):
assert isclose(r1, r2)
assert len(results) == len(tasks)
# test map with callback
for tasks in all_tasks:
results = pool.map(_function, tasks, callback=_callback)
for r1, r2 in zip(results, [_function(x) for x in tasks]):
assert isclose(r1, r2)
assert len(results) == len(tasks)
def test_mpi_with_dill():
pool = MPIPool(use_dill=True)
pool.wait(lambda: sys.exit(0))
all_tasks = [[random.random() for i in range(1000)]]
# test map alone
for tasks in all_tasks:
results = pool.map(_function, tasks)
for r1, r2 in zip(results, [_function(x) for x in tasks]):
assert isclose(r1, r2)
assert len(results) == len(tasks)
# test map with callback
for tasks in all_tasks:
results = pool.map(_function, tasks, callback=_callback)
for r1, r2 in zip(results, [_function(x) for x in tasks]):
assert isclose(r1, r2)
assert len(results) == len(tasks)
pool.close()
"""
# Standard library
import random
import sys
# Project
from schwimmbad.mpi import MPIPool, MPI
from schwimmbad.tests.test_pools import isclose, _function
def callback(x):
assert MPI.COMM_WORLD.Get_rank() == 0
if MPI is not None and MPI.COMM_WORLD.Get_size() > 1:
pool = MPIPool()
pool.wait(lambda: sys.exit(0))
all_tasks = [[random.random() for i in range(1000)]]
# test map alone
for tasks in all_tasks:
results = pool.map(_function, tasks)
for r1,r2 in zip(results, [_function(x) for x in tasks]):
assert isclose(r1, r2)
# test map with callback
for tasks in all_tasks:
results = pool.map(_function, tasks, callback=callback)
for r1,r2 in zip(results, [_function(x) for x in tasks]):
def test_mpi_worker_error():
with MPIPool() as pool:
tasks = [random.random() for i in range(1000)]
pool.map(worker_error, tasks) # should fail
pars.set_matter_power(redshifts=[0.], kmax=200.0)
pars.NonLinearModel.set_params(halofit_version='takahashi')
camb.set_feedback_level(level=100)
results = camb.get_results(pars)
if shellnums is None:
shellnum_min = int(
results.comoving_radial_distance(zmin) * h // shellwidth)
shellnum_max = int(
results.comoving_radial_distance(zmax) * h // shellwidth + 1)
shellnums = list(range(shellnum_min, shellnum_max + 1))
else:
shellnums = list(map(int, shellnums.split(",")))
try:
pool = MPIPool()
except:
pool = None
if is_cutsky:
args = zip(product([1, 2, 3], shellnums), repeat(snapshot_cutsky))
else:
args = zip(product([1, 2, 3], shellnums), repeat(None))
if pool is not None:
if not pool.is_master():
pool.wait()
sys.exit(0)
pool.map(generate_lightcone_shell, args)
else:
map(generate_lightcone_shell, args)
pars.set_matter_power(redshifts=[0.], kmax=200.0)
pars.NonLinearModel.set_params(halofit_version='takahashi')
camb.set_feedback_level(level=100)
results = camb.get_results(pars)
if shellnums is None:
shellnum_min = int(
results.comoving_radial_distance(zmin) * h // shellwidth)
shellnum_max = int(
results.comoving_radial_distance(zmax) * h // shellwidth + 1)
shellnums = list(range(shellnum_min, shellnum_max + 1))
else:
shellnums = list(map(int, shellnums.split(",")))
try:
pool = MPIPool()
except:
pool = None
### Loop over galtypes and shells for galaxies and randoms
args = product(shellnums, nz_pars, [nrandoms], [dir_out],
[lightcone_name_template], [output_name_template],
[random_name_template])
if pool is not None:
if not pool.is_master():
pool.wait()
sys.exit(0)
else:
pool.map(process_shell, args)
else:
map(process_shell, args)