Source code for bmtk.tests.builder.test_network_mpi
import pytest
import tempfile
import numpy as np
from bmtk.builder.builder_utils import check_properties_across_ranks
from bmtk.builder import NetworkBuilder
try:
from mpi4py import MPI
comm = MPI.COMM_WORLD
MPI_rank = comm.Get_rank()
MPI_size = comm.Get_size()
has_mpi = True
except:
MPI_rank = 0
MPI_size = 1
has_mpi = False
[docs]
def tmpdir():
tmp_dir = tempfile.mkdtemp() if MPI_rank == 0 else None
tmp_dir = comm.bcast(tmp_dir, 0)
return tmp_dir
[docs]
@pytest.mark.skipif(MPI_size < 2, reason='Can only run test using mpi')
def test_create_network_dir_mpi():
# For old MPI issue were both ranks were trying to create the output_dir at the same time.
tdir = tmpdir()
net = NetworkBuilder('v1')
net.add_nodes(N=1, model_type='biophysical')
net.build()
net.save_nodes(output_dir=tdir)
# check_properties_across_ranks
[docs]
@pytest.mark.skipif(MPI_size < 2, reason='Can only run test using mpi')
def test_check_properties_mpi():
check_properties_across_ranks({
'p1': 'p1',
'p2': 2.22,
'p3': 3,
'p4': True,
'p5': [0, 1, 2, 3],
'p6': ['a', 'b', 'c', 'd'],
'p7': np.array([1.0, 2.0, 3.0, 4.0])
})
[docs]
@pytest.mark.skipif(MPI_size < 2, reason='Can only run test using mpi')
def test_check_nonuniform_args():
# Check that the number of arguments is the same across ranks
with pytest.raises(IndexError):
props = {
'p1': 1,
'p2': 'two'
}
if MPI_rank > 0:
props['p3'] = 3.0
check_properties_across_ranks(props)
# Check the argument keys are the same
with pytest.raises(TypeError):
check_properties_across_ranks({
'p{}'.format(MPI_rank): 0
})
[docs]
@pytest.mark.skipif(MPI_size < 2, reason='Can only run test using mpi')
def test_nonuniform_str_mpi():
with pytest.raises(TypeError):
check_properties_across_ranks({
'p1': str(MPI_rank),
'p2': 2.22,
'p3': 3,
'p4': True,
'p5': [0, 1, 2, 3],
'p6': ['a', 'b', 'c', 'd'],
'p7': np.array([1.0, 2.0, 3.0, 4.0])
})
[docs]
@pytest.mark.skipif(MPI_size < 2, reason='Can only run test using mpi')
def test_nonuniform_int_mpi():
with pytest.raises(Exception):
check_properties_across_ranks({
'p1': 'p1',
'p2': 2.22,
'p3': int(MPI_rank),
'p4': True,
'p5': [0, 1, 2, 3],
'p6': ['a', 'b', 'c', 'd'],
'p7': np.array([1.0, 2.0, 3.0, 4.0])
})
[docs]
@pytest.mark.skipif(MPI_size < 2, reason='Can only run test using mpi')
def test_nonuniform_float_mpi():
with pytest.raises(TypeError):
check_properties_across_ranks({
'p1': 'p1',
'p2': float(MPI_rank),
'p3': 3,
'p4': True,
'p5': [0, 1, 2, 3],
'p6': ['a', 'b', 'c', 'd'],
'p7': np.array([1.0, 2.0, 3.0, 4.0])
})
[docs]
@pytest.mark.skipif(MPI_size < 2, reason='Can only run test using mpi')
def test_nonuniform_list_mpi():
with pytest.raises(TypeError):
check_properties_across_ranks({
'p1': 'p1',
'p2': 2.22,
'p3': 3,
'p4': True,
'p5': [0, 1, 2, 3 + MPI_rank],
'p6': ['a', 'b', 'c', 'd'],
'p7': np.array([1.0, 2.0, 3.0, 4.0])
})
[docs]
@pytest.mark.skipif(MPI_size < 2, reason='Can only run test using mpi')
def test_nonuniform_array_mpi():
with pytest.raises(TypeError):
check_properties_across_ranks({
'p1': 'p1',
'p2': float(MPI_rank),
'p3': 3,
'p4': True,
'p5': [0, 1, 2, 3 + MPI_rank],
'p6': ['a', 'b', 'c', 'd'],
'p7': np.array([1.0, 2.0, 3.0, 4.0 + MPI_rank])
})
if __name__ == '__main__':
test_check_properties_mpi()
test_check_nonuniform_args()
test_nonuniform_str_mpi()
test_nonuniform_float_mpi()
test_nonuniform_int_mpi()
test_nonuniform_list_mpi()
test_nonuniform_array_mpi()