python 多处理 - OverflowError ('cannot serialize a bytes object larger than 4GiB' )

Question

我们正在使用 multiprocessing 库 (python 3.6) 运行一个脚本，其中一个大的 pd.DataFrame 作为参数传递给一个函数:

from multiprocessing import Pool
import time 

def my_function(big_df):
    # do something time consuming
    time.sleep(50)

if __name__ == '__main__':
    with Pool(10) as p:
        res = {}
        output = {}
        for id, big_df in some_dict_of_big_dfs:
            res[id] = p.apply_async(my_function,(big_df ,))
        output = {id : res[id].get() for id in id_list}

问题是我们从 pickle 库中得到一个错误。

Reason: 'OverflowError('cannot serialize a bytes objects larger than4GiB',)'

我们知道 pickle v4 可以序列化更大的对象 question related , link , 但我们不知道如何修改 multiprocessing 使用的协议(protocol)。

有人知道怎么办吗？谢谢!!

Answer 1

显然有一个开放的issue关于这个话题，并且有一些关于这个特定的相关倡议的描述answer .我找到了一种方法来更改基于此 answer 的 multiprocessing 库中使用的默认 pickle 协议(protocol).正如评论中指出的那样，此解决方案仅适用于 Linux 和操作系统多处理库

解决方法:

你首先创建一个新的分离模块

pickle4reducer.py

from multiprocessing.reduction import ForkingPickler, AbstractReducer

class ForkingPickler4(ForkingPickler):
    def __init__(self, *args):
        if len(args) > 1:
            args[1] = 2
        else:
            args.append(2)
        super().__init__(*args)

    @classmethod
    def dumps(cls, obj, protocol=4):
        return ForkingPickler.dumps(obj, protocol)


def dump(obj, file, protocol=4):
    ForkingPickler4(file, protocol).dump(obj)


class Pickle4Reducer(AbstractReducer):
    ForkingPickler = ForkingPickler4
    register = ForkingPickler4.register
    dump = dump

然后，在您的主脚本中，您需要添加以下内容:

import pickle4reducer
import multiprocessing as mp
ctx = mp.get_context()
ctx.reducer = pickle4reducer.Pickle4Reducer()

with mp.Pool(4) as p:
    # do something

这可能会解决溢出的问题。

但是，警告，您可能会考虑阅读 this在做任何事情之前 否则你可能会遇到和我一样的错误:

'i' format requires -2147483648 <= number <= 2147483647

(此错误的原因在上面的 link 中有很好的解释)。长话短说，multiprocessing 使用 pickle 协议(protocol)通过其所有进程发送数据，如果您已经达到 4gb 限制，那可能意味着您可能会考虑将函数重新定义为“无效”方法，而不是输入/输出方法。所有这些入站/出站数据都会增加 RAM 的使用，可能由于构造(我的情况)而效率低下，最好将所有进程指向同一个对象，而不是为每次调用创建一个新副本。

希望对您有所帮助。