Python中MPI发送接收dtype=object大型数组的问题解决与替代方案

执行命令mpirun -np 2 python3 main.py跨CPU发送dtype=object的大型numpy数组时，数组规模较小时正常运行，增大后触发如下错误：

示例代码 #

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)
    
    # 小尺寸可正常运行
    #a=1.5*np.ones(50000000,dtype='float64')
    #b=2.5*np.ones(100000000,dtype='float64')

    # 大尺寸触发报错
    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'

错误原因 #

mpi4py对dtype=object的数组会采用pickle序列化处理，而非直接使用MPI的高效二进制传输。当数组总数据量过大时，计算得到的序列化后总字节数（12000000297）超过了32位int的最大值（2^31-1=2147483647），导致类型转换时溢出报错。

解决方法与替代策略 #

策略1：拆分发送，避免使用object dtype数组 #

object dtype数组本质存储的是子数组的指针，完全可以直接分别发送每个子数组，无需包装成object数组，既解决溢出问题，又提升传输效率。

代码示例：

import numpy as np
from mpi4py import MPI

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

if rank == 0:
    a = 1.5 * np.ones(500000000, dtype='float64')
    b = 2.5 * np.ones(1000000000, dtype='float64')
    # 依次发送两个子数组
    comm.send(a, dest=1)
    comm.send(b, dest=1)
else:
    # 依次接收子数组
    a_recv = comm.recv(source=0)
    b_recv = comm.recv(source=0)
    # 按需组合成列表或数组
    receive = [a_recv, b_recv]

策略2：使用MPI直接缓冲区传输（大写Send/Recv） #

mpi4py提供的大写Send/Recv方法会直接操作numpy数组的内存缓冲区，跳过pickle序列化，不仅支持64位长度计算，不会触发溢出，传输效率也远高于Python对象的send/recv。

代码示例：

import numpy as np
from mpi4py import MPI

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

if rank == 0:
    a = 1.5 * np.ones(500000000, dtype='float64')
    b = 2.5 * np.ones(1000000000, dtype='float64')
    # 直接发送数组缓冲区
    comm.Send(a, dest=1)
    comm.Send(b, dest=1)
else:
    # 预先分配接收缓冲区（需知道数组尺寸）
    a_recv = np.empty(500000000, dtype='float64')
    b_recv = np.empty(1000000000, dtype='float64')
    comm.Recv(a_recv, source=0)
    comm.Recv(b_recv, source=0)
    receive = [a_recv, b_recv]

策略3：分块发送保留object数组（业务必需时） #

如果业务逻辑必须保留object dtype数组的形式，可以将大数组拆分为小块发送，接收端再重组。

代码示例：

import numpy as np
from mpi4py import MPI

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
CHUNK_SIZE = 100000000  # 自定义分块大小

if rank == 0:
    a = 1.5 * np.ones(500000000, dtype='float64')
    b = 2.5 * np.ones(1000000000, dtype='float64')
    arr = np.array([a, b], dtype=object)
    
    # 先发送数组中子元素的数量
    comm.send(len(arr), dest=1)
    for sub_arr in arr:
        # 发送当前子数组的总长度
        comm.send(len(sub_arr), dest=1)
        # 分块发送子数组内容
        for i in range(0, len(sub_arr), CHUNK_SIZE):
            chunk = sub_arr[i:i+CHUNK_SIZE]
            comm.send(chunk, dest=1)
else:
    arr_len = comm.recv(source=0)
    receive = []
    for _ in range(arr_len):
        sub_len = comm.recv(source=0)
        sub_arr = np.empty(sub_len, dtype='float64')
        pos = 0
        # 分块接收并重组子数组
        while pos < sub_len:
            chunk = comm.recv(source=0)
            sub_arr[pos:pos+len(chunk)] = chunk
            pos += len(chunk)
        receive.append(sub_arr)
    # 转换回object dtype数组（如果需要）
    receive = np.array(receive, dtype=object)

内容的提问来源于stack exchange，提问作者DrMittal