Python Multiprocessing Pool无法满负载利用CPU问题及优化需求

核心问题分析 #

你的代码当前存在两个关键瓶颈：

1. 阻塞式批量处理：调用p.map时主进程会完全阻塞，停止读取CSV数据，导致子进程处理完当前批次后，必须等待主进程攒够下一批2000条才能继续工作，中间出现CPU空闲窗口，后期负载下降。
2. 未处理剩余数据：代码中没有处理最后一批不足2000条的aqiRequests，会导致部分数据丢失（若未额外补全逻辑）。

具体优化措施 #

1. 改用非阻塞的p.imap_unordered替代p.map #

map是阻塞调用，必须等所有子进程处理完整个批次才返回结果；而imap_unordered是迭代式的，子进程处理完单个任务就返回结果，主进程可同时继续读取CSV并提交新任务，让CPU始终处于忙碌状态。

优化后代码片段：

from multiprocessing import Pool, cpu_count
import csv

def doBigCalculation(req):
    # 替换为你的计算逻辑，返回可写入CSV的结果行
    param1, param2 = req
    # 这里是你的资源密集型计算
    result_val = ...
    return [result_val]  # 适配CSV writer的行格式

def main():
    p = Pool(cpu_count())
    batch_size = 300  # 4核CPU推荐200-400，可根据测试调整
    aqiRequests = []

    # 一次性打开输入输出文件，减少IO开销
    with open('input.csv', newline='', buffering=1024*1024) as inputFile, \
         open('output.csv', 'w', newline='', buffering=1024*1024) as outputFile:
        reader = csv.reader(inputFile, delimiter=',', quotechar='"')
        writer = csv.writer(outputFile, delimiter=',', quotechar='"')

        # 边读数据边提交任务，同时写入结果
        for row in reader:
            param1 = float(row[0])
            param2 = float(row[1])
            # 用元组替代自定义类，减少内存占用
            aqiRequests.append((param1, param2))
            
            if len(aqiRequests) >= batch_size:
                # 非阻塞迭代获取结果，主进程可继续读数据
                for result in p.imap_unordered(doBigCalculation, aqiRequests):
                    writer.writerow(result)
                aqiRequests = []
        
        # 处理最后一批剩余数据
        if aqiRequests:
            for result in p.imap_unordered(doBigCalculation, aqiRequests):
                writer.writerow(result)
    
    p.close()
    p.join()

2. 调整批量大小，匹配CPU核心数 #

当前2000条的批量过大，会导致主进程阻塞时间过长；批量过小则会增加进程调度开销。建议将批量大小设置为CPU核心数的50-100倍（4核推荐200-400），平衡调度开销和数据准备时间。可通过测试不同数值（100、300、500）找到最优值。

3. 优化IO性能 #

• 给文件打开添加buffering参数（如1024*1024即1MB缓冲区），减少磁盘IO的系统调用次数。
• 用元组替代AQIRequest自定义类，降低每个任务的内存占用，避免内存溢出导致的swap交换拖慢速度。
• 不要频繁打开/关闭输出文件，全程保持输出文件打开状态。

4. 验证与调优 #

运行时用htop或top监控CPU使用率：

• 若仍出现负载下降，可能是磁盘IO瓶颈，可将CSV文件转移到SSD存储。
• 若内存占用过高，可进一步缩小批量大小，或在处理完批次后强制触发垃圾回收（import gc; gc.collect()）。

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