Provide MPI communicator from Python to library.
Summary
Continuation from #3644 (closed) to pass communicator from Python.
Allow multiple simulations to be managed from the Python level.
Allow Python logic to determine the partitioning of a communicator into subcommunicators for the library to use instead of COMM_WORLD.
Use cases
Immediate benefit will be to allow gmxapi ensemble simulation code to work with libgromacs_mpi
the same way it works with thread-MPI libgromacs
.
This feature is the next step towards allowing library-external code to apply an arbitrary number of MPI ranks to an arbitrary number of simulations, such as for API-driven replica exchange, "multisim", or ensemble methods.
Impact
Currently, gmxapi.mdrun
can only manage ensemble simulation work by using mpi4py
to allocate one rank per simulation. However, gmxapi is not currently able to tell an MPI-enabled GROMACS library to use specific ranks (or subcommunicators), so it only makes sense to use gmxapi
for ensemble simulations with thread-MPI GROMACS builds, limiting the amount of resources to a single node per simulation.
This change will allow resources to be allocated to ensemble simulations in terms of MPI ranks, regardless of the ratio of ranks and nodes, while also making gmxapi much more useful to MPI-enabled GROMACS builds.
Detailed description
Part 1
The essential gmxapi
code changes will be an adjusted signature in gmxapi.simulation.context
and updates to the bindings code.
Bindings code updates:
- Add bindings for a
gmxapi::createContext(const gmxapi::ResourceAssignment&)
version ofPyContext
. - Extract the
MPI_Comm
from a Pythonmpi4py.MPI.Comm
PyObject*
. - Use the interface from #3644 (closed) to provide the communicator when initializing the simulation resources.
update: completed with resolution of #4423 (closed).
Part 2
- Bindings to allow libgmxapi to report to Python whether it can use a multi-rank communicator (see #3644 (closed))
- Pure-Python logic to more flexibly map available ranks to available work.
update: in a related (follow-up) issue, we should reinvestigate patterns for handling errors while preventing hangs on other ranks. Ref modern Python concurrency patterns: capturing exceptions as part of a Future model; cancelations, shields, well defined task cancellation propagation policies, and exceptions with special meanings.
Requirements
There are still some MPI-related global variables in the library code, and some hard-coded references to MPI_COMM_WORLD that will need to be revisited and potentially cleaned up.
mpi4py
will become a build-time requirement for the Python bindings instead of just a run-time requirement.