Homec4science
Diffusion tamaas (ci-fix-static-checks) Browse

tamaas/tests/test_mpi_routines.py
998dbf9a8e3aci-fix-static-checks

Last Change
Raw File

test_mpi_routines.py
View Options

# -*- coding: utf-8 -*-
# @file
# @section LICENSE
#
# Copyright (©) 2016-2021 EPFL (École Polytechnique Fédérale de Lausanne),
# Laboratory (LSMS - Laboratoire de Simulation en Mécanique des Solides)
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from __future__ import print_function
import inspect
import sys
import re
import os
import numpy as np
from mpi4py import MPI
from pytest import mark
import tamaas
PROG = '''
import sys
import numpy as np
sys.path = ['{}'] + sys.path
import mpi_routines
from mpi4py import MPI
res = 0
try:
mpi_routines.{}()
except Exception as e:
print(e)
res = 1
comm = MPI.Comm.Get_parent()
comm.Reduce(np.array(res, 'i'), None, op=MPI.SUM, root=0)
comm.Free()
sys.exit(0)
'''
HOSTFILE_CONTENT = 'localhost slots=24\n'
@mark.skipif(not tamaas.TamaasInfo.has_mpi, reason='tamaas compiled w/o MPI')
def test_mpi_routines(tamaas_fixture):
import mpi_routines
me_path = os.path.realpath(__file__)
current_dir = os.path.dirname(me_path)
hostfile = os.path.join(current_dir, 'hostfile-mpi')
with open(hostfile, 'w') as fh:
fh.write(HOSTFILE_CONTENT)
mpi_matcher = re.compile('^mpi_')
for name, _ in inspect.getmembers(mpi_routines, inspect.isfunction):
if mpi_matcher.match(name):
try:
info = MPI.Info.Create()
info['add-hostfile'] = hostfile
print('Spawning "{}"'.format(name))
comm = MPI.COMM_SELF.Spawn(
sys.executable, args=[
'-c', PROG.format(os.path.abspath(current_dir),
name),
], maxprocs=2, info=info)
res = np.array(0, 'i')
comm.Reduce(None, res, op=MPI.SUM, root=MPI.ROOT)
except MPI.Exception as exc:
print(exc)
raise RuntimeError('Could not spawn MPI processes for "{}"'
.format(name))
assert res == 0, name
c4science · Help