Python含超大共享数据结构的for循环并行化实现咨询

问题描述 #

需要并行化LS_calculation函数中最外层的for i in range(tensor.dims[0])循环，让每个进程处理部分迭代。涉及的数据结构包括：

• 数百MB级的二维列表uold/vold/wold
• GB级的三维列表ucopy/vcopy/wcopy（维度为numThreads × uold.shape）
• 自定义类实现的超大三维tensor

原始核心代码：

def LS_calculation(ucopy, vcopy, wcopy, uold, vold, wold, tensor, a):
    for i in range(tensor.dims[0]):
        top = [ None for i in range(3) ]
        proc_id = 0
        for j in range(tensor.dims[1]):
            for k in range(tensor.dims[2]):
                tensorAt = tensor.at([i, j, k])
                top[0] = vold[a][j] * wold[a][k] * tensorAt
                ucopy[proc_id][a][i] += top[0]
                top[1] = uold[a][i] * wold[a][k] * tensorAt
                vcopy[proc_id][a][j] += top[1]
                top[2] = uold[a][i] * vold[a][j] * tensorAt
                wcopy[proc_id][a][k] += top[2]
    return

for a in range(A):
    LS_calculation(ucopy, vcopy, wcopy, uold, vold, wold, tensor, approximation)

实现方案 #

1. 重构函数：拆分迭代区间 #

将原外层循环拆分为处理子区间的函数，每个进程负责一个子区间的i值计算，同时通过进程ID定位到ucopy/vcopy/wcopy的对应分区，避免写冲突：

from multiprocessing import current_process

def process_chunk(start_i, end_i, ucopy, vcopy, wcopy, uold, vold, wold, tensor, a):
    # 获取当前进程ID（从0开始）
    proc_id = current_process()._identity[0] - 1
    for i in range(start_i, end_i):
        top = [None] * 3
        for j in range(tensor.dims[1]):
            for k in range(tensor.dims[2]):
                tensorAt = tensor.at([i, j, k])
                
                top[0] = vold[a][j] * wold[a][k] * tensorAt
                ucopy[proc_id][a][i] += top[0]
                
                top[1] = uold[a][i] * wold[a][k] * tensorAt
                vcopy[proc_id][a][j] += top[1]
                
                top[2] = uold[a][i] * vold[a][j] * tensorAt
                wcopy[proc_id][a][k] += top[2]

2. 拆分任务并启动进程池 #

使用multiprocessing.Pool配合starmap传递多参数，将i的迭代范围均分给多个进程：

import os
from multiprocessing import Pool

def LS_calculation_parallel(ucopy, vcopy, wcopy, uold, vold, wold, tensor, a, num_processes=None):
    num_total_i = tensor.dims[0]
    # 默认用CPU核心数，不超过总迭代数
    if num_processes is None:
        num_processes = min(num_total_i, os.cpu_count())
    
    # 拆分迭代区间
    chunk_size = num_total_i // num_processes
    task_chunks = []
    for p_idx in range(num_processes):
        start = p_idx * chunk_size
        # 最后一个进程处理剩余所有迭代
        end = start + chunk_size if p_idx != num_processes - 1 else num_total_i
        # 组装任务参数元组
        task_chunks.append(
            (start, end, ucopy, vcopy, wcopy, uold, vold, wold, tensor, a)
        )
    
    # 启动进程池执行任务
    with Pool(num_processes) as pool:
        pool.starmap(process_chunk, task_chunks)

3. 替换原循环调用 #

将原串行循环替换为并行版本：

for a in range(A):
    LS_calculation_parallel(ucopy, vcopy, wcopy, uold, vold, wold, tensor, a)

4. 关键优化：大内存数据共享 #

由于ucopy/vcopy/wcopy是GB级数据，直接传递会导致每个进程复制一份，引发内存溢出，必须使用共享内存：

方案A：使用NumPy共享内存（推荐，效率更高）

from multiprocessing import shared_memory
import numpy as np

# 假设ucopy是NumPy数组，先创建共享内存
def create_shared_array(arr):
    shm = shared_memory.SharedMemory(create=True, size=arr.nbytes)
    shared_arr = np.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf)
    shared_arr[:] = arr[:]  # 复制数据到共享内存
    return shm, shared_arr

