cProfile 在运行多处理 Python 代码时导致酸洗错误

Question

我有一个正常运行时运行良好的 Python 脚本：

$ python script.py <options>

我正在尝试使用 cProfile 模块分析代码：

$ python -m cProfile -o script.prof script.py <options>

当我启动上面的命令时，我收到一个关于无法pickle函数的错误：

Traceback (most recent call last):
  File "scripts/process_grid.py", line 1500, in <module>
    _compute_write_index(kwrgs)
  File "scripts/process_grid.py", line 626, in _compute_write_index
    args,
  File "scripts/process_grid.py", line 1034, in _parallel_process
    pool.map(_apply_along_axis_palmers, chunk_params)
  File "/home/james/miniconda3/envs/climate/lib/python3.6/multiprocessing/pool.py", line 266, in map
    return self._map_async(func, iterable, mapstar, chunksize).get()
  File "/home/james/miniconda3/envs/climate/lib/python3.6/multiprocessing/pool.py", line 644, in get
    raise self._value
  File "/home/james/miniconda3/envs/climate/lib/python3.6/multiprocessing/pool.py", line 424, in _handle_tasks
    put(task)
  File "/home/james/miniconda3/envs/climate/lib/python3.6/multiprocessing/connection.py", line 206, in send
    self._send_bytes(_ForkingPickler.dumps(obj))
  File "/home/james/miniconda3/envs/climate/lib/python3.6/multiprocessing/reduction.py", line 51, in dumps
    cls(buf, protocol).dump(obj)
_pickle.PicklingError: Can't pickle <function _apply_along_axis_palmers at 0x7fe05a540b70>: attribute lookup _apply_along_axis_palmers on __main__ failed

该代码使用多处理，我假设这是进行酸洗的地方。

正在使用的代码在GitHub 上。

本质上，我在进程池中映射一个函数和一个相应的参数字典：

pool.map(_apply_along_axis_palmers, chunk_params)

_apply_along_axis_palmers据我所知，该函数是“可挑选的”，因为它是在模块的顶层定义的。同样，在 cProfile 上下文之外运行时不会发生此错误，所以这可能为酸洗添加了额外的约束？

任何人都可以评论为什么会发生这种情况，和/或我如何纠正这个问题？

Answer 1

您在这里遇到的问题是，通过使用-mcProfile，模块__main__是cProfile（代码的实际入口点），而不是您的脚本。cProfile 试图通过确保解决这一问题时，你的脚本运行时，它看到__name__的"__main__"，所以它知道它正在运行的脚本，而不是导入为模块，但sys.modules['__main__']仍然是cProfile模块。

问题是，pickle通过酸洗它们的限定名称来处理酸洗函数（加上一些样板，首先说明它是一个函数）。并且为了确保它能够在往返过程中存活下来，它总是仔细检查合格的名称是否可以在sys.modules. 因此，当您这样做时pickle.dumps(_apply_along_axis_palmers)（在这种情况下通过将其作为映射器函数显式传递或隐式传递），在_apply_along_axis_palmers主脚本中定义where 时，它会仔细检查是否sys.modules['__main__']._apply_along_axis_palmers存在。但它不cProfile._apply_along_axis_palmers存在，因为不存在。

我不知道对此有什么好的解决方案。我能想到的最好方法是手动修复sys.modules以使其正确公开您的模块及其内容。我还没有完全测试过这个，所以可能会有一些怪癖，但我发现的一个解决方案是更改一个名为mymodule.py以下形式的模块：

# imports...
# function/class/global defs...

if __name__ == '__main__':
    main()  # Or series of statements

到：

# imports...
import sys
# function/class/global defs...

if __name__ == '__main__':
    import cProfile
    # if check avoids hackery when not profiling
    # Optional; hackery *seems* to work fine even when not profiling, it's just wasteful
    if sys.modules['__main__'].__file__ == cProfile.__file__:
        import mymodule  # Imports you again (does *not* use cache or execute as __main__)
        globals().update(vars(mymodule))  # Replaces current contents with newly imported stuff
        sys.modules['__main__'] = mymodule  # Ensures pickle lookups on __main__ find matching version
    main()  # Or series of statements

从那以后，sys.modules['__main__']指的是你自己的模块，而不是cProfile，所以事情似乎有效。cProfile尽管如此，似乎仍然可以工作，并且酸洗会按预期找到您的功能。只有真正的成本是重新导入你的模块，但如果你做了足够多的实际工作，重新导入的成本应该相当小。