Python 中的 multiprocessing 用于 CPU 密集型任务

1. 概述

当 Python 代码需要大量计算时，比如处理图像、计算校验和或数值模拟，你会发现使用线程几乎不会更快。原因是 GIL（全局解释器锁）。本文简要解释这个问题。随后指导你使用 multiprocessing 模块和 ProcessPoolExecutor 将工作分散到多个进程。读完后，你将知道如何并行化 CPU 密集型任务、为其编写测试并避免常见错误。

所有示例使用 Python 3.12+ 并且只需要标准库，无需额外安装。

2. 为什么线程不能帮助

2.1. 一个 CPU 密集型任务

我们使用一个用纯 Python 编写的计数质数的函数。这是典型的 CPU 密集任务：

# cpu_task.py
import math

def count_primes(start: int, end: int) -> int:
    count = 0
    for n in range(max(start, 2), end):
        if all(n % d for d in range(2, math.isqrt(n) + 1)):
            count += 1
    return count

该函数遍历 [start, end) 中的每个数字，并尝试除以从 2 到其平方根的数字。该函数不读取文件，不调用网络。它只进行计算。

2.2. 测量时间：顺序和线程

# bench_threads.py
import time
from concurrent.futures import ThreadPoolExecutor
from cpu_task import count_primes

RANGES = [(i, i + 125_000) for i in range(0, 1_000_000, 125_000)]

def run_sequential() -> int:
    return sum(count_primes(a, b) for a, b in RANGES)

def run_threads() -> int:
    starts, ends = zip(*RANGES)
    with ThreadPoolExecutor(max_workers=8) as executor:
        return sum(executor.map(count_primes, starts, ends))

if __name__ == "__main__":
    for fn in (run_sequential, run_threads):
        t0 = time.perf_counter()
        result = fn()
        print(f"{fn.__name__}: {result} primes, {time.perf_counter() - t0:.2f}s")

RANGES 将区间 [0, 1_000_000) 分成 8 部分。run_threads 将每个部分分配给一个线程。8 核机器上的结果（你的机器上的数据会有所不同）：

run_sequential: 78498 primes, 6.41s
run_threads: 78498 primes, 6.58s

线程版本甚至更慢一点。在标准 CPython 中，任何时刻只有一个线程能执行 Python 字节码，因为它必须持有 GIL。八个线程只是轮流使用一个核心，还增加了线程切换的开销。

注意： 从 Python 3.13 开始，有一个无 GIL 的构建版本（python3.13t）。这个版本仍处于实验阶段，许多 C 扩展库还不支持它。对于大多数现有项目，使用进程仍然是最可靠的方法。

3. multiprocessing 基础

每个进程都有自己的解释器和自己的 GIL。因此，N 个进程可以在 N 个核心上真正并行运行。

3.1. 手动创建 Process

# basic_process.py
from multiprocessing import Process, Queue
from cpu_task import count_primes

def worker(start: int, end: int, results: Queue) -> None:
    results.put(count_primes(start, end))

if __name__ == "__main__":
    results = Queue()
    ranges = [(0, 500_000), (500_000, 1_000_000)]
    processes = [Process(target=worker, args=(a, b, results)) for a, b in ranges]

    for p in processes:
        p.start()
    total = sum(results.get() for _ in processes)
    for p in processes:
        p.join()

    print("Total primes:", total)

解释各部分：

• Process(target=..., args=...) 描述一个子进程。它不会运行直到你调用 start()。
• 子进程不与父进程共享内存。要发送结果回来，你必须使用通信通道。这里是 Queue。
• 程序在 join() 之前调用 results.get()。如果先 join() 而队列中还有未读取的大数据，程序可能会发生死锁。

Total primes: 78498

这种方法灵活但冗长。在实践中你很少自己管理 Process。

3.2. 使用 Pool

Pool 创建一个预先存在的 worker 组，并自动将工作分配给它们：

# pool_example.py
from multiprocessing import Pool
from cpu_task import count_primes

RANGES = [(i, i + 125_000) for i in range(0, 1_000_000, 125_000)]

if __name__ == "__main__":
    with Pool(processes=8) as pool:
        counts = pool.starmap(count_primes, RANGES)
    print(counts)
    print("Total:", sum(counts))

