使用mpi4py散射pandas dataframe并对大数组执行数学计算的问题

核心问题分析 #

• 你错误地尝试读取chunks变量拿当前进程的任务，而非scatter返回的data变量，非0进程的chunks本身就是None，必然报错
• scatter给每个进程的是列对列表（不是单个列对），需要遍历列表逐个计算
• 不需要单独发送列对应的数据，只要保证所有进程持有相同的df即可，测试场景下固定随机种子就能实现，真实场景可以让所有进程读取共享存储的同一份数据，远高效于广播大数据
• 你第二个逐次Send列对的思路效率极低，scatter批量分发任务的原思路是正确的，不需要改

可运行的正确实现 #

import pandas as pd
import numpy as np
import itertools
import time
from mpi4py import MPI

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

# 固定随机种子，保证所有进程生成的df完全一致，不需要传输整个df
np.random.seed(42)
start_time = time.time()

# 计算函数示例，你可以替换成自己的相干计算逻辑
def myFunc(X,Y):
    return np.corrcoef(X, Y)[0,1]

# 所有进程生成相同的测试df，真实场景替换为读取共享存储文件即可
df = pd.DataFrame(data=np.random.random((500, 34)), columns=[f'Col_{i:03}' for i in range(1, 35)])
columns = df.columns.to_list()

if rank == 0:
    # 生成所有列对
    all_col_pairs = list(itertools.combinations(columns, 2))
    # 均分任务到各个进程
    chunks = [[] for _ in range(size)]
    for idx, pair in enumerate(all_col_pairs):
        chunks[idx % size].append(pair)
else:
    chunks = None

# 分发任务：每个进程拿到自己要计算的列对列表
local_pairs = comm.scatter(chunks, root=0)
print(f'进程{rank}分到{len(local_pairs)}个列对计算任务')

# 本地计算
local_res = {}
for col_a, col_b in local_pairs:
    x = df[col_a].to_numpy()
    y = df[col_b].to_numpy()
    local_res[f"{col_a}_&_{col_b}"] = myFunc(x, y)

# 汇总所有进程的结果到主进程
all_res_list = comm.gather(local_res, root=0)

if rank == 0:
    # 合并所有结果
    total_res = {}
    for res_dict in all_res_list:
        total_res.update(res_dict)
    # 排序输出
    result_df = pd.DataFrame.from_dict(total_res, orient='index', columns=['real_coh'])
    result_sorted = result_df.sort_index()
    print("并行计算完成，结果维度：", result_sorted.shape)
    print("总耗时：", time.time() - start_time)
    # 后续可以加result_sorted.to_csv()等保存逻辑

补充说明 #

如果你的实际场景中df过大，且无法让所有进程各自加载，你可以用广播numpy数组的方式传输数据：

if rank == 0:
    df_np = df.to_numpy()
else:
    df_np = np.empty((500,34), dtype=np.float64)
# 广播整个数组，比传输pandas对象高效很多
comm.Bcast(df_np, root=0)
# 所有进程都可以通过列名对应索引取数据，列名可以提前单独广播一次

执行命令还是你用的mpirun -np 4 python3 file_name.py即可。

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