Python 3.9多进程读取文件时如何获取返回值并按序构造数据

使用Python 3.9实现多进程读取tif格式文件以提升读取速度，当前代码存在两个核心问题：

• 无法获取子进程的返回数据
• 最终拼接所有读取结果时需要保证和原文件列表的顺序一致
  原有实现代码如下：

# read files from file list in the given indices
def read_files(files, folder_path):
    raw_data = []
    # loops through all tif files in the given folder and parses the data.
    for file in files:
        if file[-3:] == "tif":
            curr_frame = Image.open(os.path.join(folder_path, file))
            raw_data.append(np.array(curr_frame))
    return np.asarray(raw_data).astype(np.float64)


def run_processes(folder_path=None):
    if folder_path is None:
        global PATH
        folder_path = PATH
    files = os.listdir(folder_path)

    start = time.time()
    processes = []
    num_files_per = int(len(files) / os.cpu_count())
    for i in range(os.cpu_count()):
        processes.append(Process(target=read_files, args=(files[(i*num_files_per):((i+1)*num_files_per)], folder_path)))
    for process in processes:
        process.start()
    for process in processes:
        process.join()
    end = time.time()
    print(f"Multi: {end - start}")

• 原生Process启动的子进程拥有独立内存空间，子进程内函数的返回值不会自动同步到主进程，没有跨进程通信机制的前提下主进程无法拿到子进程的计算结果
• 原有文件分片逻辑存在缺陷：当文件总数无法被CPU核心数整除时，最后剩余的文件不会被分配到任何子进程，会出现数据遗漏
• 缺少Windows平台多进程的启动保护，在Windows系统下运行会出现递归启动进程的报错

使用标准库multiprocessing.Pool进程池实现，它自带的starmap方法会自动按传入参数的顺序收集子进程返回值，天然满足顺序要求，同时进程池会自动均匀分配任务，不需要手动计算分片大小，也不会遗漏文件。
修复后的完整代码如下：

import os
import time
import numpy as np
from PIL import Image
from multiprocessing import Pool, cpu_count


# 单文件读取逻辑，任务粒度更细，负载更均匀
def read_single_file(file, folder_path):
    if file[-3:] == "tif":
        curr_frame = Image.open(os.path.join(folder_path, file))
        return np.array(curr_frame).astype(np.float64)
    return None


def run_processes(folder_path=None):
    if folder_path is None:
        global PATH
        folder_path = PATH
    files = os.listdir(folder_path)
    # 组装参数列表，保持文件原有顺序
    task_args = [(file, folder_path) for file in files]

    start = time.time()
    # 启动进程池，默认使用全部CPU核心
    with Pool(processes=cpu_count()) as pool:
        # starmap方法自动按任务传入顺序返回结果，天然保序
        results = pool.starmap(read_single_file, task_args)
    # 过滤掉非tif文件返回的None值，拼接为最终数组
    raw_data = np.asarray([res for res in results if res is not None])
    end = time.time()
    print(f"Multi: {end - start}")
    return raw_data


if __name__ == "__main__":
    # 跨平台运行多进程必须加主模块入口保护
    data = run_processes()

方案说明 #

• 进程池的starmap方法会严格按照传入任务的顺序返回结果，不需要额外做排序处理
• 把任务粒度从「一批文件」改成「单个文件」，可以避免不同批次文件大小差异导致的某一个进程运行时间过长、其他进程空等的问题，整体负载更均衡
• with语句会自动处理进程池的关闭、资源回收逻辑，不需要手动调用join/close
• 入口判断if __name__ == "__main__":兼容Windows、macOS、Linux全平台运行要求

如果读取的单文件体积较大，可以把pool.starmap换成pool.istarmap，内存占用会更低，同样可以保证返回顺序。

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