如何对imap_unordered耗时的Python代码进行性能分析？

问题背景 #

这段Linux环境下的Python代码会生成一个100MB的numpy数组，随后通过imap_unordered执行多进程任务，但实际仅做简单的数值平方计算。代码运行缓慢且每次调用_square函数的耗时基本一致，每次调用会输出一个"."，需要对其性能进行分析定位瓶颈。

待分析代码 #

from multiprocessing import Pool
import numpy as np
from pympler.asizeof import asizeof


class ParallelProcessor:
    def __init__(self, num_processes=None):
        self.vals = np.random.random((3536, 3536))
        print("Size of array in bytes", asizeof(self.vals))

    def _square(self, x):
        print(".", end="", flush=True)
        return x * x

    def process(self, data):
        """
        Processes the data in parallel using the square method.
        :param data: An iterable of items to be squared.
        :return: A list of squared results.
        """
        with Pool(1) as pool: 
            for result in pool.imap_unordered(self._square, data):
                # print(result)
                pass 

if __name__ == "__main__":
    # Create an instance of the ParallelProcessor
    processor = ParallelProcessor()

    # Input data
    data = range(1000)

    # Run the processing in parallel
    processor.process(data)

性能分析方法 #

1. 排查进程间序列化开销 #

当使用多进程调用类实例方法时，Python会将整个ParallelProcessor实例（包括100MB的vals数组）序列化后传递给子进程，这会产生巨大的通信开销。

• 验证方法：将_square改为独立函数（脱离类），或者把vals改为类变量，重新运行代码对比耗时。如果耗时大幅降低，说明序列化大对象是核心瓶颈。
• 辅助验证：用pickle.dumps(processor)测试序列化整个实例的耗时和大小，直观看到开销规模。

2. 单函数耗时拆解 #

精准统计_square每次调用的实际耗时，区分是函数计算本身还是进程通信的开销：

• 修改_square函数加入时间统计：

  import time
  def _square(self, x):
      start = time.perf_counter()
      print(".", end="", flush=True)
      res = x * x
      print(f"\nCall took: {time.perf_counter() - start:.6f}s", end="")
      return res
  

• 观察输出的耗时，如果耗时主要集中在函数执行前/后，说明是进程通信的开销；如果是函数内部，则需要排查计算逻辑（本例中计算逻辑极简单，大概率是前者）。

3. 系统级监控工具 #

用Linux原生工具定位系统层面的瓶颈：

• htop/top：观察子进程的CPU、内存占用，看是否存在内存拷贝导致的CPU等待，或者进程频繁启停的开销。
• strace：跟踪系统调用，执行strace -p <进程PID>，查看是否有大量read/write调用（对应进程间通信IO），或内存分配相关系统调用耗时过长。
• perf：采样分析进程时间分布，执行perf record -g -p <进程PID>，之后用perf report查看函数调用栈的耗时占比，精准定位时间消耗环节。

4. 多进程API对比验证 #

当前代码用了Pool(1)（单进程）和imap_unordered，可以换用其他多进程API对比，排查API本身的开销：

• 换成Pool.map或者Pool.apply，对比耗时差异。
• 改用concurrent.futures.ProcessPoolExecutor实现相同逻辑，看是否有性能差异，判断是否是imap_unordered的迭代开销导致。

5. 内存使用分析 #

用内存分析工具查看进程内存变化，确认是否存在不必要的内存拷贝：

• 用pympler的Tracker跟踪实例在主进程和子进程中的内存占用，看子进程是否加载了整个100MB的vals数组。
• 用memory_profiler装饰process方法，查看内存使用的变化曲线，定位内存开销大的环节。

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