Python: AttributeError: Can't pickle local object 'writeBuf.<locals>.write'

Python: AttributeError: Can't pickle local object 'writeBuf.<locals>.write'

我对Python一点都不熟悉,一般都是Ruby或者JS。但是我需要在运行 Python 的系统上编写基准测试脚本。我想要做的是创建一个小脚本来获取文件大小和线程数并写入一个随机缓冲区。这是我折腾了 2 个小时后得到的结果:

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 工作所必需的。

什么本来只是一个小脚本变成了一个巨大的考验,谁能指出正确的方法?

理论上,python 不能 pickle 函数。 (详情见Can't pickle Function

实际上,python 会选择函数的名称和模块,以便传递函数。但是,在您的情况下,您尝试传递的函数是 writeBuf.

返回的局部变量

改为:

  1. 删除 writeBuf 包装器。
  2. 不要使用 write 函数的闭包(bufdirectory),而是给 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)