starmap 将 RANGES 中的每个元组拆分为 start, end 参数。结果按输入的确切顺序返回：

[10746, 9895, 9622, 9480, 9637, 9554, 9454, 9610]
Total: 78498

如果不传递 processes，Pool 将使用 os.cpu_count() 个 worker。with 块在结束时会清理这些 worker。

4. ProcessPoolExecutor：更现代的 API

concurrent.futures.ProcessPoolExecutor 与 ThreadPoolExecutor 有相同的接口。因此你几乎可以通过仅改变类名来从线程切换到进程：

# parallel.py
import os
from concurrent.futures import ProcessPoolExecutor
from cpu_task import count_primes

def split_range(start: int, end: int, parts: int) -> list[tuple[int, int]]:
    step = (end - start + parts - 1) // parts
    return [(s, min(s + step, end)) for s in range(start, end, step)]

def parallel_count_primes(start: int, end: int, workers: int | None = None) -> int:
    workers = workers or os.cpu_count() or 1
    chunks = split_range(start, end, workers)
    starts, ends = zip(*chunks)
    with ProcessPoolExecutor(max_workers=workers) as executor:
        return sum(executor.map(count_primes, starts, ends))

if __name__ == "__main__":
    import time
    t0 = time.perf_counter()
    print(parallel_count_primes(0, 1_000_000, workers=8))
    print(f"{time.perf_counter() - t0:.2f}s")

78498
1.27s

时间从大约 6.4 秒降低到 1.3 秒。速度不能精确提高 8 倍有两个原因。首先，各个段的工作量不均匀：较大数字的段检查需要更长时间。其次，启动进程和传输数据也需要时间。

如果你想在每个结果完成时立即处理它，请使用 submit 结合 as_completed：

from concurrent.futures import ProcessPoolExecutor, as_completed
from cpu_task import count_primes

if __name__ == "__main__":
    with ProcessPoolExecutor() as executor:
        futures = {executor.submit(count_primes, a, a + 250_000): a
                   for a in range(0, 1_000_000, 250_000)}
        for future in as_completed(futures):
            print(f"开始位置 {futures[future]} 的段：{future.result()} 个质数")

如果 worker 抛出异常，future.result() 会在父进程中重新抛出该异常。因此你可以像平常一样用 try/except 捕捉错误。

5. 许多小任务：chunksize 和 imap_unordered

每次向 worker 发送一个任务时，Python 必须 pickle 参数、通过管道发送、然后 unpickle 结果。如果你有 100,000 个很小的任务，这个开销可能大于计算本身。你可以通过用 chunksize 将任务分组成批处理来解决这个问题：

# chunks.py
from multiprocessing import Pool

def is_prime(n: int) -> bool:
    if n < 2:
        return False
    return all(n % d for d in range(2, int(n ** 0.5) + 1))

if __name__ == "__main__":
    numbers = range(1_000_000)
    with Pool() as pool:
        total = sum(pool.imap_unordered(is_prime, numbers, chunksize=10_000))
    print(total)

• chunksize=10_000 使得每次发送包含 10,000 个数字而不是 1 个。进程间的通信次数大大减少。
• imap_unordered 以惰性方式返回结果且不保持顺序。当你只需要聚合时（例如求和），且想节省内存时，这种方式很合适。
• ProcessPoolExecutor.map 也接受 chunksize 参数。其默认值为 1，所以当有许多小任务时，你需要调整它。

6. 为每个 worker 初始化资源一次

假设每个任务需要一个耗时的查找表来创建。如果在每次函数调用时都重新创建表会很浪费。initializer 参数允许在 worker 启动时恰好运行一次函数：

# initializer_example.py
from concurrent.futures import ProcessPoolExecutor

_price_table: dict[str, float] = {}

def load_prices(currency: str) -> None:
    global _price_table
    rate = 25_000 if currency == "VND" else 1
    _price_table = {f"SKU-{i}": i * 1.5 * rate for i in range(100_000)}

def order_total(skus: list[str]) -> float:
    return sum(_price_table[s] for s in skus)

