import configparser
import os
import sys
from abc import ABC, abstractmethod
from ase.parallel import MPI4PY as ASEMPI4Py
from ase.parallel import DummyMPI as ASEDummyMPI
from ase.parallel import broadcast as ase_broadcast
from ase.parallel import paropen as ase_paropen
homepath = os.environ["HOME"]
configp = configparser.ConfigParser()
configp.read(homepath + "/.directupsampling/config")
try:
config = configp["parallel"]
except KeyError:
config = {}
[docs]
def broadcast(*args, **kwargs):
return ase_broadcast(*args, **kwargs)
[docs]
def paropen(*args, **kwargs):
return ase_paropen(*args, **kwargs)
[docs]
def parprint(*args, **kwargs):
if world.rank == 0:
print(*args, **kwargs)
[docs]
def parcall(function, *args, comm=None, **kwargs):
"""Call only on comm master rank.
kwarg :comm: is reserved and not passed to :function:,
otherwise args and kwargs are passed.
"""
if comm is None:
comm = world
if comm.rank == 0:
return function(*args, **kwargs)
[docs]
def parlog(message, logger, level="info", comm=None):
"""For logger call only on master rank."""
return parcall(
getattr(logger, level),
message,
comm=comm,
)
def _get_comm():
"""Get the correct MPI world object."""
if "mpi4py" in sys.modules or config.get("always_use_mpi"):
return MPI4Py()
return DummyMPI()
[docs]
class BaseMPI(ABC):
[docs]
@abstractmethod
def scatter(self, list_, root=0): ...
[docs]
@abstractmethod
def gather(self, value, root=0): ...
[docs]
@abstractmethod
def allgather(self, value): ...
[docs]
@abstractmethod
def broadcast(self, value, root=0): ...
[docs]
class DummyMPI(BaseMPI, ASEDummyMPI):
rank = 0
size = 1
[docs]
def scatter(self, list_, root=0):
assert root == 0
assert len(list_) == 1
return list_[0]
[docs]
def gather(self, value, root=0):
assert root == 0
return [value]
[docs]
def allgather(self, value):
return [value]
[docs]
def broadcast(self, value, root=0):
assert root == 0
return value
[docs]
class MPI4Py(BaseMPI, ASEMPI4Py):
[docs]
def scatter(self, list_, root=0):
return self.comm.scatter(list_, root=root)
[docs]
def gather(self, value, root=0):
return self.comm.gather(value, root=root)
[docs]
def allgather(self, value):
return self.comm.allgather(value)
[docs]
def broadcast(self, value, root=0):
return self.comm.bcast(value, root=root)
[docs]
class MPI:
"""Wrapper for MPI world object.
Decides at runtime (after all imports) which one to use:
* MPI4Py
* a dummy implementation for serial runs
"""
def __init__(self):
self.comm = None
def __getattr__(self, name):
if self.comm is None:
self.comm = _get_comm()
return getattr(self.comm, name)
world = MPI()