类方法中运行multiprocessing遇pickle错误，求解决方案

问题场景 #

在类方法中使用multiprocessing.Pool调用外部任务函数时，触发以下错误：

TypeError: cannot pickle '_thread.RLock' object

任务函数负责调用OpenAI的GPT-4o接口，调用时通过zip(repeat(client), prompts)将主进程中的OpenAI客户端实例传递给子进程。

错误原因 #

OpenAI客户端实例内部包含_thread.RLock（线程锁）这类不可被pickle序列化的对象，而multiprocessing.Pool在向子进程传递参数时，必须对所有参数进行序列化操作，因此无法传递客户端实例。

可行解决方案 #

方案1：在子进程中初始化客户端 #

避免在主进程传递客户端实例，改为传递API密钥，让每个子进程自行初始化客户端：

修改任务函数：

def run_a_request(api_key, prompt, stream=False):
    from openai import OpenAI
    # 子进程内创建客户端
    client = OpenAI(api_key=api_key)
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ],
        temperature=0.1,
        stream=stream
    )
    if not stream:
        response = response.choices[0].message.content
    return response

修改类方法中的调用逻辑：

def run_intermediate_response(self, subgraph, components, question):
    # ... 原有生成prompts的逻辑不变 ...

    pool = Pool(5)
    # 传递API密钥而非客户端实例
    result = pool.map(run_a_request, zip(repeat("你的API密钥"), prompts))

    print(result)
    return 1

方案2：利用进程池初始化函数创建全局客户端 #

通过Pool的initializer参数，在每个子进程启动时初始化客户端，子进程内通过全局变量复用该客户端：

先定义全局变量和初始化函数：

# 全局变量，用于子进程存储客户端实例
_global_client = None

def init_worker(api_key):
    global _global_client
    from openai import OpenAI
    _global_client = OpenAI(api_key=api_key)

修改任务函数：

def run_a_request(prompt, stream=False):
    global _global_client
    response = _global_client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ],
        temperature=0.1,
        stream=stream
    )
    if not stream:
        response = response.choices[0].message.content
    return response

修改类方法中的调用逻辑：

def run_intermediate_response(self, subgraph, components, question):
    # ... 原有生成prompts的逻辑不变 ...

    # 初始化进程池时传入初始化函数和API密钥
    pool = Pool(5, initializer=init_worker, initargs=("你的API密钥",))
    # 仅传递prompt参数即可
    result = pool.map(run_a_request, prompts)

    print(result)
    return 1

方案3：改用异步客户端（适合IO密集场景） #

如果你的场景是高并发IO请求，可以改用OpenAI异步客户端结合asyncio，避免多进程序列化问题：

import asyncio
from openai import AsyncOpenAI

async def run_a_request_async(api_key, prompt, stream=False):
    client = AsyncOpenAI(api_key=api_key)
    response = await client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1,
        stream=stream
    )
    if not stream:
        response = response.choices[0].message.content
    return response

async def batch_requests(api_key, prompts):
    tasks = [run_a_request_async(api_key, prompt) for prompt in prompts]
    return await asyncio.gather(*tasks)

# 在类方法中调用：
def run_intermediate_response(self, subgraph, components, question):
    # ... 原有生成prompts的逻辑不变 ...

    result = asyncio.run(batch_requests("你的API密钥", prompts))
    print(result)
    return 1

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