"""Module for wrapping MPI functions.
Most functions will do nothing if mpi4py is not present and will simply
replicate expected behaviour. This allows for MPI and non-MPI code to be
run without any changes (In theory...).
"""
import enum
import logging
import typing as t
from functools import lru_cache
from functools import wraps
import numpy as np
from .types import AnyValType
T = t.TypeVar("T", bound=AnyValType)
[docs]
class Ops(str, enum.Enum):
"""Available MPI operations."""
SUM = "sum"
PROD = "prod"
MAX = "max"
MIN = "min"
[docs]
def has_mpi() -> bool:
"""Checks if mpi4py is installed.
Returns
-------
bool:
True if mpi4py is installed, False otherwise
"""
try:
import mpi4py # noqa: F401
except ImportError:
return False
return True
[docs]
def convert_op(operation: Ops) -> t.Any:
"""Converts string to MPI operation.
Parameters
----------
operation : Ops
Operation to convert.
Returns
-------
t.Any
MPI operation.
"""
from mpi4py import MPI
return getattr(MPI, str(operation.upper()))
[docs]
@lru_cache(maxsize=10)
def shared_comm() -> t.Any:
"""Return shared memory communicator.
Returns the process id within a node. Used for shared memory.
"""
from mpi4py import MPI
return MPI.COMM_WORLD.Split_type(MPI.COMM_TYPE_SHARED)
[docs]
@lru_cache(maxsize=10)
def nprocs() -> int:
"""Gets number of processes or returns 1 if mpi is not installed.
Returns
-------
int:
Rank of process or 1 if MPI is not installed
"""
try:
from mpi4py import MPI
except ImportError:
return 1
comm = MPI.COMM_WORLD
return comm.Get_size()
[docs]
def allgather(value: T) -> t.List[T]:
"""Gathers all values from all processes.
Returns value if mpi is not installed.
Parameters
----------
value:
Value to gather
Returns
-------
list:
List of gathered values or list containing value if mpi
is not installed
"""
try:
from mpi4py import MPI
except ImportError:
return [value]
comm = MPI.COMM_WORLD
data = value
data = comm.allgather(data)
return data
[docs]
def allreduce(value: T, op: Ops) -> T:
"""Reduces all values from all processes.
Returns value if mpi is not installed.
Parameters
----------
value:
Value to reduce
op: :class:`~taurex.mpi.Ops`
Operation to perform
Returns
-------
result:
Reduced value or value if mpi is not installed
"""
try:
from mpi4py import MPI
except ImportError:
return value
comm = MPI.COMM_WORLD
data = value
data = comm.allreduce(value, op=convert_op(op))
return data
[docs]
def broadcast(array: T, rank: t.Optional[int] = 0) -> T:
"""Broadcasts array from rank or returns array if mpi is not installed.
Parameters
----------
array:
Array to broadcast
rank: int, optional
Rank to broadcast from, default is 0
Returns
-------
array:
Broadcasted array or array if mpi is not installed
"""
import numpy as np
try:
from mpi4py import MPI
except ImportError:
return array
comm = MPI.COMM_WORLD
if isinstance(array, np.ndarray):
data = None
if get_rank() == rank:
data = np.copy(array)
else:
data = np.zeros_like(array)
comm.Bcast(data, root=rank)
else:
data = comm.bcast(array, root=rank)
return data
[docs]
@lru_cache(maxsize=10)
def get_rank(comm: t.Any = None) -> int:
"""Gets rank or returns 0 if mpi is not installed.
Parameters
----------
comm: int, optional
MPI communicator, default is MPI_COMM_WORLD
Returns
-------
int:
Rank of process in communitor or 0 if MPI is not installed
"""
try:
from mpi4py import MPI
except ImportError:
return 0
comm = comm or MPI.COMM_WORLD
rank = comm.Get_rank()
return rank
[docs]
def barrier(comm: t.Any = None) -> None:
"""Waits for all processes to finish.
Does nothing if mpi4py not present
Parameters
----------
comm: int, optional
MPI communicator, default is MPI_COMM_WORLD
"""
try:
from mpi4py import MPI
except ImportError:
return
comm = comm or MPI.COMM_WORLD
comm.Barrier()
[docs]
def only_master_rank(f) -> t.Callable:
"""A decorator to ensure only the master MPI rank can run it.
Parameters
----------
f: callable
Function to decorate
Returns
-------
callable
Decorated function
"""
@wraps(f)
def wrapper(*args, **kwargs):
if get_rank() == 0:
return f(*args, **kwargs)
return wrapper
[docs]
@lru_cache(maxsize=10)
def shared_rank() -> int:
"""Gets rank within shared memory communicator (MPI only)."""
return shared_comm().Get_rank()
[docs]
def materialize(
arr: t.Optional[t.Any],
) -> t.Optional[np.ndarray]:
"""Turn a dataset-like source into a plain numpy array.
Anything that is already an array (or ``None``) is returned as-is.
Used on the code paths where shared memory is not used, so that
callers handing in e.g. an :class:`h5py.Dataset` still get an
in-memory array back.
"""
if arr is None or isinstance(arr, np.ndarray):
return arr
return np.asarray(arr[()])
_SHARED_FILL_CHUNK_BYTES = 32 * 1024 * 1024
"""Maximum amount of data read at a time when filling shared memory."""
[docs]
def fill_shared(shared_array: np.ndarray, arr: t.Any) -> None:
"""Copy ``arr`` into an already-allocated shared array.
In-memory arrays are copied in one go. Dataset-like sources are read
along their leading axis in bounded chunks so that the shared-memory
root rank never has to materialise a second full private copy of the
whole array just to fill the shared segment.
"""
if isinstance(arr, np.ndarray):
np.copyto(shared_array, arr)
return
leading = arr.shape[0]
if leading == 0:
return
row_bytes = arr.size // leading * arr.dtype.itemsize
step = max(1, _SHARED_FILL_CHUNK_BYTES // max(row_bytes, 1))
for start in range(0, leading, step):
stop = min(start + step, leading)
shared_array[start:stop] = arr[start:stop]
[docs]
def allocate_as_shared( # noqa: C901
arr: t.Any,
logger: t.Optional[logging.Logger] = None,
force_shared: t.Optional[bool] = False,
):
"""Converts a numpy array into an MPI shared memory.
This allow for things like opacities to be loaded only
once per node when using MPI. Only activates if mpi4py
installed and when enabled via the ``mpi_use_shared`` input::
[Global]
mpi_use_shared = True
or ``force_shared=True`` otherwise does nothing and
returns the same array back
When using MPI shared memory, only the shared-memory root rank
(rank 0 in the shared communicator) needs to pass the actual data
array. Non-root ranks can pass ``None`` to avoid allocating a
private copy of the array in their local memory. The shape, dtype,
and size are automatically broadcast from the root rank.
On the root rank ``arr`` may also be an on-disk dataset (anything
exposing ``shape``, ``dtype``, ``size`` and ``__getitem__``, such as
an :class:`h5py.Dataset`). It is then read into the shared segment in
bounded chunks, so the root never holds a second full private copy of
the data. Such sources are materialised into a plain array if sharing
turns out to be inactive.
Parameters
----------
arr: numpy array, dataset-like or None
Array to convert to shared memory. Can be None on non-root
ranks of the shared communicator.
logger: :class:`~taurex.log.logger.Logger`
Logger object to print outputs
force_shared: bool
Force conversion to shared memory
Returns
-------
array:
If enabled and MPI present, shared memory version of array
otherwise the original array
"""
try:
from mpi4py import MPI
except ImportError:
return materialize(arr)
from taurex.cache import GlobalCache
if GlobalCache()["mpi_use_shared"] or force_shared:
if logger is not None:
logger.info("Moving to shared memory")
comm = shared_comm()
sh_rank = shared_rank()
sh_size = comm.Get_size()
if logger is not None:
logger.info(
"Shared-memory group: rank %s / %s (world rank %s)",
sh_rank,
sh_size,
MPI.COMM_WORLD.Get_rank(),
)
if sh_size <= 1:
if logger is not None:
logger.warning(
"Shared-memory group has only %s rank(s)! "
"MPI shared memory is NOT reducing duplication — "
"all ranks have private copies. Check that ranks "
"are on the same NUMA node (socket) and that your "
"MPI implementation supports COMM_TYPE_SHARED.",
sh_size,
)
return materialize(arr)
# Determine shape, dtype, and size across ranks.
# On non-root ranks, arr may be None to avoid private allocation.
if sh_rank == 0:
if arr is None:
raise ValueError(
"Shared-memory root rank (0) must provide the data array"
)
shape = arr.shape
dtype = arr.dtype
itemsize = arr.dtype.itemsize
nbytes = arr.size * itemsize
nbytes_alloc = nbytes
else:
shape = None
dtype = None
itemsize = None
nbytes = 0
nbytes_alloc = 0
# Broadcast metadata from root to all ranks in shared communicator
nbytes = comm.bcast(nbytes, root=0)
itemsize = comm.bcast(itemsize, root=0)
shape = comm.bcast(shape, root=0)
dtype = comm.bcast(dtype, root=0)
# CRITICAL: Only rank 0 allocates shared memory; other ranks pass 0.
# MPI_Win_Allocate_shared sums sizes across ALL ranks in the
# communicator, so non-root ranks MUST pass 0 to avoid multiplying
# the shared memory footprint by the number of ranks per node.
window = MPI.Win.Allocate_shared(nbytes_alloc, itemsize, comm=comm)
buf, actual_itemsize = window.Shared_query(0)
if actual_itemsize != itemsize:
raise Exception(
f"Shared memory itemsize {actual_itemsize} != "
f"array itemsize {itemsize}"
)
shared_array = np.ndarray(buffer=buf, dtype=dtype, shape=shape)
if sh_rank == 0:
fill_shared(shared_array, arr)
comm.Barrier()
return shared_array
else:
return materialize(arr)