使用mpi4py与scipy.minimize时遭遇死锁问题求助

我尝试实现并行化函数调用的功能，采用类的方式编写代码，核心逻辑如下（简化版）：

from scipy.optimize import minimize
from mpi4py import MPI
import numpy as np

class Solver():
    def __init__(self, SampleTimes, InitialArray):
        self.comm = MPI.COMM_WORLD
        self.rank = self.comm.Get_rank()
        self.size = self.comm.Get_size()
        
        self.SampleTimes  = SampleTimes
        self.InitialArray = InitialArray
        
        self.Finished = False


    def f(x_):
        return <some scalar result>

    def Objective(self, x):

        self.x = self.comm.bcast(x if self.rank == 0 else None, root=0)
        
        tstep_select = np.array_split(self.SampleTimes, self.size)[self.rank]
        
        local_results = []
        for t in tstep_select:
            result = self.f(self.x) # Some function of x 
            local_results.append( (t, result) )    
        
        all_results = self.comm.gather(local_results, root=0)

        if self.rank==0:
            all_results = [item for sublist in all_results for item in sublist]
            all_results = np.array(all_results)
            all_results = all_results[all_results[:,0].argsort()]

            scalar = np.trapz(all_results[:,1], all_results[:,0])
            
            return -scalar

    def Maximize(self,):
        
        if self.rank == 0:
            self.Max = minimize(self.Objective, self.InitialArray)
            self.Finished = self.comm.bcast(True, root=0)
            return self.Max
            
        else:
            while not self.Finished:
                self.Objective(None)


if __name__=='__main__':
    t_eval = np.linspace(0, 100, 100)
    Instance = Solver(SampleTimes=t_eval)
    print(Instance.Maximize())

我清楚只有rank 0进程需要调用minimizer，其他进程要参与Objective函数的计算——这也是Maximize方法里if/else分支的设计目的。我原本认为可以给非0进程的Objective传入None，因为self.x会通过rank 0的广播来设置。但怀疑问题可能出在scipy的minimize要求每次调用Objective都返回单个值，当前非0进程的Objective返回None，不过minimizer只在rank 0执行，所以不确定这是不是问题所在。

编辑（代码v2） #

我移除了"Finished"变量，但仍然出现死锁：

class Solver():
    def __init__(self, SampleTimes, InitialArray):
        self.comm = MPI.COMM_WORLD
        self.rank = self.comm.Get_rank()
        self.size = self.comm.Get_size()
        
        self.SampleTimes  = SampleTimes
        self.InitialArray = InitialArray
      
    def f(x_):
        return <some scalar result>

    def Objective(self, x):

        self.x = self.comm.bcast(x if self.rank == 0 else None, root=0)
        
        tstep_select = np.array_split(self.SampleTimes, self.size)[self.rank]
        
        local_results = []
        for t in tstep_select:
            result = self.f(self.x) # Some function of x 
            local_results.append( (t, result) )    
        
        all_results = self.comm.gather(local_results, root=0)

        if self.rank==0:
            all_results = [item for sublist in all_results for item in sublist]
            all_results = np.array(all_results)
            all_results = all_results[all_results[:,0].argsort()]

            scalar = np.trapz(all_results[:,1], all_results[:,0])
            
            return -scalar

    def Maximize(self,):
        
        if self.rank == 0:
            self.Max = minimize(self.Objective, self.InitialArray)
            return self.Max
            
        else: self.Objective(None)


if __name__=='__main__':
    t_eval = np.linspace(0, 100, 100)
    Instance = Solver(SampleTimes=t_eval)
    print(Instance.Maximize())

死锁原因分析 #

1. 非0进程执行次数不匹配
   Scipy的minimize会多次调用Objective函数（每一轮迭代都要调用），但v2版本中非0进程只调用一次Objective(None)，执行完就退出。后续rank 0发起的广播、gather操作找不到其他进程配合，直接触发死锁。
   第一个版本的Finished逻辑也有问题：rank 0要等minimize完全执行完才广播终止信号，但非0进程在循环中调用Objective时，minimize还在运行，此时rank 0会正常参与通信；但minimize结束后，非0进程的循环还会继续调用Objective，这时候rank 0不再发起通信，导致非0进程卡在bcast步骤，引发死锁。

2. 函数定义错误
   代码中的f方法缺少self参数，调用self.f(self.x)时会抛出异常，中断进程执行，进而导致通信配对失败，触发死锁。

3. 实例变量的状态混乱
   用self.x存储广播后的变量，多进程环境下实例变量的状态无法同步，可能导致非0进程读取到错误的x值，引发计算错误或通信异常。

解决办法 #

修正后的完整代码

from scipy.optimize import minimize
from mpi4py import MPI
import numpy as np

class Solver():
    def __init__(self, SampleTimes, InitialArray):
        self.comm = MPI.COMM_WORLD
        self.rank = self.comm.Get_rank()
        self.size = self.comm.Get_size()
        
        self.SampleTimes  = SampleTimes
        self.InitialArray = InitialArray
      
    def f(self, x_):
        # 替换为实际的标量计算逻辑
        return np.sum(x_)

    def Objective(self, x):
        # 统一广播逻辑：rank0传入当前优化的x，其他进程自动接收
        x_bcast = self.comm.bcast(x, root=0)
        
        # 静态拆分计算任务
        tstep_select = np.array_split(self.SampleTimes, self.size)[self.rank]
        
        local_results = []
        for t in tstep_select:
            result = self.f(x_bcast)
            local_results.append( (t, result) )    
        
        # 收集所有结果到rank0
        all_results = self.comm.gather(local_results, root=0)

        if self.rank == 0:
            # 合并结果并计算积分
            all_results = [item for sublist in all_results for item in sublist]
            all_results = np.array(all_results)
            all_results = all_results[all_results[:,0].argsort()]
            scalar = np.trapz(all_results[:,1], all_results[:,0])
            return -scalar
        # 非0进程无需返回值，直接结束
        return None

    def Maximize(self):
        if self.rank == 0:
            # rank0启动优化流程
            self.Max = minimize(self.Objective, self.InitialArray)
            # 优化完成后广播None作为终止信号
            self.comm.bcast(None, root=0)
            return self.Max
        else:
            # 非0进程持续监听任务，直到收到终止信号
            while True:
                # 接收rank0的广播：None表示任务结束
                x = self.comm.bcast(None, root=0)
                if x is None:
                    break
                # 执行当前迭代的计算任务
                self.Objective(x)
            return None


if __name__=='__main__':
    t_eval = np.linspace(0, 100, 100)
    # 补充初始化数组参数（原代码遗漏）
    init_arr = np.array([1.0, 2.0])
    Instance = Solver(SampleTimes=t_eval, InitialArray=init_arr)
    result = Instance.Maximize()
    # 仅rank0打印结果
    if Instance.rank == 0:
        print(result)

关键修正点

• 修复f方法参数：添加self参数，解决调用时的异常问题。
• 非0进程循环监听：让非0进程进入持续循环，接收rank0每一轮迭代的广播任务，直到收到None终止信号，确保每一次通信都能配对。
• 终止信号机制：rank0在优化完成后广播None，通知非0进程退出循环，避免无意义的等待。
• 避免实例变量存储x：使用局部变量x_bcast存储广播后的参数，避免多进程下的状态混乱。
• 补充遗漏参数：修复原代码中Solver实例化时缺少InitialArray的问题。

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