yuqi-zheng
· also available in English

Python 多进程零拷贝：shared_memory 与 Protocol 5 的选择

pythonperformance

本文是 Python 零拷贝序列化系列第二篇。第一篇 深入介绍 PEP 574 和 PickleBuffer API 的内部机制。

Python 多进程通过 pickle 序列化数据在进程间传递。传几个整数没问题；传一个 800MB 的 DataFrame，标准序列化耗时可达 12 秒——计算还没开始，时间已经花光了。

消除这一开销有两个工具：Pickle Protocol 5 和 multiprocessing.shared_memory。它们解决的是不同的问题。

为什么大数据序列化慢

multiprocessing 用 pickle 序列化参数，通过管道或套接字发送字节流。一个 NumPy 数组的旧路径：

  1. __reduce_ex__ 分配数组数据的 bytes 副本
  2. pickle.dumps 将其拷贝到 pickle 流
  3. 流写入管道
  4. Worker 从管道读入缓冲区
  5. pickle.loads 将缓冲区拷贝到新数组

一个数组经历三到四次拷贝。CPU 时间与峰值内存均随数组大小线性增长。

工具一：Pickle Protocol 5（PEP 574）

Protocol 5 将大型二进制数据与元数据流分离。数据以 PickleBuffer（零拷贝视图）传递，而非嵌入 pickle 字节流。

import pickle
import numpy as np

a = np.zeros((10_000, 10_000))  # Լ 800MB

buffers = []
data = pickle.dumps(a, protocol=5, buffer_callback=buffers.append)
# buffers 持有 a 的内存视图，尚无拷贝

b = pickle.loads(data, buffers=buffers)
# b 与 a 共享内存，仍无拷贝

在多进程场景中，Worker 预先分配目标缓冲区（在共享内存或其他位置），作为 buffers 传入反序列化，数组数据直接落入目标，跳过中间拷贝。

NumPy 和 Apache Arrow 已实现带 PickleBuffer 支持的 __reduce_ex__，使用 protocol=5 即可免费获得近零拷贝传输。

适用场景：向 Worker 一次性传递大型对象，不需要写回。

工具二：multiprocessing.shared_memory

shared_memory 将同一块物理内存页映射到多个进程中。没有 pickle，没有管道，没有拷贝——Worker 直接读写数组。

import numpy as np
from multiprocessing import shared_memory, Process

# 主进程：分配并填充共享内存
arr = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
shm = shared_memory.SharedMemory(create=True, size=arr.nbytes)
shared_arr = np.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf)
shared_arr[:] = arr[:]

def worker(name, shape, dtype):
    shm = shared_memory.SharedMemory(name=name)
    work_arr = np.ndarray(shape, dtype=dtype, buffer=shm.buf)
    work_arr[0] += 1   # 对主进程立即可见
    shm.close()

p = Process(target=worker, args=(shm.name, arr.shape, arr.dtype))
p.start()
p.join()

print(shared_arr[0])  # 2.0 —— Worker 的写入已可见

shm.close()
shm.unlink()  # 必须调用一次以释放 OS 资源

跨越进程边界的只有 shm.name——一个短字符串。数组本身从不移动。

适用场景：多个 Worker 需要反复读写同一份大型数据——共享特征矩阵、滚动价格缓冲区、一次加载多次查询的推理模型。

对比

shared_memoryPickle Protocol 5
拷贝开销无，直接内存映射近零，缓冲区视图 + 少量元数据
编程复杂度高，需手动管理生命周期、同步、资源释放低，库支持即透明获得
写回Worker 可写，主进程立即可见单向传输，不支持写回
并发需显式 Lock / Semaphore不适用，每个 Worker 独立副本
风险忘记 unlink() 导致资源泄漏；未同步导致数据竞争版本依赖（需 Protocol 5 环境）
最佳场景持久共享状态，多次访问大型对象一次性传递给 Worker

shared_memory 资源泄漏防范

shm.unlink() 删除 OS 级别的共享内存段。若创建进程崩溃前未调用，该段将持续存在直到重启（Linux 上可在 /dev/shm 中看到）。

用 try/finally 保证清理：

shm = shared_memory.SharedMemory(create=True, size=arr.nbytes)
try:
    # ... 使用 shm ...
finally:
    shm.close()
    shm.unlink()

附加共享内存的 Worker 只需调用 shm.close()，不调用 shm.unlink()——只有创建者负责释放。

并发写入的同步

多个进程同时写同一共享数组，在没有同步机制的情况下会导致数据损坏：

from multiprocessing import Lock

lock = Lock()

def worker(name, shape, dtype, lock):
    shm = shared_memory.SharedMemory(name=name)
    arr = np.ndarray(shape, dtype=dtype, buffer=shm.buf)
    with lock:
        arr[0] += 1
    shm.close()

若各 Worker 操作数组的不重叠区间（分片），则无需加锁——每个 Worker 独占自己的分片。

场景决策指南

场景推荐方案
800MB DataFrame → Worker，返回结果Protocol 5（单向，无需写回）
模型权重一次加载，多 Worker 查询shared_memory（只读，无需加锁）
滚动价格缓冲区，生产者写，消费者读shared_memory + 写时加锁
小型对象（< 几 MB）标准 pickle，开销可忽略
跨机器传输两者均不适用，使用 Arrow IPC、gRPC 或 Ray 对象存储

第三方库参考

框架特点
joblib在 Parallel 中对 NumPy 数组透明使用 shared_memory
ray基于 Plasma 的全局对象存储，本地零拷贝，跨机器 Protocol 5
dask分布式调度 + 惰性计算，支持超内存计算
npshmex对 ProcessPoolExecutor + shared_memory 的轻量封装

Ray 值得单独指出：它使用 Plasma（Apache Arrow 的共享内存存储），本地 Worker 间传递对象完全零拷贝；跨机器传输则使用带外缓冲区，与 Protocol 5 原理相同。

总结

  • 标准 pickle 会将大型数组拷贝 3～4 次，两种工具都能消除大部分甚至全部拷贝开销。
  • Protocol 5 是低摩擦选项：升级 NumPy，指定 protocol=5，一次性 Worker 参数近零拷贝，无需更改业务代码。
  • shared_memory 是真正的零拷贝，但需要手动管理生命周期和并发写入时的同步，适合持久共享状态场景。
  • 大多数科学计算场景下，joblib 或 Ray 已经封装好了复杂性，直接使用即可。

参考资料