MPI (taurex.mpi)

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…).

class Ops(value)

Bases: str, Enum

Available MPI operations.

MAX = 'max'

MIN = 'min'

PROD = 'prod'

SUM = 'sum'

allgather(value: T) → List[T]

Gathers all values from all processes.

Returns value if mpi is not installed.

Parameters:

value – Value to gather

Returns:

List of gathered values or list containing value if mpi is not installed

Return type:

list

allocate_as_shared(arr: Any, logger: Logger | None = None, force_shared: bool | None = 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 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 (Logger) – Logger object to print outputs

• force_shared (bool) – Force conversion to shared memory

Returns:

If enabled and MPI present, shared memory version of array otherwise the original array

Return type:

array

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 (Ops) – Operation to perform

Returns:

Reduced value or value if mpi is not installed

Return type:

result

barrier(comm: Any | None = 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

broadcast(array: T, rank: int | None = 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:

Broadcasted array or array if mpi is not installed

Return type:

array

convert_op(operation: Ops) → Any

Converts string to MPI operation.

Parameters:

operation (Ops) – Operation to convert.

Returns:

MPI operation.

Return type:

t.Any

fill_shared(shared_array: ndarray, arr: 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.

get_rank(comm: 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:

Rank of process in communitor or 0 if MPI is not installed

Return type:

int

has_mpi() → bool

Checks if mpi4py is installed.

Returns:

True if mpi4py is installed, False otherwise

Return type:

bool

materialize(arr: Any | None) → ndarray | None

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 h5py.Dataset still get an in-memory array back.

nprocs() → int

Gets number of processes or returns 1 if mpi is not installed.

Returns:

Rank of process or 1 if MPI is not installed

Return type:

int

only_master_rank(f) → Callable

A decorator to ensure only the master MPI rank can run it.

Parameters:

f (callable) – Function to decorate

Returns:

Decorated function

Return type:

callable

shared_comm() → Any

Return shared memory communicator.

Returns the process id within a node. Used for shared memory.

shared_rank() → int

Gets rank within shared memory communicator (MPI only).