如何在 python 中使用 dtype=object mpi 发送和接收大型数组？

Question

以下是演示我的问题的玩具代码。我运行它作为

mpirun -np 2 python3 main.py

我想跨 CPU 发送大型数组 (dtype=object)。该代码适用于较小的数组，但不适用于较大的数组。错误消息显示在代码下方。

import numpy as np
from mpi4py import MPI

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
N = comm.Get_size()

if rank == 0:
    print(N)
    
    arr = np.zeros(2,dtype=object)
    
    #This works
    #a=1.5*np.ones(50000000,dtype='float64')
    #b=2.5*np.ones(100000000,dtype='float64')

    #This doesn't work
    a=1.5*np.ones(500000000,dtype='float64')
    b=2.5*np.ones(1000000000,dtype='float64')

    arr[0] = a
    arr[1] = b

    comm.send(arr,dest=1)
        
else:
    receive = comm.recv(source=0)

错误信息：

2
Traceback (most recent call last):
  File "/home/shikhar/Documents/Examples/mpi/main.py", line 19, in <module>
    comm.send(arr,dest=1)
  File "mpi4py/MPI/Comm.pyx", line 1406, in mpi4py.MPI.Comm.send
  File "mpi4py/MPI/msgpickle.pxi", line 211, in mpi4py.MPI.PyMPI_send
  File "mpi4py/MPI/msgpickle.pxi", line 147, in mpi4py.MPI.pickle_dump
  File "mpi4py/MPI/msgbuffer.pxi", line 50, in mpi4py.MPI.downcast
OverflowError: integer 12000000297 does not fit in 'int'

如何纠正上述问题，或者如果人们可以针对这种情况提出替代策略，我们将不胜感激？

Answer 1

按照this page和this（由@9769953建议），代码可以运行。我们需要使用

mpi4py.util.pkl5

。

import numpy as np
from mpi4py import MPI
from mpi4py.util import pkl5

comm = pkl5.Intracomm(MPI.COMM_WORLD)
rank = comm.Get_rank()
N = comm.Get_size()

if rank == 0:
    print(N)
    
    arr = np.zeros(2,dtype=object)

    a=1.5*np.ones(500000000,dtype='float64')
    b=2.5*np.ones(1000000000,dtype='float64')

    arr[0] = a
    arr[1] = b

    comm.send(arr,dest=1)
        
else:
    receive = comm.recv(source=0)