跨函数共享内存映射numpy数组时避免段错误的问题求助

问题现象 #

通过multiprocessing.shared_memory实现零拷贝共享numpy ndarray时，出现段错误：在加载数组的get_np()函数内部操作数组（比如np.sum()）完全正常，但将数组返回给主函数后再执行同样的操作就会崩溃。调用copy()可以临时解决问题，但会产生内存拷贝，违背了零拷贝的初衷。

核心原因 #

问题出在**SharedMemory对象的生命周期管理**上：

• 在get_np()函数中，shm是局部变量，当函数执行完毕返回numpy数组时，这个局部的shm对象会被Python的垃圾回收机制自动调用close()方法，导致共享内存的映射被释放。
• 返回的numpy数组只是一个指向共享内存的视图，它本身不持有共享内存的引用。底层内存已经失效，后续对数组的操作自然会触发段错误。

解决方案（零拷贝） #

需要让SharedMemory对象的生命周期与numpy数组保持一致，确保数组被使用期间，共享内存映射不会被释放。以下是两种可行方案：

方案一：同时返回SharedMemory对象和数组 #

在get_np()中同时返回数组和对应的shm对象，在主函数中持有shm的引用，直到不再需要数组时再关闭共享内存。

修改后的代码：

from multiprocessing import shared_memory
import numpy as np

NP_SHARED_NAME = "sharednp"
ARRAY_SHAPE = (1000, )
NP_DATA_TYPE = np.int64

def generate_integer_array(size):
    return np.random.randint(-100, 100, size)

def release_shared(name):
    shm = shared_memory.SharedMemory(name=name, create=False)
    shm.close()
    shm.unlink()

def put_np(data: np.ndarray):
    np_data_type = data.dtype
    np_shape = data.shape
    d_size = np.dtype(np_data_type).itemsize * np.prod(np_shape)
    shm = shared_memory.SharedMemory(create=True, size=d_size, name=NP_SHARED_NAME)
    dst = np.ndarray(shape=np_shape, dtype=np_data_type, buffer=shm.buf)
    dst[:] = data[:]
    return shm.name

def get_np(name):
    shm = shared_memory.SharedMemory(name=name)
    np_array = np.ndarray(ARRAY_SHAPE, dtype=NP_DATA_TYPE, buffer=shm.buf)
    print(np.sum(np_array))
    # 同时返回共享内存对象和数组，避免局部变量被回收
    return np_array, shm

if __name__ == '__main__':
    data = generate_integer_array(1000)
    sm_name = put_np(data)
    array, shm = get_np(sm_name)
    print(np.sum(array))
    # 先关闭当前进程的共享内存映射，再销毁共享内存
    shm.close()
    release_shared(sm_name)

方案二：将SharedMemory对象附加到数组属性 #

把shm对象绑定到numpy数组的自定义属性上，让数组的引用保持shm对象存活，直到数组被回收。

修改后的代码：

from multiprocessing import shared_memory
import numpy as np

NP_SHARED_NAME = "sharednp"
ARRAY_SHAPE = (1000, )
NP_DATA_TYPE = np.int64

def generate_integer_array(size):
    return np.random.randint(-100, 100, size)

def release_shared(name):
    shm = shared_memory.SharedMemory(name=name, create=False)
    shm.close()
    shm.unlink()

def put_np(data: np.ndarray):
    np_data_type = data.dtype
    np_shape = data.shape
    d_size = np.dtype(np_data_type).itemsize * np.prod(np_shape)
    shm = shared_memory.SharedMemory(create=True, size=d_size, name=NP_SHARED_NAME)
    dst = np.ndarray(shape=np_shape, dtype=np_data_type, buffer=shm.buf)
    dst[:] = data[:]
    return shm.name

def get_np(name):
    shm = shared_memory.SharedMemory(name=name)
    np_array = np.ndarray(ARRAY_SHAPE, dtype=NP_DATA_TYPE, buffer=shm.buf)
    print(np.sum(np_array))
    # 将shm对象附加到数组的自定义属性，保持引用
    np_array._shm = shm
    return np_array

if __name__ == '__main__':
    data = generate_integer_array(1000)
    sm_name = put_np(data)
    array = get_np(sm_name)
    print(np.sum(array))
    # 先关闭数组关联的共享内存映射，再销毁
    array._shm.close()
    release_shared(sm_name)

说明 #

两种方案都能实现零拷贝共享，核心是确保SharedMemory对象在数组被使用期间不会被垃圾回收。在多进程场景中，每个访问共享内存的进程都需要创建自己的SharedMemory对象，并且在使用完毕后关闭对应的映射，最后由创建共享内存的进程负责销毁（unlink()）。

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