如何在FastAPI启动时初始化不可Pickle对象并供多进程池函数调用

问题描述 #

在FastAPI项目中使用多进程池处理任务时遇到以下问题：

• 当clientObject作为全局变量提前声明时，代码可正常运行
• 改为通过FastAPI的start_up事件将客户端存入app.state后，main()函数能正常访问app.state.clientObject，但多进程池中的get_response函数报错：'State' object has no attribute 'clientObject'
• 由于clientObject不可被Pickle序列化，无法直接作为参数传入子进程函数

原全局变量版本核心代码：

from multiprocessing import Pool
from functools import partial
from fastapi import FastAPI

clientObject = package.client() # 对象不可被Pickle序列化
    
app = FastAPI()

# 省略wrapper、get_response、main及接口定义

修改为app.state存储后的核心代码（报错版本）：

@app.on_event("start_up")
def start_connection():
    app.state.clientObject = package.client()

def get_response(data_point): 
     data_point = some_other_processes(data_point)
     ans = app.state.clientObject.process(data_point) # 此处报错
     return ans

原因分析 #

多进程池的子进程通过父进程fork生成，FastAPI的start_up事件仅在父进程中执行，子进程的app.state并未经过初始化流程，因此在子进程中访问app.state.clientObject会触发属性不存在的错误。同时因为客户端对象不可序列化，无法通过进程间通信传递给子进程。

解决方案 #

使用multiprocessing.Pool的initializer参数，在每个子进程启动时单独初始化客户端对象，避免跨进程传递非可序列化对象的问题。优化后的代码如下：

def initialize_workers():
    global clientObject
    clientObject = package.client()

def get_response(data_point):
    global clientObject
    data_point = some_other_processes(data_point)
    ans = clientObject.process(data_point)
    return ans

def main(raw_data):
    list_data_point = preprocess(raw_data)
    with Pool(initializer=initialize_workers) as pool:
        wrapped_workload = partial(wrapper, 
                                   get_response,
                                   3)
        results = pool.map(wrapped_workload, list_data_point)
        pool.close()
        pool.join()
    return results

方案说明 #

• initializer=initialize_workers指定了子进程启动时必须执行的初始化函数
• 每个子进程都会独立调用initialize_workers，在自身进程空间内创建clientObject全局变量，完全避免跨进程传递对象的需求
• 子进程的get_response直接访问自身进程内的全局客户端对象，不会出现序列化失败或对象未初始化的问题

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