Python多进程读取npy文件并合并效率低，请求排查代码问题

我有大量独立的npy文件，希望通过多进程读取并合并为单个npy文件以提升处理速度。以下是我编写的代码，但运行效率与单进程读取几乎一致。后续在集群环境测试发现，磁盘I/O速度可达1GB/s，但代码仅能达到200MB/s，怀疑并非磁盘I/O限制，请求帮忙排查代码是否存在问题。

原代码 #

import numpy as np
import os, re
from itertools import repeat
from multiprocessing import Process, Manager, Pool
import multiprocessing


def compressFile(inputs, targets, lock, rootdir, filename):
    with lock:
        data = np.load(os.path.join(rootdir, filename))
        inputs.append(data)
        targets.append('a' in filename) # a second list to store something about the file name


if __name__ == '__main__':

    rootdir = 'directory'

    with Manager() as manager, Pool(multiprocessing.cpu_count()) as pool:

        lock = manager.Lock()
        inputs = manager.list()
        targets = manager.list()
        file_list = [f for f in os.listdir(rootdir) if '.npy' in f]
        pool.starmap(compressFile, zip(repeat(inputs), repeat(targets), repeat(lock), repeat(rootdir), file_list))

        np.save(os.path.join(rootdir, 'inputs.npy'), np.asarray(inputs))
        np.save(os.path.join(rootdir, 'targets.npy'), np.asarray(targets))

问题分析 #

代码效率低下的核心原因是全局锁的滥用：

• 你用lock把整个文件读取和共享列表追加操作完全包裹，导致所有进程必须串行执行——每个进程都要等前一个进程释放锁才能开始读取，完全失去了多进程并行的意义，速度自然和单进程无差异。
• 另外，Manager.list这类跨进程共享对象本身存在额外的通信开销，进一步拖慢了整体速度。

优化方案 #

正确的思路是让每个进程独立读取文件并返回结果，主进程统一收集所有结果后再合并保存，完全不需要共享锁和共享列表：

import numpy as np
import os
from multiprocessing import Pool
import multiprocessing


def process_file(rootdir, filename):
    # 每个进程独立读取文件，返回数据和目标标记
    data = np.load(os.path.join(rootdir, filename))
    target = 'a' in filename
    return data, target


if __name__ == '__main__':
    rootdir = 'directory'
    file_list = [f for f in os.listdir(rootdir) if '.npy' in f]

    # 用Pool并行处理，每个进程返回各自的结果
    with Pool(multiprocessing.cpu_count()) as pool:
        # 用starmap传递参数，每个任务对应一个文件
        results = pool.starmap(process_file, zip(repeat(rootdir), file_list))

    # 拆分结果为inputs和targets
    inputs = np.array([res[0] for res in results])
    targets = np.array([res[1] for res in results])

    # 保存合并后的文件
    np.save(os.path.join(rootdir, 'inputs.npy'), inputs)
    np.save(os.path.join(rootdir, 'targets.npy'), targets)

优化点说明 #

1. 移除全局锁：每个进程独立读取文件，充分利用磁盘并行I/O能力，真正实现多进程加速。
2. 避免跨进程共享对象：子进程直接返回结果，主进程统一收集，消除了共享列表的通信开销。
3. 简化逻辑：去掉不必要的Manager和lock，代码更简洁高效。

修改后应该能充分利用集群的磁盘I/O带宽，大幅提升读取合并的速度。

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