mpi4py中接收数据后调用extend导致代码冻结的问题排查

我用mpi4py实现多处理器数据通信，定义了嵌套列表neighbors，每个元素对应一个数据子区域，列表内的整数值是分配给该子区域的MPI进程编号。例如将数据划分为3个子区域、每个子区域分配1个MPI进程时，neighbors=[[0],[1],[2]]。

通常各数据子区域独立运行，偶尔需要在子区域间传递数据片段。每个处理器会遍历其相邻数据子区域及对应的进程，执行非阻塞发送（isend）和阻塞接收（recv）操作。注释掉bankr.extend(received)语句时程序运行正常，3进程场景下会输出预期的收发日志；但取消该语句注释、尝试将接收的数据追加到一个大列表时，代码在输出部分日志后冻结。

阻塞接收部分代码 #

bankr = []
# for each neighbor (subsection of data)
for name in neighbors:
    # for each processor asigned to that neighbor
    for source in neighbors[name]:
        # blocking test for a message:
        if MPI.COMM_WORLD.probe(source=source):
            received = MPI.COMM_WORLD.recv(source=source)
            print('rank ', mcdc["mpi_rank"], "received shape ", received.shape)
            # bankr.extend(received)

正常运行时的输出 #

rank  0 sent message to  1
rank  0 received shape  (8,)
rank  1 sent message to  2
rank  1 sent message to  0
rank  1 received shape  (49,)
rank  1 received shape  (53,)
rank  2 sent message to  1
rank  2 received shape  (9,)

冻结时的输出 #

rank  0 sent message to  1
rank  0 received shape  (8,)

问题分析与解决建议 #

核心问题 #

1. extend操作的耗时阻塞：从输出看received是numpy数组，bankr.extend(received)会遍历数组的每个元素逐个追加到Python列表。如果数组规模较大，这个操作会占用大量CPU时间，导致当前进程无法及时处理后续MPI通信，引发进程间死锁。比如rank 0在执行extend时卡住，无法响应rank 1的消息，而rank 1可能在等待对应操作完成，最终导致整个程序冻结。

2. MPI通信的时序冲突：代码中用probe阻塞检查指定source的消息，再执行recv。当extend操作占用过多时间时，其他进程已发送的消息无法被及时接收，非阻塞发送（isend）的完成依赖接收方的确认，未被确认的发送会导致发送进程无法继续后续操作，进而引发全局死锁。

解决建议 #

• 改用numpy数组存储数据：避免用Python列表逐个追加元素，预先分配足够大的numpy数组，通过切片赋值或np.concatenate合并接收的数据，大幅降低数据合并的耗时。
• 延迟数据合并操作：先将所有接收的数据暂存到一个数组列表中，等所有MPI通信操作完成后，再一次性合并数据，避免通信过程中被耗时操作打断。
• 优化MPI接收逻辑：去掉不必要的probe检查，直接使用recv；或改用非阻塞接收（irecv）配合wait/waitall，实现通信与数据处理的重叠，提升效率。
• 确保消息全接收：检查所有发送的消息都被对应进程接收，避免因遗漏消息导致进程一直等待而死锁。

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