if __name__ == "__main__":
    orders = [["SKU-1", "SKU-2"], ["SKU-10", "SKU-20", "SKU-30"]]
    with ProcessPoolExecutor(max_workers=2, initializer=load_prices,
                             initargs=("VND",)) as executor:
        print(list(executor.map(order_total, orders)))

[112500.0, 2250000.0]

变量 _price_table 是在每个 worker 中的全局变量。父进程看不到这些更改。这是一种安全的方式，使每个 worker 都有自己的资源，而不必随每个任务发送查找表。

7. 为并行代码编写测试

你可以像测试普通代码一样用 pytest 测试使用进程的代码。条件是 worker 函数必须在模块级别定义：

# test_parallel.py
from cpu_task import count_primes
from parallel import parallel_count_primes, split_range

def test_count_primes_known_value():
    assert count_primes(0, 100) == 25

def test_split_range_covers_whole_interval():
    chunks = split_range(0, 10, 3)
    assert chunks == [(0, 4), (4, 8), (8, 10)]

def test_parallel_matches_sequential():
    expected = count_primes(0, 50_000)
    assert parallel_count_primes(0, 50_000, workers=4) == expected

pip install pytest
pytest -q test_parallel.py

...                                                    [100%]
3 passed in 0.62s

最重要的测试是 test_parallel_matches_sequential。它确保分段不会遗漏或重复计算任何元素。这是自己编写 split_range 时常见的错误。

8. 常见错误和何时不应该使用

忘记 if __name__ == "__main__":。 使用 spawn 启动方法（在 Windows 和 macOS 上的默认方法）或 forkserver（从 Python 3.14 开始在 Linux 上的默认方法），子进程会重新导入主模块。如果创建池的代码在模块级别，每个子进程都会创建一个新池。然后你会遇到 RuntimeError 或进程无限创建。

传递 lambda 或嵌套函数。 发送给 worker 的所有东西都必须是可 pickle 的：

from concurrent.futures import ProcessPoolExecutor

if __name__ == "__main__":
    with ProcessPoolExecutor() as executor:
        print(list(executor.map(lambda x: x * x, range(5))))

_pickle.PicklingError: Can't pickle <function <lambda> at 0x...>: attribute lookup <lambda> on __main__ failed

你应该在模块级别用 def 定义函数。如果需要固定某些参数，使用 functools.partial。

传递过大的数据。 一个包含数百万元素的列表传递给每个任务会被多次 pickle。复制时间可能超过计算部分。有几种解决方法：只发送索引或文件路径、使用 initializer，或将 multiprocessing.shared_memory 用于数组类型的数据。

依赖 fork 的行为。 在旧版 Linux 上能工作的代码，因为 fork（子进程继承全局变量），在 macOS、Windows 或 Python 3.14 上可能会出错。为了获得一致的行为，应明确选择启动方法：

import multiprocessing as mp

if __name__ == "__main__":
    mp.set_start_method("spawn")

注意： 不应该为 I/O 密集型任务（如 HTTP 调用或数据库查询）使用 multiprocessing。在等待 I/O 时，线程仍会释放 GIL，因此 ThreadPoolExecutor 或 asyncio 要轻量得多。当每个任务只需要几微秒而没有分组，或当计算已经在释放 GIL 的 C 库中时（NumPy、大数据哈希函数等），也不应该使用 multiprocessing。在这些情况下，线程或向量化通常就足够了。

创建过多 worker。 设置 worker 数量大于 CPU 核心数不会使 CPU 密集型代码更快。它只会浪费 RAM，因为每个进程都是一个完整的解释器。从 os.cpu_count() 开始是合理的。

9. 结论

GIL 使得线程无法加速 CPU 密集型的纯 Python 代码。每个进程都有自己的 GIL，所以将工作分散到多个进程是充分利用所有 CPU 核心的简单方法。在大多数情况下，ProcessPoolExecutor 是最简洁的选择。当你需要 imap_unordered 或 starmap 时，Pool 很有用。而手动 Process 仅适用于特殊情况。

使用 multiprocessing 时，你需要记住以下几点：

• 始终有 if __name__ == "__main__":。
• 仅传递可 pickle 的函数和数据。
• 当有许多小任务时增加 chunksize。
• 为重型资源使用 initializer。
• 编写一个测试来比较并行结果与顺序结果，以确保分段正确。