# 为ucopy/vcopy/wcopy创建共享内存
ucopy_shm, ucopy_shared = create_shared_array(np.array(ucopy))
vcopy_shm, vcopy_shared = create_shared_array(np.array(vcopy))
wcopy_shm, wcopy_shared = create_shared_array(np.array(wcopy))

# 修改process_chunk函数，通过共享内存名称获取数组
def process_chunk(start_i, end_i, ucopy_shm_name, ucopy_shape, ucopy_dtype,
                  vcopy_shm_name, vcopy_shape, vcopy_dtype,
                  wcopy_shm_name, wcopy_shape, wcopy_dtype,
                  uold, vold, wold, tensor, a):
    # 连接共享内存
    ucopy_shm = shared_memory.SharedMemory(name=ucopy_shm_name)
    ucopy = np.ndarray(ucopy_shape, dtype=ucopy_dtype, buffer=ucopy_shm.buf)
    # vcopy、wcopy同理
    vcopy_shm = shared_memory.SharedMemory(name=vcopy_shm_name)
    vcopy = np.ndarray(vcopy_shape, dtype=vcopy_dtype, buffer=vcopy_shm.buf)
    wcopy_shm = shared_memory.SharedMemory(name=wcopy_shm_name)
    wcopy = np.ndarray(wcopy_shape, dtype=wcopy_dtype, buffer=wcopy_shm.buf)
    
    # 后续计算逻辑和之前一致
    proc_id = current_process()._identity[0] -1
    ...
    
    # 关闭共享内存连接（无需unlink，主进程负责清理）
    ucopy_shm.close()
    vcopy_shm.close()
    wcopy_shm.close()

# 修改任务组装逻辑，传递共享内存参数
task_chunks.append(
    (start, end,
     ucopy_shm.name, ucopy_shared.shape, ucopy_shared.dtype,
     vcopy_shm.name, vcopy_shared.shape, vcopy_shared.dtype,
     wcopy_shm.name, wcopy_shared.shape, wcopy_shared.dtype,
     uold, vold, wold, tensor, a)
)

# 计算完成后清理共享内存
ucopy_shm.close()
ucopy_shm.unlink()
vcopy_shm.close()
vcopy_shm.unlink()
wcopy_shm.close()
wcopy_shm.unlink()

方案B：使用multiprocessing.RawArray（纯列表场景）

如果必须使用Python列表，可将列表展平为RawArray（共享内存数组），再在子进程中恢复为三维结构：

from multiprocessing import RawArray
import ctypes
import numpy as np

# 将三维列表转为RawArray
def list_to_rawarray(three_d_list, dtype=ctypes.c_double):
    flat_list = [val for proc in three_d_list for a_slice in proc for val in a_slice]
    raw_arr = RawArray(dtype, flat_list)
    shape = (len(three_d_list), len(three_d_list[0]), len(three_d_list[0][0]))
    return raw_arr, shape

# 转换ucopy/vcopy/wcopy
ucopy_raw, ucopy_shape = list_to_rawarray(ucopy)
vcopy_raw, vcopy_shape = list_to_rawarray(vcopy)
wcopy_raw, wcopy_shape = list_to_rawarray(wcopy)

# 修改process_chunk函数
def process_chunk(start_i, end_i, ucopy_raw, ucopy_shape,
                  vcopy_raw, vcopy_shape,
                  wcopy_raw, wcopy_shape,
                  uold, vold, wold, tensor, a):
    # 将RawArray转为NumPy数组再转为列表视图
    ucopy = np.ctypeslib.as_array(ucopy_raw).reshape(ucopy_shape).tolist()
    vcopy = np.ctypeslib.as_array(vcopy_raw).reshape(vcopy_shape).tolist()
    wcopy = np.ctypeslib.as_array(wcopy_raw).reshape(wcopy_shape).tolist()
    
    # 后续计算逻辑不变
    ...

注意事项 #

• 自定义tensor类必须支持序列化（可被pickle处理），或通过共享内存存储其数据，否则子进程无法访问。
• 进程数不宜超过CPU核心数，避免上下文切换开销。
• 本场景中每个进程写自己的proc_id分区，无数据冲突，无需加锁。

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