adrn / schwimmbad

A common interface to processing pools.
MIT License
115 stars 18 forks source link

Overcome the 2GB limit in MPI pickling #41

Open inonchiu opened 3 years ago

inonchiu commented 3 years ago

The MPI returns an error if the memory size is too large to be picked during the MPI. This issue has been reported before (see this). It raises an error like this.

File "/***/lib/python3.7/site-packages/schwimmbad/mpi.py", line 196, in map
    status=status)
File "mpi4py/MPI/Comm.pyx", line 1173, in mpi4py.MPI.Comm.recv
File "mpi4py/MPI/msgpickle.pxi", line 302, in mpi4py.MPI.PyMPI_recv
File "mpi4py/MPI/msgpickle.pxi", line 263, in mpi4py.MPI.PyMPI_recv_match
File "mpi4py/MPI/msgpickle.pxi", line 139, in mpi4py.MPI.Pickle.alloc
SystemError: Negative size passed to PyBytes_FromStringAndSize

Recently, mpi4py has provided a workaround in mpi4py.util (see this). To overcome this memory limit in MPI, I added two lines in schwimmbad/mpi.py as follows, and it worked. In the function _import_mpi,

def _import_mpi(quiet=False, use_dill=False):
    global MPI 
    try:
        from mpi4py import MPI as _MPI
        # see https://mpi4py.readthedocs.io/en/stable/mpi4py.util.pkl5.html
        from mpi4py.util import pkl5 # NEW, use pickle5
        _MPI.COMM_WORLD = pkl5.Intracomm(_MPI.COMM_WORLD) # NEW, use comm in pkl5
       ...

It will be great to enhance this in a future release of schwimmbad.