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()
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]):
assert isclose(r1, r2)
pool.close()