Metaflow中使用multiprocessing调用fasttext模型遇序列化错误求助

你在Metaflow框架中使用Python的multiprocessing结合fasttext模型执行预测任务时，遇到了TypeError: cannot pickle 'fasttext_pybind.fasttext' object错误。根本原因是fasttext的模型对象基于C++ pybind实现，无法被Python的pickle机制序列化，而你尝试通过multiprocessing.Manager传递模型对象、用pickle.dumps序列化模型的思路，都绕不开序列化这个核心障碍。

• 使用multiprocessing.Manager().list()存储fasttext模型：Manager创建的共享对象需要在进程间传递，底层依赖pickle序列化，而fasttext模型对象不支持序列化，直接触发报错。
• 尝试pickle.dumps(model_abn)：fasttext的pybind对象本身未实现pickle序列化接口，直接序列化必然失败。

最优思路是避免在进程间传递模型对象，改为让每个子进程独立加载模型。利用multiprocessing.Pool的initializer参数，在子进程启动时统一加载模型，既解决序列化问题，又保证执行效率。

修改后的完整代码 #

import pickle
import os
import boto3
import multiprocessing
import time

from metaflow import batch, conda, FlowSpec, step, conda_base
from util import pip_install_module


@conda_base(libraries={'scikit-learn': '0.23.1', 'numpy': '1.22.4', 'pandas': '1.5.1', 'fasttext': '0.9.2'}) 
class BatchInference(FlowSpec):
    pip_install_module("python-dev-tools", "2023.3.24")

    @batch(cpu=10, memory=120000)
    @step
    def start(self):
        self.df_input = ['af', 'febrt' ,'fefv fd we' ,'fe hth dw hytht' ,' dfegrtg hg df reg']
        self.next(self.predict)

    @batch(cpu=10, memory=120000)
    @step
    def predict(self):
        import fasttext
        fasttext.FastText.eprint = lambda x: None

        # 子进程初始化函数：每个子进程启动时加载模型到全局变量
        def init_worker():
            global model_abn
            # 确保模型文件路径正确，若从S3下载需先将模型拉取到本地
            model_abn = fasttext.load_model('fasttext_model.bin')

        print('启动子进程并加载模型')
        time_start = time.time()

        # 创建进程池，指定初始化函数，子进程启动时自动加载模型
        pool = multiprocessing.Pool(initializer=init_worker)
        # 直接传递文本数据，子进程通过全局变量访问模型
        results = pool.map(self.predict_abn, self.df_input)

        pool.close()
        pool.join()

        time_end = time.time()
        print(f"Time elapsed: {round(time_end - time_start, 2)}s")
        self.next(self.end)

    @step
    def end(self):
        print("Predictions evaluated successfully")

    def predict_abn(self, text):
        # 直接使用子进程内的全局模型对象
        return model_abn.predict(text, k=1)


if __name__ == '__main__':
    BatchInference()

关键修改说明 #

1. 移除multiprocessing.Manager：彻底放弃进程间传递模型的思路，避免序列化需求。
2. 添加init_worker初始化函数：每个子进程启动时独立加载模型到全局变量，子进程内部直接访问该变量执行预测。
3. 简化predict_abn方法：不再接收模型参数，减少进程间通信开销。
4. 进程池指定initializer：确保子进程启动时自动执行模型加载逻辑，无需手动传递模型。

额外优化建议 #

• 若模型体积较大，可设置multiprocessing.Pool(initializer=init_worker, maxtasksperchild=100)，限制每个子进程处理的任务数，避免内存持续增长。
• 若模型存储在S3，需在predict步骤开头先将模型下载到本地临时路径，确保init_worker能正确加载文件。

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