如何用Ray在Python自定义函数中并行化嵌套for循环

问题描述 #

需要并行化code.py中for p in range(5)的循环段，该循环包含400多行代码，不想重写全部逻辑，且项目依赖Ray库。代码结构如下：

import numpy as np
from tqdm import tqdm
import time
import os
from Auxiliar_py.functions import func1
... # 省略其他导入与代码

constant1=10
constant2=20
loops=fits.open('loops.fits')
loops_list=len(loops)
...

def mainFunction(arg1,arg2,arg3,...):
    var01=constant1+1
    var02=func1(var01)
    ...
    attempt=0
    max_attempts=5
    for i in tqdm(range(len(loops))):
         while attempt < max_attempts:
                try:
                    for p in range(5):           # 目标并行化循环
                        var11=func1(var01)
                        var12=func1(var11)
                        var13=func2(var01,var12)
                        ...                      # 400多行核心代码
                    break  
                except Exception as e:
                    print(f"An error occurred: {e}")
                    attempt += 1
                    print(f"Attempt {attempt} out of {max_attempts}")
    var21=var11+20
    var22=func3(var13)
    ....

可行方案：最小改动实现Ray并行化 #

不需要重写全部代码，只需将循环内的核心逻辑封装为Ray远程函数，仅需整理输入参数和返回值，步骤如下：

1. 初始化Ray #

在代码开头或mainFunction内初始化Ray（若未初始化）：

import ray
ray.init()  # 可根据需求添加参数，如指定CPU核心数

2. 封装循环核心逻辑为远程函数 #

将for p in range(5)内的400多行代码复制到一个新函数中，用@ray.remote装饰，并传入所有循环依赖的外部变量（如var01、constant1等），最后返回后续代码需要用到的变量（如var11、var12、var13）：

@ray.remote(max_retries=0)  # 先关闭自动重试，后续手动处理整体重试逻辑
def process_p(p, var01, constant1, ...):  # 传入所有循环内用到的外部变量
    # 直接复制原循环内的400多行代码
    var11 = func1(var01)
    var12 = func1(var11)
    var13 = func2(var01, var12)
    ...  # 原400多行代码
    return var11, var12, var13  # 返回后续逻辑需要的变量

3. 替换原循环为并行调用 #

将原for p in range(5)循环替换为Ray并行任务提交与结果收集，同时保留原有的重试逻辑：

def mainFunction(arg1,arg2,arg3,...):
    var01=constant1+1
    var02=func1(var01)
    ...
    for i in tqdm(range(len(loops))):
        attempt=0  # 将attempt变量移到i循环内，避免跨i迭代累积
        max_attempts=5
        while attempt < max_attempts:
            try:
                # 提交5个并行任务
                futures = [process_p.remote(p, var01, constant1, ...) for p in range(5)]
                # 等待所有任务完成并获取结果
                results = ray.get(futures)
                
                # 根据原逻辑处理结果：若原循环是保留最后一次p的结果，取最后一个返回值
                var11, var12, var13 = results[-1]
                # 若需要合并所有p的结果，可遍历results处理，比如：
                # all_var11 = [res[0] for res in results]
                
                break  
            except Exception as e:
                print(f"An error occurred: {e}")
                attempt += 1
                print(f"Attempt {attempt} out of {max_attempts}")
    var21=var11+20
    var22=func3(var13)
    ....

4. 注意事项 #

• 变量依赖检查：确保process_p函数包含所有循环内用到的外部变量（如全局常量、var01等），避免因变量未序列化导致的错误。
• 序列化问题：若传入的变量（如fits对象）无法被Ray序列化，需提前提取可序列化的数据（如数组、数值）传入函数。
• 重试逻辑调整：原代码的重试是针对整个for p循环，上述方案保留了该逻辑——若任意一个并行任务失败，将重试整个并行任务组。若需要针对单个任务重试，可修改@ray.remote(max_retries=4)让Ray自动重试单个失败任务。
• 资源控制：若process_p占用大量资源，可通过@ray.remote(num_cpus=1)指定每个任务的CPU核心数，避免资源耗尽。

替代优化：并行化外层for i循环 #

如果for i in range(len(loops))的迭代之间也无依赖，可进一步并行化该外层循环，整体性能提升会更明显，方法类似：将每个i对应的逻辑封装为远程函数，批量提交任务。

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