Python:AttributeError:无法pickle本地对象“writeBuf.<locals>.write”

Question

我对Python一点也不熟悉，我通常使用Ruby或JS。但我需要在运行Python的系统上编写一个基准测试脚本。我要做的是创建一个小脚本，获取文件大小和线程数，并编写一个随机缓冲区。这是我摆弄了两个小时后得到的：

from multiprocessing import Pool
import os, sys

def writeBuf(buf):
    def write(n):
        f = open(os.path.join(directory, 'n' + str(n)), 'w')
        try:
            f.write(buf)
            f.flush()
            os.fsync(f.fileno)
        finally:
            f.close()
    return write

if __name__ == '__main__':
    targetDir = sys.argv[1]
    numThreads = int(sys.argv[2])
    numKiloBytes = int(sys.argv[3])
    numFiles = int(102400 / numKiloBytes)

    buf = os.urandom(numKiloBytes * 1024)

    directory = os.path.join(targetDir, str(numKiloBytes) + 'k')
    if not os.path.exists(directory):
        os.makedirs(directory)

    with Pool(processes=numThreads) as pool:
        pool.map(writeBuf(buf), range(numFiles))

但它抛出了错误：AttributeError: Can't pickle local object 'writeBuf.<locals>.write'

我以前尝试过在不使用闭包的情况下使用write，但在尝试在__name__ == '__main__'部分中定义函数时出错。省略if也会导致错误，我读到Pool需要它才能工作。

什么应该只是一个小小的剧本变成一个巨大的考验，谁能指点我正确的方向？

Answer 1

理论上，python不能pickle函数。（有关详细信息，请参见Can't pickle Function）

实际上，python对函数的名称和模块进行pickle，以便传递函数。在您的例子中，您试图传递的函数是由writeBuf返回的局部变量。

相反：

1. 移除writeBuf包装。
2. 不要使用write函数的闭包（buf和directory），而是将所需的一切作为参数给出write。

结果：

def write(args):
    directory, buf, n = args

    with open(os.path.join(directory, 'n' + str(n)), 'w') as f:
        # might as well use with-statements ;)
        f.write(buf)
        f.flush()
        os.fsync(f.fileno)

if __name__ == '__main__':
    ...

    with Pool(processes=numThreads) as pool:
        nargs = [(directory, buf, n) for n in range(numFiles)]
        pool.map(write, nargs)